update release

Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>
Fix Qwen3-TTS gradio demo (#2372 )
2026-03-31 20:39:22 +08:00 · 2026-03-31 11:42:09 +00:00 · 2026-03-31 09:51:35 +00:00 · 2026-03-31 16:15:47 +08:00 · 2026-03-31 14:59:00 +08:00 · 2026-03-31 14:30:12 +08:00
107 changed files with 5966 additions and 2428 deletions
@@ -64,15 +64,22 @@ while true; do
 done

 echo "--- Pulling container"
-image_name="public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:${BUILDKITE_COMMIT}-rocm-omni"
+## Temporary change to use AMD Docker Hub to store the vllm-ci image
+# to bypass the rate limit issue with ECR Public Gallery.
+# TODO: @tjtanaa point back to ECR Public Gallery
+# once the amd agents are configured to use ECR Public Gallery.
+# image_name="public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:${BUILDKITE_COMMIT}-rocm-omni"
+image_name="rocm/vllm-ci:${BUILDKITE_COMMIT}-rocm-omni"
 container_name="rocm_${BUILDKITE_COMMIT}_$(tr -dc A-Za-z0-9 < /dev/urandom | head -c 10; echo)"

-# Install AWS CLI to authenticate to ECR Public Gallery to get higher rate limit for pulling images
-sudo apt-get update && sudo apt-get install -y awscli
-# Use safe docker login helper to prevent race conditions
-source "$(dirname "${BASH_SOURCE[0]}")/../docker_login_ecr_public.sh"
-safe_docker_login_ecr_public
-# Pull the container from ECR Public Gallery
+# TODO: @tjtanaa uncomment this once the amd agents are configured to use ECR Public Gallery.
+# # Install AWS CLI to authenticate to ECR Public Gallery to get higher rate limit for pulling images
+# sudo apt-get update && sudo apt-get install -y awscli
+##  Use safe docker login helper to prevent race conditions
+# source "$(dirname "${BASH_SOURCE[0]}")/../docker_login_ecr_public.sh"
+# safe_docker_login_ecr_public
+
+## Pull the container from AMD Docker Hub

 docker pull "${image_name}"

@@ -126,8 +133,6 @@ if [[ $commands == *"--shard-id="* ]]; then
        -e VLLM_ROCM_USE_AITER=0 \
        -e HIP_VISIBLE_DEVICES="${GPU}" \
        -e HF_TOKEN \
-        -e AWS_ACCESS_KEY_ID \
-        -e AWS_SECRET_ACCESS_KEY \
        -v "${HF_CACHE}:${HF_MOUNT}" \
        -e "HF_HOME=${HF_MOUNT}" \
        -e "PYTHONPATH=${MYPYTHONPATH}" \
@@ -160,8 +165,6 @@ else
          -e MIOPEN_DEBUG_CONV_GEMM=0 \
          -e VLLM_ROCM_USE_AITER=0 \
          -e HF_TOKEN \
-          -e AWS_ACCESS_KEY_ID \
-          -e AWS_SECRET_ACCESS_KEY \
          -v "${HF_CACHE}:${HF_MOUNT}" \
          -e "HF_HOME=${HF_MOUNT}" \
          -e "PYTHONPATH=${MYPYTHONPATH}" \
@@ -134,15 +134,33 @@ steps:
    - timeout 10m pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
    - timeout 20m pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "advanced_model" --run-level "advanced_model"

- label: "Qwen3-TTS E2E Test"
+- label: "Qwen3-TTS CustomVoice E2E Test"
  agent_pool: mi325_2
  depends_on: amd-build
  mirror_hardwares: [amdproduction]
  grade: Blocking
  commands:
-    - export VLLM_LOGGING_LEVEL=DEBUG
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - timeout 20m pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py
+    - |
+      timeout 20m bash -c '
+        export VLLM_LOGGING_LEVEL=DEBUG
+        export VLLM_WORKER_MULTIPROC_METHOD=spawn
+        export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
+        pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m "advanced_model" --run-level "advanced_model" && pytest -s -v tests/e2e/offline_inference/test_qwen3_tts_customvoice.py
+      '
+
+- label: "Qwen3-TTS Base E2E Test"
+  agent_pool: mi325_2
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - |
+      timeout 20m bash -c '
+        export VLLM_LOGGING_LEVEL=DEBUG
+        export VLLM_WORKER_MULTIPROC_METHOD=spawn
+        export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
+        pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py -m "advanced_model" --run-level "advanced_model" && pytest -s -v tests/e2e/offline_inference/test_qwen3_tts_base.py
+      '

 - label: "Diffusion Image Edit Test"
  agent_pool: mi325_1
@@ -124,7 +124,8 @@ steps:
  commands:
    - export VLLM_LOGGING_LEVEL=DEBUG
    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - timeout 20m pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py
+    - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
+    - timeout 20m pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m "core_model" --run-level "core_model"

 - label: "Diffusion Image Edit Test"
  agent_pool: mi325_1
@@ -208,6 +208,52 @@ steps:
          volumes:
            - "/fsx/hf_cache:/fsx/hf_cache"

+  - label: "Qwen3-TTS CustomVoice E2E Test"
+    depends_on: upload-merge-pipeline
+    commands:
+      - |
+        timeout 20m bash -c '
+          export VLLM_LOGGING_LEVEL=DEBUG
+          export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
+          pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m "advanced_model" --run-level "advanced_model" && pytest -s -v tests/e2e/offline_inference/test_qwen3_tts_customvoice.py
+        '
+    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"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Qwen3-TTS Base E2E Test"
+    depends_on: upload-merge-pipeline
+    commands:
+      - |
+        timeout 20m bash -c '
+          export VLLM_LOGGING_LEVEL=DEBUG
+          export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
+          pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py -m "advanced_model" --run-level "advanced_model" && pytest -s -v tests/e2e/offline_inference/test_qwen3_tts_base.py
+        '
+    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"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
  - label: "Omni Model Test with H100"
    timeout_in_minutes: 30
    depends_on: upload-merge-pipeline
@@ -50,13 +50,15 @@ steps:
                  path: /mnt/hf-cache
                  type: DirectoryOrCreate

-  - label: ":full_moon: Omni Model Test"
-    timeout_in_minutes: 60
+  - label: ":full_moon: Omni Model Test with L4"
+    timeout_in_minutes: 90
    depends_on: upload-nightly-pipeline
    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
    commands:
      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - pytest -s -v tests/examples/online_serving/test_qwen2_5_omni.py -m "advanced_model" --run-level "advanced_model"
+      - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
+      - pytest -s -v tests/examples/ -m "advanced_model and L4 and omni" --run-level "advanced_model"
+      - pytest -s -v tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and L4 and omni" --run-level "advanced_model"
    agents:
      queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
    plugins:
@@ -217,29 +219,6 @@ steps:
                  path: /mnt/hf-cache
                  type: DirectoryOrCreate

-  - label: ":full_moon: Qwen3-TTS Non-Async-Chunk E2E Test"
-    timeout_in_minutes: 30
-    depends_on: upload-nightly-pipeline
-    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
-    commands:
-      - |
-        huggingface-cli download Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice
-        export VLLM_WORKER_MULTIPROC_METHOD=spawn
-        pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py::TestQwen3TTSNoAsyncChunk
-    agents:
-      queue: "gpu_4_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: ":full_moon: Omni Model Perf Test & Test Case Statistics"
    key: nightly-omni-performance
    timeout_in_minutes: 180
@@ -248,6 +227,7 @@ steps:
    commands:
      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
      - export BENCHMARK_DIR=tests/dfx/perf/results
+      - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
      - pytest -s -v tests/dfx/perf/scripts/run_benchmark.py
      - buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
      - python tools/nightly/buildkite_testcase_statistics.py -o tests/dfx/perf/results/buildkite_testcase_statistics.html
@@ -375,6 +355,48 @@ steps:
                  path: /mnt/hf-cache
                  type: DirectoryOrCreate

+  - label: ":full_moon: Wan22 I2V Accuracy Test with H100"
+    key: nightly-wan22-i2v-accuracy
+    timeout_in_minutes: 180
+    depends_on: upload-nightly-pipeline
+    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
+    commands:
+      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - pytest -s -v tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py --run-level advanced_model
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:c392ce21e9cf9ea65c52b866447793db10e0261c
+                resources:
+                  limits:
+                    nvidia.com/gpu: 2
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: hf-cache
+                    mountPath: /root/.cache/huggingface
+                env:
+                  - name: HF_HOME
+                    value: /root/.cache/huggingface
+                  - 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: Qwen-Image Diffusion Perf Test with H100"
    key: nightly-qwen-image-performance
    timeout_in_minutes: 300
@@ -216,44 +216,42 @@ steps:
          volumes:
            - "/fsx/hf_cache:/fsx/hf_cache"

-  # - label: "Omni Model Test with H100"
-  #   depends_on: upload-ready-pipeline
-  #   commands:
-  #     - |
-  #       timeout 20m bash -c '
-  #         export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  #         export VLLM_TEST_CLEAN_GPU_MEMORY="1"
-  #         # - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
-  #         pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "core_model" --run-level "core_model"
-  #       '
-  #   agents:
-  #     queue: "mithril-h100-pool"
-  #   plugins:
-  #     - kubernetes:
-  #         podSpec:
-  #           containers:
-  #             - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-  #               resources:
-  #                 limits:
-  #                   nvidia.com/gpu: 2
-  #               volumeMounts:
-  #                 - name: devshm
-  #                   mountPath: /dev/shm
-  #                 - name: hf-cache
-  #                   mountPath: /root/.cache/huggingface
-  #               env:
-  #                 - name: HF_HOME
-  #                   value: /root/.cache/huggingface
-  #           nodeSelector:
-  #             node.kubernetes.io/instance-type: gpu-h100-sxm
-  #           volumes:
-  #             - name: devshm
-  #               emptyDir:
-  #                 medium: Memory
-  #             - name: hf-cache
-  #               hostPath:
-  #                 path: /mnt/hf-cache
-  #                 type: DirectoryOrCreate
+  - label: "Omni Model Test with H100"
+    depends_on: upload-ready-pipeline
+    commands:
+      - |
+        timeout 20m bash -c '
+          export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "core_model" --run-level "core_model"
+        '
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                resources:
+                  limits:
+                    nvidia.com/gpu: 2
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: hf-cache
+                    mountPath: /root/.cache/huggingface
+                env:
+                  - name: HF_HOME
+                    value: /root/.cache/huggingface
+            nodeSelector:
+              node.kubernetes.io/instance-type: gpu-h100-sxm
+            volumes:
+              - name: devshm
+                emptyDir:
+                  medium: Memory
+              - name: hf-cache
+                hostPath:
+                  path: /mnt/hf-cache
+                  type: DirectoryOrCreate

  - label: "MiMo-Audio E2E Test with H100"
    depends_on: upload-ready-pipeline
@@ -303,13 +301,13 @@ steps:
    commands:
      - |
        timeout 20m bash -c '
-          huggingface-cli download Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice
          export VLLM_LOGGING_LEVEL=DEBUG
          export VLLM_WORKER_MULTIPROC_METHOD=spawn
-          pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py -k "not NoAsyncChunk"
+          export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
+          pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m "core_model" --run-level "core_model"
        '
    agents:
-      queue: "gpu_4_queue"
+      queue: "gpu_1_queue"
    plugins:
      - docker#v5.2.0:
          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
@@ -3,7 +3,7 @@
   Last synced: 2025-12-15
   Modifications: Removed unused CUDA/NVIDIA logic, keeping only AMD tests
 #}
-{% set docker_image_amd = "public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT-rocm-omni" %}
+{% set docker_image_amd = "rocm/vllm-ci:$BUILDKITE_COMMIT-rocm-omni" %}
 {% set default_working_dir = "/app/vllm-omni" %}

  - group: "AMD Tests"
@@ -13,7 +13,6 @@
        depends_on: ~
        soft_fail: false
        commands:
-          - "source .buildkite/scripts/docker_login_ecr_public.sh && safe_docker_login_ecr_public"
          - "docker build -f docker/Dockerfile.rocm -t {{ docker_image_amd }} --target test --build-arg ARG_PYTORCH_ROCM_ARCH=gfx942 --progress plain ."
          - "docker push {{ docker_image_amd }}"
        key: "amd-build"
@@ -22,13 +21,15 @@
        retry:
          automatic:
            - exit_status: -1  # Agent was lost
-              limit: 1
+              limit: 2
            - exit_status: -10  # Agent was lost
-              limit: 1
+              limit: 2
+            - exit_status: 128  # Git connectivity issues
+              limit: 2
            - exit_status: 1  # Machine occasionally fail
              limit: 1
        agents:
-          queue: cpu_queue_premerge
+          queue: amd-cpu

    {% for step in steps %}
    {% if step.mirror_hardwares and mirror_hw in step.mirror_hardwares %}
@@ -16,7 +16,7 @@ Easy, fast, and cheap omni-modality model serving for everyone
 ---

 *Latest News* 🔥
-
+- [2026/03] We released [0.18.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.18.0) - strengthens the core runtime through a large entrypoint refactor and scheduler/runtime cleanups, expands unified quantization and diffusion execution, broadens multimodal model coverage, and improves production readiness across audio, omni, image, video, RL, and multi-platform deployments.
 - [2026/03] Check out our first public [project deepdive](https://youtu.be/sgwNfsNnR9I) at the vLLM Hong Kong Meetup!
 - [2026/03] **[vllm-omni-skills](https://github.com/hsliuustc0106/vllm-omni-skills)** is a community-driven collection of AI assistant skills that help developers work with vLLM-Omni more effectively. These skills can be used with popular agentic AI coding assistants like **Cursor IDE**, **Claude**, **Codex**, and more.
 - [2026/02] We released [0.16.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.16.0) - A major alignment + capability release that rebases onto **upstream vLLM v0.16.0** and significantly expands performance, distributed execution, and production readiness across **Qwen3-Omni / Qwen3-TTS**, **Bagel**, **MiMo-Audio**, **GLM-Image** and the **Diffusion (DiT) image/video stack**—while also improving platform coverage (CUDA / ROCm / NPU / XPU), CI quality, and documentation.
@@ -54,20 +54,28 @@ nav:
  - Features:
    - Sleep Mode: features/sleep_mode.md
    - Diffusion Features:
-      - Acceleration Overview: user_guide/diffusion_acceleration.md
-      - TeaCache: user_guide/diffusion/teacache.md
-      - Cache-DiT: user_guide/diffusion/cache_dit_acceleration.md
+      - Overview: user_guide/diffusion_features.md
+      - Feature Compatibility: user_guide/feature_compatibility.md
+      - Cache Acceleration:
+        - TeaCache: user_guide/diffusion/cache_acceleration/teacache.md
+        - Cache-DiT: user_guide/diffusion/cache_acceleration/cache_dit.md
      - Quantization:
        - Overview: user_guide/diffusion/quantization/overview.md
        - FP8: user_guide/diffusion/quantization/fp8.md
        - Int8: user_guide/diffusion/quantization/int8.md
        - GGUF: user_guide/diffusion/quantization/gguf.md
-      - Step Execution: user_guide/diffusion/step_execution.md
-      - Parallelism Acceleration: user_guide/diffusion/parallelism_acceleration.md
+      - Parallelism:
+        - Overview: user_guide/diffusion/parallelism/overview.md
+        - CFG Parallel: user_guide/diffusion/parallelism/cfg_parallel.md
+        - Expert Parallel: user_guide/diffusion/parallelism/expert_parallel.md
+        - Hybrid Sharded Data Parallel: user_guide/diffusion/parallelism/hsdp.md
+        - Sequence Parallel: user_guide/diffusion/parallelism/sequence_parallel.md
+        - Tensor Parallel: user_guide/diffusion/parallelism/tensor_parallel.md
+        - VAE Patch Parallel: user_guide/diffusion/parallelism/vae_patch_parallel.md
      - CPU Offloading: user_guide/diffusion/cpu_offload_diffusion.md
      - LoRA: user_guide/diffusion/lora.md
-      - Hybrid Sharded Data Parallel: design/feature/hsdp.md
      - Custom Pipeline: features/custom_pipeline.md
+      - Step Execution: user_guide/diffusion/step_execution.md
    - ComfyUI: features/comfyui.md
 - Developer Guide:
  - General:
@@ -92,6 +100,8 @@ nav:
      - design/feature/cfg_parallel.md
      - design/feature/sequence_parallel.md
      - design/feature/tensor_parallel.md
+      - design/feature/vae_parallel.md
+      - design/feature/hsdp.md
      - design/feature/cache_dit.md
      - design/feature/teacache.md
      - design/feature/async_chunk_design.md
@@ -20,7 +20,4 @@ For introduction, please check [Introduction for stage config](./stage_configs.m

 ## Optimization Features

- **[TeaCache Configuration](../user_guide/diffusion/teacache.md)** - Enable TeaCache adaptive caching for DiT models to achieve 1.5x-2.0x speedup with minimal quality loss
- **[Cache-DiT Configuration](../user_guide/diffusion/cache_dit_acceleration.md)** - Enable Cache-DiT as cache acceleration backends for DiT models
- **[Parallelism Configuration](../user_guide/diffusion/parallelism_acceleration.md)** - Enable parallelism (e.g., sequence parallelism) for for DiT models
- **[CPU Offloading](../user_guide/diffusion/cpu_offload_diffusion.md)** - Enable CPU offloading (model-level and layerwise) for for DiT models
+- **[Diffusion Features Overview](../user_guide/diffusion_features.md)** - Complete overview of all diffusion model features and supported models
@@ -548,103 +548,19 @@ L4 level testing is a comprehensive quality audit before a version release. It e
 -   ***Execution Environment***: ***GPU*** server clusters to meet the resource demands of performance testing.
 -   ***Script Example***:

-???+ example "Test Examples"
+??? example "Test Examples: Documentation Example Tests"

-    When adding L4-level ***documentation example Tests***, please pay attention to the following guides.
+    --8<-- "docs/contributing/ci/test_examples/l4_doc_example_tests.inc.md"

-    --8<-- "docs/contributing/ci/test_examples/doc_example_tests.inc.md"
+??? example "Test Examples: Performance Tests"

-    When you want to add L4-level ***performance test*** cases, you can refer to the following format for case addition in tests/dfx/perf/tests/test.json:
+    --8<-- "docs/contributing/ci/test_examples/l4_performance_tests.inc.md"

-    ```JSON
-    {
-        "test_name": "test_qwen3_omni",
-        "server_params": {
-            "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
-            "stage_config_name": "qwen3_omni.yaml"
-        },
-        "benchmark_params": [
-            {
-                "dataset_name": "random",
-                "num_prompts": [10, 20],
-                "request_rate": [0.5, 1],
-                "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": 100000,
-                    "mean_audio_ttfp_ms": 100000,
-                    "mean_audio_rtf": 100000
-                }
-            }
-        ]
-    }
-    ```
+??? example "Test Examples: Functionality Tests"

-    **Parameter Explanation**
+    --8<-- "docs/contributing/ci/test_examples/l4_functionality_tests.inc.md"

-    *Overview*
-
-    | Field            | Required | Description                                                     |
-    | ---------------- | -------- | --------------------------------------------------------------- |
-    | test_name        | Yes      | Unique identifier for the test case                             |
-    | server_params    | Yes      | Server-side configuration parameters                            |
-    | benchmark_params | Yes      | Benchmark running parameters (supports multiple configurations) |
-
-    **server_params Configuration**
-
-    *Basic Parameters*
-
-    | Parameter         | Required | Example                            | Description                   |
-    | ----------------- | -------- | ---------------------------------- | ----------------------------- |
-    | model             | Yes      | "Qwen/Qwen3-Omni-30B-A3B-Instruct" | Model name or path            |
-    | stage_config_name | Yes      | "qwen3_omni.yaml"                  | Stage configuration file name |
-
-    *Dynamic Configuration (update/delete)*
-
-    Supports incremental modifications based on the basic configuration:
-
-    | Operation | Description                          |
-    | --------- | ------------------------------------ |
-    | update    | Update or add configuration items    |
-    | delete    | Delete specified configuration items |
-
-    ***Example***:
-
-    ```
-    "update": {
-        "async_chunk": true,  // Enable asynchronous chunk processing
-        "stage_args": {
-            "0": {
-                "engine_args.custom_process_next_stage_input_func": "vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_async_chunk"
-            }
-        }
-    },
-    "delete": {
-        "stage_args": {
-            "2": ["custom_process_input_func"]  // Delete this configuration for stage 2
-        }
-    }
-    ```
-
-    **benchmark_params Configuration**
-
-    You can add any benchmark running parameters you need here. For all optional parameters, refer to the [benchmark documentation](https://github.com/vllm-project/vllm-omni/blob/main/docs/cli/bench/serve.md). General modifications are as follows:
-
-    1.  Change the ---xxx-xx-xx running parameters to xxx_xx_xx format and fill them as keys in the JSON file.
-    2.  For boolean variables in the running parameters, modify them to forms such as ignore_eos: true/false and fill them into the JSON file.
-    3.  Add the baseline parameter to specify the required validation values, ensuring the validation metric names match those in the result.json generated by the benchmark.
-    4.  The qps and concurrency modes are mutually exclusive. For detailed explanations, see the table below:
-
-    | Parameter       | Type        | Required | Example/Values  | Description                                                                                                                                                                                                                                                          |
-    | --------------- | ----------- | -------- | --------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
-    | num_prompts     | int / array | Yes      | 10,[10, 20, 30] | Number of requests. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of qps or max_concurrency, e.g., [10,10,10]. If an array is used, its length must match the number of qps or max_concurrency. |
-    | request_rate    | int / array | No       | 1, [1, 2, 3]    | Queries per second. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of num_prompts, e.g., [1,1,1]. If an array is used, its length must match the number of num_prompts.                          |
-    | max_concurrency | int / array | No       | 1, [1, 2, 3]    | Queries per second. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of num_prompts, e.g., [1,1,1]. If an array is used, its length must match the number of num_prompts.                          |
-    </details>
-
-    -   -   ***Run Command***: (Specific commands would depend on the performance testing tool and configuration defined in `nightly.json`).
+-   ***Run Command***: (Specific commands would depend on the performance testing tool and configuration defined in `nightly.json`).

 ## Chapter 4: L5 Level Testing - Stability and Reliability Testing

@@ -1,10 +1,106 @@
 # CI Failures

-What should I do when a CI job fails on my PR, but I don't think my PR caused the failure?
+## Overview of CI Checks

+When you open a PR against vLLM-Omni, several CI checks run automatically:

-## Common Case 1:
+| Check | Platform | What it does |
+| ----- | -------- | ------------ |
+| **pre-commit** | GitHub Actions | Runs linting (Ruff), formatting, spell-checking (typos), and YAML validation. |
+| **Build Wheel** | GitHub Actions | Builds Python wheels for Python 3.11 and 3.12 on Ubuntu. Skipped for docs-only or Markdown-only changes (controlled by `paths-ignore` in the workflow). |
+| **DCO** | GitHub | Verifies every commit has a `Signed-off-by` line. |
+| **docs/readthedocs.org:vllm-omni** | Read the Docs | Builds the MkDocs documentation site. |
+| **buildkite/vllm-omni** | Buildkite | Runs GPU-based tests on NVIDIA CUDA hardware (L4, H100). |
+| **buildkite/vllm-omni-amd** | Buildkite | Runs GPU-based tests on AMD ROCm hardware (MI325). |
+| **buildkite/vllm-omni-intel** | Buildkite | Runs GPU-based tests on Intel XPU hardware (Intel Arc BMG). |

-🚨 Error: The command was interrupted by a signal: signal: terminated
+## Step 1: Identify the Failing Check

-Reason: The test is terminated due to exceed the time limits. Sometimes rebuild the test will pass.
+Click the **Details** link next to the failing check on your PR to open the build log. The most common failures fall into these categories:
+
+### pre-commit failures
+
+These are typically formatting or linting issues introduced by your PR. Fix them locally:
+
+```bash
+uv pip install pre-commit
+pre-commit run --all-files
+```
+
+Then commit the fixes and push.
+
+### DCO failures
+
+Every commit must include a `Signed-off-by` line. If you forgot, amend your commits:
+
+```bash
+git commit --amend -s
+git push --force-with-lease
+```
+
+For multiple commits, use an interactive rebase to add the sign-off to each one.
+
+### Read the Docs failures
+
+The documentation build uses MkDocs with `fail_on_warning: true`, so even a minor warning (not just errors) will cause the build to fail. To reproduce locally:
+
+```bash
+uv pip install -e ".[docs]"
+mkdocs build --strict
+```
+
+Common causes include broken cross-references, invalid admonition syntax, or missing files referenced by `--8<--` includes.
+
+### Buildkite failures
+
+Buildkite runs GPU tests in Docker containers. These are the most complex checks and can fail for reasons unrelated to your PR (infrastructure issues, flaky tests, etc.). See the sections below for how to investigate.
+
+## Step 2: Check if the Failure Is a Known Issue
+
+Before spending time debugging, check whether the failure already exists on the `main` branch:
+
+1. **Look at the Buildkite build log** — the test name and error message are usually enough to identify the issue.
+2. **Check recent CI runs on `main`** — if the same test is failing there, the failure is not caused by your PR.
+3. **Search existing issues** — look for open issues in the [vllm-omni issue tracker](https://github.com/vllm-project/vllm-omni/issues) with the test name or error message.
+
+If the failure is already tracked, leave a comment on your PR noting that the failure is pre-existing and link the issue.
+
+## Step 3: Investigate the Failure
+
+If the failure appears to be new, investigate whether your changes caused it.
+
+### Reading Buildkite Logs
+
+1. Click **Details** next to the Buildkite check on your PR.
+2. Find the failing step in the pipeline (e.g., "Diffusion Model Test", "Simple Unit Test").
+3. Expand the step to see the full test output with the traceback.
+
+### Running Tests Locally
+
+For instructions on running tests locally (including specific test files, functions, and markers), see the [Running Tests](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/test_guide/#running-tests) section in the Test Guide.
+
+## Step 4: Raise an Issue or Fix It
+
+### If the failure is pre-existing (not caused by your PR)
+
+1. **Raise a new issue** if one doesn't already exist, using the title format:
+   `[CI Failure]: [job-name] - [test-path]`
+2. Include the error message, relevant log excerpts, and the commit hash where the failure occurs (e.g., "Still failing on main as of commit `abc1234`").
+3. Leave a comment on your PR linking to the issue and noting that the failure is unrelated to your changes.
+
+### If the failure is caused by your PR
+
+1. Fix the issue in your branch and push the update.
+2. If the fix is non-trivial, consider adding a test to prevent regression.
+
+## Common Failure Patterns
+
+| Symptom | Likely Cause | Fix |
+| ------- | ------------ | --- |
+| `ruff` or formatting errors | Code style violation | Run `pre-commit run --all-files` |
+| `Signed-off-by` missing | DCO check | Amend commits with `git commit --amend -s` |
+| MkDocs build warning | Broken docs reference | Run `mkdocs build --strict` locally |
+| `OOM` or `CUDA out of memory` | Test exceeds GPU memory | Check if your changes increased memory usage; use `--vae_use_slicing` / `--vae_use_tiling` for diffusion tests |
+| Import errors | Missing or changed dependency | Check `pyproject.toml` and make sure dependencies are correct |
+| Timeout (step exceeded N minutes) | Test is too slow or hangs | Profile the test; check for infinite loops or deadlocks |
+| `Agent lost` in Buildkite | Infrastructure issue (not your fault) | Re-trigger the build; comment on your PR |
@@ -0,0 +1,46 @@
+**Scope**
+
+For diffusion models, the L4 functionality test covers all or common *diffusion features* that are supported by this model, including several [parallelism acceleration methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/parallelism_acceleration/), [CPU offloading](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cpu_offload_diffusion/), [TeaCache](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/teacache/) and [Cache-DiT](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cache_dit_acceleration/) cache backends, [quantization methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/quantization/overview/).
+
+**Test Case Design**
+
+For a *high priority* model (currently listed in [issue #1832](https://github.com/vllm-project/vllm-omni/issues/1832)), we use several test cases, each with multiple features turned on, so that each supported feature is tested in at least one test case. This is to relieve the GPU workload on the CI machine. The suggested test case combination is as follows:
+
+- If the model can fit into 4 L4 GPU (with quantization and tensor parallel always on) (20GB RAM each)
+    - (1 GPU) TeaCache + Layerwise CPU offloading + GGUF
+    - (4 GPUs) CacheDiT + Ulysses=2 + TP=2 + VAE=2 + FP8
+    - (4 GPUs) CacheDiT + Ring=2 + HSDP=2 + VAE=2 + GGUF
+    - (4 GPUs) TeaCache + CFG=2 + TP=2 + VAE=2 + FP8
+- Otherwise, consider 2 H100 GPU environment (80GB RAM each) with the following tests
+    - (1 GPU) TeaCache + Layerwise CPU offloading + GGUF
+    - (2 GPUs) CacheDiT + Ulysses=2 + FP8
+    - (2 GPUs) CacheDiT + Ring=2 + GGUF
+    - (2 GPUs) TeaCache + CFG=2 + FP8
+    - (2 GPUs) CacheDiT + TP=2 + VAE=2 + FP8
+    - (2 GPUs) CacheDiT + HSDP=2 + VAE=2 + GGUF
+- If 2 H100 GPU cannot handle the model either (e.g., HunyuanImage 3.0)
+    - Still design tests and feature combinations that can best fit real-world scenario.
+    - Do not include it in CI (or exclude it from the CI's filtering criteria). Instead, relevant PR authors are suggested to run these tests locally.
+- Fallback plan
+    - If the model does not support layerwise CPU offloading, replace the corresponding test case with module-wise offloading
+    - If the model only supports specific or no caching feature, use this option or remove this option in all test cases.
+    - If the model only supports specific or no quantization feature, use this option or remove this option in all test cases.
+    - If the model does not support certain other features, remove this option from that test case. If, consequently, the coverage of this modified test case completely overlaps with others, remove this test case.
+
+For a *normal priority* model, further reduce the number of test cases.
+
+- Only write one or two test cases for the most common feature combinations.
+- The author can explore themselves to see which feature combination balances output quality and performance. Alternatively, the author can refer to any example code in the PR that adds the model, or the example code in the PR that adds a feature (if the code involves this model of interest).
+
+Currently all the features are available in online serving mode. Hence, only need to add `tests/e2e/online_serving/test_{model}_expansion.py`.
+
+**Code Style**
+
+- Validation: test that the multimodal output files of your model have the correct shapes. `OpenAIClientHandler.send_diffusion_request` should have taken care of this.
+- Test marks: always add `advanced_model` and `diffusion`. Add GPU-related marks if needed. Ref: [Markers for Tests](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/tests_markers/).
+- To maximize code reuse, you may refer to
+    - `tests/conftest.py` for `omni_server` (running server in subprocess) and `openai_client` fixtures (sending requests and validating output), `generate_synthetic_image` and `assert_XXX_valid` helper.
+    - `tests/utils.py` for `@hardware_test(...)` and `hardware_marks`.
+    - [Parametrizing tests (pytest doc)](https://docs.pytest.org/en/stable/example/parametrize.html) to reuse test function implementation for different cases.
+- Doc: add a concise docstring for each test function.
+- Reference L4 test implementation: [tests/e2e/online_serving/test_qwen_image_edit_expansion.py](https://github.com/vllm-project/vllm-omni/blob/main/tests/e2e/online_serving/test_qwen_image_edit_expansion.py).
@@ -0,0 +1,89 @@
+When you want to add L4-level ***performance test*** cases, you can refer to the following format for case addition in tests/dfx/perf/tests/test.json:
+
+```JSON
+{
+    "test_name": "test_qwen3_omni",
+    "server_params": {
+        "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+        "stage_config_name": "qwen3_omni.yaml"
+    },
+    "benchmark_params": [
+        {
+            "dataset_name": "random",
+            "num_prompts": [10, 20],
+            "max_concurrency": [1, 4],
+            "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": [500, 800],
+                "mean_audio_ttfp_ms": [2000, 3500],
+                "mean_audio_rtf": [0.25, 0.35]
+            }
+        }
+    ]
+}
+```
+
+**Parameter Explanation**
+
+*Overview*
+
+| Field            | Required | Description                                                     |
+| ---------------- | -------- | --------------------------------------------------------------- |
+| test_name        | Yes      | Unique identifier for the test case                             |
+| server_params    | Yes      | Server-side configuration parameters                            |
+| benchmark_params | Yes      | Benchmark running parameters (supports multiple configurations) |
+
+**`server_params` Configuration**
+
+*Basic Parameters*
+
+| Parameter         | Required | Example                            | Description                   |
+| ----------------- | -------- | ---------------------------------- | ----------------------------- |
+| model             | Yes      | "Qwen/Qwen3-Omni-30B-A3B-Instruct" | Model name or path            |
+| stage_config_name | Yes      | "qwen3_omni.yaml"                  | Stage configuration file name |
+
+*Dynamic Configuration (update/delete)*
+
+Supports incremental modifications based on the basic configuration:
+
+| Operation | Description                          |
+| --------- | ------------------------------------ |
+| update    | Update or add configuration items    |
+| delete    | Delete specified configuration items |
+
+**Example**:
+
+```
+"update": {
+    "async_chunk": true,  // Enable asynchronous chunk processing
+    "stage_args": {
+        "0": {
+            "engine_args.custom_process_next_stage_input_func": "vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_async_chunk"
+        }
+    }
+},
+"delete": {
+    "stage_args": {
+        "2": ["custom_process_input_func"]  // Delete this configuration for stage 2
+    }
+}
+```
+
+**`benchmark_params` Configuration**
+
+You can add any benchmark running parameters you need here. For all optional parameters, refer to the [benchmark documentation](https://github.com/vllm-project/vllm-omni/blob/main/docs/cli/bench/serve.md). General modifications are as follows:
+
+1.  Change the --xxx-xx-xx running parameters to xxx_xx_xx format and fill them as keys in the JSON file.
+2.  For boolean variables in the running parameters, modify them to forms such as ignore_eos: true/false and fill them into the JSON file.
+3.  Optionally add a `baseline` object (see **Baseline thresholds** below). If you omit `baseline` or leave it empty, the performance test still runs but does not assert metric thresholds from this field.
+4.  The qps and concurrency modes are recommended to be mutually exclusive. For detailed explanations, see the table below:
+
+| Parameter       | Type        | Required | Example/Values  | Description                                                                                                                                                                                                                                                          |
+| --------------- | ----------- | -------- | --------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| num_prompts     | int / array | Yes      | 10,[10, 20, 30] | Number of requests. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of qps or max_concurrency, e.g., [10,10,10]. If an array is used, its length must match the number of qps or max_concurrency. |
+| request_rate    | float / array | No  | 0.5, [0.5, 1, inf] | Queries per second. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of num_prompts, e.g., [1,1,1]. If an array is used, its length must match the number of num_prompts.                          |
+| max_concurrency | int / array | No       | 1, [1, 2, 3]    | Maximum concurrent in-flight requests. Same array / expansion rules as `request_rate` (mutually exclusive with QPS mode).                                                                                                                                                                                                             |
+| baseline        | object      | No       | see above       | Optional per-metric thresholds; keys must match benchmark output fields. Scalar, list (per sweep step), or object (keyed by concurrency or QPS string).  
@@ -653,25 +653,7 @@ For a fair comparison, keep the same **prompt**, **seed**, **resolution**, **num

 To ensure project maintainability and sustainable development, please submit test code (unit tests, system tests, or end-to-end tests) alongside their code changes.

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

 ---

@@ -256,7 +256,7 @@ cache_config={
 }
 ```

-Check the [user guide for cache_dit](../../user_guide/diffusion/cache_dit_acceleration.md) for more adjustable parameters.
+Check the [user guide for cache_dit](../../user_guide/diffusion/cache_acceleration/cache_dit.md) for more adjustable parameters.

 ---

@@ -34,7 +34,7 @@ vLLM-omni provides `CFGParallelMixin` that encapsulates all CFG parallel logic.
 | Method | Purpose | Automatic Behavior |
 |--------|---------|-------------------|
 | [`predict_noise_maybe_with_cfg()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Predict noise with CFG | Detects parallel mode, distributes computation, gathers results |
-| [`scheduler_step_maybe_with_cfg()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Step scheduler with sync | Rank 0 steps, broadcasts latents to all ranks |
+| [`scheduler_step_maybe_with_cfg()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Step scheduler | All ranks step locally (no broadcast needed) |
 | [`combine_cfg_noise()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Combine positive/negative | Applies CFG formula with optional normalization |
 | [`predict_noise()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Forward pass wrapper | Override for custom transformer calls |
 | [`cfg_normalize_function()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Normalize CFG output | Override for custom normalization |
@@ -47,7 +47,7 @@ vLLM-omni provides `CFGParallelMixin` that encapsulates all CFG parallel logic.
  - Rank 0 computes positive prompt prediction
  - Rank 1 computes negative prompt prediction
  - Results are gathered via `all_gather()`
-  - Combined on rank 0 using CFG formula
+  - All ranks compute CFG combine locally (deterministic, identical results)

 - **Sequential mode** (when `cfg_world_size == 1`):
  - Single rank computes both positive and negative predictions
@@ -55,14 +55,7 @@ vLLM-omni provides `CFGParallelMixin` that encapsulates all CFG parallel logic.

 `scheduler_step_maybe_with_cfg()` ensures consistent latent states across all ranks:

- **CFG-Parallel mode**:
-  - Only rank 0 performs the scheduler step (applies noise prediction to update latents)
-  - Updated latents are broadcast to all other ranks via `broadcast()`
-  - All ranks maintain synchronized latent states for the next iteration
-
- **Sequential mode**:
-  - Single rank directly performs the scheduler step
-  - No synchronization needed
+- All ranks compute the scheduler step locally — no broadcast needed because `predict_noise_maybe_with_cfg` already ensures all ranks have identical noise predictions after `all_gather` + local combine.

 ---

@@ -177,6 +170,70 @@ class LongCatImagePipeline(nn.Module, CFGParallelMixin):
        # return noise_pred
 ```

+
+### Override `combine_cfg_noise()` for Multi-Output Models
+
+When `predict_noise()` returns a tuple (e.g., video + audio), the default `combine_cfg_noise()` applies CFG to every element. Override it to apply different logic per element — for example, CFG on video but positive-only on audio:
+
+```python
+class MyVideoAudioPipeline(nn.Module, CFGParallelMixin):
+    def combine_cfg_noise(self, positive_noise_pred, negative_noise_pred, scale, normalize):
+        (video_pos, audio_pos) = positive_noise_pred
+        (video_neg, audio_neg) = negative_noise_pred
+        video_combined = super().combine_cfg_noise(video_pos, video_neg, scale, normalize)
+        return (video_combined, audio_pos)  # audio: positive only, no CFG
+```
+
+This also requires `predict_noise()` to return a tuple (see [Override predict_noise](#override-predict_noise-for-custom-transformer-calls) above).
+
+### Implement a Composite Scheduler for Multi-Output Models
+
+When each output has its own denoising schedule, implement a composite scheduler that dispatches to per-output schedulers. Assign it to `self.scheduler` so the default `scheduler_step()` works without override.
+
+**Complete example (video + audio with separate schedulers and diffuse loop):**
+
+```python
+class VideoAudioScheduler:
+    """Composite scheduler dispatching to video and audio schedulers."""
+    def __init__(self, video_scheduler, audio_scheduler):
+        self.video_scheduler = video_scheduler
+        self.audio_scheduler = audio_scheduler
+
+    def step(self, noise_pred, t, latents, return_dict=False, generator=None):
+        video_out = self.video_scheduler.step(noise_pred[0], t[0], latents[0], return_dict=False, generator=generator)[0]
+        audio_out = self.audio_scheduler.step(noise_pred[1], t[1], latents[1], return_dict=False, generator=generator)[0]
+        return ((video_out, audio_out),)
+
+class MyVideoAudioPipeline(nn.Module, CFGParallelMixin):
+    def __init__(self, ...):
+        self.scheduler = VideoAudioScheduler(video_sched, audio_sched)
+
+    def predict_noise(self, **kwargs):
+        video_pred, audio_pred = self.transformer(**kwargs)
+        return (video_pred, audio_pred)
+
+    def combine_cfg_noise(self, positive_noise_pred, negative_noise_pred, scale, normalize):
+        # ... (as above)
+
+    def diffuse(self, video_latents, audio_latents, timesteps_video, timesteps_audio, ...):
+        for t_v, t_a in zip(timesteps_video, timesteps_audio):
+            positive_kwargs = {...}  
+            negative_kwargs = {...} if do_true_cfg else None
+
+            video_pred, audio_pred = self.predict_noise_maybe_with_cfg(
+                do_true_cfg=do_true_cfg, true_cfg_scale=self.guidance_scale,
+                positive_kwargs=positive_kwargs, negative_kwargs=negative_kwargs,
+            )
+            video_latents, audio_latents = self.scheduler_step_maybe_with_cfg(
+                (video_pred, audio_pred), (t_v, t_a),
+                (video_latents, audio_latents), do_true_cfg=do_true_cfg,
+                generator=generator,
+            )
+        return video_latents, audio_latents
+```
+
+> **Note:** If you use a non-deterministic scheduler, e.g., DDPM, please set `self.scheduler_step_maybe_with_cfg(..., generator=torch.Generator(device).manual_seed(seed))` explicitly to control the randomness of scheduler step among ranks.
+
 ---

 ## Testing
@@ -369,7 +369,7 @@ images = omni.generate(
 2. **Compare performance** - Measure speedup vs baseline (expect 1.5x-2.0x)
 3. **Verify output quality** - Visually compare cached vs uncached outputs (should be nearly identical)

-See more detailed examples in [user guide for teacache](../../user_guide/diffusion/teacache.md).
+See more detailed examples in [user guide for teacache](../../user_guide/diffusion/cache_acceleration/teacache.md).

 ---

@@ -29,6 +29,11 @@ uv pip install vllm --torch-backend=auto
 uv pip install vllm-omni
 ```

+To run Gradio demos, also install the optional extras:
+```bash
+uv pip install 'vllm-omni[demo]'
+```
+
 # --8<-- [end:pre-built-wheels]

 # --8<-- [start:build-wheel-from-source]
@@ -50,6 +55,11 @@ cd vllm-omni
 uv pip install -e .
 ```

+To run Gradio demos, install with optional extras:
+```bash
+uv pip install -e '.[demo]'
+```
+
 <details><summary>(Optional) Installation of vLLM from source</summary>
 If you want to check, modify or debug with source code of vLLM, install the library from source with the following instructions:

@@ -0,0 +1,285 @@
+# Cache-DiT Guide
+
+
+## Table of Content
+
+- [Overview](#overview)
+- [Quick Start](#quick-start)
+- [Example Script](#example-script)
+- [Acceleration Methods](#acceleration-methods)
+- [Configuration Parameters](#configuration-parameters)
+- [Best Practices](#best-practices)
+- [Troubleshooting](#troubleshooting)
+- [Summary](#summary)
+- [Additional Resources](#additional-resources)
+
+---
+
+## Overview
+
+Cache-DiT accelerates diffusion transformer models through intelligent caching mechanisms, providing significant speedup with minimal quality loss. It supports multiple acceleration techniques that can be combined for optimal performance:
+
+- **DBCache**: Dual Block Cache for reducing redundant computations
+- **TaylorSeer**: Taylor expansion-based forecasting for faster inference
+- **SCM**: Step Computation Masking for selective step computation
+
+See supported models list in [Supported Models](../../diffusion_features.md#supported-models).
+
+---
+
+## Quick Start
+
+### Basic Usage
+
+Enable cache-dit acceleration by simply setting `cache_backend="cache_dit"`:
+
+```python
+from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+omni = Omni(
+    model="Qwen/Qwen-Image",
+    cache_backend="cache_dit",  # Enable Cache-DiT with defaults
+)
+
+outputs = omni.generate(
+    "a beautiful landscape",
+    OmniDiffusionSamplingParams(num_inference_steps=50),
+)
+```
+
+**Note**: When `cache_config` is not provided, Cache-DiT uses optimized default values. See the [Configuration Parameters](#configuration-parameters) section for details.
+
+### Custom Configuration
+
+To customize cache-dit settings, provide a `cache_config` dictionary, for example:
+
+```python
+omni = Omni(
+    model="Qwen/Qwen-Image",
+    cache_backend="cache_dit",
+    cache_config={
+        "Fn_compute_blocks": 1,
+        "Bn_compute_blocks": 0,
+        "max_warmup_steps": 4,
+        "residual_diff_threshold": 0.12,
+    },
+)
+```
+
+---
+
+## Example Script
+
+### Offline Inference
+
+Use the example script under `examples/offline_inference/text_to_image`:
+
+```bash
+cd examples/offline_inference/text_to_image
+python text_to_image.py \
+    --model Qwen/Qwen-Image \
+    --prompt "a cup of coffee on the table" \
+    --cache-backend cache_dit \
+    --num-inference-steps 50
+```
+
+See the [text_to_image.py](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/text_to_image/text_to_image.py) for detailed configuration options.
+
+The script uses cache-dit acceleration with a hybrid configuration combining DBCache, SCM, and TaylorSeer:
+
+```python
+omni = Omni(
+    model="Qwen/Qwen-Image",
+    cache_backend="cache_dit",
+    cache_config={
+        # Scheme: Hybrid DBCache + SCM + TaylorSeer
+        "Fn_compute_blocks": 1,  # Optimized for single-transformer models
+        "Bn_compute_blocks": 0,  # Number of backward compute blocks
+        "max_warmup_steps": 4,  # Maximum warmup steps (works for few-step models)
+        "residual_diff_threshold": 0.24,  # Higher threshold for more aggressive caching
+        "max_continuous_cached_steps": 3,  # Limit to prevent precision degradation
+        # TaylorSeer parameters [cache-dit only]
+        "enable_taylorseer": False,  # Disabled by default (not suitable for few-step models)
+        "taylorseer_order": 1,  # TaylorSeer polynomial order
+        # SCM (Step Computation Masking) parameters [cache-dit only]
+        "scm_steps_mask_policy": None,  # SCM mask policy: None (disabled), "slow", "medium", "fast", "ultra"
+        "scm_steps_policy": "dynamic",  # SCM steps policy: "dynamic" or "static"
+    }
+)
+```
+
+You can customize the configuration by modifying the `cache_config` dictionary to use only specific methods (e.g., DBCache only, DBCache + SCM, etc.) based on your quality and speed requirements.
+
+For image-to-image tasks, use the example script under `examples/offline_inference/image_to_image`:
+
+```bash
+cd examples/offline_inference/image_to_image
+python image_edit.py \
+    --model Qwen/Qwen-Image-Edit \
+    --prompt "make the sky more colorful" \
+    --image path/to/input/image.jpg \
+    --cache-backend cache_dit \
+    --num-inference-steps 50 \
+    --cache-dit-max-continuous-cached-steps 3 \
+    --cache-dit-residual-diff-threshold 0.24 \
+    --cache-dit-enable-taylorseer
+```
+
+See the [image_edit.py](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/image_to_image/image_edit.py) for detailed configuration options.
+
+### Online Serving
+
+```bash
+# Default configuration (recommended)
+vllm serve Qwen/Qwen-Image --omni --port 8091 --cache-backend cache_dit
+
+# Custom configuration
+vllm serve Qwen/Qwen-Image --omni --port 8091 \
+  --cache-backend cache_dit \
+  --cache-config '{"Fn_compute_blocks": 1, "residual_diff_threshold": 0.12}'
+```
+
+---
+
+## Acceleration Methods
+
+For comprehensive illustration, please view Cache-DiT [User Guide](https://cache-dit.readthedocs.io/en/latest/user_guide/OVERVIEWS/).
+
+### 1. DBCache (Dual Block Cache)
+
+DBCache intelligently caches intermediate transformer block outputs when the residual differences between consecutive steps are small, reducing redundant computations without sacrificing quality.
+
+**Example Configuration**:
+
+```python
+cache_config={
+    "Fn_compute_blocks": 8,           # Use first 8 blocks for difference computation
+    "Bn_compute_blocks": 0,           # No additional fusion blocks
+    "max_warmup_steps": 8,            # Cache after 8 warmup steps
+    "residual_diff_threshold": 0.12,  # Lower threshold for faster inference
+    "max_cached_steps": -1,           # No limit on cached steps
+}
+```
+
+**Performance Tips**:
+
+- Default `Fn_compute_blocks=1` works well for most cases. Some models (e.g., [FLUX.2-klein](https://github.com/wtomin/vllm-omni/blob/main/vllm_omni/diffusion/cache/cache_dit_backend.py#L363)) use a larger value for `Fn_compute_blocks` for a balanced performance.
+- Increase `residual_diff_threshold` (e.g., 0.12-0.15) for faster inference with slight quality trade-off, or decrease from default 0.24 for higher quality.
+- Default `max_warmup_steps=4` is optimized for few-step models. Increase to 6-8 for more steps if needed.
+
+### 2. TaylorSeer
+
+TaylorSeer uses Taylor expansion to forecast future hidden states, allowing the model to skip some computation steps while maintaining quality.
+
+**Example Configuration**:
+
+```python
+cache_config={
+    "enable_taylorseer": True,
+    "taylorseer_order": 1,  # First-order Taylor expansion
+}
+```
+
+**Performance Tips**:
+
+- TaylorSeer is **not suitable for few-step distilled models**.
+- Use `taylorseer_order=1` for most cases (good balance of speed and quality).
+- Combine with DBCache for maximum acceleration.
+- Higher orders (2-3) may improve quality but reduce speed gains.
+
+### 3. SCM (Step Computation Masking)
+
+SCM allows you to specify which steps must be computed and which can use cached results, similar to LeMiCa/EasyCache style acceleration.
+
+`scm_steps_mask_policy` options (number of compute steps out of 28):
+
+| Policy | Compute Steps | Speed | Quality |
+|--------|--------------|-------|---------|
+| `None` (default) | All | Baseline | Best |
+| `"slow"` | 18 / 28 | Moderate | High |
+| `"medium"` | 15 / 28 | Balanced | Good |
+| `"fast"` | 11 / 28 | Fast | Moderate |
+| `"ultra"` | 8 / 28 | Fastest | Lower |
+
+**Example Configuration**:
+
+```python
+cache_config={
+    "scm_steps_mask_policy": "medium",  # Balanced speed/quality
+    "scm_steps_policy": "dynamic",      # Use dynamic cache
+}
+```
+
+**Performance Tips**:
+
+- SCM is disabled by default. Enable it by setting a policy value if you need additional acceleration.
+- Start with `"medium"` policy and adjust based on quality requirements.
+- Use `"fast"` or `"ultra"` for maximum speed when quality can be slightly compromised.
+- `"dynamic"` policy generally provides better quality than `"static"`.
+- SCM mask is automatically regenerated when `num_inference_steps` changes during inference.
+
+---
+
+## Configuration Parameters
+
+In `cache_config` passed to `Omni` constructor, it accepts the arguments of `DBCacheConfig` ([Cache-DiT API Reference](https://cache-dit.readthedocs.io/en/latest/user_guide/CACHE_API/)). Key parameters are listed below:
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `Fn_compute_blocks` | int | 1 | First n blocks for difference computation (optimized for single-transformer models) |
+| `Bn_compute_blocks` | int | 0 | Last n blocks for fusion |
+| `max_warmup_steps` | int | 4 | Steps before caching starts (optimized for few-step distilled models) |
+| `max_cached_steps` | int | -1 | Max cached steps (-1 = unlimited) |
+| `max_continuous_cached_steps` | int | 3 | Max consecutive cached steps (prevents precision degradation) |
+| `residual_diff_threshold` | float | 0.24 | Residual difference threshold (higher for more aggressive caching) |
+| `num_inference_steps` | int \| None | None | Initial inference steps for SCM mask generation (optional, auto-refreshed during inference) |
+| `enable_taylorseer` | bool | False | Enable TaylorSeer acceleration (not suitable for few-step distilled models) |
+| `taylorseer_order` | int | 1 | Taylor expansion order |
+| `scm_steps_mask_policy` | str \| None | None | SCM mask policy (None, "slow", "medium", "fast", "ultra") |
+| `scm_steps_policy` | str | "dynamic" | SCM computation policy ("dynamic" or "static") |
+
+---
+
+## Best Practices
+
+### When to Use
+
+**Good for:**
+
+- Production deployments requiring fast inference
+- Diffusion transformer models (DiT architecture)
+- Scenarios where 1.5x-3x speedup is valuable
+
+**Not for:**
+
+- Non-DiT architectures (use model-specific acceleration instead)
+- Models already using few-step distillation (< 10 steps)
+
+---
+
+## Troubleshooting
+
+### Common Issue 1: Quality Degradation
+
+**Symptoms**: Generated images have visible artifacts or lower quality
+
+**Solution**:
+```python
+# Reduce aggressiveness - use more conservative settings
+cache_config={
+    "residual_diff_threshold": 0.20,  # Lower threshold (closer to default 0.24)
+    "Fn_compute_blocks": 8,            # Use more blocks for better decisions
+    "max_warmup_steps": 6,             # Longer warmup
+    "scm_steps_mask_policy": "slow",   # More compute steps
+}
+```
+
+---
+
+## Summary
+
+Using Cache-DiT acceleration:
+
+1. ✅ **Enable Cache-DiT** - Set `cache_backend="cache_dit"` to get 1.5x-3x speedup with optimized defaults
+2. ✅ **(Optional) Customize** - Adjust `cache_config` parameters for specific speed/quality trade-offs
@@ -0,0 +1,194 @@
+# TeaCache Guide
+
+
+## Table of Content
+
+- [Overview](#overview)
+- [Quick Start](#quick-start)
+- [Example Script](#example-script)
+- [Configuration Parameters](#configuration-parameters)
+- [Best Practices](#best-practices)
+- [Troubleshooting](#troubleshooting)
+- [Summary](#summary)
+
+---
+
+## Overview
+
+TeaCache accelerates diffusion model inference by caching transformer computations when consecutive timesteps are similar, providing **1.5x-2.0x speedup** with minimal quality loss. It dynamically decides whether to reuse cached outputs based on input similarity, making it ideal for production deployments where inference speed matters without sacrificing generation quality.
+
+See supported models list in [Supported Models](../../diffusion_features.md#supported-models).
+
+---
+
+## Quick Start
+
+
+
+### Basic Usage
+
+
+```python
+from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+omni = Omni(
+    model="Qwen/Qwen-Image",
+    cache_backend="tea_cache",
+)
+
+outputs = omni.generate(
+    "A cat sitting on a windowsill",
+    OmniDiffusionSamplingParams(num_inference_steps=50),
+)
+```
+
+### Custom Configuration
+
+```python
+omni = Omni(
+    model="Qwen/Qwen-Image",
+    cache_backend="tea_cache",
+    cache_config={
+        "rel_l1_thresh": 0.2,  # Controls speed/quality tradeoff
+    },
+)
+```
+
+### Using Environment Variable
+
+You can also enable TeaCache via environment variable:
+
+```bash
+export DIFFUSION_CACHE_BACKEND=tea_cache
+```
+
+Then initialize without explicitly setting `cache_backend`:
+
+```python
+from vllm_omni import Omni
+
+omni = Omni(
+    model="Qwen/Qwen-Image",
+    cache_config={"rel_l1_thresh": 0.2}
+)
+```
+
+---
+
+## Example Script
+
+### Offline Inference
+
+Use python script under `examples/offline_inference/text_to_image/` or `examples/offline_inference/image_to_image/` with CLI:
+
+```bash
+# Text-to-image example
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model Qwen/Qwen-Image \
+  --cache-backend tea_cache
+
+# Image-to-image example
+python examples/offline_inference/image_to_image/image_edit.py \
+  --model Qwen/Qwen-Image-Edit \
+  --image input.png \
+  --prompt "Edit description" \
+  --cache-backend tea_cache \
+  --tea-cache-rel-l1-thresh 0.25
+```
+
+See the [text_to_image.py](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/text_to_image/text_to_image.py) or [image_edit.py](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/image_to_image/image_edit.py) for detailed configuration options.
+
+### Online Serving
+
+```bash
+# Default configuration
+vllm serve Qwen/Qwen-Image --omni --port 8091 --cache-backend tea_cache
+
+# Custom configuration
+vllm serve Qwen/Qwen-Image --omni --port 8091 \
+  --cache-backend tea_cache \
+  --cache-config '{"rel_l1_thresh": 0.2}'
+```
+
+---
+
+## Configuration Parameters
+
+In `OmniDiffusionConfig`
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `rel_l1_thresh` | float | `0.2` | Similarity threshold for cache reuse. Lower values prioritize quality (less caching), higher values prioritize speed (more caching). Suggested range: 0.1-0.8 |
+| `coefficients` | list[float] \| None | `None` | Polynomial coefficients for rescaling L1 distance. Must contain exactly 5 elements if provided. If `None`, uses model-specific defaults based on transformer type. |
+
+Users can find the default model coefficients in [`vllm_omni/diffusion/cache/teacache/config.py`](https://github.com/vllm-project/vllm-omni/blob/main/vllm_omni/diffusion/cache/teacache/config.py), for example:
+
+```python
+_MODEL_COEFFICIENTS = {
+    # Qwen-Image transformer coefficients from ComfyUI-TeaCache
+    # Tuned specifically for Qwen's dual-stream transformer architecture
+    # Used for all Qwen-Image Family pipelines, in general
+    "QwenImageTransformer2DModel": [
+        -4.50000000e02,
+        2.80000000e02,
+        -4.50000000e01,
+        3.20000000e00,
+        -2.00000000e-02,
+    ],
+    ...
+}
+```
+
+---
+
+## Best Practices
+
+### When to Use
+
+**Good for:**
+
+- Production deployments requiring faster inference, tolerant of minimal quality loss
+- Scenarios where 1.5-2x speedup is valuable
+- Useful for single-card acceleration
+
+**Not for:**
+
+- Maximum quality requirements where no degradation is acceptable
+- Very short inference runs (< 20 steps) where caching overhead may outweigh benefits
+
+
+---
+
+## Troubleshooting
+
+### Common Issue 1: Quality Degradation
+
+**Symptoms**: Generated images show artifacts, reduced detail, or inconsistent quality compared to non-cached results
+
+**Solution**:
+
+```python
+# Lower the threshold for more conservative caching
+cache_config={"rel_l1_thresh": 0.1}
+```
+
+### Common Issue 2: Limited Speedup
+
+**Symptoms**: Actual speedup is less than expected (< 1.3x)
+
+**Solutions**:
+1. Increase the threshold to enable more aggressive caching:
+   ```python
+   cache_config={"rel_l1_thresh": 0.8}
+   ```
+2. Ensure you're using sufficient inference steps (35+ recommended)
+3. Check that your model architecture is supported (see Supported Models section)
+
+---
+
+
+## Summary
+
+1. ✅ **Enable TeaCache** - Set `cache_backend="tea_cache"` to get 1.5x-2.0x speedup with optimized defaults
+2. ✅ **(Optional) Customize** - Adjust thresholds and polynomial coefficients for specific speed/quality trade-offs
@@ -1,228 +0,0 @@
-# Cache-DiT Acceleration Guide
-
-This guide explains how to use cache-dit acceleration in vLLM-Omni to speed up diffusion model inference.
-
-## Overview
-
-Cache-dit is a library that accelerates diffusion transformer models through intelligent caching mechanisms. It supports multiple acceleration techniques that can be combined for optimal performance:
-
- **DBCache**: Dual Block Cache for reducing redundant computations
- **TaylorSeer**: Taylor expansion-based forecasting for faster inference
- **SCM**: Step Computation Masking for selective step computation
-
-## Quick Start
-
-### Basic Usage
-
-Enable cache-dit acceleration by simply setting `cache_backend="cache_dit"`. Cache-dit will use its recommended default parameters:
-
-```python
-from vllm_omni.entrypoints.omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-# Simplest way: just enable cache-dit with default parameters
-omni = Omni(
-    model="Qwen/Qwen-Image",
-    cache_backend="cache_dit",
-)
-
-images = omni.generate(
-    "a beautiful landscape",
-    OmniDiffusionSamplingParams(num_inference_steps=50),
-)
-```
-
-**Default Parameters**: When `cache_config` is not provided, cache-dit uses optimized default values. See the [Configuration Reference](#configuration-reference) section for a complete list of all parameters and their default values.
-
-### Custom Configuration
-
-To customize cache-dit settings, provide a `cache_config` dictionary, for example:
-
-```python
-omni = Omni(
-    model="Qwen/Qwen-Image",
-    cache_backend="cache_dit",
-    cache_config={
-        "Fn_compute_blocks": 1,
-        "Bn_compute_blocks": 0,
-        "max_warmup_steps": 4,
-        "residual_diff_threshold": 0.12,
-    },
-)
-```
-
-## Online Serving (OpenAI-Compatible)
-
-Enable Cache-DiT for online serving by passing `--cache-backend cache_dit` when starting the server:
-
-```bash
-# Use Cache-DiT default (recommended) parameters
-vllm serve Qwen/Qwen-Image --omni --port 8091 --cache-backend cache_dit
-```
-
-To customize Cache-DiT settings for online serving, pass a JSON string via `--cache-config`:
-
-```bash
-vllm serve Qwen/Qwen-Image --omni --port 8091 \
-  --cache-backend cache_dit \
-  --cache-config '{"Fn_compute_blocks": 1, "Bn_compute_blocks": 0, "max_warmup_steps": 4, "residual_diff_threshold": 0.12}'
-```
-
-## Acceleration Methods
-
-For comprehensive illustration, please view cache-dit [User_Guide](https://cache-dit.readthedocs.io/en/latest/user_guide/OVERVIEWS/)
-
-### 1. DBCache (Dual Block Cache)
-
-DBCache intelligently caches intermediate transformer block outputs when the residual differences between consecutive steps are small, reducing redundant computations without sacrificing quality.
-
-**Key Parameters**:
-
- `Fn_compute_blocks` (int, default: 1): Number of **first n** transformer blocks used to compute stable feature differences. Higher values provide more accurate caching decisions but increase computation.
- `Bn_compute_blocks` (int, default: 0): Number of **last n** transformer blocks used for additional fusion. These blocks act as an auto-scaler for approximate hidden states.
- `max_warmup_steps` (int, default: 4): Number of initial steps where caching is disabled to ensure the model learns sufficient features before caching begins. Optimized for few-step distilled models.
- `residual_diff_threshold` (float, default: 0.24): Threshold for residual difference. Higher values lead to faster performance but may reduce precision. Default uses a relatively higher threshold for more aggressive caching.
- `max_cached_steps` (int, default: -1): Maximum number of cached steps. Set to -1 for unlimited caching.
- `max_continuous_cached_steps` (int, default: 3): Maximum number of consecutive cached steps. Limits consecutive caching to prevent precision degradation.
-
-**Example Configuration**:
-
-```python
-cache_config={
-    "Fn_compute_blocks": 8,      # Use first 8 blocks for difference computation
-    "Bn_compute_blocks": 0,       # No additional fusion blocks
-    "max_warmup_steps": 8,        # Cache after 8 warmup steps
-    "residual_diff_threshold": 0.12,  # Higher threshold for faster inference
-    "max_cached_steps": -1,        # No limit on cached steps
-}
-```
-
-**Performance Tips**:
-
- Default `Fn_compute_blocks=1` works well for most cases. Increase to 8-12 for larger models or when more accuracy is needed
- Increase `residual_diff_threshold` (e.g., 0.12-0.15) for faster inference with slight quality trade-off, or decrease from default 0.24 for higher quality
- Default `max_warmup_steps=4` is optimized for few-step models. Increase to 6-8 for more steps if needed
-
-### 2. TaylorSeer
-
-TaylorSeer uses Taylor expansion to forecast future hidden states, allowing the model to skip some computation steps while maintaining quality.
-
-**Key Parameters**:
-
- `enable_taylorseer` (bool, default: False): Enable TaylorSeer acceleration
- `taylorseer_order` (int, default: 1): Order of Taylor expansion. Higher orders provide better accuracy but require more computation.
-
-**Example Configuration**:
-
-```python
-cache_config={
-    "enable_taylorseer": True,
-    "taylorseer_order": 1,  # First-order Taylor expansion
-}
-```
-
-**Performance Tips**:
-
- Use `taylorseer_order=1` for most cases (good balance of speed and quality)
- Combine with DBCache for maximum acceleration
- Higher orders (2-3) may improve quality but reduce speed gains
-
-### 3. SCM (Step Computation Masking)
-
-SCM allows you to specify which steps must be computed and which can use cached results, similar to LeMiCa/EasyCache style acceleration.
-
-**Key Parameters**:
-
- `scm_steps_mask_policy` (str | None, default: None): Predefined mask policy. Options:
-  - `None`: SCM disabled (default)
-  - `"slow"`: More compute steps, higher quality (18 compute steps out of 28)
-  - `"medium"`: Balanced (15 compute steps out of 28)
-  - `"fast"`: More cache steps, faster inference (11 compute steps out of 28)
-  - `"ultra"`: Maximum speed (8 compute steps out of 28)
- `scm_steps_policy` (str, default: "dynamic"): Policy for cached steps:
-  - `"dynamic"`: Use dynamic cache for masked steps (recommended)
-  - `"static"`: Use static cache for masked steps
-
-**Example Configuration**:
-
-```python
-cache_config={
-    "scm_steps_mask_policy": "medium",  # Balanced speed/quality
-    "scm_steps_policy": "dynamic",      # Use dynamic cache
-}
-```
-
-**Performance Tips**:
-
- SCM is disabled by default (`scm_steps_mask_policy=None`). Enable it by setting a policy value if you need additional acceleration
- Start with `"medium"` policy and adjust based on quality requirements
- Use `"fast"` or `"ultra"` for maximum speed when quality can be slightly compromised
- `"dynamic"` policy generally provides better quality than `"static"`
- SCM mask is automatically regenerated when `num_inference_steps` changes during inference
-
-## Configuration Reference
-
-### DiffusionCacheConfig Parameters
-
-| Parameter | Type | Default | Description |
-|-----------|------|---------|-------------|
-| `Fn_compute_blocks` | int | 1 | First n blocks for difference computation (optimized for single-transformer models) |
-| `Bn_compute_blocks` | int | 0 | Last n blocks for fusion |
-| `max_warmup_steps` | int | 4 | Steps before caching starts (optimized for few-step distilled models) |
-| `max_cached_steps` | int | -1 | Max cached steps (-1 = unlimited) |
-| `max_continuous_cached_steps` | int | 3 | Max consecutive cached steps (prevents precision degradation) |
-| `residual_diff_threshold` | float | 0.24 | Residual difference threshold (higher for more aggressive caching) |
-| `num_inference_steps` | int \| None | None | Initial inference steps for SCM mask generation (optional, auto-refreshed during inference) |
-| `enable_taylorseer` | bool | False | Enable TaylorSeer acceleration (not suitable for few-step distilled models) |
-| `taylorseer_order` | int | 1 | Taylor expansion order |
-| `scm_steps_mask_policy` | str \| None | None | SCM mask policy (None, "slow", "medium", "fast", "ultra") |
-| `scm_steps_policy` | str | "dynamic" | SCM computation policy ("dynamic" or "static") |
-
-## Example: Accelerate Text-to-Image Generation with CacheDiT
-
-See `examples/offline_inference/text_to_image/text_to_image.py` for a complete working example with cache-dit acceleration.
-
-```bash
-# Enable cache-dit with hybrid acceleration
-cd examples/offline_inference/text_to_image
-python text_to_image.py \
-    --model Qwen/Qwen-Image \
-    --prompt "a cup of coffee on the table" \
-    --cache-backend cache_dit \
-    --num-inference-steps 50
-```
-
-
-The script uses cache-dit acceleration with a hybrid configuration combining DBCache, SCM, and TaylorSeer:
-
-```python
-omni = Omni(
-    model="Qwen/Qwen-Image",
-    cache_backend="cache_dit",
-    cache_config={
-        # Scheme: Hybrid DBCache + SCM + TaylorSeer
-        # DBCache
-        "Fn_compute_blocks": 8,
-        "Bn_compute_blocks": 0,
-        "max_warmup_steps": 4,
-        "residual_diff_threshold": 0.12,
-        # TaylorSeer
-        "enable_taylorseer": True,
-        "taylorseer_order": 1,
-        # SCM
-        "scm_steps_mask_policy": "fast",  # Set to None to disable SCM
-        "scm_steps_policy": "dynamic",
-    },
-)
-```
-
-You can customize the configuration by modifying the `cache_config` dictionary to use only specific methods (e.g., DBCache only, DBCache + SCM, etc.) based on your quality and speed requirements.
-
-To test another model, you can modify `--model` with the target model identifier like `Tongyi-MAI/Z-Image-Turbo` and update `cache_config` according the model architecture (e.g., number of transformer blocks).
-
-
-## Additional Resources
-
- [Cache-DiT User Guide](https://cache-dit.readthedocs.io/en/latest/user_guide/OVERVIEWS/)
- [Cache-DiT Benchmark](https://cache-dit.readthedocs.io/en/latest/benchmark/HYBRID_CACHE/)
- [DBCache Technical Details](https://cache-dit.readthedocs.io/en/latest/user_guide/CACHE_API/)
@@ -0,0 +1,169 @@
+# CFG-Parallel Guide
+
+
+## Table of Content
+
+- [Overview](#overview)
+- [Quick Start](#quick-start)
+- [Example Script](#example-script)
+- [Configuration Parameters](#configuration-parameters)
+- [Best Practices](#best-practices)
+- [Troubleshooting](#troubleshooting)
+- [Summary](#summary)
+
+---
+
+## Overview
+
+CFG-Parallel accelerates diffusion models by distributing positive and negative classifier-free guidance (CFG) passes across different GPUs, providing ~1.8x speedup when CFG is enabled. It's ideal for image editing tasks that require guidance scales greater than 1.0.
+
+See supported models list in [Supported Models](../../diffusion_features.md#supported-models).
+
+---
+
+## Quick Start
+
+### Basic Usage
+
+Simplest working example:
+
+```python
+from vllm_omni import Omni
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from PIL import Image
+
+omni = Omni(
+    model="Qwen/Qwen-Image-Edit",
+    parallel_config=DiffusionParallelConfig(cfg_parallel_size=2),  # Enable CFG-Parallel
+)
+
+input_image = Image.open("input.png").convert("RGB")
+outputs = omni.generate(
+    {
+        "prompt": "turn this cat to a dog",
+        "negative_prompt": "low quality, blurry",
+        "multi_modal_data": {"image": input_image},
+    },
+    OmniDiffusionSamplingParams(
+        true_cfg_scale=4.0,
+        num_inference_steps=50,
+    ),
+)
+```
+
+---
+
+## Example Script
+
+### Offline Inference
+
+Use python script under `examples/offline_inference/image_to_image/image_edit.py`:
+
+```bash
+cd examples/offline_inference/image_to_image/
+python image_edit.py \
+  --model "Qwen/Qwen-Image-Edit" \
+  --image "input.png" \
+  --prompt "turn this cat to a dog" \
+  --negative-prompt "low quality, blurry" \
+  --cfg-scale 4.0 \
+  --output "edited_image.png" \
+  --cfg-parallel-size 2
+```
+
+### Online Serving
+
+Enable CFG-Parallel in online serving:
+
+```bash
+# Default configuration
+vllm serve Qwen/Qwen-Image-Edit --omni --port 8091 --cfg-parallel-size 2
+
+```
+
+---
+
+## Configuration Parameters
+
+In `DiffusionParallelConfig`
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `cfg_parallel_size` | int | 1 | Number of GPUs for CFG parallelism. Set to 2 to enable CFG-Parallel (rank 0 for positive, rank 1 for negative branch) |
+
+
+!!! info
+    Most models support `cfg_parallel_size=2` (positive branch on rank 0, negative branch on rank 1). **Bagel** is an exception: it supports `cfg_parallel_size=3`, which adds a third branch on rank 2 for full three-way CFG parallelism.
+
+
+---
+
+## Best Practices
+
+### When to Use
+
+**Good for:**
+
+- Tasks requiring classifier-free guidance
+- Multi-GPU setups (at least 2 GPUs available)
+- Combining with other parallelism methods (sequence/tensor parallel)
+
+**Not for:**
+
+- Single GPU setups
+- Models that don't support CFG-Parallel (check [supported models](../../diffusion_features.md#supported-models))
+- Workloads without negative prompts or classifier-free guidance
+- Very short inference runs (< 10 steps) where parallelism overhead may outweigh benefits
+
+### Expected Performance
+
+| Configuration | Speedup | Quality | Use Case |
+|--------------|---------|---------|----------|
+| CFG-Parallel (2 GPUs) | 1.5~1.8x | No degradation | Large model, VRAM limited |
+
+---
+
+## Troubleshooting
+
+### Common Issue 1: No Speedup with CFG-Parallel
+
+**Symptoms**: CFG-Parallel enabled but no performance improvement
+
+**Solutions**:
+
+1. **Ensure CFG scale is set correctly:**
+```python
+# Bad: No CFG effect
+sampling_params = OmniDiffusionSamplingParams(num_inference_steps=50)
+
+# Good: CFG-Parallel will work
+sampling_params = OmniDiffusionSamplingParams(
+    num_inference_steps=50,
+    true_cfg_scale=4.0  # Must be > 1.0
+)
+```
+
+2. **Add negative prompt:**
+```python
+outputs = omni.generate(
+    {
+        "prompt": "beautiful landscape",
+        "negative_prompt": "low quality, blurry",  # Required for best results
+        "multi_modal_data": {"image": input_image}
+    },
+    sampling_params
+)
+```
+
+3. **Check model support:**
+   - Verify your model in [supported models](../../diffusion_features.md#supported-models)
+   - Some models don't support CFG-Parallel
+
+---
+
+## Summary
+
+1. ✅ **Enable CFG-Parallel** - Set `cfg_parallel_size=2` in `DiffusionParallelConfig` to get speedup when using CFG
+2. ✅ **Set CFG Scale** - Ensure `true_cfg_scale > 1.0` in `OmniDiffusionSamplingParams` for CFG-Parallel to take effect
+3. ✅ **Check Model Support** - Verify your model supports CFG-Parallel in [supported models](../../diffusion_features.md#supported-models)
@@ -0,0 +1,87 @@
+# Expert Parallelism Guide
+
+
+## Table of Content
+
+- [Overview](#overview)
+- [Quick Start](#quick-start)
+- [Configuration Parameters](#configuration-parameters)
+- [Best Practices](#best-practices)
+- [Summary](#summary)
+
+---
+
+## Overview
+
+Unlike Tensor Parallelism which shards every layer's weights, Expert Parallelism (EP) only shards the MoE expert MLP blocks. This significantly reduces the memory footprint of MoE models (e.g., HunyuanImage3.0) while maintaining constant dense-equivalent compute efficiency.
+
+During the forward pass, a gating mechanism routes tokens to their designated experts, requiring all-to-all communication to dispatch tokens to the correct ranks and combine results.
+
+See supported models list in [Supported Models](../../diffusion_features.md#supported-models).
+
+!!! note "EP Size Constraint"
+    The effective EP size equals `tp × sp × cfg × dp`. At least one of TP/SP/CFG/DP must be set when EP is enabled.
+
+---
+
+## Quick Start
+
+### Basic Usage
+
+```python
+from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+
+omni = Omni(
+    model="tencent/HunyuanImage-3.0",
+    parallel_config=DiffusionParallelConfig(
+        tensor_parallel_size=8,
+        enable_expert_parallel=True,
+    ),
+)
+
+outputs = omni.generate(
+    "A brown and white dog is running on the grass",
+    OmniDiffusionSamplingParams(
+        num_inference_steps=50,
+        width=1024,
+        height=1024,
+    ),
+)
+```
+
+---
+
+## Configuration Parameters
+
+In `DiffusionParallelConfig`:
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `enable_expert_parallel` | bool | False | Enable Expert Parallelism for MoE models |
+
+EP size is derived automatically as `tp × sp × cfg × dp` — configure at least one of those to set the EP degree.
+
+---
+
+## Best Practices
+
+### When to Use
+
+**Good for:**
+
+- MoE models (e.g., HunyuanImage3.0) with numbers of experts
+- Memory-constrained multi-GPU setups where only expert blocks need sharding
+
+**Not for:**
+
+- Dense models (no MoE layers) — EP has no effect
+- Single GPU setups
+
+---
+
+## Summary
+
+1. ✅ **Enable EP** - Set `enable_expert_parallel=True` in `DiffusionParallelConfig` for MoE models
+2. ✅ **Set parallelism degree** - At least one of `tensor_parallel_size` / `ulysses_degree` / `cfg_parallel_size` must be > 1 to define the EP size
@@ -0,0 +1,149 @@
+# HSDP Guide
+
+
+## Table of Content
+
+- [Overview](#overview)
+- [Quick Start](#quick-start)
+- [Example Script](#example-script)
+- [Configuration Parameters](#configuration-parameters)
+- [Best Practices](#best-practices)
+- [Summary](#summary)
+
+---
+
+## Overview
+
+HSDP (Hybrid Sharded Data Parallel) shards model weights across GPUs to reduce per-GPU memory usage. This enables inference of large models (e.g., Wan2.2 14B) on GPUs with limited memory.
+
+Unlike Tensor Parallelism which splits computation, HSDP uses PyTorch's FSDP2 to shard and redistribute weights at runtime. Each GPU only holds a fraction of the model weights, and weights are gathered on-demand during forward passes.
+
+See supported models list in [Supported Models](../../diffusion_features.md#supported-models).
+
+**Operating Modes:**
+
+- **Standalone Mode**: HSDP alone without other parallelism. Must specify `hsdp_shard_size` explicitly.
+- **Combined Mode**: HSDP overlays on top of other parallelism (Ulysses-SP, CFG-Parallel). HSDP dimensions must match world_size.
+
+---
+
+## Quick Start
+
+### Basic Usage
+
+Simplest working example (standalone HSDP, shard across 4 GPUs):
+
+```python
+from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+
+omni = Omni(
+    model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
+    parallel_config=DiffusionParallelConfig(
+        use_hsdp=True,
+        hsdp_shard_size=4,  # Shard across 4 GPUs
+    ),
+)
+
+outputs = omni.generate(
+    "A cat playing piano",
+    OmniDiffusionSamplingParams(num_inference_steps=50),
+)
+```
+
+### Combined with Sequence Parallel
+
+```python
+omni = Omni(
+    model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
+    parallel_config=DiffusionParallelConfig(
+        ulysses_degree=4,  # Sequence parallel
+        use_hsdp=True,     # HSDP overlays on SP
+    ),
+)
+```
+
+---
+
+## Example Script
+
+### Offline Inference
+
+Use Python script under `examples/offline_inference/image_to_video/`:
+
+```bash
+# Standalone HSDP: shard across 4 GPUs
+python examples/offline_inference/image_to_video/image_to_video.py \
+    --model Wan-AI/Wan2.2-T2V-A14B-Diffusers \
+    --use-hsdp \
+    --hsdp-shard-size 4
+
+# Combined HSDP + Sequence Parallel
+python examples/offline_inference/image_to_video/image_to_video.py \
+    --model Wan-AI/Wan2.2-T2V-A14B-Diffusers \
+    --ulysses-degree 4 \
+    --use-hsdp
+```
+
+### Online Serving
+
+**Standalone HSDP** (shard model across 4 GPUs):
+
+```bash
+vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --port 8091 \
+    --use-hsdp --hsdp-shard-size 4
+```
+
+**Combined with Sequence Parallel**:
+
+```bash
+vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --port 8091 \
+    --use-hsdp --usp 4
+```
+
+---
+
+## Configuration Parameters
+
+In `DiffusionParallelConfig`:
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `use_hsdp` | bool | False | Enable HSDP |
+| `hsdp_shard_size` | int | -1 | Number of GPUs to shard weights across. `-1` = auto (requires other parallelism > 1) |
+| `hsdp_replicate_size` | int | 1 | Number of replica groups. Each group holds a full sharded copy |
+
+**Constraints:**
+
+- `hsdp_replicate_size × hsdp_shard_size == world_size`
+- HSDP cannot be used with Tensor Parallelism (`tensor_parallel_size` must be 1)
+
+---
+
+## Best Practices
+
+### When to Use
+
+**Good for:**
+
+- Very large models (e.g., Wan2.2 14B)
+- Multi-GPU setups where memory reduction is the primary goal
+- Combining with Sequence Parallelism for large video models
+
+**Not for:**
+
+- Models that fit comfortably in single-GPU memory
+- Use cases requiring Tensor Parallelism (HSDP and TP are mutually exclusive)
+
+### Adding HSDP Support to New Models
+
+For detailed instructions on adding HSDP support to new models, see the [HSDP Contributing Guide](../../../design/feature/hsdp.md).
+
+---
+
+## Summary
+
+1. ✅ **Enable HSDP** - Set `use_hsdp=True` and `hsdp_shard_size` to reduce per-GPU memory for large models
+2. ✅ **Combine with SP** - Use together with `ulysses_degree` for video models requiring both memory reduction and sequence parallelism
+3. ⚠️ **Incompatible with TP** - `tensor_parallel_size` must be 1 when HSDP is enabled
@@ -0,0 +1,16 @@
+# Parallelism Acceleration Guide
+
+This guide covers the parallelism methods in vLLM-Omni for speeding up diffusion model inference and reducing per-device memory requirements.
+
+## Supported Methods
+
+| Method | Description |
+|--------|-------------|
+| **[Tensor Parallelism](tensor_parallel.md)** | Shards DiT weights across GPUs to reduce per-GPU memory |
+| **[Sequence Parallelism](sequence_parallel.md)** | Splits sequence dimension across GPUs (Ulysses-SP, Ring-Attention, or hybrid) for high-resolution images and videos |
+| **[CFG-Parallel](cfg_parallel.md)** | Runs CFG positive/negative branches on separate GPUs for ~1.8x speedup on guided generation |
+| **[VAE Patch Parallelism](vae_patch_parallel.md)** | Distributes VAE decode spatially across GPUs to reduce peak VAE memory |
+| **[HSDP](hsdp.md)** | Shards full model weights via PyTorch FSDP2 to enable large-model inference on memory-constrained GPUs |
+| **[Expert Parallelism](expert_parallel.md)** | Shards MoE expert blocks across GPUs for MoE models (e.g. HunyuanImage3.0) |
+
+See [Supported Models](../../diffusion_features.md#supported-models) for per-model compatibility.
@@ -0,0 +1,233 @@
+# Sequence Parallelism Guide
+
+
+## Table of Content
+
+- [Overview](#overview)
+- [Quick Start](#quick-start)
+- [Example Script](#example-script)
+- [Best Practices](#best-practices)
+- [Troubleshooting](#troubleshooting)
+- [Summary](#summary)
+
+---
+
+## Overview
+
+Sequence parallelism splits the input along the sequence dimension across multiple GPUs, allowing each device to process only a portion of the sequence. vLLM-Omni provides 1.5x-3.6x speedup for large images and videos using DeepSpeed Ulysses, Ring-Attention, or hybrid approaches. Use sequence parallelism when generating high-resolution images/videos that don't fit on a single GPU or require faster inference.
+
+See supported models list in [Diffusion Features - Supported Models](../../diffusion_features.md#supported-models).
+
+**Supported Methods:**
+
+- **DeepSpeed Ulysses Sequence Parallel (Ulysses-SP)** ([paper](https://arxiv.org/pdf/2309.14509)): Uses all-to-all communication for subset of attention heads per device
+- **Ring-Attention** ([paper](https://arxiv.org/abs/2310.01889)): Uses ring-based P2P communication with sharded sequence dimension throughout
+- **Hybrid Ulysses + Ring**: Combines both for larger scale parallelism (`ulysses_degree × ring_degree`)
+
+---
+
+## Quick Start
+
+### Basic Usage - Ulysses-SP
+
+Simplest working example with Ulysses Sequence Parallel:
+
+```python
+from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+
+omni = Omni(
+    model="Qwen/Qwen-Image",
+    parallel_config=DiffusionParallelConfig(ulysses_degree=2)  # Enable Ulysses-SP
+)
+
+outputs = omni.generate(
+    "A cat sitting on a windowsill",
+    OmniDiffusionSamplingParams(num_inference_steps=50, width=1024, height=1024),
+)
+```
+
+!!! note "Experimental UAA mode"
+    `ulysses_mode="advanced_uaa"` is an experimental extension to Ulysses-SP. It lets Ulysses attention handle arbitrary sequence lengths and arbitrary attention head counts without relying on `attention_mask`-based token padding.
+
+    In hybrid Ulysses + Ring mode, Ring still requires every rank in the same ring group to observe the same post-Ulysses sequence length. If that condition is not met, vLLM-Omni raises a validation error instead of entering the ring kernel with inconsistent shapes.
+
+To enable the experimental UAA mode, use a model/configuration that requires it. For example, `Tongyi-MAI/Z-Image-Turbo` has 30 attention heads, so `ulysses_degree=4` requires UAA because 30 is not divisible by 4:
+
+```python
+omni = Omni(
+    model="Tongyi-MAI/Z-Image-Turbo",
+    parallel_config=DiffusionParallelConfig(
+        ulysses_degree=4,
+        ulysses_mode="advanced_uaa",
+    ),
+)
+```
+
+### Alternative Methods
+
+**Ring-Attention** (better for very long sequences):
+
+```python
+omni = Omni(
+    model="Qwen/Qwen-Image",
+    parallel_config=DiffusionParallelConfig(ring_degree=2)  # Enable Ring-Attention
+)
+```
+
+**Hybrid Ulysses + Ring** (for larger scale):
+
+```python
+omni = Omni(
+    model="Qwen/Qwen-Image",
+    parallel_config=DiffusionParallelConfig(ulysses_degree=2, ring_degree=2)  # 4 GPUs total
+)
+```
+
+---
+
+## Example Script
+
+### Offline Inference
+
+Use Python script under `examples/offline_inference/text_to_image/text_to_image.py`:
+
+**Ulysses-SP:**
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+    --model Qwen/Qwen-Image \
+    --prompt "A cat sitting on a windowsill" \
+    --ulysses-degree 2 \
+    --width 1024 --height 1024
+```
+
+**Ring-Attention:**
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+    --model Qwen/Qwen-Image \
+    --prompt "A cat sitting on a windowsill" \
+    --ring-degree 2 \
+    --width 1024 --height 1024
+```
+
+**Hybrid Ulysses + Ring:**
+
+```bash
+# Hybrid: 2 Ulysses × 2 Ring = 4 GPUs total
+python examples/offline_inference/text_to_image/text_to_image.py \
+    --model Qwen/Qwen-Image \
+    --prompt "A cat sitting on a windowsill" \
+    --ulysses-degree 2 --ring-degree 2 \
+    --width 1024 --height 1024
+```
+
+### Online Serving
+
+**Ulysses-SP:**
+
+```bash
+# Text-to-image (requires >= 2 GPUs)
+vllm serve Qwen/Qwen-Image --omni --port 8091 --usp 2
+```
+
+**Ulysses-SP with UAA mode** (for models with non-divisible head counts):
+
+```bash
+vllm serve Tongyi-MAI/Z-Image-Turbo --omni --port 8091 --usp 4 --ulysses-mode advanced_uaa
+```
+
+**Ring-Attention:**
+
+```bash
+# Text-to-image (requires >= 2 GPUs)
+vllm serve Qwen/Qwen-Image --omni --port 8091 --ring 2
+```
+
+**Hybrid Ulysses + Ring:**
+
+```bash
+# Text-to-image (requires >= 4 GPUs)
+vllm serve Qwen/Qwen-Image --omni --port 8091 --usp 2 --ring 2
+```
+
+---
+
+## Configuration Parameters
+
+In `DiffusionParallelConfig`:
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `ulysses_degree` | int | 1 | Number of GPUs for Ulysses-SP. Uses all-to-all communication. |
+| `ring_degree` | int | 1 | Number of GPUs for Ring-Attention. Uses P2P ring communication. |
+| `ulysses_mode` | str | `"default"` | Ulysses attention mode. Set to `"advanced_uaa"` to handle arbitrary sequence lengths and head counts without padding. |
+
+**Notes:**
+- Total sequence parallel size equals to `ulysses_degree × ring_degree`
+- Degrees must evenly divide the sequence length for optimal performance (or use `ulysses_mode="advanced_uaa"` for Ulysses-SP)
+
+
+## Best Practices
+
+### When to Use
+
+**Good for:**
+
+- Large images (1024x1024 or higher) or videos
+- Fast inter-GPU communication, larger bandwidth (e.g., NVLink)
+
+**Not for:**
+
+- Small images (<1024px) - overhead exceeds benefit, use single GPU with cache instead
+
+
+---
+
+## Troubleshooting
+
+### Common Issue 1: Performance Not Scaling
+
+**Symptoms**: Adding GPUs doesn't improve speed proportionally, or higher parallelism degree is slower
+
+**Diagnosis:**
+```bash
+# Check GPU topology
+nvidia-smi topo -m
+
+```
+
+**Solutions:**
+
+1. Check inter-GPU communication - NVLink is better than PCIe
+2. Reduce parallelism degree if over-parallelized:
+```python
+# If 4 GPUs is slower than 2
+parallel_config=DiffusionParallelConfig(ulysses_degree=2)
+```
+3. Try to switch between Ring-Attention and Ulysses-SP
+
+- Ring-Attention has advantages, like communication-computation overlap, but the block-wise loop overhead is relatively higher, especially for short sequences
+- Ulysses-SP: can benefit from larger bandwidth (such as NVLink), with two major constraints, the sequence length should be divisible by usp size, and the number of heads should be divisible by usp size (or use `ulysses_mode="advanced_uaa"`)
+
+
+### Common Issue 2: Out of Memory (OOM)
+
+**Symptoms**: CUDA OOM errors or process crashes with memory errors
+
+**Solutions:**
+
+1. Increase parallelism degree to split sequence more:
+```python
+parallel_config=DiffusionParallelConfig(ulysses_degree=4)  # From 2
+```
+2. Combine with other parallelism method, e.g., tensor parallel, and memory optimization methods, e.g., cpu offloading.
+
+
+## Summary
+
+1. ✅ **Enable Sequence Parallelism** - Set `ulysses_degree` or `ring_degree` for long sequence generation
+2. ✅ **UAA mode** - Use `ulysses_mode="advanced_uaa"` when head count is not divisible by `ulysses_degree`
+3. ✅ **Troubleshooting** - Check GPU topology with `nvidia-smi topo -m`, reduce degree if performance doesn't scale
@@ -0,0 +1,151 @@
+# Tensor Parallelism Guide
+
+
+## Table of Content
+
+- [Overview](#overview)
+- [Quick Start](#quick-start)
+- [Example Script](#example-script)
+- [Configuration Parameters](#configuration-parameters)
+- [Best Practices](#best-practices)
+- [Troubleshooting](#troubleshooting)
+- [Summary](#summary)
+
+---
+
+## Overview
+
+Tensor Parallelism (TP) shards some model weights across multiple GPUs, usually the Linear layers. This enables running large models that don't fit on a single GPU. It's essential for memory-constrained setups or very large models.
+
+See supported models list in [Supported Models](../../diffusion_features.md#supported-models).
+
+!!! note "TP Limitations for Diffusion Models"
+    We currently implement Tensor Parallelism (TP) only for the DiT (Diffusion Transformer) blocks. This is because the `text_encoder` component in vLLM-Omni uses the original Transformers implementation, which does not yet support TP.
+
+    - Good news: The text_encoder typically has minimal impact on overall inference performance.
+    - Bad news: When TP is enabled, every TP process retains a full copy of the text_encoder weights, leading to significant GPU memory waste.
+
+    We are actively refactoring this design to address this. For details and progress, please refer to [Issue #771](https://github.com/vllm-project/vllm-omni/issues/771).
+
+---
+
+## Quick Start
+
+
+### Basic Usage
+
+Simplest working example:
+
+```python
+from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+
+omni = Omni(
+    model="Tongyi-MAI/Z-Image-Turbo",
+    parallel_config=DiffusionParallelConfig(tensor_parallel_size=2),  # Enable TP
+)
+
+outputs = omni.generate(
+    "a cat reading a book",
+    OmniDiffusionSamplingParams(num_inference_steps=9),
+)
+```
+
+---
+
+## Example Script
+
+### Offline Inference
+
+Use Python script under `examples/offline_inference`, and enable TP:
+
+```bash
+# Text-to-Image with Qwen-Image
+python examples/offline_inference/text_to_image/text_to_image.py \
+    --model Qwen/Qwen-Image \
+    --tensor-parallel-size 2
+
+# Image Editing with Qwen-Image-Edit
+python examples/offline_inference/image_to_image/image_edit.py \
+    --model Qwen/Qwen-Image-Edit \
+    --image input.png \
+    --prompt "Edit description" \
+    --tensor-parallel-size 2
+```
+
+### Online Serving
+
+You can enable tensor parallelism in online serving via `--tensor-parallel-size`:
+
+```bash
+# Text-to-Image with Qwen-Image on 2 GPUs
+vllm serve Qwen/Qwen-Image --omni --port 8091 \
+    --tensor-parallel-size 2
+
+# Text-to-Image with Z-Image (TP=2 only)
+vllm serve Tongyi-MAI/Z-Image-Turbo --omni --port 8091 \
+    --tensor-parallel-size 2
+```
+
+---
+
+## Configuration Parameters
+
+In `DiffusionParallelConfig`:
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `tensor_parallel_size` | int | 1 | Number of GPUs to shard model weights across. Must divide number of heads. |
+
+
+---
+
+## Best Practices
+
+### When to Use
+
+**Good for:**
+
+- Large models that don't fit on a single GPU, especially for models with large DiT blocks (transformer layers)
+- Memory-constrained environments
+
+**Not for:**
+
+- When maximum throughput is needed and memory is sufficient
+- Models with incompatible dimensions (e.g., Z-Image `num_heads=30`, which now supports `tensor_parallel_size=2`)
+
+
+## Troubleshooting
+
+### Common Issue 1: Out of Memory (OOM)
+
+**Symptoms**: CUDA OOM errors during model loading or inference, process crashes with memory errors
+
+**Solution**:
+```python
+# Step 1: Enable TP with smallest degree
+parallel_config=DiffusionParallelConfig(tensor_parallel_size=2)
+
+# Step 2: If still OOM, increase TP degree
+parallel_config=DiffusionParallelConfig(tensor_parallel_size=4)
+
+```
+
+### Common Issue 2: Divisibility Error
+
+**Symptoms**: Error like "Model dimension X not divisible by tensor_parallel_size Y"
+
+**Solutions**:
+1. Check model-specific constraints (e.g., Z-Image only supports TP=2)
+2. Use a smaller TP size that divides model dimensions
+3. Consult [Supported Models](../../diffusion_features.md#supported-models) for compatible TP sizes
+
+
+---
+
+## Summary
+
+1. ✅ **Enable TP** - Set `--tensor-parallel-size` to reduce per-GPU memory
+2. ✅ **Increase TP size** - Only increase if OOM persists
+3. ⚠️ **Text encoder not sharded** - Known limitation
@@ -0,0 +1,200 @@
+# VAE Patch Parallelism Guide
+
+
+## Table of Content
+
+- [Overview](#overview)
+- [Quick Start](#quick-start)
+- [Example Script](#example-script)
+- [Configuration Parameters](#configuration-parameters)
+- [Best Practices](#best-practices)
+- [Troubleshooting](#troubleshooting)
+- [Summary](#summary)
+
+---
+
+## Overview
+
+VAE Patch Parallelism distributes the VAE (Variational AutoEncoder) decode/encode computation across multiple GPUs by splitting the latent space into spatial tiles or patches. Each GPU processes a subset of tiles in parallel, significantly reducing peak memory consumption during the VAE decode stage while maintaining output quality.
+
+This is particularly useful for:
+- **High-resolution image generation** where VAE decode can become a memory bottleneck
+- **Memory-constrained environments** where the VAE decode activation peak exceeds available VRAM
+- **Multi-GPU setups** where you want to leverage distributed resources for the VAE stage
+
+See supported models list in [Supported Models](../../diffusion_features.md#supported-models).
+
+
+VAE Patch Parallelism uses two strategies based on image size:
+
+| Strategy | Use Case | How It Works | Overlap Handling | Output Quality |
+|----------|----------|--------------|------------------|----------------|
+| **Tiled Decode** | Large images (triggers VAE tiling) | Distributes existing VAE tiling computation across ranks. Each rank decodes a subset of overlapping tiles. | Uses VAE's native `blend_v` and `blend_h` functions to seamlessly merge overlapping regions | Bit-identical (same logic as single-GPU tiling) |
+| **Patch Decode** | Small images (no VAE tiling) | Splits latent into spatial patches with halos. Each rank decodes one patch with boundary context. | Halo regions provide edge context; core regions are directly stitched without blending | Near-identical (diff < 0.5%, visually imperceptible) |
+
+
+VAE Patch Parallelism **reuses the DiT process group** (`dit_group`) and does not initialize a separate ProcessGroup. This means:
+
+- **Shared ranks**: VAE patch parallelism uses the same GPU ranks as DiT parallelism (Tensor Parallel, Sequence Parallel, etc.)
+- **Combined usage**: VAE patch parallelism is typically used together with other parallelism methods
+- **Configuration alignment**: The `vae_patch_parallel_size` should be no greater than the size of your DiT process group
+
+---
+
+## Quick Start
+
+### Basic Usage
+
+Simplest working example:
+
+```python
+from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+
+# TP=2 for DiT, VAE patch parallel also uses these 2 GPUs
+omni = Omni(
+    model="Tongyi-MAI/Z-Image-Turbo",
+    parallel_config=DiffusionParallelConfig(
+        tensor_parallel_size=2,          # Enable tensor parallelism for DiT
+        vae_patch_parallel_size=2,       # Enable VAE patch parallelism
+    ),
+    vae_use_tiling=True,  # Required for VAE patch parallelism
+)
+
+outputs = omni.generate(
+    "a futuristic city at sunset, high resolution, 8k",
+    OmniDiffusionSamplingParams(
+        num_inference_steps=9,
+        height=1152,  # High resolution benefits from VAE patch parallel
+        width=1152,
+    ),
+)
+```
+
+---
+
+## Example Script
+
+### Offline Inference
+
+Use Python script under `examples/offline_inference/text_to_image/`:
+
+```bash
+# Text-to-Image with Z-Image
+python examples/offline_inference/text_to_image/text_to_image.py \
+    --model Tongyi-MAI/Z-Image-Turbo \
+    --prompt "a futuristic city at sunset" \
+    --height 1152 \
+    --width 1152 \
+    --tensor-parallel-size 2 \
+    --vae-patch-parallel-size 2 \
+    --vae-use-tiling
+```
+
+### Online Serving
+
+You can enable VAE patch parallelism in online serving via `--vae-patch-parallel-size`:
+
+```bash
+# Text-to-Image with Z-Image, TP=2 + VAE patch parallel=2
+vllm serve Tongyi-MAI/Z-Image-Turbo --omni --port 8091 \
+    --tensor-parallel-size 2 \
+    --vae-patch-parallel-size 2 \
+    --vae-use-tiling
+```
+
+---
+
+## Configuration Parameters
+
+In `DiffusionParallelConfig`:
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `vae_patch_parallel_size` | int | 1 | Number of GPUs for VAE patch/tile parallelism. Set to 2 or higher to enable. Should typically match `tensor_parallel_size` as they share the same process group. |
+
+Additional requirements:
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `vae_use_tiling` | bool | False | Must be set to `True` when using VAE patch parallelism. |
+
+!!! note "Automatic VAE Tiling"
+    When `vae_patch_parallel_size > 1` and the model has a distributed VAE (`DistributedVaeMixin`), the system automatically sets `vae_use_tiling=True` if not already enabled.
+
+---
+
+## Best Practices
+
+### When to Use
+
+**Good for:**
+
+- High-resolution image generation and long video generation
+- Memory-constrained setups where VAE decode causes OOM
+- Multi-GPU environments
+
+**Not for:**
+
+- Low-resolution images/videos where VAE decode is not a bottleneck
+- Single GPU setups should use vae tiling decode, but not parallel vae tiling decode
+- Models that do not support vae patch parallel
+
+---
+
+## Troubleshooting
+
+### Common Issue 1: Model Not Support VAE Patch Parallel
+
+**Symptoms**:
+```
+WARNING: vae_patch_parallel_size=2 is set but VAE patch parallelism is NOT enabled for xxxPipeline; ignoring.
+```
+
+**Root Cause**: VAE Patch Parallelism requires the model's VAE to implement `DistributedVaeMixin`. At startup, `vllm_omni/diffusion/registry.py` checks whether the instantiated pipeline has a `.vae` attribute that is an instance of `DistributedVaeMixin`. If it does not, the setting is silently ignored:
+
+```python
+vae_pp_size = od_config.parallel_config.vae_patch_parallel_size
+is_distributed_vae = hasattr(model, "vae") and isinstance(model.vae, DistributedVaeMixin)
+if vae_pp_size > 1 and not is_distributed_vae:
+    logger.warning(
+        "vae_patch_parallel_size=%d is set but VAE patch parallelism is NOT enabled for %s; ignoring.",
+        vae_pp_size,
+        od_config.model_class_name,
+    )
+```
+
+**Solutions**:
+
+1. **Use a supported model** (recommended): check [Supported Models](../../diffusion_features.md#supported-models) for the VAE-Patch-Parallel column.
+
+2. To add support for a new model, implement `DistributedVaeMixin` on its VAE class (contributions are welcome).
+
+
+### Common Issue 2: `vae_patch_parallel_size` Exceeds DiT Process Group Size
+
+**Symptoms**: Shows warning message, and vae patch parallel size is resized to DiT process group size
+
+**Root Cause**: VAE Patch Parallelism reuses the DiT process group.
+
+**Recommendation**: Always set `vae_patch_parallel_size` to be no greater than your DiT process group size.
+
+Note that the size of DiT process group size equals to:
+```text
+dit_parallel_size = data_parallel_size
+                  × cfg_parallel_size
+                  × sequence_parallel_size
+                  × pipeline_parallel_size
+                  × tensor_parallel_size
+
+```
+_sequence_parallel_size = ulysses_degree × ring_degree_
+
+---
+
+## Summary
+
+1. ✅ **Enable VAE Patch Parallelism** - Set `vae_patch_parallel_size`， `vae_use_tiling=True` in `DiffusionParallelConfig` to reduce VAE decode peak memory
+2. ✅ **Use Long Sequence** - VAE patch parallelism benefits are most apparent at long sequence decoding
+3. ✅ **Combine with other parallelism methods** - Suggest to use together with Tensor Parallel or CFG-Parallel for maximum memory savings
@@ -1,477 +0,0 @@
-# Parallelism Acceleration Guide
-
-This guide includes how to use parallelism methods in vLLM-Omni to speed up diffusion model inference as well as reduce the memory requirement on each device.
-
-## Overview
-
-The following parallelism methods are currently supported in vLLM-Omni:
-
-1. DeepSpeed Ulysses Sequence Parallel (DeepSpeed Ulysses-SP) ([arxiv paper](https://arxiv.org/pdf/2309.14509)): Ulysses-SP splits the input along the sequence dimension and uses all-to-all communication to allow each device to compute only a subset of attention heads.
-
-2. [Ring-Attention](#ring-attention) - splits the input along the sequence dimension and uses ring-based P2P communication to accumulate attention results, keeping the sequence dimension sharded
-
-3. Classifier-Free-Guidance Parallel (CFG-Parallel): CFG-Parallel runs the positive/negative prompts of classifier-free guidance (CFG) on different devices, then merges on a single device to perform the scheduler step.
-
-4. [Tensor Parallelism](#tensor-parallelism): Tensor parallelism shards model weights across devices. This can reduce per-GPU memory usage. Note that for diffusion models we currently shard the majority of layers within the DiT.
-
-5. [VAE Patch Parallelism](#vae-patch-parallelism): VAE patch parallelism shards VAE decode spatially across ranks. This can reduce the peak memory of VAE decode and (depending on resolution and communication overhead) speed up VAE decode.
-
-6. [HSDP](#hsdp): Hybrid Sharded Data Parallel shards model weights across GPUs using PyTorch FSDP2. This reduces per-GPU memory usage, enabling inference of large models on GPUs with limited memory.
-
-7. [Expert Parallel](#expert-parallelism): Expert Parallelism shards the Experts of a Mixture-of-Experts (MoE) layer across multiple devices. During the forward, a gating mechanism routes tokens to their designated experts, necessitating cross-cards communication(all-to-all) to dispatch tokens to the correct ranks and combine the results. This parallelism allows for massive scaling of model parameters without a proportional increase in the computational load per device.
-
-The following table shows which models are currently supported by parallelism method:
-
-### ImageGen
-
-| Model                    | Model Identifier                     | Ulysses-SP | Ring-SP | CFG-Parallel | Tensor-Parallel | VAE-Patch-Parallel | Expert-Parallel | HSDP |
-|--------------------------|--------------------------------------|:----------:|:-------:|:------------:|:---------------:|:------------------:|:---------------:|:----:|
-| **LongCat-Image**        | `meituan-longcat/LongCat-Image`      |     ✅      |    ✅    |      ❌       |        ✅        |         ❌          |      N/A        | ❌   |
-| **LongCat-Image-Edit**   | `meituan-longcat/LongCat-Image-Edit` |     ✅      |    ✅    |      ❌       |        ✅        |         ❌          |      N/A        | ❌   |
-| **Ovis-Image**           | `OvisAI/Ovis-Image`                  |     ❌      |    ❌    |      ❌       |        ❌        |         ❌          |      N/A        | ❌   |
-| **Qwen-Image**           | `Qwen/Qwen-Image`                    |     ✅      |    ✅    |      ✅       |        ✅        |         ✅          |      N/A        | ❌   |
-| **Qwen-Image-Edit**      | `Qwen/Qwen-Image-Edit`               |     ✅      |    ✅    |      ✅       |        ✅        |         ❌          |      N/A        | ❌   |
-| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509`          |     ✅      |    ✅    |      ✅       |        ✅        |         ❌          |      N/A        | ❌   |
-| **Qwen-Image-Layered**   | `Qwen/Qwen-Image-Layered`            |     ✅      |    ✅    |      ✅       |        ✅        |         ❌          |      N/A        | ❌   |
-| **Z-Image**              | `Tongyi-MAI/Z-Image-Turbo`           |     ✅      |    ✅    |      ❌       |  ✅ (TP=2 only)  |         ✅          |      N/A        | ❌   |
-| **Stable-Diffusion3.5**  | `stabilityai/stable-diffusion-3.5`   |     ❌      |    ❌    |      ❌       |        ✅        |         ✅          |      N/A        | ❌   |
-| **FLUX.2-klein**         | `black-forest-labs/FLUX.2-klein-4B`  |     ✅      |    ✅    |      ❌       |        ✅        |         ❌          |      N/A        | ✅   |
-| **FLUX.1-dev**           | `black-forest-labs/FLUX.1-dev`       |     ❌      |    ❌    |      ✅       |        ✅        |         ❌          |      N/A        | ✅   |
-| **FLUX.2-dev**           | `black-forest-labs/FLUX.2-dev`       |     ❌      |    ❌    |      ❌       |        ✅        |         ❌          |      N/A        | ✅   |
-| **HunyuanImage3.0**      | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` |     ❌      |    ❌    |      ❌       |        ✅        |         ❌          |      ✅        | ❌   |
-| **Bagel**                  | `ByteDance-Seed/BAGEL-7B-MoT` |     ✅      |    ✅    |      ✅       |        ✅        |         ❌          |      N/A        | ❌   |
-| **DreamID-Omni**           | `XuGuo699/DreamID-Omni`       |     ❌      |    ❌    |      ✅       |        ❌        |         ❌          |      N/A        | ❌   |
-| **FLUX.1-Kontext-dev**   | `black-forest-labs/FLUX.1-Kontext-dev`       |     ❌      |    ❌    |      ❌       |        ✅        |         ❌          |      N/A        | ✅   |
-| **OmniGen2**               | `OmniGen2/OmniGen2`           |     ❌      |    ❌    |      ❌       |        ✅        |         ❌          |      N/A        | ❌   |
-
-!!! note "TP Limitations for Diffusion Models"
-    We currently implement Tensor Parallelism (TP) only for the DiT (Diffusion Transformer) blocks. This is because the `text_encoder` component in vLLM-Omni uses the original Transformers implementation, which does not yet support TP.
-
-    - Good news: The text_encoder typically has minimal impact on overall inference performance.
-    - Bad news: When TP is enabled, every TP process retains a full copy of the text_encoder weights, leading to significant GPU memory waste.
-
-    We are actively refactoring this design to address this. For details and progress, please refer to [Issue #771](https://github.com/vllm-project/vllm-omni/issues/771).
-
-
-!!! note "Why Z-Image is TP=2 only"
-    Z-Image Turbo is currently limited to `tensor_parallel_size` of **1 or 2** due to model shape divisibility constraints.
-    For example, the model has `n_heads=30` and a final projection out dimension of `64`, so valid TP sizes must divide both 30 and 64; the only common divisors are **1 and 2**.
-
-### VideoGen
-
-| Model | Model Identifier | Ulysses-SP | Ring-Attention | Tensor-Parallel | HSDP | VAE-Patch-Parallel |
-|-------|------------------|:----------:|:--------------:|:---------------:|:----:| :----:|
-| **Wan2.1** | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers` | ✅ | ✅ | ✅ | ✅ | ✅ |
-| **Wan2.1** | `Wan-AI/Wan2.1-T2V-14B-Diffusers` | ✅ | ✅ | ✅ | ✅ | ✅ |
-| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ✅ | ✅ | ✅ | ✅ | ✅ |
-| **LTX-2** | `Lightricks/LTX-2` | ✅ | ✅ | ✅ | ❌ | ❌ |
-
-### Tensor Parallelism
-
-Tensor parallelism splits model parameters across GPUs. In vLLM-Omni, tensor parallelism is configured via `DiffusionParallelConfig.tensor_parallel_size`.
-
-#### Offline Inference
-
-```python
-from vllm_omni import Omni
-from vllm_omni.diffusion.data import DiffusionParallelConfig
-
-omni = Omni(
-    model="Tongyi-MAI/Z-Image-Turbo",
-    parallel_config=DiffusionParallelConfig(tensor_parallel_size=2),
-)
-
-outputs = omni.generate(
-    "a cat reading a book",
-    OmniDiffusionSamplingParams(
-        num_inference_steps=9,
-        width=512,
-        height=512,
-    ),
-)
-```
-
-### VAE Patch Parallelism
-
-VAE patch parallelism distributes the VAE decode workload across multiple ranks by splitting the latent spatially. It is configured via `DiffusionParallelConfig.vae_patch_parallel_size` and can be combined with other parallelism methods (e.g., TP).
-
-!!! note "Enablement and feature gate"
-    - VAE patch parallelism is currently **enabled only for validated pipelines** (check [ImageGen](#imagegen) and [VideoGen](#videogen) for more information).
-    - If `vae_patch_parallel_size > 1` is set for a validated pipeline, vLLM-Omni will automatically enable `vae_use_tiling` as a safety gate. (We use `vae_use_tiling` because it indicates the VAE supports diffusers tiling parameters like `tile_latent_min_size` and `tile_overlap_factor`.)
-
-#### Offline Inference
-
-```python
-from vllm_omni import Omni
-from vllm_omni.diffusion.data import DiffusionParallelConfig
-
-omni = Omni(
-    model="Tongyi-MAI/Z-Image-Turbo",
-    parallel_config=DiffusionParallelConfig(
-        tensor_parallel_size=2,
-        vae_patch_parallel_size=2,
-    ),
-    vae_use_tiling=True,
-)
-
-outputs = omni.generate(
-    prompt="a cat reading a book",
-    num_inference_steps=9,
-    width=1024,
-    height=1024,
-)
-```
-
-#### How it works (method selection)
-
-VAE patch parallelism automatically selects between two internal decode methods based on whether diffusers tiling would kick in:
-
- `_distributed_tiled_decode`: Used when the latent spatial size exceeds `vae.tile_latent_min_size` (i.e., diffusers tiled decode). Each rank decodes a subset of tiles; rank0 gathers and runs the same overlap+blend+stitch logic as diffusers. This matches the single-rank diffusers tiled output.
-
- `_distributed_patch_decode`: Used when diffusers tiling would not kick in. Each rank decodes a grid patch expanded with a latent-space halo; then rank0 gathers the cropped core patches and stitches them into the full image. This path has no blending and can introduce small numerical differences compared to the non-parallel decode.
-
-### Sequence Parallelism
-
-#### Ulysses-SP
-
-!!! note "Experimental UAA mode"
-    `ulysses_mode="advanced_uaa"` is an experimental extension to Ulysses-SP. It lets Ulysses attention handle arbitrary sequence lengths and arbitrary attention head counts without relying on `attention_mask`-based token padding.
-
-    In hybrid Ulysses + Ring mode, Ring still requires every rank in the same ring group to observe the same post-Ulysses sequence length. If that condition is not met, vLLM-Omni raises a validation error instead of entering the ring kernel with inconsistent shapes.
-
-##### Offline Inference
-
-An example of offline inference script using [Ulysses-SP](https://arxiv.org/pdf/2309.14509) is shown below:
-```python
-from vllm_omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-from vllm_omni.diffusion.data import DiffusionParallelConfig
-ulysses_degree = 2
-
-omni = Omni(
-    model="Qwen/Qwen-Image",
-    parallel_config=DiffusionParallelConfig(ulysses_degree=2)
-)
-
-outputs = omni.generate(
-    "A cat sitting on a windowsill",
-    OmniDiffusionSamplingParams(num_inference_steps=50, width=2048, height=2048),
-)
-```
-
-See `examples/offline_inference/text_to_image/text_to_image.py` for a complete working example.
-
-To enable the experimental UAA mode explicitly, use a model/configuration that actually requires it. For example, `Tongyi-MAI/Z-Image-Turbo` has 30 attention heads, so `ulysses_degree=4` requires UAA because 30 is not divisible by 4:
-
-```python
-omni = Omni(
-    model="Tongyi-MAI/Z-Image-Turbo",
-    parallel_config=DiffusionParallelConfig(
-        ulysses_degree=4,
-        ulysses_mode="advanced_uaa",
-    ),
-)
-```
-
-##### Online Serving
-
-You can enable Ulysses-SP in online serving for diffusion models via `--usp`:
-
-```bash
-# Text-to-image (requires >= 2 GPUs)
-vllm serve Qwen/Qwen-Image --omni --port 8091 --usp 2
-
-# Experimental UAA mode for a model with 30 attention heads
-vllm serve Tongyi-MAI/Z-Image-Turbo --omni --port 8091 --usp 4 --ulysses-mode advanced_uaa
-```
-
-##### Benchmarks
-!!! note "Benchmark Disclaimer"
-    These benchmarks are provided for **general reference only**. The configurations shown use default or common parameter settings and have not been exhaustively optimized for maximum performance. Actual performance may vary based on:
-
-    - Specific model and use case
-    - Hardware configuration
-    - Careful parameter tuning
-    - Different inference settings (e.g., number of steps, image resolution)
-
-
-To measure the parallelism methods, we run benchmarks with **Qwen/Qwen-Image** model generating images (**2048x2048** as long sequence input) with 50 inference steps. The hardware devices are NVIDIA H800 GPUs. `sdpa` is the attention backends.
-
-| Configuration | Ulysses degree |Generation Time | Speedup |
-|---------------|----------------|---------|---------|
-| **Baseline (diffusers)** | - | 112.5s | 1.0x |
-| Ulysses-SP  |  2  |  65.2s | 1.73x |
-| Ulysses-SP  |  4  | 39.6s | 2.84x |
-| Ulysses-SP  |  8  | 30.8s | 3.65x |
-
-#### Ring-Attention
-
-Ring-Attention ([arxiv paper](https://arxiv.org/abs/2310.01889)) splits the input along the sequence dimension and uses ring-based P2P communication to accumulate attention results. Unlike Ulysses-SP which uses all-to-all communication, Ring-Attention keeps the sequence dimension sharded throughout the computation and circulates Key/Value blocks through a ring topology.
-
-##### Offline Inference
-
-An example of offline inference script using Ring-Attention is shown below:
-```python
-from vllm_omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-from vllm_omni.diffusion.data import DiffusionParallelConfig
-ring_degree = 2
-
-omni = Omni(
-    model="Qwen/Qwen-Image",
-    parallel_config=DiffusionParallelConfig(ring_degree=2)
-)
-
-outputs = omni.generate(
-    "A cat sitting on a windowsill",
-    OmniDiffusionSamplingParams(num_inference_steps=50, width=2048, height=2048),
-)
-```
-
-See `examples/offline_inference/text_to_image/text_to_image.py` for a complete working example.
-
-
-##### Online Serving
-
-You can enable Ring-Attention in online serving for diffusion models via `--ring`:
-
-```bash
-# Text-to-image (requires >= 2 GPUs)
-vllm serve Qwen/Qwen-Image --omni --port 8091 --ring 2
-```
-
-##### Benchmarks
-!!! note "Benchmark Disclaimer"
-    These benchmarks are provided for **general reference only**. The configurations shown use default or common parameter settings and have not been exhaustively optimized for maximum performance. Actual performance may vary based on:
-
-    - Specific model and use case
-    - Hardware configuration
-    - Careful parameter tuning
-    - Different inference settings (e.g., number of steps, image resolution)
-
-
-To measure the parallelism methods, we run benchmarks with **Qwen/Qwen-Image** model generating images (**1024x1024** as long sequence input) with 50 inference steps. The hardware devices are NVIDIA A100 GPUs. `flash_attn` is the attention backends.
-
-| Configuration | Ring degree |Generation Time | Speedup |
-|---------------|----------------|---------|---------|
-| **Baseline (diffusers)** | - | 45.2s | 1.0x |
-| Ring-Attention  |  2  |  29.9s | 1.51x |
-| Ring-Attention  |  4  | 23.3s | 1.94x |
-
-
-#### Hybrid Ulysses + Ring
-
-You can combine both Ulysses-SP and Ring-Attention for larger scale parallelism. The total sequence parallel size equals `ulysses_degree × ring_degree`.
-
-!!! note "Experimental UAA in hybrid mode"
-    `ulysses_mode="advanced_uaa"` can also be used with hybrid Ulysses + Ring, but this does not remove Ring's shape requirement. Every rank in the same ring group must still have the same post-Ulysses sequence length.
-
-##### Offline Inference
-
-```python
-from vllm_omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-from vllm_omni.diffusion.data import DiffusionParallelConfig
-
-# Hybrid: 2 Ulysses × 2 Ring = 4 GPUs total
-omni = Omni(
-    model="Qwen/Qwen-Image",
-    parallel_config=DiffusionParallelConfig(
-        ulysses_degree=2,
-        ring_degree=2,
-    )
-)
-
-outputs = omni.generate(
-    "A cat sitting on a windowsill",
-    OmniDiffusionSamplingParams(num_inference_steps=50, width=2048, height=2048),
-)
-```
-
-##### Online Serving
-
-```bash
-# Text-to-image (requires >= 4 GPUs)
-vllm serve Qwen/Qwen-Image --omni --port 8091 --usp 2 --ring 2
-```
-
-##### Benchmarks
-!!! note "Benchmark Disclaimer"
-    These benchmarks are provided for **general reference only**. The configurations shown use default or common parameter settings and have not been exhaustively optimized for maximum performance. Actual performance may vary based on:
-
-    - Specific model and use case
-    - Hardware configuration
-    - Careful parameter tuning
-    - Different inference settings (e.g., number of steps, image resolution)
-
-
-To measure the parallelism methods, we run benchmarks with **Qwen/Qwen-Image** model generating images (**1024x1024** as long sequence input) with 50 inference steps. The hardware devices are NVIDIA A100 GPUs. `flash_attn` is the attention backends.
-
-| Configuration | Ulysses degree | Ring degree | Generation Time | Speedup |
-|---------------|----------------|-------------|-----------------|---------|
-| **Baseline (diffusers)** | - | - | 45.2s | 1.0x |
-| Hybrid Ulysses + Ring  |  2  |  2  |  24.3s | 1.87x |
-
-
-### CFG-Parallel
-
-#### Offline Inference
-
-CFG-Parallel is enabled through `DiffusionParallelConfig(cfg_parallel_size=2)`, which runs one rank for the positive branch and one rank for the negative branch.
-
-An example of offline inference using CFG-Parallel (image-to-image) is shown below:
-
-```python
-from vllm_omni import Omni
-from vllm_omni.diffusion.data import DiffusionParallelConfig
-
-image_path = "path_to_image.png"
-omni = Omni(
-    model="Qwen/Qwen-Image-Edit",
-    parallel_config=DiffusionParallelConfig(cfg_parallel_size=2),
-)
-input_image = Image.open(image_path).convert("RGB")
-
-outputs = omni.generate(
-    {
-        "prompt": "turn this cat to a dog",
-        "negative_prompt": "low quality, blurry",
-        "multi_modal_data": {"image": input_image},
-    },
-    OmniDiffusionSamplingParams(
-        true_cfg_scale=4.0,
-        num_inference_steps=50,
-    ),
-)
-```
-
-Notes:
-
- CFG-Parallel is only effective when a `negative_prompt` is provided AND a guidance scale (or `cfg_scale`) is greater than 1.
-
-See `examples/offline_inference/image_to_image/image_edit.py` for a complete working example.
-```bash
-cd examples/offline_inference/image_to_image/
-python image_edit.py \
-  --model "Qwen/Qwen-Image-Edit" \
-  --image "qwen_image_output.png" \
-  --prompt "turn this cat to a dog" \
-  --negative-prompt "low quality, blurry" \
-  --cfg-scale 4.0 \
-  --output "edited_image.png" \
-  --cfg-parallel-size 2
-```
-
-#### Online Serving
-
-You can enable CFG-Parallel in online serving for diffusion models via `--cfg-parallel-size`:
-
-```bash
-vllm serve Qwen/Qwen-Image-Edit --omni --port 8091 --cfg-parallel-size 2
-```
-
-### HSDP
-
-HSDP (Hybrid Sharded Data Parallel) shards model weights across GPUs to reduce per-GPU memory usage. This enables inference of large models (e.g., Wan2.2 14B) on GPUs with limited memory.
-
-Unlike Tensor Parallelism which splits computation, HSDP uses PyTorch's FSDP2 to shard and redistribute weights at runtime. Each GPU only holds a fraction of the model weights, and weights are gathered on-demand during forward passes.
-
-#### Configuration
-
-HSDP is configured via `DiffusionParallelConfig`:
-
-| Parameter | Type | Default | Description |
-|-----------|------|---------|-------------|
-| `use_hsdp` | bool | False | Enable HSDP |
-| `hsdp_shard_size` | int | -1 | Number of GPUs to shard weights across. -1 = auto (requires other parallelism > 1) |
-| `hsdp_replicate_size` | int | 1 | Number of replica groups. Each group holds a full sharded copy |
-
-**Constraints:**
-
- `hsdp_replicate_size × hsdp_shard_size == world_size`
- HSDP cannot be used with Tensor Parallelism (`tensor_parallel_size` must be 1)
-
-#### Operating Modes
-
-HSDP can work in two modes:
-
- **Standalone Mode**: HSDP alone without other parallelism. Must specify `hsdp_shard_size` explicitly.
- **Combined Mode**: HSDP overlays on top of other parallelism (Ulysses Sequence Parallel, CFG Parallel). HSDP dimensions must match world_size.
-
-#### Offline Inference
-
-**Standalone HSDP** (shard across 4 GPUs, no other parallelism):
-
-```python
-from vllm_omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-from vllm_omni.diffusion.data import DiffusionParallelConfig
-
-omni = Omni(
-    model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
-    parallel_config=DiffusionParallelConfig(
-        use_hsdp=True,
-        hsdp_shard_size=4,  # Shard across 4 GPUs
-    ),
-)
-
-outputs = omni.generate(
-    "A cat playing piano",
-    OmniDiffusionSamplingParams(num_inference_steps=50),
-)
-```
-
-**Combined HSDP + Sequence Parallel**:
-
-```python
-omni = Omni(
-    model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
-    parallel_config=DiffusionParallelConfig(
-        ulysses_degree=4,  # Sequence parallel
-        use_hsdp=True,     # HSDP overlays on SP
-    ),
-)
-```
-
-#### Online Serving
-
-**Standalone HSDP** (shard model across 4 GPUs):
-
-```bash
-vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --port 8091 --use-hsdp --hsdp-shard-size 4
-```
-
-**Combined with Sequence Parallel**:
-
-```bash
-vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --port 8091 --use-hsdp --usp 4
-```
-
-#### Adding HSDP Support to New Models
-
-For detailed instructions on adding HSDP support to new models, see the [HSDP Contributing Guide](../../design/feature/hsdp.md).
-
-### Expert Parallelism
-
-Unlike Tensor Parallelism which shards every layer's weights, EP only shards the MoE expert MLP blocks. This significantly reduces the memory footprint of MoE models (e.g., HunyuanImage3.0) while maintaining constant dense-equivalent compute efficiency. Expert Parallelism is enabled via `DiffusionParallelConfig.enable_expert_parallel`. And `self.ep = tp * sp * cfg * dp` for now, so at least one of TP/SP/CFG/DP should set when EP enabled.
-
-#### Offline Inference
-
-```python
-from vllm_omni import Omni
-from vllm_omni.diffusion.data import DiffusionParallelConfig
-
-omni = Omni(
-    model="tencent/HunyuanImage-3.0",
-    parallel_config=DiffusionParallelConfig(tensor_parallel_size=8, enable_expert_parallel=True),
-)
-
-outputs = omni.generate(
-    "A brown and white dog is running on the grass",
-    OmniDiffusionSamplingParams(
-        num_inference_steps=50,
-        width=1024,
-        height=1024,
-    ),
-)
-```
@@ -65,6 +65,7 @@ The available `ignored_layers` names depend on the model architecture (e.g., `to
 | Flux | `black-forest-labs/FLUX.1-dev` | All layers | None |
 | HunyuanImage-3 | `tencent/HunyuanImage3` | All layers | None |
 | HunyuanVideo-1.5 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `720p_t2v`, `480p_i2v` | All layers | None |
+| Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | All layers | None |

 ## Combining with Other Features

@@ -46,7 +46,7 @@ its stepwise request state machine. For normal diffusion inference, leave it
 disabled unless your workflow depends on this mode.

 If you are looking for general diffusion speedups, see
-[Diffusion Acceleration Overview](../diffusion_acceleration.md).
+[Diffusion Features Overview](../diffusion_features.md).

 ## Troubleshooting

@@ -1,145 +0,0 @@
-# TeaCache Configuration Guide
-
-TeaCache speeds up diffusion model inference by caching transformer computations when consecutive timesteps are similar. This typically provides **1.5x-2.0x speedup** with minimal quality loss.
-
-## Quick Start
-
-Enable TeaCache by setting `cache_backend` to `"tea_cache"`:
-
-```python
-from vllm_omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-# Simple configuration - model_type is automatically extracted from pipeline.__class__.__name__
-omni = Omni(
-    model="Qwen/Qwen-Image",
-    cache_backend="tea_cache",
-    cache_config={
-        "rel_l1_thresh": 0.2  # Optional, defaults to 0.2
-    }
-)
-outputs = omni.generate(
-    "A cat sitting on a windowsill",
-    OmniDiffusionSamplingParams(
-        num_inference_steps=50,
-    ),
-)
-```
-
-### Using Environment Variable
-
-You can also enable TeaCache via environment variable:
-
-```bash
-export DIFFUSION_CACHE_BACKEND=tea_cache
-```
-
-Then initialize without explicitly setting `cache_backend`:
-
-```python
-from vllm_omni import Omni
-
-omni = Omni(
-    model="Qwen/Qwen-Image",
-    cache_config={"rel_l1_thresh": 0.2}  # Optional
-)
-```
-
-## Online Serving (OpenAI-Compatible)
-
-Enable TeaCache for online serving by passing `--cache-backend tea_cache` when starting the server:
-
-```bash
-vllm serve Qwen/Qwen-Image --omni --port 8091 \
-  --cache-backend tea_cache \
-  --cache-config '{"rel_l1_thresh": 0.2}'
-```
-
-## Configuration Parameters
-
-### `rel_l1_thresh` (float, default: `0.2`)
-
-Controls the balance between speed and quality. Lower values prioritize quality, higher values prioritize speed.
-
-**Recommended values:**
-
- `0.2` - **~1.5x speedup** with minimal quality loss (recommended)
- `0.4` - **~1.8x speedup** with slight quality loss
- `0.6` - **~2.0x speedup** with noticeable quality loss
- `0.8` - **~2.25x speedup** with significant quality loss
-
-## Examples
-
-### Python API
-
-```python
-from vllm_omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-omni = Omni(
-    model="Qwen/Qwen-Image",
-    cache_backend="tea_cache",
-    cache_config={"rel_l1_thresh": 0.2}
-)
-outputs = omni.generate(
-    "A cat sitting on a windowsill",
-    OmniDiffusionSamplingParams(
-        num_inference_steps=50,
-    ),
-)
-```
-
-## Performance Tuning
-
-Start with the default `rel_l1_thresh=0.2` and adjust based on your needs:
-
- **Maximum quality**: Use `0.1-0.2`
- **Balanced**: Use `0.2-0.4` (recommended)
- **Maximum speed**: Use `0.6-0.8` (may reduce quality)
-
-## Troubleshooting
-
-### Quality Degradation
-
-If you notice quality issues, lower the threshold:
-
-```python
-cache_config={"rel_l1_thresh": 0.1}  # More conservative caching
-```
-
-## Supported Models
-
-### ImageGen
-
-<style>
-th {
-  white-space: nowrap;
-  min-width: 0 !important;
-}
-</style>
-
-| Architecture | Models | Example HF Models |
-|--------------|--------|-------------------|
-| `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` |
-| `QwenImageEditPipeline` | Qwen-Image-Edit | `Qwen/Qwen-Image-Edit` |
-| `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2509 | `Qwen/Qwen-Image-Edit-2509` |
-| `QwenImageLayeredPipeline` | Qwen-Image-Layered | `Qwen/Qwen-Image-Layered` |
-| `BagelForConditionalGeneration` | BAGEL (DiT-only) | `ByteDance-Seed/BAGEL-7B-MoT` |
-
-### VideoGen
-
-No VideoGen models are supported by TeaCache yet.
-
-### Coming Soon
-
-<style>
-th {
-  white-space: nowrap;
-  min-width: 0 !important;
-}
-</style>
-
-| Architecture | Models | Example HF Models |
-|--------------|--------|-------------------|
-| `FluxPipeline` | Flux | - |
-| `CogVideoXPipeline` | CogVideoX | - |
@@ -1,380 +0,0 @@
-# Diffusion Acceleration Overview
-
-vLLM-Omni supports various acceleration methods to speed up diffusion model inference with minimal quality degradation. These include **cache methods** that intelligently cache intermediate computations to avoid redundant work across diffusion timesteps, **parallelism methods** that distribute the computation across multiple devices, and **quantization methods** that reduce memory footprint while preserving accuracy.
-
-## Supported Acceleration Methods
-
-vLLM-Omni currently supports two main cache acceleration backends:
-
-1. **[TeaCache](diffusion/teacache.md)** - Hook-based adaptive caching that caches transformer computations when consecutive timesteps are similar
-2. **[Cache-DiT](diffusion/cache_dit_acceleration.md)** - Library-based acceleration using multiple techniques:
-    - **DBCache** (Dual Block Cache): Caches intermediate transformer block outputs based on residual differences
-    - **TaylorSeer**: Uses Taylor expansion-based forecasting for faster inference
-    - **SCM** (Step Computation Masking): Selectively computes steps based on adaptive masking
-
-Both methods can provide significant speedups (typically **1.5x-2.0x**) while maintaining high output quality.
-
-vLLM-Omni also supports quantization methods:
-
-3. **[Quantization](diffusion/quantization/overview.md)** - Reduces DiT linear layers from BF16 to FP8 or Int8, providing ~1.28x speedup with minimal quality loss. Supports per-layer skip for sensitive layers.
-
-vLLM-Omni also supports parallelism methods for diffusion models, including:
-
-1. [Ulysses-SP](diffusion/parallelism_acceleration.md#ulysses-sp) - splits the input along the sequence dimension and uses all-to-all communication to allow each device to compute only a subset of attention heads.
-
-2. [Ring-Attention](diffusion/parallelism_acceleration.md#ring-attention) - splits the input along the sequence dimension and uses ring-based P2P communication to accumulate attention results, keeping the sequence dimension sharded.
-
-3. [CFG-Parallel](diffusion/parallelism_acceleration.md#cfg-parallel) - runs the positive/negative prompts of classifier-free guidance (CFG) on different devices, then merges on a single device to perform the scheduler step.
-
-4. [Tensor Parallelism](diffusion/parallelism_acceleration.md#tensor-parallelism) - shards DiT weights across devices to reduce per-GPU memory usage.
-
-5. [VAE Patch Parallelism](diffusion/parallelism_acceleration.md#vae-patch-parallelism) - shards VAE decode/encode spatially across ranks to reduce VAE peak memory (and can speed up VAE decode).
-
-6. [HSDP](diffusion/parallelism_acceleration.md#hsdp) - Hybrid Sharded Data Parallel shards model weights across GPUs to reduce per-GPU memory usage, enabling inference of large models on limited GPU memory.
-
-## Quick Comparison
-
-### Cache Methods
-
-| Method | Configuration | Description | Best For |
-|--------|--------------|-------------|----------|
-| **TeaCache** | `cache_backend="tea_cache"` | Simple, adaptive caching with minimal configuration | Quick setup, balanced speed/quality |
-| **Cache-DiT** | `cache_backend="cache_dit"` | Advanced caching with multiple techniques (DBCache, TaylorSeer, SCM) | Maximum acceleration, fine-grained control |
-
-### Quantization Methods
-
-| Method | Configuration | Description | Best For |
-|--------|--------------|-------------|----------|
-| **FP8** | `quantization="fp8"` | FP8 W8A8 on Ada/Hopper, weight-only on older GPUs | Memory reduction, inference speedup |
-| **Int8** | `quantization="int8"` | Int8 W8A8 | Memory reduction, inference speedup |
-
-## Supported Models
-
-The following table shows which models are currently supported by each acceleration method:
-
-### ImageGen
-
-| Model | Model Identifier | TeaCache | Cache-DiT | Ulysses-SP | Ring-Attention | CFG-Parallel | Tensor-Parallel | VAE-Patch-Parallel |
-|-------|------------------|:--------:|:---------:|:----------:|:--------------:|:------------:|:---------------:|:------------------:|
-| **LongCat-Image** | `meituan-longcat/LongCat-Image` | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ |
-| **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ |
-| **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ |
-| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
-| **Qwen-Image-2512** | `Qwen/Qwen-Image-2512` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
-| **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
-| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
-| **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
-| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ | ❌ | ❌ | ❌ | ✅ (TP=2 only) | ✅ |
-| **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ |
-| **Bagel** | `ByteDance-Seed/BAGEL-7B-MoT` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
-| **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ |
-| **NextStep-1.1** | `stepfun-ai/NextStep-1.1` | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ |
-| **FLUX.2-klein** | `black-forest-labs/FLUX.2-klein-4B` | ✅ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ |
-| **FLUX.2-dev** | `black-forest-labs/FLUX.2-dev` | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ |
-| **FLUX.1-Kontext-dev** | `black-forest-labs/FLUX.1-Kontext-dev` | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ |
-| **GLM-Image** | `zai-org/GLM-Image` | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ |
-
-### VideoGen
-
-| Model | Model Identifier | TeaCache | Cache-DiT | Ulysses-SP | Ring-Attention | CFG-Parallel | HSDP | VAE-Patch-Parallel |
-|-------|------------------|:--------:|:---------:|:----------:|:--------------:|:------------:|:----:|:----:|
-| **Wan2.1-T2V** | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers` | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
-| **Wan2.1-T2V** | `Wan-AI/Wan2.1-T2V-14B-Diffusers` | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
-| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
-| **LTX-2** | `Lightricks/LTX-2` | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-| **DreamID-Omni** | `XuGuo699/DreamID-Omni` | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ |
-
-### Quantization
-
-| Model | Model Identifier | FP8 | Int8 |
-|-------|------------------|:---:|:---:|
-| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ |
-| **Qwen-Image-2512** | `Qwen/Qwen-Image-2512` | ✅ | ✅ |
-| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ |
-
-
-### AudioGen
-
-| Model                  | Model Identifier                         | TeaCache | Cache-DiT | Ulysses-SP | Ring-Attention | CFG-Parallel |
-|------------------------|------------------------------------------|:--------:|:---------:|:----------:|:--------------:|:------------:|
-| **Stable-Audio-Open**  | `stabilityai/stable-audio-open-1.0`      |    ✅    |           |            |                |              |
-
-## Performance Benchmarks
-
-The following benchmarks were measured on **Qwen/Qwen-Image** and **Qwen/Qwen-Image-Edit** models generating 1024x1024 images with 50 inference steps:
-
-!!! note "Benchmark Disclaimer"
-    These benchmarks are provided for **general reference only**. The configurations shown use default or common parameter settings and have not been exhaustively optimized for maximum performance. Actual performance may vary based on:
-
-    - Specific model and use case
-    - Hardware configuration
-    - Careful parameter tuning
-    - Different inference settings (e.g., number of steps, image resolution)
-
-    For optimal performance in your specific scenario, we recommend experimenting with different parameter configurations as described in the detailed guides below.
-
-| Model | Cache Backend | Cache Config | Generation Time | Speedup | Notes |
-|-------|---------------|--------------|----------------|---------|-------|
-| **Qwen/Qwen-Image** | None  | None | 20.0s | 1.0x | Baseline (diffusers) |
-| **Qwen/Qwen-Image** | TeaCache | `rel_l1_thresh=0.2` | 10.47s | **1.91x** | Recommended default setting |
-| **Qwen/Qwen-Image** | Cache-DiT | DBCache + TaylorSeer (Fn=1, Bn=0, W=8, TaylorSeer order=1) | 10.8s | **1.85x** | - |
-| **Qwen/Qwen-Image** | Cache-DiT | DBCache + TaylorSeer + SCM (Fn=8, Bn=0, W=4, TaylorSeer order=1, SCM fast) | 14.0s | **1.43x** | - |
-| **Qwen/Qwen-Image-Edit** | None | No acceleration | 51.5s | 1.0x | Baseline (diffusers) |
-| **Qwen/Qwen-Image-Edit** | Cache-DiT | Default (Fn=1, Bn=0, W=4, TaylorSeer disabled, SCM disabled) | 21.6s | **2.38x** | - |
-
-To measure the parallelism methods, we run benchmarks with **Qwen/Qwen-Image** model generating images (**2048x2048** as long sequence input) with 50 inference steps. The hardware devices are NVIDIA H800 GPUs. `sdpa` is the attention backends.
-
-| Configuration | Ulysses degree |Generation Time | Speedup |
-|---------------|----------------|---------|---------|
-| **Baseline (diffusers)** | - | 112.5s | 1.0x |
-| Ulysses-SP  |  2  |  65.2s | 1.73x |
-| Ulysses-SP  |  4  | 39.6s | 2.84x |
-| Ulysses-SP  |  8  | 30.8s | 3.65x |
-
-## Quick Start
-
-### Using TeaCache
-
-```python
-from vllm_omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-omni = Omni(
-    model="Qwen/Qwen-Image",
-    cache_backend="tea_cache",
-    cache_config={"rel_l1_thresh": 0.2}  # Optional, defaults to 0.2
-)
-
-outputs = omni.generate(
-    "A cat sitting on a windowsill",
-    OmniDiffusionSamplingParams(
-        num_inference_steps=50,
-    ),
-)
-```
-
-### Using Cache-DiT
-
-```python
-from vllm_omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-omni = Omni(
-    model="Qwen/Qwen-Image",
-    cache_backend="cache_dit",
-    cache_config={
-        "Fn_compute_blocks": 1,
-        "Bn_compute_blocks": 0,
-        "max_warmup_steps": 8,
-        "enable_taylorseer": True,
-        "taylorseer_order": 1,
-    }
-)
-
-outputs = omni.generate(
-    "A cat sitting on a windowsill",
-    OmniDiffusionSamplingParams(
-        num_inference_steps=50,
-    ),
-)
-```
-
-### Using Ulysses-SP
-
-Run text-to-image:
-```python
-from vllm_omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-from vllm_omni.diffusion.data import DiffusionParallelConfig
-ulysses_degree = 2
-
-omni = Omni(
-    model="Qwen/Qwen-Image",
-    parallel_config=DiffusionParallelConfig(ulysses_degree=ulysses_degree)
-)
-
-outputs = omni.generate(
-    "A cat sitting on a windowsill",
-    OmniDiffusionSamplingParams(num_inference_steps=50, width=2048, height=2048),
-)
-```
-
-
-Run image-to-image:
-```python
-from vllm_omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-from vllm_omni.diffusion.data import DiffusionParallelConfig
-ulysses_degree = 2
-
-omni = Omni(
-    model="Qwen/Qwen-Image-Edit",
-    parallel_config=DiffusionParallelConfig(ulysses_degree=ulysses_degree)
-)
-
-outputs = omni.generate(
-    {
-        "prompt": "turn this cat to a dog",
-        "multi_modal_data": {"image": input_image}
-    },
-    OmniDiffusionSamplingParams(num_inference_steps=50),
-)
-```
-
-### Using Ring-Attention
-
-Run text-to-image:
-```python
-from vllm_omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-from vllm_omni.diffusion.data import DiffusionParallelConfig
-ring_degree = 2
-
-omni = Omni(
-    model="Qwen/Qwen-Image",
-    parallel_config=DiffusionParallelConfig(ring_degree=2)
-)
-
-outputs = omni.generate(
-    "A cat sitting on a windowsill",
-    OmniDiffusionSamplingParams(num_inference_steps=50, width=2048, height=2048),
-)
-```
-
-### Using Tensor Parallelism
-
-```python
-from vllm_omni import Omni
-from vllm_omni.diffusion.data import DiffusionParallelConfig
-
-omni = Omni(
-    model="Tongyi-MAI/Z-Image-Turbo",
-    parallel_config=DiffusionParallelConfig(tensor_parallel_size=2),
-)
-
-outputs = omni.generate(
-    prompt="a cat reading a book",
-    num_inference_steps=9,
-    width=512,
-    height=512,
-)
-```
-
-### Using VAE Patch Parallelism
-
-```python
-from vllm_omni import Omni
-from vllm_omni.diffusion.data import DiffusionParallelConfig
-
-omni = Omni(
-    model="Tongyi-MAI/Z-Image-Turbo",
-    parallel_config=DiffusionParallelConfig(
-        tensor_parallel_size=2,
-        vae_patch_parallel_size=2,
-    ),
-)
-
-outputs = omni.generate(
-    prompt="a cat reading a book",
-    num_inference_steps=9,
-    width=1024,
-    height=1024,
-)
-```
-
-### Using HSDP
-
-HSDP (Hybrid Sharded Data Parallel) shards model weights across GPUs to reduce per-GPU memory usage. This enables inference of large models (e.g., Wan2.2 14B) on GPUs with limited memory.
-
-```python
-from vllm_omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-from vllm_omni.diffusion.data import DiffusionParallelConfig
-
-omni = Omni(
-    model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
-    parallel_config=DiffusionParallelConfig(
-        use_hsdp=True,
-        hsdp_replicate_size=1,  # Number of replica groups
-        hsdp_shard_size=8,      # Number of GPUs to shard across (or -1 for auto)
-    ),
-)
-
-outputs = omni.generate(
-    "A cat playing piano",
-    OmniDiffusionSamplingParams(num_inference_steps=50),
-)
-```
-
-### Using CFG-Parallel
-
-Run image-to-image:
-
-CFG-Parallel splits the CFG positive/negative branches across GPUs. Use it when you set a non-trivial `true_cfg_scale`.
-
-```python
-from vllm_omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-from vllm_omni.diffusion.data import DiffusionParallelConfig
-cfg_parallel_size = 2
-
-omni = Omni(
-    model="Qwen/Qwen-Image-Edit",
-    parallel_config=DiffusionParallelConfig(cfg_parallel_size=cfg_parallel_size)
-)
-
-outputs = omni.generate(
-    {
-        "prompt": "turn this cat to a dog",
-        "multi_modal_data": {"image": input_image}
-    },
-    OmniDiffusionSamplingParams(num_inference_steps=50, true_cfg_scale=4.0),
-)
-```
-
-### Using FP8 Quantization
-
-```python
-from vllm_omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-omni = Omni(
-    model="<your-model>",
-    quantization="fp8",
-)
-
-outputs = omni.generate(
-    "A cat sitting on a windowsill",
-    OmniDiffusionSamplingParams(num_inference_steps=50),
-)
-```
-
-### Using Int8 Quantization
-
-```python
-from vllm_omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-omni = Omni(
-    model="<your-model>",
-    quantization="int8",
-)
-
-outputs = omni.generate(
-    "A cat sitting on a windowsill",
-    OmniDiffusionSamplingParams(num_inference_steps=50),
-)
-```
-
-## Documentation
-
-For detailed information on each acceleration method:
-
- **[TeaCache Guide](diffusion/teacache.md)** - Complete TeaCache documentation, configuration options, and best practices
- **[Cache-DiT Acceleration Guide](diffusion/cache_dit_acceleration.md)** - Comprehensive Cache-DiT guide covering DBCache, TaylorSeer, SCM, and configuration parameters
- **[Quantization Guide](diffusion/quantization/overview.md)** - Quantization for DiT models with per-layer control
- **[Tensor Parallelism](diffusion/parallelism_acceleration.md#tensor-parallelism)** - Guidance on how to enable TP for diffusion models.
- **[Sequence Parallelism](diffusion/parallelism_acceleration.md#sequence-parallelism)** - Guidance on how to set sequence parallelism with configuration.
- **[CFG-Parallel](diffusion/parallelism_acceleration.md#cfg-parallel)** - Guidance on how to set CFG-Parallel to run positive/negative branches across ranks.
- **[VAE Patch Parallelism](diffusion/parallelism_acceleration.md#vae-patch-parallelism)** - Guidance on how to reduce VAE memory via patch/tile parallelism.
- **[HSDP](diffusion/parallelism_acceleration.md#hsdp)** - Hybrid Sharded Data Parallel for memory-efficient inference of large models.
@@ -0,0 +1,190 @@
+# Diffusion Advanced Features
+
+## Table of Contents
+
+- [Overview](#overview)
+- [Supported Features](#supported-features)
+- [Supported Models](#supported-models)
+- [Feature Compatibility](#feature-compatibility)
+- [Learn More](#learn-more)
+
+## Overview
+
+vLLM-Omni supports various advanced features for diffusion models:
+
+- Acceleration: **cache methods**, **parallelism methods**
+- Memory optimization: **cpu offloading**, **quantization**
+- Extensions: **LoRA inference**
+
+## Supported Features
+
+### Acceleration
+
+#### Lossy Acceleration
+
+Cache methods trade minimal quality for significant speedup. Quality loss is typically imperceptible with proper tuning.
+
+| Method | Description | Best For |
+|--------|-------------|----------|
+| **[TeaCache](diffusion/cache_acceleration/teacache.md)** | Adaptive caching using modulated inputs | Quick setup, balanced quality/speed on single GPU |
+| **[Cache-DiT](diffusion/cache_acceleration/cache_dit.md)** | Multiple caching techniques: DBCache, TaylorSeer, SCM | Fine-grained control, tunable quality-speed tradeoff |
+
+
+#### Lossless Acceleration
+
+Parallelism methods distribute computation across GPUs without quality loss (mathematically equivalent to single-GPU).
+
+| Method | Description | Best For |
+|--------|-------------|----------|
+| **[Ulysses-SP](diffusion/parallelism/sequence_parallel.md)** | Sequence parallelism via all-to-all communication | High-resolution images (>1536px) or long videos with 2-8 GPUs |
+| **[Ring-Attention](diffusion/parallelism/sequence_parallel.md)** | Sequence parallelism via ring-based communication | Videos, very long sequences, memory-constrained, with 2-8 GPUs |
+| **[CFG-Parallel](diffusion/parallelism/cfg_parallel.md)** | Splits CFG positive/negative branches across devices | Image editing with CFG guidance (true_cfg_scale > 1) on 2 GPUs |
+| **[Tensor Parallelism](diffusion/parallelism/tensor_parallel.md)** | Shards model weights across devices | Large models that don't fit in single GPU, with 2+ GPUs |
+| **[HSDP](diffusion/parallelism/hsdp.md)** | Weight sharding via FSDP2, redistributed on-demand at runtime | Very large models (14B+) on limited VRAM, combinable with SP |
+| **[Expert Parallelism](diffusion/parallelism/expert_parallel.md)** | Shards MoE expert MLP blocks across devices | MoE diffusion models (e.g., HunyuanImage3.0) |
+
+**Note:** Some acceleration methods can be combined together for optimized performance. See [Feature Compatibility Table](#feature-compatibility) and [Feature Compatibility Tutorial](feature_compatibility.md) for detailed configuration examples.
+
+### Memory Optimization
+
+Memory optimization methods help reduce GPU memory usage, enabling inference on resource-constrained hardware or larger models.
+
+| Method | Description | Best For |
+|--------|-------------|----------|
+| **[CPU Offload](diffusion/cpu_offload_diffusion.md)** | Offloads model components to CPU memory | Limited VRAM, large models on consumer GPUs |
+| **[Quantization](diffusion/quantization/overview.md)** | Reduces DiT layers from BF16 to FP8/INT8/etc. | Limited VRAM, minimal accuracy loss    |
+| **[VAE Patch Parallelism](diffusion/parallelism/vae_patch_parallel.md)** | Distributes VAE decode tiling across GPUs | High-resolution generation with reduced VAE memory peak |
+
+### Extensions
+
+Extension methods add specialized capabilities to diffusion models beyond standard inference.
+
+| Method | Description | Best For |
+|--------|-------------|----------|
+| **[LoRA Inference](diffusion/lora.md)** | Enables inference with Low-Rank Adaptation (LoRA) adapters weights | Reinforcement learning extensions |
+
+
+### Quantization Methods
+
+| Method | Configuration | Description | Best For |
+|--------|--------------|-------------|----------|
+| **[FP8](diffusion/quantization/fp8.md)** | `quantization="fp8"` | FP8 W8A8 on Ada/Hopper, weight-only on older GPUs | Memory reduction, inference speedup |
+| **[INT8](diffusion/quantization/int8.md)** | `quantization="int8"` | INT8 weight-only, no calibration or pre-quantized checkpoint needed | Memory reduction, broad GPU compatibility |
+| **[GGUF](diffusion/quantization/gguf.md)** | `quantization="gguf"` | Native GGUF transformer-only weights (Q4, Q8, etc.) | Memory reduction on consumer GPUs |
+
+## Supported Models
+
+The following tables show which models support each feature:
+
+- **🔀SP (Ulysses & Ring)**: Includes both Ulysses-SP and Ring-Attention methods
+- ✅ = Fully supported
+- ❌ = Not supported
+
+> Notes:
+
+> 1. CPU Offload has two methods: Module-wise (default for models with DiT + text encoder) and Layerwise. The tables below show **Layerwise support** only.
+> 2. The **💾Quantization** column is collapsed for readability. See [Quantization Overview](diffusion/quantization/overview.md) for per-method (FP8, GGUF, …) and per-model support details.
+
+### ImageGen
+
+| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization |
+|-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|
+| **Bagel** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
+| **FLUX.1-dev** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ |
+| **FLUX.2-klein** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ |
+| **FLUX.1-Kontext-dev** | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
+| **FLUX.2-dev** | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
+| **GLM-Image** | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
+| **HunyuanImage3** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ |
+| **LongCat-Image** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
+| **LongCat-Image-Edit** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
+| **MammothModa2(T2I)** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
+| **Nextstep_1(T2I)** | ❓ | ❓ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
+| **OmniGen2** | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
+| **Ovis-Image** | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
+| **Qwen-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ |
+| **Qwen-Image-2512** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ |
+| **Qwen-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ |
+| **Qwen-Image-Edit-2509** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ |
+| **Qwen-Image-Layered** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ |
+| **Stable-Diffusion3.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
+| **Z-Image** | ✅ | ✅ | ✅ | ❓ | ✅ (TP=2 only) | ❌ | ❌ | ✅ | ✅ |
+
+> Notes:
+> 1. Nextstep_1(T2I) does not support cache acceleration methods such as TeaCache or Cache-DiT.
+> 2. `Tongyi-MAI/Z-Image-Turbo` is a distilled model with minimal NFEs; CFG-Parallel is not necessary.
+
+### VideoGen
+
+| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization |
+|-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|
+| **Wan2.2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
+| **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
+| **Helios** | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
+| **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
+
+### AudioGen
+
+| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization |
+|-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|
+| **Stable-Audio-Open** | ❌ | ❌ | ❓ | ❓ | ❌ | ❌ | ❌ | ❌ | ✅ |
+
+
+## Feature Compatibility
+
+**Legend:**
+
+- ✅: Functionality is supported
+- ❌: No support plan
+- ❓: Not verified yet and Not Recommended
+
+|  | ⚡TeaCache | ⚡Cache-DiT | 🔀Ulysses-SP | 🔀Ring-Attn | 🔀CFG-Parallel | 🔀Tensor Parallel | 🔀HSDP | 🔀Expert Parallel | 💾CPU Offloading (Layerwise) | 💾CPU Offloading (Module-wise) | 💾VAE Patch Parallel | 💾FP8 Quant | 🔧LoRA Inference |
+|---|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|
+| **⚡TeaCache** | | | | | | | | | | | | | |
+| **⚡Cache-DiT** | ❌ | | | | | | | | | | | | |
+| **🔀Ulysses-SP** | ✅ | ✅ | | | | | | | | | | | |
+| **🔀Ring-Attn** | ✅ | ✅ | ✅ | | | | | | | | | | |
+| **🔀CFG-Parallel** | ✅ | ✅ | ✅ | ✅ | | | | | | | | | |
+| **🔀Tensor Parallel** | ✅ | ✅ | ✅ | ✅ | ✅ | | | | | | | | |
+| **🔀HSDP** | ❓ | ❓ | ❓ | ❓ | ❓ | ❌ | | | | | | | |
+| **🔀Expert Parallel** | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | | | | | | |
+| **💾CPU Offloading (Layerwise)** | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | | | | | |
+| **💾CPU Offloading (Module-wise)** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❓ | ❓ | ❌ | | | | |
+| **💾VAE Patch Parallel** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | | | |
+| **💾FP8 Quant** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❓ | ❓ | ✅ | ✅ | ✅ | | |
+| **🔧LoRA Inference** | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | |
+
+!!! info
+
+    1. Tensor Parallel and HSDP are not compatible.
+    2. TeaCache and Cache-DiT are not compatible.
+    3. CPU Offloading (Layerwise) and CPU Offloading (Module-wise) are not compatible.
+    4. CPU Offloading (Layerwise) supports single-card for now.
+    5. Using FP8-Quant as an example of qunatization methods.
+
+
+## Learn More
+
+**Cache Acceleration:**
+
+- **[TeaCache Configuration Guide](diffusion/cache_acceleration/teacache.md)** - Parameter tuning, performance tips, troubleshooting
+- **[Cache-DiT Advanced Guide](diffusion/cache_acceleration/cache_dit.md)** - DBCache, TaylorSeer, SCM techniques and optimization
+
+**Parallelism Methods:**
+
+- **[Parallelism Overview](diffusion/parallelism/overview.md)** - Tensor Parallelism, Sequence Parallelism, CFG Parallelism, HSDP, and Expert Parallelism
+
+**Memory Optimization:**
+
+- **[CPU Offload Guide](diffusion/cpu_offload_diffusion.md)** - Offload model components to CPU, reduce GPU memory usage
+- **[VAE Patch Parallelism Guide](diffusion/parallelism/vae_patch_parallel.md)** - Distribute VAE decode tiling across GPUs for high-resolution images
+- **[Quantization Overview](diffusion/quantization/overview.md)** - Overview of quantization methods for diffusion models
+
+**Extensions:**
+
+- **[LoRA Inference Guide](diffusion/lora.md)** - Low-Rank Adaptation for style customization and fine-tuning
+
+**Advanced Topics:**
+
+- **[Feature Compatibility](feature_compatibility.md)** - How to combine multiple features for maximum performance
@@ -256,6 +256,15 @@ lora_adapter/

 ## Web UI Demo

+!!! note "Gradio is an optional dependency"
+    The Gradio demo requires the `[demo]` extras. Install them first:
+
+    ```bash
+    pip install 'vllm-omni[demo]'
+    ```
+
+    Or, if installing from source: `pip install -e '.[demo]'`
+
 Launch the Gradio demo:

 ```bash
@@ -11,6 +11,15 @@ Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serv

 ## Gradio Demo

+!!! note "Gradio is an optional dependency"
+    The Gradio demo requires the `[demo]` extras. Install them first:
+
+    ```bash
+    pip install 'vllm-omni[demo]'
+    ```
+
+    Or, if installing from source: `pip install -e '.[demo]'`
+
 An interactive Gradio demo is available with text-to-speech synthesis, voice cloning, and streaming support.

 ```bash
@@ -122,6 +122,15 @@ python openai_chat_client.py --input input1.png input2.png --prompt "Combine the

 ### Method 4: Using Gradio Demo

+!!! note "Gradio is an optional dependency"
+    The Gradio demo requires the `[demo]` extras. Install them first:
+
+    ```bash
+    pip install 'vllm-omni[demo]'
+    ```
+
+    Or, if installing from source: `pip install -e '.[demo]'`
+
 ```bash
 python gradio_demo.py
 # Visit http://localhost:7861
@@ -149,6 +149,15 @@ python openai_chat_completion_client_for_multimodal_generation.py \

 ## Run Local Web UI Demo

+!!! note "Gradio is an optional dependency"
+    The Gradio demo requires the `[demo]` extras. Install them first:
+
+    ```bash
+    pip install 'vllm-omni[demo]'
+    ```
+
+    Or, if installing from source: `pip install -e '.[demo]'`
+
 This Web UI demo allows users to interact with the model through a web browser.

 ### Running Gradio Demo
@@ -154,7 +154,7 @@ print(response.choices[1].message.audio)    # Audio response
 ```

 ## Streaming Output
-If you want to enable streaming output, please set the argument as below. The final output will be obtained just after generated by corresponding stage. Now we only support text streaming output. Other modalities can output normally.
+If you want to enable streaming output, please set the argument as below. The final output will be obtained just after generated by corresponding stage. We support both text streaming output and audio streaming output. Other modalities can output normally.
 ```bash
 python openai_chat_completion_client_for_multimodal_generation.py \
    --query-type use_image \
@@ -163,6 +163,15 @@ python openai_chat_completion_client_for_multimodal_generation.py \

 ## Run Local Web UI Demo

+!!! note "Gradio is an optional dependency"
+    The Gradio demo requires the `[demo]` extras. Install them first:
+
+    ```bash
+    pip install 'vllm-omni[demo]'
+    ```
+
+    Or, if installing from source: `pip install -e '.[demo]'`
+
 This Web UI demo allows users to interact with the model through a web browser.

 ### Running Gradio Demo
@@ -19,6 +19,15 @@ Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/

 ## Gradio Demo

+!!! note "Gradio is an optional dependency"
+    The Gradio demo requires the `[demo]` extras. Install them first:
+
+    ```bash
+    pip install 'vllm-omni[demo]'
+    ```
+
+    Or, if installing from source: `pip install -e '.[demo]'`
+
 Two interactive Gradio demos are available, both supporting all 3 task types:

 | Demo     | File                     | Transport    | Streaming Quality                                  |
@@ -105,6 +105,15 @@ python openai_chat_client.py --prompt "A beautiful landscape painting" --output

 ### Method 4: Using Gradio Demo

+!!! note "Gradio is an optional dependency"
+    The Gradio demo requires the `[demo]` extras. Install them first:
+
+    ```bash
+    pip install 'vllm-omni[demo]'
+    ```
+
+    Or, if installing from source: `pip install -e '.[demo]'`
+
 ```bash
 python gradio_demo.py
 # Visit http://localhost:7860
@@ -0,0 +1,218 @@
+# Feature Compatibility
+
+This guide explains the compatibility matrix of different diffusion features in vLLM-Omni. You can use cache methods together with parallelism methods and other features to achieve optimal speed and efficiency.
+
+## Overview
+
+vLLM-Omni supports combining:
+
+- **Cache methods** (TeaCache, Cache-DiT) with **Parallelism methods** (Ulysses-SP, Ring-Attention, CFG-Parallel, Tensor Parallelism)
+- **Multiple parallelism methods** together (e.g., Ulysses-SP + Ring-Attention, CFG-Parallel + Sequence Parallelism)
+- **LoRA adapters** with most acceleration features
+- **CPU offloading** with other memory optimization features
+
+See the feature compatibility matrix in [Table](diffusion_features.md#feature-compatibility)
+
+## Common Combinations
+
+### 1. Cache + Sequence Parallelism (Recommended)
+
+Best for: **Large images (>1536px) or videos**
+
+Combines cache acceleration with sequence parallelism for maximum speedup on single-device-challenging workloads.
+
+**Using TeaCache + Ulysses-SP:**
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model Qwen/Qwen-Image \
+  --prompt "A beautiful mountain landscape" \
+  --cache-backend tea_cache \
+  --ulysses-degree 2
+```
+
+**Using Cache-DiT + Ring-Attention:**
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model Qwen/Qwen-Image \
+  --prompt "A futuristic city" \
+  --cache-backend cache_dit \
+  --ring-degree 2
+```
+
+### 2. Cache + CFG-Parallel
+
+Best for: **Image editing with Classifier-Free Guidance**
+
+Accelerates both the diffusion process and CFG computation.
+
+```bash
+python examples/offline_inference/image_to_image/image_edit.py \
+  --model Qwen/Qwen-Image-Edit \
+  --prompt "make it sunset" \
+  --negative-prompt "low quality, blurry" \
+  --image input.png \
+  --cache-backend cache_dit \
+  --cfg-parallel-size 2 \
+  --cfg-scale 4.0
+```
+
+### 3. CFG-Parallel + Sequence Parallelism
+
+Best for: **Large resolution image editing with CFG**
+
+Combines both CFG branch splitting and sequence parallelism for maximum GPU utilization.
+
+**CFG-Parallel + Ulysses-SP:**
+
+```bash
+python examples/offline_inference/image_to_image/image_edit.py \
+  --model Qwen/Qwen-Image-Edit \
+  --prompt "transform into autumn scene" \
+  --negative-prompt "low quality" \
+  --image input.png \
+  --cache-backend cache_dit \
+  --cfg-parallel-size 2 \
+  --ulysses-degree 2 \
+  --cfg-scale 4.0
+```
+
+### 4. Hybrid Ulysses + Ring + Vae tiling
+
+Best for: **Very large images or videos on multiple devices**
+
+Combines Ulysses-SP (all-to-all) with Ring-Attention (ring P2P) for scalable parallelism.
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model Qwen/Qwen-Image \
+  --prompt "Epic fantasy landscape" \
+  --cache-backend cache_dit \
+  --ulysses-degree 2 \
+  --ring-degree 2 \
+  --num-inference-steps 50 \
+  --width 2048 \
+  --height 2048 \
+  --vae-use-tiling
+```
+
+### 5. Cache + Tensor Parallelism
+
+Best for: **Large models that don't fit in single GPU memory**
+
+Reduces per-GPU memory usage while maintaining cache acceleration.
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model Tongyi-MAI/Z-Image-Turbo \
+  --prompt "A cat reading a book" \
+  --cache-backend tea_cache \
+  --tensor-parallel-size 2 \
+  --num-inference-steps 9 \
+```
+
+## Online Serving
+
+### Cache + Sequence Parallelism
+
+```bash
+# TeaCache + Ulysses-SP
+vllm serve Qwen/Qwen-Image --omni --port 8091 \
+  --cache-backend tea_cache \
+  --cache-config '{"rel_l1_thresh": 0.2}' \
+  --usp 2
+
+# Cache-DiT + Ring-Attention
+vllm serve Qwen/Qwen-Image --omni --port 8091 \
+  --cache-backend cache_dit \
+  --cache-config '{"Fn_compute_blocks": 1, "max_warmup_steps": 8}' \
+  --ring 2
+```
+
+### Cache + CFG-Parallel
+
+```bash
+vllm serve Qwen/Qwen-Image-Edit --omni --port 8091 \
+  --cache-backend cache_dit \
+  --cfg-parallel-size 2
+```
+
+### Multiple Parallelism Methods
+
+```bash
+# CFG-Parallel + Ulysses-SP (4 GPUs total)
+vllm serve Qwen/Qwen-Image-Edit --omni --port 8091 \
+  --cache-backend cache_dit \
+  --cfg-parallel-size 2 \
+  --usp 2
+
+# Hybrid Ulysses + Ring (4 GPUs total)
+vllm serve Qwen/Qwen-Image --omni --port 8091 \
+  --cache-backend cache_dit \
+  --usp 2 \
+  --ring 2
+```
+
+
+## Limitations
+
+### Incompatibilities
+
+- **TeaCache + Cache-DiT**: These two cache methods cannot be used together. Only one cache backend can be active at a time. Attempting to enable both will result in an error.
+
+### Partial Support
+
+- **Tensor Parallelism — Text Encoder Not Sharded**: TP currently only shards the DiT blocks. Each TP rank retains a **full copy of the text encoder weights**, leading to significant GPU memory overhead proportional to TP degree. Tracked in [Issue #771](https://github.com/vllm-project/vllm-omni/issues/771).
+
+- **CPU Offloading — Two Modes Are Mutually Exclusive**: Model-level offload (`enable_cpu_offload`) and layerwise offload (`enable_layerwise_offload`) cannot be used simultaneously. If both are set, layerwise takes priority and model-level is silently ignored.
+
+- **CPU Offloading — VAE stays on GPU**: Both offloading strategies keep the VAE on GPU at all times. For high-resolution generation, VAE decode can still cause OOM. Mitigate by combining with `vae_use_tiling=True` or VAE Patch Parallelism.
+
+- **VAE Patch Parallelism — DistributedVaeExecutor Required**: VAE Patch Parallelism is only enabled for models that have `DistributedVaeExecutor`. Unsupported models will silently ignore `vae_patch_parallel_size`, and use sequential vae tiling instead.
+
+### Configuration Constraints
+
+- **GPU Count Must Match Parallel Degrees**: Total GPU count must satisfy:
+  ```
+  total_gpus = ulysses_degree × ring_degree × cfg_parallel_size × tensor_parallel_size
+  ```
+  Any mismatch will cause a configuration error at startup.
+
+- **VAE Patch Parallel Size ≤ DiT Process Group Size**: `vae_patch_parallel_size` reuses the DiT process group and cannot exceed it. Larger values are automatically clamped with a warning.
+
+- **Model-Specific TP Constraints**: Some models impose divisibility constraints on TP size. For example, Z-Image Turbo (`num_heads=30`) only supports `tensor_parallel_size=2`. Check [Supported Models](diffusion_features.md#supported-models) for per-model constraints.
+
+## Troubleshooting
+
+### Performance Not Scaling
+
+**Symptoms:** Adding more GPUs doesn't improve speed proportionally
+
+**Solutions:**
+1. Check GPU communication bandwidth (use `nvidia-smi topo -m`)
+2. Reduce parallelism degree if communication overhead is high
+3. For very long sequences, prefer Ring-Attention over Ulysses-SP
+4. Ensure batch size is large enough to saturate GPUs
+
+### Out of Memory with Parallelism
+
+**Symptoms:** OOM errors when combining methods
+
+**Solutions:**
+1. Enable Tensor Parallelism to shard weights
+2. Reduce resolution or batch size
+3. Combine with memory efficient methods, such as cpu offloading
+
+### Configuration Errors
+
+**Symptoms:** Errors about invalid parallel configuration
+
+**Solutions:**
+1. Verify total GPU count matches: `ulysses × ring × cfg × tp`
+2. Check model supports all enabled methods
+3. Ensure divisibility constraints (e.g., Z-Image TP=1 or 2 only)
+
+## See Also
+
+- [Diffusion Acceleration Overview](diffusion_features.md) - Main acceleration guide
@@ -129,8 +129,17 @@ The first run loads model weights. Subsequent runs are faster:
 --stage-init-timeout 900  # Increase timeout for slow storage
 ```

+### `Transformers does not recognize this architecture` Error
+
+Your have to upgrade `transformers` package to `5.3.0` or above:
+
+```
+pip install --upgrade transformers
+```
+
 ## Requirements

 - vLLM-Omni with GLM-Image support
 - CUDA-capable GPU (recommended: H100/A100 with 80GB)
 - GLM-Image model weights
+- `transformers` v5.3.0 or above
@@ -178,6 +178,12 @@ def parse_args() -> argparse.Namespace:
        action="store_true",
        help="Enable diffusion pipeline profiler to display stage durations.",
    )
+    parser.add_argument(
+        "--quantization",
+        type=str,
+        default=None,
+        help="Quantization method (e.g. 'fp8').",
+    )

    # Memory & parallelism
    parser.add_argument("--vae-use-slicing", action="store_true", help="Enable VAE slicing.")
@@ -213,6 +219,7 @@ def main():
        parallel_config=parallel_config,
        enforce_eager=args.enforce_eager,
        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
+        quantization=args.quantization,
    )

    # Validate I2V / V2V arguments
@@ -1,7 +1,10 @@
 import argparse
 from functools import lru_cache

-import gradio as gr
+try:
+    import gradio as gr
+except ImportError:
+    raise ImportError("gradio is required to run this demo. Install it with: pip install 'vllm-omni[demo]'") from None
 import torch

 from vllm_omni.entrypoints.omni import Omni
@@ -64,8 +64,8 @@ def parse_args() -> argparse.Namespace:
    parser.add_argument(
        "--guidance-scale",
        type=float,
-        default=1.0,
-        help="Classifier-free guidance scale.",
+        default=4.0,
+        help="Classifier-free guidance scale. HunyuanImage3 recommends 4.0-5.0.",
    )
    parser.add_argument("--height", type=int, default=1024, help="Height of generated image.")
    parser.add_argument("--width", type=int, default=1024, help="Width of generated image.")
@@ -16,7 +16,10 @@ import argparse
 import base64
 import io

-import gradio as gr
+try:
+    import gradio as gr
+except ImportError:
+    raise ImportError("gradio is required to run this demo. Install it with: pip install 'vllm-omni[demo]'") from None
 import httpx
 import numpy as np
 import soundfile as sf
@@ -196,3 +196,9 @@ gpu_memory_utilization: 0.5
 ```

 - The first request may be slow due to model warmup. Subsequent requests will be faster.
+
+- If you encounter `Transformers does not recognize this architecture` error, your have to upgrade `transformers` package to `5.3.0` or above:
+
+```
+pip install --upgrade transformers
+```
@@ -10,7 +10,10 @@ import argparse
 import base64
 from io import BytesIO

-import gradio as gr
+try:
+    import gradio as gr
+except ImportError:
+    raise ImportError("gradio is required to run this demo. Install it with: pip install 'vllm-omni[demo]'") from None
 import requests
 from PIL import Image

@@ -6,7 +6,10 @@ import random
 from pathlib import Path
 from typing import Any

-import gradio as gr
+try:
+    import gradio as gr
+except ImportError:
+    raise ImportError("gradio is required to run this demo. Install it with: pip install 'vllm-omni[demo]'") from None
 import numpy as np
 import soundfile as sf
 import torch
@@ -211,7 +211,7 @@ print(response.choices[1].message.audio)
 Supported speaker names depend on the model (e.g. `Ethan`, `Chelsie`, `Aiden`). Omit `speaker` to use the default.

 ## Streaming Output
-If you want to enable streaming output, please set the argument as below. The final output will be obtained just after generated by corresponding stage. Now we only support text streaming output. Other modalities can output normally.
+If you want to enable streaming output, please set the argument as below. The final output will be obtained just after generated by corresponding stage. We support both text streaming output and audio streaming output. Other modalities can output normally.
 ```bash
 python examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py \
    --query-type use_image \
@@ -6,7 +6,10 @@ import random
 from pathlib import Path
 from typing import Any

-import gradio as gr
+try:
+    import gradio as gr
+except ImportError:
+    raise ImportError("gradio is required to run this demo. Install it with: pip install 'vllm-omni[demo]'") from None
 import numpy as np
 import soundfile as sf
 import torch
@@ -24,7 +24,10 @@ import io
 import json
 import logging

-import gradio as gr
+try:
+    import gradio as gr
+except ImportError:
+    raise ImportError("gradio is required to run this demo. Install it with: pip install 'vllm-omni[demo]'") from None
 import httpx
 import numpy as np
 import soundfile as sf
@@ -122,7 +125,7 @@ PLAYER_HTML = """
 def _build_player_js(sample_rate: int) -> str:
    """Build the JavaScript that powers the AudioWorklet player."""
    return f"""
-() => {{
+    <script>
    const SR = {sample_rate};
    const WC = {json.dumps(WORKLET_JS)};
    let ctx = null, node = null, abort = null, gen = false, st = {{}};
@@ -281,7 +284,7 @@ def _build_player_js(sample_rate: int) -> str:
            }}
        }}
    }};
-}}
+    </script>
 """


@@ -432,10 +435,7 @@ def create_app(api_base: str):
    )

    with gr.Blocks(
-        css=css,
        title="Qwen3-TTS Demo",
-        js=_build_player_js(PCM_SAMPLE_RATE),
-        theme=theme,
    ) as demo:
        gr.HTML(f"""
        <div style="display:flex; align-items:center; gap:16px; margin-bottom:8px;">
@@ -770,7 +770,14 @@ def create_app(api_base: str):

        demo.queue()

-    return gr.mount_gradio_app(fastapi_app, demo, path="/")
+    return gr.mount_gradio_app(
+        fastapi_app,
+        demo,
+        path="/",
+        css=css,
+        theme=theme,
+        head=_build_player_js(PCM_SAMPLE_RATE),
+    )


 def parse_args():
@@ -3,7 +3,10 @@
 import base64
 import io

-import gradio as gr
+try:
+    import gradio as gr
+except ImportError:
+    raise ImportError("gradio is required to run this demo. Install it with: pip install 'vllm-omni[demo]'") from None
 import httpx
 import numpy as np
 import soundfile as sf
@@ -100,7 +103,7 @@ def build_payload(

    if task_type == "CustomVoice":
        if voice:
-            payload["speaker"] = voice
+            payload["voice"] = voice
        if instructions and instructions.strip():
            payload["instructions"] = instructions.strip()

@@ -10,7 +10,10 @@ import argparse
 import base64
 from io import BytesIO

-import gradio as gr
+try:
+    import gradio as gr
+except ImportError:
+    raise ImportError("gradio is required to run this demo. Install it with: pip install 'vllm-omni[demo]'") from None
 import requests
 from PIL import Image

@@ -20,7 +20,10 @@ from datetime import datetime
 from pathlib import Path
 from typing import Any

-import gradio as gr
+try:
+    import gradio as gr
+except ImportError:
+    raise ImportError("gradio is required to run this demo. Install it with: pip install 'vllm-omni[demo]'") from None
 import httpx
 import numpy as np
 import soundfile as sf
@@ -53,9 +53,14 @@ dev = [
    "av", # for ComfyUI tests
    "openpyxl>=3.0.0", # for nightly CI
    "pyttsx3>=2.99",
+    "opencc>=1.2.0",
    "mistune>=3.2.0", # for example tests
 ]

+demo = [
+    "gradio>=6.7.0",
+]
+
 docs = [
    "mkdocs>=1.5.0",
    "mkdocs-api-autonav",
@@ -4,7 +4,6 @@ librosa>=0.11.0
 resampy>=0.4.3
 diffusers>=0.36.0
 accelerate==1.12.0
-gradio==5.50
 soundfile>=0.13.1
 cache-dit==1.3.0
 tqdm>=4.66.0
@@ -18,3 +17,4 @@ prettytable>=3.8.0
 aenum==3.1.16
 pyzmq>=25.0.0
 janus>=1.0.0
+pydub
@@ -118,10 +118,6 @@ def benchmark_params(request, omni_server):
    if param_index >= len(all_params):
        raise ValueError(f"No benchmark parameters found for index {param_index} in test: {test_name}")

-    if all_params[param_index]["dataset_name"] == "random-mm":
-        # TODO: Due to known issues, skip the random-mm dataset.
-        pytest.skip("Skipping parameter for random-mm dataset.")
-
    current = param_index + 1
    total = len(all_params)
    print(f"\n  Running benchmark {current}/{total} for {test_name}")
@@ -132,15 +128,70 @@ def benchmark_params(request, omni_server):
    }


-def assert_result(result, params, num_prompt):
+def _resolve_baseline_value(
+    baseline_raw: Any,
+    *,
+    sweep_index: int | None,
+    max_concurrency: Any = None,
+    request_rate: Any = None,
+) -> Any:
+    """Pick the baseline threshold for this sweep step.
+
+    Supported shapes per metric:
+    - **Scalar** — same threshold for every concurrency / QPS.
+    - **List** — aligned with ``max_concurrency`` / ``request_rate`` sweep order; use ``sweep_index``.
+    - **Dict** — keyed by concurrency or rate, e.g. ``{"1": 500, "4": 800}`` (keys are strings in JSON).
+
+    For dict lookup, ``max_concurrency`` is preferred when both are set (concurrency sweep).
+    """
+    if baseline_raw is None:
+        # If no baseline is set, the maximum value will be used.
+        return 100000
+    if isinstance(baseline_raw, dict):
+        if max_concurrency is not None:
+            for key in (max_concurrency, str(max_concurrency)):
+                if key in baseline_raw:
+                    return baseline_raw[key]
+        if request_rate is not None:
+            for key in (request_rate, str(request_rate)):
+                if key in baseline_raw:
+                    return baseline_raw[key]
+        raise KeyError(
+            f"baseline dict has no key for max_concurrency={max_concurrency!r} "
+            f"or request_rate={request_rate!r}; keys={list(baseline_raw.keys())!r}"
+        )
+    if isinstance(baseline_raw, (list, tuple)):
+        return baseline_raw[sweep_index]
+    return baseline_raw
+
+
+def assert_result(
+    result,
+    params,
+    num_prompt,
+    *,
+    sweep_index: int | None = None,
+    max_concurrency: Any = None,
+    request_rate: Any = None,
+) -> None:
    assert result["completed"] == num_prompt, "Request failures exist"
    baseline_data = params.get("baseline", {})
-    for metric_name, baseline_value in baseline_data.items():
+    for metric_name, baseline_raw in baseline_data.items():
        current_value = result[metric_name]
+        baseline_value = _resolve_baseline_value(
+            baseline_raw,
+            sweep_index=sweep_index,
+            max_concurrency=max_concurrency,
+            request_rate=request_rate,
+        )
        if "throughput" in metric_name:
-            assert current_value >= baseline_value, f"{metric_name}: {current_value} < {baseline_value}"
+            if current_value <= baseline_value:
+                print(
+                    f"ERROR: Throughput test results were below baseline: {metric_name}: {current_value} > {baseline_value}"
+                )
        else:
-            assert current_value <= baseline_value, f"{metric_name}: {current_value} > {baseline_value}"
+            if current_value >= baseline_value:
+                print(f"ERROR: Test results exceeded baseline: {metric_name}: {current_value} < {baseline_value}")


@pytest.mark.parametrize("omni_server", test_params, indirect=True)
@@ -195,8 +246,8 @@ def test_performance_benchmark(omni_server, benchmark_params):
        elif not isinstance(value, bool):
            args.extend([arg_name, str(value)])

-    # QPS test
-    for qps, num_prompt in zip(qps_list, num_prompt_list):
+    # QPS test (sweep_index aligns with qps_list / num_prompt_list for this loop)
+    for i, (qps, num_prompt) in enumerate(zip(qps_list, num_prompt_list)):
        args = args + ["--request-rate", str(qps), "--num-prompts", str(num_prompt)]
        result = run_benchmark(
            args=args,
@@ -205,10 +256,16 @@ def test_performance_benchmark(omni_server, benchmark_params):
            dataset_name=dataset_name,
            num_prompt=num_prompt,
        )
-        assert_result(result, params, num_prompt=num_prompt)
+        assert_result(
+            result,
+            params,
+            num_prompt=num_prompt,
+            sweep_index=i,
+            request_rate=qps,
+        )

-    # concurrency test
-    for concurrency, num_prompt in zip(max_concurrency_list, num_prompt_list):
+    # concurrency test (sweep_index aligns with max_concurrency_list for separate thresholds per concurrency)
+    for i, (concurrency, num_prompt) in enumerate(zip(max_concurrency_list, num_prompt_list)):
        args = args + ["--max-concurrency", str(concurrency), "--num-prompts", str(num_prompt), "--request-rate", "inf"]
        result = run_benchmark(
            args=args,
@@ -217,4 +274,10 @@ def test_performance_benchmark(omni_server, benchmark_params):
            dataset_name=dataset_name,
            num_prompt=num_prompt,
        )
-        assert_result(result, params, num_prompt=num_prompt)
+        assert_result(
+            result,
+            params,
+            num_prompt=num_prompt,
+            sweep_index=i,
+            max_concurrency=concurrency,
+        )
@@ -25,9 +25,9 @@
                "ignore_eos": true,
                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
                "baseline": {
-                    "mean_ttft_ms": 100000,
-                    "mean_audio_ttfp_ms": 100000,
-                    "mean_audio_rtf": 100000
+                    "mean_ttft_ms": [1000, 3000, 5000],
+                    "mean_audio_ttfp_ms": [8000, 10000, 13000],
+                    "mean_audio_rtf": [0.2, 0.25, 0.45]
                }
            },
            {
@@ -39,10 +39,10 @@
                    40,
                    100
                ],
-                "max_concurrency": [
-                    1,
-                    4,
-                    10
+                "request_rate": [
+                    0.1,
+                    0.3,
+                    0.5
                ],
                "random_input_len": 100,
                "random_output_len": 100,
@@ -62,9 +62,31 @@
                },
                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
                "baseline": {
-                    "mean_ttft_ms": 100000,
-                    "mean_audio_ttfp_ms": 100000,
-                    "mean_audio_rtf": 100000
+                    "mean_ttft_ms": [2000, 4000, 6000],
+                    "mean_audio_ttfp_ms": [10000, 13000, 15000],
+                    "mean_audio_rtf": [0.25, 0.35, 0.45]
+                }
+            },
+            {
+                "dataset_name": "random",
+                "backend": "openai-chat-omni",
+                "endpoint": "/v1/chat/completions",
+                "num_prompts": [
+                    4,
+                    16
+                ],
+                "max_concurrency": [
+                    1,
+                    4
+                ],
+                "random_input_len": 2500,
+                "random_output_len": 900,
+                "ignore_eos": true,
+                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
+                "baseline": {
+                    "mean_ttft_ms": [1000, 3000],
+                    "mean_audio_ttfp_ms": [30000, 60000],
+                    "mean_audio_rtf": [0.35, 0.45]
                }
            }
        ]
@@ -113,9 +135,9 @@
                "ignore_eos": true,
                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
                "baseline": {
-                    "mean_ttft_ms": 100000,
-                    "mean_audio_ttfp_ms": 100000,
-                    "mean_audio_rtf": 100000
+                    "mean_ttft_ms": [1000, 3000, 5000],
+                    "mean_audio_ttfp_ms": [1000, 3000, 5000],
+                    "mean_audio_rtf": [0.2, 0.35, 0.6]
                }
            },
            {
@@ -127,10 +149,10 @@
                    40,
                    100
                ],
-                "max_concurrency": [
-                    1,
-                    4,
-                    10
+                "request_rate": [
+                    0.1,
+                    0.3,
+                    0.5
                ],
                "random_input_len": 100,
                "random_output_len": 100,
@@ -150,9 +172,31 @@
                },
                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
                "baseline": {
-                    "mean_ttft_ms": 100000,
-                    "mean_audio_ttfp_ms": 100000,
-                    "mean_audio_rtf": 100000
+                    "mean_ttft_ms": [2000, 4000, 6000],
+                    "mean_audio_ttfp_ms": [2000, 4000, 6000],
+                    "mean_audio_rtf": [0.25, 0.4, 0.7]
+                }
+            },
+            {
+                "dataset_name": "random",
+                "backend": "openai-chat-omni",
+                "endpoint": "/v1/chat/completions",
+                "num_prompts": [
+                    4,
+                    16
+                ],
+                "max_concurrency": [
+                    1,
+                    4
+                ],
+                "random_input_len": 2500,
+                "random_output_len": 900,
+                "ignore_eos": true,
+                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
+                "baseline": {
+                    "mean_ttft_ms": [1000, 3000],
+                    "mean_audio_ttfp_ms": [1000, 3000],
+                    "mean_audio_rtf": [0.35, 0.45]
                }
            }
        ]
@@ -183,8 +227,8 @@
                },
                "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
                "baseline": {
-                    "mean_audio_ttfp_ms": 100000,
-                    "mean_audio_rtf": 100000
+                    "mean_audio_ttfp_ms": [6000, 6000],
+                    "mean_audio_rtf": [0.3, 0.3]
                }
            }
        ]
@@ -209,12 +209,9 @@ def _test_cfg_parallel_worker(
            cfg_normalize=test_config["cfg_normalize"],
        )

-    # Only rank 0 has valid output in CFG parallel mode
-    if cfg_rank == 0:
-        assert noise_pred is not None
-        result_queue.put(noise_pred.cpu())
-    else:
-        assert noise_pred is None
+    # CFG parallel returns the combined prediction on every rank.
+    assert noise_pred is not None
+    result_queue.put((cfg_rank, noise_pred.cpu()))

    destroy_distributed_env()

@@ -348,7 +345,18 @@ def test_predict_noise_maybe_with_cfg(cfg_parallel_size: int, dtype: torch.dtype

    # Get results from queues
    baseline_output = baseline_queue.get()
-    cfg_parallel_output = cfg_parallel_queue.get()
+    cfg_parallel_outputs = [cfg_parallel_queue.get() for _ in range(cfg_parallel_size)]
+    cfg_parallel_outputs.sort(key=lambda item: item[0])
+    cfg_parallel_output = cfg_parallel_outputs[0][1]
+
+    for cfg_rank, rank_output in cfg_parallel_outputs[1:]:
+        torch.testing.assert_close(
+            rank_output,
+            cfg_parallel_output,
+            rtol=0,
+            atol=0,
+            msg=f"CFG parallel ranks produced different outputs (rank 0 vs rank {cfg_rank})",
+        )

    # Verify shapes match
    assert baseline_output.shape == cfg_parallel_output.shape, (
@@ -7,6 +7,49 @@ import pytest
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]


+class TestSetForwardContextNumTokens:
+    """Test _set_forward_context_num_tokens defensive fix for vLLM 0.18.0."""
+
+    def test_sets_num_tokens_when_context_available(self, mocker):
+        """num_tokens should be set on ForwardContext when available."""
+        import vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe as hunyuan_moe
+
+        mock_ctx = mocker.MagicMock()
+        del mock_ctx.in_profile_run  # simulate missing attr
+        mocker.patch.object(hunyuan_moe._vllm_fc, "is_forward_context_available", return_value=True)
+        mocker.patch.object(hunyuan_moe._vllm_fc, "get_forward_context", return_value=mock_ctx)
+
+        hunyuan_moe._set_forward_context_num_tokens(1024)
+
+        assert mock_ctx.num_tokens == 1024
+        assert mock_ctx.in_profile_run is False
+
+    def test_sets_in_profile_run_only_if_missing(self, mocker):
+        """in_profile_run should not be overwritten if already set."""
+        import vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe as hunyuan_moe
+
+        mock_ctx = mocker.MagicMock()
+        mock_ctx.in_profile_run = True  # already set
+        mocker.patch.object(hunyuan_moe._vllm_fc, "is_forward_context_available", return_value=True)
+        mocker.patch.object(hunyuan_moe._vllm_fc, "get_forward_context", return_value=mock_ctx)
+
+        hunyuan_moe._set_forward_context_num_tokens(512)
+
+        assert mock_ctx.num_tokens == 512
+        assert mock_ctx.in_profile_run is True  # not overwritten
+
+    def test_noop_when_context_unavailable(self, mocker):
+        """Should do nothing when ForwardContext is not available."""
+        import vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe as hunyuan_moe
+
+        mocker.patch.object(hunyuan_moe._vllm_fc, "is_forward_context_available", return_value=False)
+        mock_get = mocker.patch.object(hunyuan_moe._vllm_fc, "get_forward_context")
+
+        hunyuan_moe._set_forward_context_num_tokens(256)
+
+        mock_get.assert_not_called()
+
+
 class TestHunyuanFusedMoEPlatformDispatch:
    """Test platform dispatch via platform qualname hooks."""

@@ -47,6 +47,19 @@ def pytest_addoption(parser):
        help="Balanced sample count per GEdit task group",
    )
    group.addoption("--accuracy-workers", action="store", type=int, default=1, help="Worker count for accuracy benches")
+    group.addoption(
+        "--wan22-i2v-image-source",
+        action="store",
+        default=None,
+        help="Image source for Wan2.2 I2V accuracy tests. Can be local path or remote URL.",
+    )
+    group.addoption(
+        "--wan22-i2v-online-timeout-seconds",
+        action="store",
+        type=int,
+        default=1200,
+        help="Online serving timeout in seconds for Wan2.2 I2V accuracy tests.",
+    )


 def _hf_cache_root() -> Path:
@@ -142,6 +155,17 @@ def accuracy_workers(request: pytest.FixtureRequest) -> int:
    return int(request.config.getoption("accuracy_workers"))


+@pytest.fixture(scope="session")
+def wan22_i2v_image_source(request: pytest.FixtureRequest) -> str | None:
+    value = request.config.getoption("wan22_i2v_image_source")
+    return str(value) if value else None
+
+
+@pytest.fixture(scope="session")
+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 gebench_samples_per_type(request: pytest.FixtureRequest) -> int:
    return int(request.config.getoption("gebench_samples_per_type"))
@@ -0,0 +1,165 @@
+from __future__ import annotations
+
+import argparse
+import base64
+import json
+from io import BytesIO
+from pathlib import Path
+
+import requests
+import torch
+from diffusers import UniPCMultistepScheduler, WanImageToVideoPipeline
+from diffusers.pipelines.wan import pipeline_wan_i2v as wan_i2v_module
+from diffusers.utils import export_to_video, load_image
+from PIL import Image
+
+from tests.e2e.accuracy.wan22_i2v.wan22_i2v_video_similarity_common import BOUNDARY_RATIO
+
+
+def _parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Run Wan2.2 I2V diffusers offline generation.")
+    parser.add_argument("--model", required=True)
+    parser.add_argument("--image-source", required=True)
+    parser.add_argument("--prompt", required=True)
+    parser.add_argument("--negative-prompt", required=True)
+    parser.add_argument("--size", required=True)
+    parser.add_argument("--fps", type=int, required=True)
+    parser.add_argument("--num-frames", type=int, required=True)
+    parser.add_argument("--guidance-scale", type=float, required=True)
+    parser.add_argument("--guidance-scale-2", type=float, required=True)
+    parser.add_argument("--flow-shift", type=float, required=True)
+    parser.add_argument("--num-inference-steps", type=int, required=True)
+    parser.add_argument("--seed", type=int, required=True)
+    parser.add_argument("--output", required=True)
+    parser.add_argument("--metadata-output", required=True)
+    return parser.parse_args()
+
+
+def _parse_size(size: str) -> tuple[int, int]:
+    width_str, height_str = size.lower().split("x", 1)
+    return int(width_str), int(height_str)
+
+
+class _IdentityFtfy:
+    @staticmethod
+    def fix_text(text: str) -> str:
+        return text
+
+
+def _ensure_wan_ftfy_fallback() -> None:
+    if not hasattr(wan_i2v_module, "ftfy"):
+        wan_i2v_module.ftfy = _IdentityFtfy()
+
+
+def _offline_cuda_device() -> torch.device:
+    return torch.device("cuda:0")
+
+
+def _load_input_image(source: str) -> Image.Image:
+    if source.startswith("data:image"):
+        _, encoded = source.split(",", 1)
+        image = Image.open(BytesIO(base64.b64decode(encoded)))
+        image.load()
+        return image.convert("RGB")
+
+    source_path = Path(source)
+    if source_path.exists():
+        image = Image.open(source_path)
+        image.load()
+        return image.convert("RGB")
+
+    image = load_image(source)
+    if isinstance(image, Image.Image):
+        image.load()
+        return image.convert("RGB")
+
+    response = requests.get(source, timeout=60)
+    response.raise_for_status()
+    image = Image.open(BytesIO(response.content))
+    image.load()
+    return image.convert("RGB")
+
+
+def _resize_to_target(image: Image.Image, *, width: int, height: int) -> Image.Image:
+    return image.resize((width, height), Image.Resampling.LANCZOS)
+
+
+def _configure_scheduler(pipe: WanImageToVideoPipeline, *, flow_shift: float) -> None:
+    pipe.scheduler = UniPCMultistepScheduler.from_config(
+        pipe.scheduler.config,
+        flow_shift=flow_shift,
+    )
+
+
+def _write_metadata(
+    path: Path,
+    *,
+    args: argparse.Namespace,
+    width: int,
+    height: int,
+    frame_count: int,
+) -> None:
+    payload = {
+        "model": args.model,
+        "image_source": args.image_source,
+        "size": args.size,
+        "width": width,
+        "height": height,
+        "fps": args.fps,
+        "num_frames": args.num_frames,
+        "actual_frame_count": frame_count,
+        "guidance_scale": args.guidance_scale,
+        "guidance_scale_2": args.guidance_scale_2,
+        "boundary_ratio": BOUNDARY_RATIO,
+        "flow_shift": args.flow_shift,
+        "num_inference_steps": args.num_inference_steps,
+        "seed": args.seed,
+        "world_size": 1,
+    }
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
+
+
+def main() -> int:
+    args = _parse_args()
+    device = _offline_cuda_device()
+    torch.cuda.set_device(device)
+    _ensure_wan_ftfy_fallback()
+
+    pipe = WanImageToVideoPipeline.from_pretrained(args.model, torch_dtype=torch.bfloat16)
+    pipe.register_to_config(boundary_ratio=BOUNDARY_RATIO)
+    _configure_scheduler(pipe, flow_shift=args.flow_shift)
+    pipe.to(device)
+    pipe.set_progress_bar_config(disable=False)
+
+    input_image = _load_input_image(args.image_source)
+    width, height = _parse_size(args.size)
+    resized_image = _resize_to_target(input_image, width=width, height=height)
+
+    generator = torch.Generator(device=device.type).manual_seed(args.seed)
+    frames = pipe(
+        image=resized_image,
+        prompt=args.prompt,
+        negative_prompt=args.negative_prompt,
+        height=height,
+        width=width,
+        num_frames=args.num_frames,
+        guidance_scale=args.guidance_scale,
+        guidance_scale_2=args.guidance_scale_2,
+        num_inference_steps=args.num_inference_steps,
+        generator=generator,
+    ).frames[0]
+
+    output_path = Path(args.output)
+    metadata_path = Path(args.metadata_output)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    export_to_video(frames, str(output_path), fps=args.fps)
+    _write_metadata(metadata_path, args=args, width=width, height=height, frame_count=len(frames))
+
+    if hasattr(pipe, "maybe_free_model_hooks"):
+        pipe.maybe_free_model_hooks()
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
@@ -0,0 +1,672 @@
+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 pathlib import Path
+from typing import Any
+from urllib.parse import urlparse
+
+import pytest
+import requests
+import torch
+from diffusers import UniPCMultistepScheduler
+from PIL import Image
+
+from tests.conftest import OmniServerParams
+from tests.e2e.accuracy.wan22_i2v.run_wan22_i2v_diffusers_cp import (
+    _configure_scheduler,
+    _ensure_wan_ftfy_fallback,
+    _IdentityFtfy,
+    _offline_cuda_device,
+    _resize_to_target,
+)
+from tests.e2e.accuracy.wan22_i2v.wan22_i2v_video_similarity_common import (
+    FLOW_SHIFT,
+    FPS,
+    GUIDANCE_SCALE,
+    GUIDANCE_SCALE_2,
+    HEIGHT,
+    MODEL_NAME,
+    NEGATIVE_PROMPT,
+    NUM_FRAMES,
+    NUM_INFERENCE_STEPS,
+    PROMPT,
+    PSNR_THRESHOLD,
+    RABBIT_IMAGE_URL,
+    SEED,
+    SIZE,
+    SSIM_THRESHOLD,
+    WIDTH,
+)
+from tests.utils import hardware_test
+
+
+def test_parse_video_metadata_extracts_dimensions_and_fps() -> None:
+    payload = {
+        "streams": [
+            {
+                "width": 832,
+                "height": 480,
+                "avg_frame_rate": "16/1",
+                "nb_read_frames": "93",
+            }
+        ]
+    }
+
+    metadata = _parse_video_metadata(payload)
+
+    assert metadata["width"] == 832
+    assert metadata["height"] == 480
+    assert metadata["fps"] == 16.0
+    assert metadata["frame_count"] == 93
+
+
+def test_parse_ssim_summary_extracts_all_score() -> None:
+    output = """
+    [Parsed_ssim_0 @ 000001] SSIM Y:0.971903 (15.512007) U:0.965077 (14.569044) V:0.962414 (14.252637) All:0.968311 (15.035654)
+    """
+
+    assert _parse_ssim_score(output) == 0.968311
+
+
+def test_parse_psnr_summary_extracts_average_score() -> None:
+    output = """
+    [Parsed_psnr_0 @ 000001] PSNR y:32.670157 u:31.844621 v:31.513839 average:32.148004 min:31.563744 max:33.201457
+    """
+
+    assert _parse_psnr_score(output) == 32.148004
+
+
+def test_build_diffusers_command_uses_python_runner_path(tmp_path: Path) -> None:
+    runner_path = tmp_path / "run_wan22_i2v_diffusers_cp.py"
+    command = _build_diffusers_command(
+        runner_path=runner_path,
+        image_source=RABBIT_IMAGE_URL,
+        output_path=tmp_path / "offline.mp4",
+        metadata_path=tmp_path / "offline.json",
+    )
+
+    assert command[:2] == [
+        sys.executable,
+        str(runner_path),
+    ]
+    assert "--output" in command
+    assert "--metadata-output" in command
+
+
+def test_resolve_image_source_prefers_existing_local_path(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+    image_path = tmp_path / "rabbit.png"
+    Image.new("RGB", (8, 8), color=(255, 128, 64)).save(image_path)
+
+    assert _resolve_image_source(str(image_path)) == str(image_path.resolve())
+
+
+def test_build_online_image_reference_uses_data_url_for_local_path(tmp_path: Path) -> None:
+    image_path = tmp_path / "rabbit.png"
+    Image.new("RGB", (4, 2), color=(10, 20, 30)).save(image_path)
+
+    reference = _build_online_image_reference(str(image_path))
+
+    assert reference.startswith("data:image/png;base64,")
+    encoded = reference.split(",", 1)[1]
+    assert len(b64decode(encoded)) > 0
+
+
+def test_send_video_request_with_timeout_uses_requested_timeout(monkeypatch: pytest.MonkeyPatch) -> None:
+    class DummyResponse:
+        def raise_for_status(self) -> None:
+            return None
+
+        def json(self) -> dict[str, str]:
+            return {"id": "video_123"}
+
+    class DummyClient:
+        def __init__(self) -> None:
+            self.run_level = "core_model"
+            self.wait_args: tuple[str, int, int] | None = None
+
+        def _build_url(self, path: str) -> str:
+            return f"http://localhost:8000{path}"
+
+        def _wait_until_video_completed(
+            self,
+            video_id: str,
+            poll_interval_seconds: int = 2,
+            timeout_seconds: int = 300,
+        ) -> None:
+            self.wait_args = (video_id, poll_interval_seconds, timeout_seconds)
+
+        def _download_video_content(self, video_id: str) -> bytes:
+            return b"video-bytes"
+
+    monkeypatch.setattr(requests, "post", lambda *args, **kwargs: DummyResponse())
+
+    client = DummyClient()
+    result = _send_video_request_with_timeout(
+        client,
+        {
+            "form_data": {"prompt": "test"},
+            "image_reference": RABBIT_IMAGE_URL,
+        },
+        timeout_seconds=1200,
+    )
+
+    assert result == b"video-bytes"
+    assert client.wait_args == ("video_123", 2, 1200)
+
+
+def test_online_timeout_defaults_to_1200() -> None:
+    assert _online_timeout_seconds(None) == 1200
+
+
+def test_artifact_dir_is_under_repo_result_folder(tmp_path: Path) -> None:
+    artifact_dir = _artifact_dir(str(tmp_path / "rabbit.png"))
+
+    assert artifact_dir.parent == Path(__file__).parent / "result"
+    assert artifact_dir.name.startswith("rabbit-")
+
+
+def test_offline_cuda_device_uses_indexed_cuda_device() -> None:
+    assert _offline_cuda_device() == torch.device("cuda:0")
+
+
+def test_ensure_wan_ftfy_fallback_sets_identity(monkeypatch: pytest.MonkeyPatch) -> None:
+    from diffusers.pipelines.wan import pipeline_wan_i2v as wan_i2v_module
+
+    monkeypatch.delattr(wan_i2v_module, "ftfy", raising=False)
+    _ensure_wan_ftfy_fallback()
+
+    assert hasattr(wan_i2v_module, "ftfy")
+    assert isinstance(wan_i2v_module.ftfy, _IdentityFtfy)
+    assert wan_i2v_module.ftfy.fix_text("abc") == "abc"
+
+
+def test_resize_to_target_matches_requested_dimensions() -> None:
+    image = Image.new("RGB", (100, 50), color=(1, 2, 3))
+
+    resized = _resize_to_target(image, width=832, height=480)
+
+    assert resized.size == (832, 480)
+
+
+def test_configure_scheduler_switches_to_unipc() -> None:
+    class DummyPipe:
+        def __init__(self) -> None:
+            self.scheduler = UniPCMultistepScheduler(num_train_timesteps=1000, flow_shift=1.0)
+
+    pipe = DummyPipe()
+    _configure_scheduler(pipe, flow_shift=5.0)
+
+    assert isinstance(pipe.scheduler, UniPCMultistepScheduler)
+    assert pipe.scheduler.config.flow_shift == 5.0
+
+
+REPO_ROOT = Path(__file__).resolve().parents[4]
+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(
+            model=MODEL_NAME,
+            server_args=[
+                "--usp",
+                "2",
+                "--use-hsdp",
+                "--hsdp-shard-size",
+                "2",
+            ],
+            use_omni=True,
+        ),
+        id="wan22_i2v_usp2_hsdp2",
+    )
+]
+
+
+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,
+    image_source: str,
+    output_path: Path,
+    metadata_path: Path,
+) -> list[str]:
+    return [
+        sys.executable,
+        str(runner_path),
+        "--model",
+        MODEL_NAME,
+        "--image-source",
+        image_source,
+        "--prompt",
+        PROMPT,
+        "--negative-prompt",
+        NEGATIVE_PROMPT,
+        "--size",
+        SIZE,
+        "--fps",
+        str(FPS),
+        "--num-frames",
+        str(NUM_FRAMES),
+        "--guidance-scale",
+        str(GUIDANCE_SCALE),
+        "--guidance-scale-2",
+        str(GUIDANCE_SCALE_2),
+        "--flow-shift",
+        str(FLOW_SHIFT),
+        "--num-inference-steps",
+        str(NUM_INFERENCE_STEPS),
+        "--seed",
+        str(SEED),
+        "--output",
+        str(output_path),
+        "--metadata-output",
+        str(metadata_path),
+    ]
+
+
+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:
+    configured = configured or RABBIT_IMAGE_URL
+    candidate = Path(configured)
+    if candidate.exists():
+        return str(candidate.resolve())
+    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}"
+
+
+def _artifact_paths(image_source: str) -> tuple[Path, Path, Path]:
+    artifact_dir = _artifact_dir(image_source)
+    artifact_dir.mkdir(parents=True, exist_ok=True)
+    return (
+        artifact_dir / "online.mp4",
+        artifact_dir / "offline.mp4",
+        artifact_dir / "offline_metadata.json",
+    )
+
+
+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,
+    openai_client,
+    image_source: str,
+    online_timeout_seconds: int,
+) -> Path:
+    online_path, _, _ = _artifact_paths(image_source)
+    request_config = {
+        "model": omni_server.model,
+        "form_data": {
+            "prompt": PROMPT,
+            "negative_prompt": NEGATIVE_PROMPT,
+            "size": SIZE,
+            "fps": FPS,
+            "num_frames": NUM_FRAMES,
+            "guidance_scale": GUIDANCE_SCALE,
+            "guidance_scale_2": GUIDANCE_SCALE_2,
+            "flow_shift": FLOW_SHIFT,
+            "num_inference_steps": NUM_INFERENCE_STEPS,
+            "seed": SEED,
+        },
+        "image_reference": _build_online_image_reference(image_source),
+    }
+    online_video_bytes = _send_video_request_with_timeout(
+        openai_client,
+        request_config,
+        timeout_seconds=_online_timeout_seconds(online_timeout_seconds),
+    )
+    online_path.write_bytes(online_video_bytes)
+    return online_path
+
+
+def _generate_offline_video(*, image_source: str) -> tuple[Path, Path]:
+    _, offline_path, offline_metadata_path = _artifact_paths(image_source)
+    command = _build_diffusers_command(
+        runner_path=RUNNER_PATH,
+        image_source=image_source,
+        output_path=offline_path,
+        metadata_path=offline_metadata_path,
+    )
+    subprocess.run(
+        command,
+        cwd=REPO_ROOT,
+        env=_runner_env(),
+        check=True,
+        timeout=VIDEO_TIMEOUT_SECONDS,
+    )
+    return offline_path, offline_metadata_path
+
+
+@pytest.mark.advanced_model
+@pytest.mark.benchmark
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+def test_wan22_i2v_diffusers_offline_generates_video(
+    wan22_i2v_image_source: str | None,
+) -> None:
+    if not torch.cuda.is_available():
+        pytest.skip("Wan2.2 I2V diffusers offline test requires CUDA.")
+
+    _probe_binary("ffprobe")
+    if not RUNNER_PATH.exists():
+        raise AssertionError(f"Offline diffusers runner does not exist: {RUNNER_PATH}")
+
+    image_source = _resolve_image_source(wan22_i2v_image_source)
+    _validate_image_source(image_source)
+    offline_path, offline_metadata_path = _generate_offline_video(image_source=image_source)
+    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
+
+
+@pytest.mark.advanced_model
+@pytest.mark.benchmark
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "H100"}, num_cards=2)
+@pytest.mark.parametrize("omni_server", SERVER_CASES, indirect=True)
+def test_wan22_i2v_online_serving_generates_video(
+    omni_server,
+    openai_client,
+    wan22_i2v_image_source: str | None,
+    wan22_i2v_online_timeout_seconds: int,
+) -> None:
+    if not torch.cuda.is_available() or torch.cuda.device_count() < 2:
+        pytest.skip("Wan2.2 I2V similarity e2e test requires >= 2 CUDA GPUs.")
+
+    _probe_binary("ffprobe")
+    image_source = _resolve_image_source(wan22_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,
+        online_timeout_seconds=wan22_i2v_online_timeout_seconds,
+    )
+    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
+
+
+@pytest.mark.advanced_model
+@pytest.mark.benchmark
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "H100"}, num_cards=2)
+def test_wan22_i2v_serving_matches_diffusers_video_similarity(
+    wan22_i2v_image_source: str | None,
+) -> None:
+    if not torch.cuda.is_available() or torch.cuda.device_count() < 2:
+        pytest.skip("Wan2.2 I2V similarity e2e test requires >= 2 CUDA GPUs.")
+
+    _probe_binary("ffmpeg")
+    _probe_binary("ffprobe")
+    if not RUNNER_PATH.exists():
+        raise AssertionError(f"Offline diffusers runner does not exist: {RUNNER_PATH}")
+
+    image_source = _resolve_image_source(wan22_i2v_image_source)
+    _validate_image_source(image_source)
+    online_path, offline_path, offline_metadata_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() or not offline_metadata_path.exists():
+        pytest.skip(f"Missing offline artifacts from prerequisite test: {offline_path}, {offline_metadata_path}")
+
+    assert online_path.exists(), f"Expected online video artifact at {online_path}"
+    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}"
+
+    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 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"
+    )
+    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}"
+    )
@@ -0,0 +1,24 @@
+from __future__ import annotations
+
+RABBIT_IMAGE_URL = "https://vllm-public-assets.s3.us-west-2.amazonaws.com/omni-assets/rabbit.png"
+MODEL_NAME = "Wan-AI/Wan2.2-I2V-A14B-Diffusers"
+SIZE = "832x480"
+WIDTH = 832
+HEIGHT = 480
+FPS = 16
+NUM_FRAMES = 93
+GUIDANCE_SCALE = 3.5
+GUIDANCE_SCALE_2 = 3.5
+FLOW_SHIFT = 5.0
+NUM_INFERENCE_STEPS = 50
+SEED = 42
+BOUNDARY_RATIO = 0.875
+SSIM_THRESHOLD = 0.94
+PSNR_THRESHOLD = 28.0
+
+PROMPT = """一只棕色野兔的正面特写镜头，采用低角度仰拍视角，营造亲密而庄严的视觉冲击。兔子一双圆润漆黑的大眼睛直视镜头深处，眼神中交织着野生动物的警觉与一丝难以言喻的温柔好奇，仿佛在与观者建立跨越物种的静默对话。它毛色呈现层次丰富的棕褐渐变，从浅奶油色腹部过渡到深棕背部，每根毛发纹理清晰可辨，在侧光下泛着丝绸般的光泽。细长洁白的胡须共有三对，随呼吸节奏微微颤动，偶尔因捕捉气流信息而轻轻摇摆。
+两只标志性的长耳完全竖立，耳廓外侧覆盖短密棕毛，内侧则露出粉嫩的血管网络，薄如蝉翼的皮肤下血液流动隐约可见，耳朵以细微幅度不时转动，精准定位声源方向。背景是一片澄澈的蔚蓝天空，形态蓬松的白色积云以缓慢速度横向漂移，云影在兔子头顶交替变化，光线随之明暗流转。晴朗天气的明媚阳光从画面左上方45度角倾泻而下，在兔脸右侧形成柔和的伦勃朗式阴影，强化了面部立体感和皮毛质感。
+兔子湿润的黑鼻子持续进行每秒三至四次的快速抽动，这是它们感知化学信号的本能动作，粉色三瓣嘴随之轻启，露出正在反刍的洁白门齿，下颌以稳定节奏左右研磨。摄影采用大光圈浅景深，焦点牢牢锁定在兔子双眼连线所在的焦平面，背景天空和远景绿色植被虚化成圆润的彩色光斑，前景几根嫩绿草叶闯入画面边缘，以缓慢弧线随风摇曳，暗示着和煦的春日微风。
+整个场景弥漫着宁静致远的田园诗意，色彩温暖饱和，充满生命力。兔子在持续五秒的对视后，以典型 lagomorph 特征完成一次完整的瞬膜眨眼——第三眼睑从内侧横向滑过眼球，继而缓缓歪头向右十五度，这个行为在动物行为学中代表认知加工和好奇心表达，耳朵随之向同一方向倾斜，最终恢复正视姿态，胡须舒展，完成这段短暂而珍贵的自然纪录。"""
+
+NEGATIVE_PROMPT = """色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"""
@@ -0,0 +1,83 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+E2E offline tests for Qwen3-TTS Base model with text input and audio output.
+
+Async_chunk disable, cuda_graph disabled (no_async_chunk stage config).
+CUDA graph is disabled by setting engine_args.enforce_eager=true via modify_stage_config().
+Same structure as test_qwen3_omni (models, stage_configs, test_params, parametrize omni_runner).
+"""
+
+import os
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
+
+from pathlib import Path
+
+import pytest
+
+from tests.conftest import modify_stage_config
+from tests.utils import hardware_test
+
+MODEL = "Qwen/Qwen3-TTS-12Hz-0.6B-Base"
+REF_AUDIO_URL = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-TTS-Repo/clone_2.wav"
+REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."
+
+
+def get_cuda_graph_config():
+    path = modify_stage_config(
+        get_stage_config(),
+        updates={
+            "stage_args": {
+                0: {
+                    "engine_args.enforce_eager": "true",
+                },
+                1: {"engine_args.enforce_eager": "true"},
+            },
+        },
+    )
+    return path
+
+
+def get_stage_config(name: str = "qwen3_tts_no_async_chunk.yaml"):
+    """Get the no_async_chunk stage config path (async_chunk disable, cuda_graph disabled)."""
+    return str(Path(__file__).parent.parent.parent.parent / "vllm_omni" / "model_executor" / "stage_configs" / name)
+
+
+# Same structure as test_qwen3_omni: models, stage_configs, test_params
+tts_server_params = [
+    pytest.param(
+        (MODEL, get_cuda_graph_config()),
+        id="no_cuda_graph",
+    )
+]
+
+
+def get_prompt():
+    """Text prompt for text-to-audio (same role as get_question in test_qwen3_omni)."""
+    return "Hello, this is a test for text to audio."
+
+
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("omni_runner", tts_server_params, indirect=True)
+def test_text_to_audio_001(omni_runner, omni_runner_handler) -> None:
+    """
+    Test text input processing and audio output via offline Omni runner.
+    Deploy Setting: qwen3_tts_no_async_chunk.yaml + enforce_eager=true
+    Input Modal: text
+    Output Modal: audio
+    Input Setting: stream=False
+    Extra Setting: task_type=Base, voice=clone, ref_audio/ref_text provided
+    Datasets: few requests
+    """
+    request_config = {
+        "input": get_prompt(),
+        "task_type": "Base",
+        "voice": "clone",
+        "ref_audio": REF_AUDIO_URL,
+        "ref_text": REF_TEXT,
+    }
+    omni_runner_handler.send_audio_speech_request(request_config)
@@ -0,0 +1,74 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+E2E offline tests for Qwen3-TTS CustomVoice model with text input and audio output.
+
+Async_chunk disable, cuda_graph disabled (no_async_chunk stage config).
+CUDA graph is disabled by setting engine_args.enforce_eager=true via modify_stage_config().
+Same structure as test_qwen3_omni (models, stage_configs, test_params, parametrize omni_runner).
+"""
+
+import os
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
+
+from pathlib import Path
+
+import pytest
+
+from tests.conftest import modify_stage_config
+from tests.utils import hardware_test
+
+MODEL = "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice"
+
+
+def get_cuda_graph_config():
+    path = modify_stage_config(
+        get_stage_config(),
+        updates={
+            "stage_args": {
+                0: {
+                    "engine_args.enforce_eager": "true",
+                },
+                1: {"engine_args.enforce_eager": "true"},
+            },
+        },
+    )
+    return path
+
+
+def get_stage_config(name: str = "qwen3_tts_no_async_chunk.yaml"):
+    """Get the no_async_chunk stage config path (async_chunk disable, cuda_graph disabled)."""
+    return str(Path(__file__).parent.parent.parent.parent / "vllm_omni" / "model_executor" / "stage_configs" / name)
+
+
+# Same structure as test_qwen3_omni: models, stage_configs, test_params
+tts_server_params = [
+    pytest.param(
+        (MODEL, get_cuda_graph_config()),
+        id="no_cuda_graph",
+    )
+]
+
+
+def get_prompt():
+    """Text prompt for text-to-audio (same role as get_question in test_qwen3_omni)."""
+    return "Hello, this is a test for text to audio."
+
+
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("omni_runner", tts_server_params, indirect=True)
+def test_text_to_audio_001(omni_runner, omni_runner_handler) -> None:
+    """
+    Test text input processing and audio output via offline Omni runner.
+    Deploy Setting: qwen3_tts_no_async_chunk.yaml + enforce_eager=true
+    Input Modal: text
+    Output Modal: audio
+    Input Setting: stream=False
+    Datasets: few requests
+    """
+    request_config = {"input": get_prompt(), "voice": "vivian"}
+    omni_runner_handler.send_audio_speech_request(request_config)
@@ -23,7 +23,7 @@ from tests.conftest import (
 )
 from tests.utils import hardware_test

-models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]
+model = "Qwen/Qwen3-Omni-30B-A3B-Instruct"

 AUDIO_KEY = ["test"]
 IMAGE_KEY = ["square", "quadrate"]
@@ -49,16 +49,32 @@ def get_chunk_config(default_path):
    return path


+def get_batch_token_config(default_path):
+    path = modify_stage_config(
+        default_path,
+        updates={
+            "stage_args": {1: {"engine_args.max_num_batched_tokens": 64}},
+        },
+    )
+    return path
+
+
 # CI stage config for 2*H100-80G GPUs
 default_path = str(Path(__file__).parent.parent / "stage_configs" / "qwen3_omni_ci.yaml")
-stage_configs = [default_path, get_chunk_config(default_path)]

 if current_omni_platform.is_xpu():
-    stage_configs = [str(Path(__file__).parent.parent / "stage_configs" / "xpu" / "qwen3_omni_ci.yaml")]
+    default_path = str(Path(__file__).parent.parent / "stage_configs" / "xpu" / "qwen3_omni_ci.yaml")

 # Create parameter combinations for model and stage config
 test_params = [
-    OmniServerParams(model=model, stage_config_path=stage_config) for model in models for stage_config in stage_configs
+    pytest.param(OmniServerParams(model=model, stage_config_path=default_path), id="default"),
+    pytest.param(OmniServerParams(model=model, stage_config_path=get_chunk_config(default_path)), id="async_chunk"),
+]
+
+test_token_params = [
+    pytest.param(
+        OmniServerParams(model=model, stage_config_path=get_batch_token_config(default_path)), id="batch_token_64"
+    )
 ]


@@ -85,6 +101,8 @@ def get_prompt(prompt_type="text_only"):
        "text_video": "What is in this video? ",
        "text_image": "What is in this image? ",
        "text_audio": "What is in this audio? ",
+        "text_audio_video": "First, what is in this audio? Then, what is in this video? ",
+        "one_word": "What is the capital of UK? Answer in one word",
    }
    return prompts.get(prompt_type, prompts["text_only"])

@@ -121,7 +139,7 @@ def test_text_to_audio_001(omni_server, openai_client) -> None:
@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
-@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+@pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True)
 def test_text_to_text_audio_001(omni_server, openai_client) -> None:
    """
    Input Modal: text
@@ -288,7 +306,7 @@ def test_video_to_text_audio_001(omni_server, openai_client) -> None:
@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
-@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+@pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True)
 def test_text_audio_to_text_audio_001(omni_server, openai_client) -> None:
    """
    Input Modal: text, audio
@@ -313,7 +331,7 @@ def test_text_audio_to_text_audio_001(omni_server, openai_client) -> None:
@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
-@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+@pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True)
 def test_text_image_to_text_audio_001(omni_server, openai_client) -> None:
    """
    Input Modal: text, image
@@ -339,7 +357,7 @@ def test_text_image_to_text_audio_001(omni_server, openai_client) -> None:
@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
-@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+@pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True)
 def test_text_video_to_text_audio_001(omni_server, openai_client) -> None:
    """
    Input Modal: text, video
@@ -367,7 +385,7 @@ def test_text_video_to_text_audio_001(omni_server, openai_client) -> None:
@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
-@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+@pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True)
 def test_mix_to_text_audio_001(omni_server, openai_client) -> None:
    """
    Input Modal: text, audio, image, video
@@ -393,3 +411,106 @@ def test_mix_to_text_audio_001(omni_server, openai_client) -> None:
        "key_words": {"audio": AUDIO_KEY, "image": IMAGE_KEY, "video": VIDEO_KEY},
    }
    openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
+
+
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_audio_in_video_001(omni_server, openai_client) -> None:
+    """
+    Input Modal: text + video (synthetic MP4 with embedded audio; ``use_audio_in_video`` uses audio from the video).
+    Output Modal: text, audio
+    Input Setting: stream=False
+    Datasets: single request
+    """
+    video_data_url = f"data:video/mp4;base64,{generate_synthetic_video(224, 224, 300, embed_audio=True)['base64']}"
+    messages = dummy_messages_from_mix_data(
+        system_prompt=get_system_prompt(),
+        video_data_url=video_data_url,
+        content_text=get_prompt("text_audio_video"),
+    )
+
+    request_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "stream": False,
+        "use_audio_in_video": True,
+        "key_words": {"video": VIDEO_KEY, "audio": AUDIO_KEY + ["beep", "electronic"]},
+    }
+    openai_client.send_omni_request(request_config)
+
+
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_audio_in_video_002(omni_server, openai_client) -> None:
+    """
+    Input Modal: text + video (synthetic MP4 with embedded audio; ``use_audio_in_video`` uses audio from the video).
+    Output Modal: text, audio
+    Input Setting: stream=True
+    Datasets: few requests
+    """
+    video_data_url = f"data:video/mp4;base64,{generate_synthetic_video(224, 224, 300, embed_audio=True)['base64']}"
+    messages = dummy_messages_from_mix_data(
+        system_prompt=get_system_prompt(),
+        video_data_url=video_data_url,
+        content_text=get_prompt("text_audio_video"),
+    )
+
+    request_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "stream": True,
+        "use_audio_in_video": True,
+        "key_words": {"video": VIDEO_KEY, "audio": AUDIO_KEY + ["beep", "electronic"]},
+    }
+
+    # Retry when assert_omni_response fails on key_words (see tests/conftest.py).
+    _keyword_assert_msg = "The output does not contain any of the keywords."
+    _max_retries = 3
+    for attempt in range(_max_retries):
+        try:
+            openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
+            break
+        except AssertionError as e:
+            if _keyword_assert_msg not in str(e) or attempt == _max_retries - 1:
+                raise
+            print(f"Keyword assertion failed, retrying {attempt + 2}/{_max_retries}: {e!r}")
+
+
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_one_word_prompt_001(omni_server, openai_client) -> None:
+    """
+    Input Modal: text only (one-word answer constraint).
+    Output Modal: text, audio (default ``modalities``); ``key_words`` only assert on text.
+    Input Setting: stream=True
+    Datasets: single request
+    """
+    messages = dummy_messages_from_mix_data(
+        system_prompt=get_system_prompt(),
+        content_text=get_prompt("one_word"),
+    )
+
+    request_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "stream": True,
+        "key_words": {"text": ["london"]},
+    }
+
+    # Retry only when assert_omni_response fails on text/audio cosine similarity (see tests/conftest.py).
+    _similarity_assert_msg = "The audio content is not same as the text"
+    _max_retries = 3
+    for attempt in range(_max_retries):
+        try:
+            openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
+            break
+        except AssertionError as e:
+            if _similarity_assert_msg not in str(e) or attempt == _max_retries - 1:
+                raise
+            print(f"Similarity assertion failed, retrying {attempt + 2}/{_max_retries}: {e!r}")
@@ -1,285 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-E2E Online tests for Qwen3-TTS model with text input and audio output.
-
-These tests verify the /v1/audio/speech endpoint works correctly with
-actual model inference, not mocks.
-"""
-
-import os
-
-os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
-os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
-
-from pathlib import Path
-
-import httpx
-import pytest
-
-from tests.conftest import OmniServer
-from tests.utils import hardware_test
-
-MODEL = "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice"
-
-
-def get_stage_config(name: str = "qwen3_tts.yaml"):
-    """Get the stage config path for Qwen3-TTS."""
-    return str(Path(__file__).parent.parent.parent.parent / "vllm_omni" / "model_executor" / "stage_configs" / name)
-
-
-@pytest.fixture(scope="class")
-def omni_server():
-    """Start vLLM-Omni server with CustomVoice model."""
-    stage_config_path = get_stage_config()
-
-    print(f"Starting OmniServer with model: {MODEL}")
-
-    with OmniServer(
-        MODEL,
-        [
-            "--stage-configs-path",
-            stage_config_path,
-            "--stage-init-timeout",
-            "120",
-            "--trust-remote-code",
-            "--enforce-eager",
-            "--disable-log-stats",
-        ],
-    ) as server:
-        print("OmniServer started successfully")
-        yield server
-        print("OmniServer stopping...")
-
-    print("OmniServer stopped")
-
-
-def make_speech_request(
-    host: str,
-    port: int,
-    text: str,
-    voice: str = "vivian",
-    language: str = "English",
-    task_type: str | None = None,
-    instructions: str | None = None,
-    timeout: float = 120.0,
-) -> httpx.Response:
-    """Make a request to the /v1/audio/speech endpoint."""
-    url = f"http://{host}:{port}/v1/audio/speech"
-    payload = {
-        "input": text,
-        "voice": voice,
-        "language": language,
-    }
-    if task_type:
-        payload["task_type"] = task_type
-    if instructions:
-        payload["instructions"] = instructions
-
-    with httpx.Client(timeout=timeout) as client:
-        return client.post(url, json=payload)
-
-
-def verify_wav_audio(content: bytes) -> bool:
-    """Verify that content is valid WAV audio data."""
-    # WAV files start with "RIFF" header
-    if len(content) < 44:  # Minimum WAV header size
-        return False
-    return content[:4] == b"RIFF" and content[8:12] == b"WAVE"
-
-
-# Minimum expected audio size for a short sentence (~1 second of 24kHz 16-bit mono WAV)
-MIN_AUDIO_BYTES = 10000
-
-
-class TestQwen3TTSCustomVoice:
-    """E2E tests for Qwen3-TTS CustomVoice model."""
-
-    @pytest.mark.core_model
-    @pytest.mark.omni
-    @hardware_test(res={"cuda": "L4"}, num_cards=4)
-    def test_speech_english_basic(self, omni_server) -> None:
-        """Test basic English TTS generation."""
-        response = make_speech_request(
-            host=omni_server.host,
-            port=omni_server.port,
-            text="Hello, how are you?",
-            voice="vivian",
-            language="English",
-        )
-
-        assert response.status_code == 200, f"Request failed: {response.text}"
-        assert response.headers.get("content-type") == "audio/wav"
-        assert verify_wav_audio(response.content), "Response is not valid WAV audio"
-        assert len(response.content) > MIN_AUDIO_BYTES, (
-            f"Audio content too small ({len(response.content)} bytes), expected at least {MIN_AUDIO_BYTES} bytes"
-        )
-
-    @pytest.mark.core_model
-    @pytest.mark.omni
-    @hardware_test(res={"cuda": "L4"}, num_cards=4)
-    def test_speech_chinese_basic(self, omni_server) -> None:
-        """Test basic Chinese TTS generation."""
-        response = make_speech_request(
-            host=omni_server.host,
-            port=omni_server.port,
-            text="你好，我是通义千问",
-            voice="vivian",
-            language="Chinese",
-        )
-
-        assert response.status_code == 200, f"Request failed: {response.text}"
-        assert response.headers.get("content-type") == "audio/wav"
-        assert verify_wav_audio(response.content), "Response is not valid WAV audio"
-        assert len(response.content) > MIN_AUDIO_BYTES, (
-            f"Audio content too small ({len(response.content)} bytes), expected at least {MIN_AUDIO_BYTES} bytes"
-        )
-
-    @pytest.mark.core_model
-    @pytest.mark.omni
-    @hardware_test(res={"cuda": "L4"}, num_cards=4)
-    def test_speech_different_voices(self, omni_server) -> None:
-        """Test TTS with different voice options."""
-        voices = ["vivian", "ryan"]
-        for voice in voices:
-            response = make_speech_request(
-                host=omni_server.host,
-                port=omni_server.port,
-                text="Testing voice selection.",
-                voice=voice,
-                language="English",
-            )
-
-            assert response.status_code == 200, f"Request failed for voice {voice}: {response.text}"
-            assert verify_wav_audio(response.content), f"Invalid WAV for voice {voice}"
-
-    @pytest.mark.core_model
-    @pytest.mark.omni
-    @hardware_test(res={"cuda": "L4"}, num_cards=4)
-    def test_speech_binary_response_not_utf8_error(self, omni_server) -> None:
-        """
-        Regression test: Verify binary audio is returned, not UTF-8 error.
-
-        This test ensures the multimodal_output property correctly retrieves
-        audio from completion outputs, preventing the "TTS model did not
-        produce audio output" error.
-        """
-        response = make_speech_request(
-            host=omni_server.host,
-            port=omni_server.port,
-            text="This should return binary audio, not a JSON error.",
-            voice="vivian",
-            language="English",
-        )
-
-        # Should NOT be a JSON error response
-        assert response.status_code == 200, f"Request failed: {response.text}"
-
-        # Verify it's binary audio, not JSON
-        try:
-            # If this succeeds and starts with {"error", it's a bug
-            text = response.content.decode("utf-8")
-            assert not text.startswith('{"error"'), f"Got error response instead of audio: {text}"
-        except UnicodeDecodeError:
-            # This is expected - binary audio can't be decoded as UTF-8
-            pass
-
-        assert verify_wav_audio(response.content), "Response is not valid WAV audio"
-
-
-class TestQwen3TTSAPIEndpoints:
-    """Test API endpoint functionality."""
-
-    @pytest.mark.core_model
-    @pytest.mark.omni
-    @hardware_test(res={"cuda": "L4"}, num_cards=4)
-    def test_list_voices_endpoint(self, omni_server) -> None:
-        """Test the /v1/audio/voices endpoint returns available voices."""
-        url = f"http://{omni_server.host}:{omni_server.port}/v1/audio/voices"
-
-        with httpx.Client(timeout=30.0) as client:
-            response = client.get(url)
-
-        assert response.status_code == 200
-        data = response.json()
-        assert "voices" in data
-        assert isinstance(data["voices"], list)
-        assert len(data["voices"]) > 0
-        # Check some expected voices are present
-        voices_lower = [v.lower() for v in data["voices"]]
-        assert "vivian" in voices_lower or "ryan" in voices_lower
-
-    @pytest.mark.core_model
-    @pytest.mark.omni
-    @hardware_test(res={"cuda": "L4"}, num_cards=4)
-    def test_models_endpoint(self, omni_server) -> None:
-        """Test the /v1/models endpoint returns loaded model."""
-        url = f"http://{omni_server.host}:{omni_server.port}/v1/models"
-
-        with httpx.Client(timeout=30.0) as client:
-            response = client.get(url)
-
-        assert response.status_code == 200
-        data = response.json()
-        assert "data" in data
-        assert len(data["data"]) > 0
-
-
-@pytest.fixture(scope="class")
-def omni_server_no_async_chunk():
-    """Start vLLM-Omni server with non-async-chunk config."""
-    stage_config_path = get_stage_config("qwen3_tts_no_async_chunk.yaml")
-
-    with OmniServer(
-        MODEL,
-        [
-            "--stage-configs-path",
-            stage_config_path,
-            "--stage-init-timeout",
-            "120",
-            "--trust-remote-code",
-            "--enforce-eager",
-            "--disable-log-stats",
-        ],
-    ) as server:
-        yield server
-
-
-class TestQwen3TTSNoAsyncChunk:
-    """E2E tests for Qwen3-TTS in non-async-chunk (full decode) mode."""
-
-    @pytest.mark.core_model
-    @pytest.mark.omni
-    @hardware_test(res={"cuda": "L4"}, num_cards=4)
-    def test_speech_english(self, omni_server_no_async_chunk) -> None:
-        """Test English TTS with non-async-chunk pipeline."""
-        response = make_speech_request(
-            host=omni_server_no_async_chunk.host,
-            port=omni_server_no_async_chunk.port,
-            text="Hello, how are you?",
-            voice="vivian",
-            language="English",
-        )
-
-        assert response.status_code == 200, f"Request failed: {response.text}"
-        assert response.headers.get("content-type") == "audio/wav"
-        assert verify_wav_audio(response.content), "Response is not valid WAV audio"
-        assert len(response.content) > MIN_AUDIO_BYTES
-
-    @pytest.mark.core_model
-    @pytest.mark.omni
-    @hardware_test(res={"cuda": "L4"}, num_cards=4)
-    def test_speech_chinese(self, omni_server_no_async_chunk) -> None:
-        """Test Chinese TTS with non-async-chunk pipeline."""
-        response = make_speech_request(
-            host=omni_server_no_async_chunk.host,
-            port=omni_server_no_async_chunk.port,
-            text="你好，我是通义千问",
-            voice="vivian",
-            language="Chinese",
-        )
-
-        assert response.status_code == 200, f"Request failed: {response.text}"
-        assert response.headers.get("content-type") == "audio/wav"
-        assert verify_wav_audio(response.content), "Response is not valid WAV audio"
-        assert len(response.content) > MIN_AUDIO_BYTES
@@ -1,10 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """
-E2E tests for Qwen3-TTS Base voice-clone model.
+E2E Online tests for Qwen3-TTS model with text input and audio output.

-Regression test for #1663: speech tokenizer loaded in bfloat16 produced
-all-silence PCM due to NaN overflow in SnakeBeta activation.
+These tests verify the /v1/audio/speech endpoint works correctly with
+actual model inference, not mocks.
 """

 import os
@@ -12,212 +12,99 @@ import os
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"

-import struct
-import tempfile
 from pathlib import Path

-import httpx
-import numpy as np
 import pytest

-from tests.conftest import (
-    OmniServer,
-    convert_audio_file_to_text,
-    cosine_similarity_text,
-)
+from tests.conftest import OmniServerParams
 from tests.utils import hardware_test

-MODEL = "Qwen/Qwen3-TTS-12Hz-1.7B-Base"
+MODEL = "Qwen/Qwen3-TTS-12Hz-0.6B-Base"

-# Official Qwen3-TTS reference audio/text from examples/offline_inference/qwen3_tts/end2end.py
 REF_AUDIO_URL = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-TTS-Repo/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."
-SYN_TEXT = "Good one. Okay, fine, I'm just gonna leave this sock monkey here. Goodbye."


-def get_stage_config():
-    return str(
-        Path(__file__).parent.parent.parent.parent / "vllm_omni" / "model_executor" / "stage_configs" / "qwen3_tts.yaml"
-    )
+def get_stage_config(name: str = "qwen3_tts.yaml"):
+    """Get the stage config path from vllm_omni model_executor stage_configs."""
+    return str(Path(__file__).parent.parent.parent.parent / "vllm_omni" / "model_executor" / "stage_configs" / name)


-@pytest.fixture(scope="module")
-def omni_server():
-    """Start vLLM-Omni server with Base voice-clone model."""
-    stage_config_path = get_stage_config()
-
-    with OmniServer(
-        MODEL,
-        [
-            "--stage-configs-path",
-            stage_config_path,
-            "--stage-init-timeout",
-            "120",
-            "--trust-remote-code",
-            "--enforce-eager",
-            "--disable-log-stats",
-        ],
-    ) as server:
-        yield server
-
-
-def make_base_speech_request(
-    host: str,
-    port: int,
-    text: str = SYN_TEXT,
-    ref_text: str = REF_TEXT,
-    ref_audio: str = REF_AUDIO_URL,
-    response_format: str = "wav",
-    timeout: float = 120.0,
-) -> httpx.Response:
-    url = f"http://{host}:{port}/v1/audio/speech"
-    payload = {
-        "model": MODEL,
-        "input": text,
-        "task_type": "Base",
-        "ref_text": ref_text,
-        "ref_audio": ref_audio,
-        "response_format": response_format,
+def get_prompt(prompt_type="text"):
+    """Text prompt for text-to-audio tests (same as test_qwen3_omni - beijing test case)."""
+    prompts = {
+        "text": "The weather is nice today, perfect for a walk in the park.",
    }
-    with httpx.Client(timeout=timeout) as client:
-        return client.post(url, json=payload)
+    return prompts.get(prompt_type, prompts["text"])


-def verify_wav_audio(content: bytes) -> bool:
-    if len(content) < 44:
-        return False
-    return content[:4] == b"RIFF" and content[8:12] == b"WAVE"
+def get_max_batch_size(size_type="few"):
+    """Batch size for concurrent requests (same as test_qwen3_omni)."""
+    batch_sizes = {"few": 5, "medium": 100, "large": 256}
+    return batch_sizes.get(size_type, 5)


-def assert_not_silence(pcm_bytes: bytes):
-    """Assert PCM16 samples are not all identical (e.g. all-silence)."""
-    samples = struct.unpack(f"<{len(pcm_bytes) // 2}h", pcm_bytes)
-    unique = set(samples)
-    assert len(unique) > 1, (
-        f"All-silence detected: {len(samples)} samples, unique values: {unique}. "
-        "See https://github.com/vllm-project/vllm-omni/issues/1663"
+tts_server_params = [
+    pytest.param(
+        OmniServerParams(
+            model=MODEL,
+            stage_config_path=get_stage_config("qwen3_tts.yaml"),
+            server_args=["--trust-remote-code", "--disable-log-stats"],
+        ),
+        id="async_chunk",
    )
+]


-MIN_AUDIO_BYTES = 10000
-MIN_HNR_DB = 1.2  # Clean voice clone > 1.2 dB; distorted < 1.0 dB
-
-
-def compute_hnr_db(pcm_samples: np.ndarray, sr: int = 24000) -> float:
-    """Compute mean Harmonic-to-Noise Ratio (dB) for speech quality.
-
-    Clean cloned speech has HNR > 1.2 dB; distorted speech (e.g. lost
-    ref_code decoder context) drops below 1.0 dB.
+@pytest.mark.advanced_model
+@pytest.mark.core_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
+def test_text_to_audio_001(omni_server, openai_client) -> None:
+    """
+    Test text input processing and audio output via OpenAI API.
+    Deploy Setting: default yaml
+    Input Modal: text
+    Output Modal: audio
+    Input Setting: stream=False
+    Datasets: few requests
    """
-    frame_len = int(0.03 * sr)  # 30ms frames
-    hop = frame_len // 2
-    hnr_values: list[float] = []

-    for start in range(0, len(pcm_samples) - frame_len, hop):
-        frame = pcm_samples[start : start + frame_len]
-        if np.max(np.abs(frame)) < 0.01:
-            continue
-        ac = np.correlate(frame, frame, mode="full")[len(frame) - 1 :]
-        ac = ac / (ac[0] + 1e-10)
-        min_lag = int(sr / 400)
-        max_lag = min(int(sr / 80), len(ac))
-        if min_lag >= max_lag:
-            continue
-        peak = float(np.max(ac[min_lag:max_lag]))
-        if 0 < peak < 1:
-            hnr_values.append(10 * np.log10(peak / (1 - peak + 1e-10)))
-
-    return float(np.mean(hnr_values)) if hnr_values else 0.0
+    request_config = {
+        "model": omni_server.model,
+        "input": get_prompt(),
+        "stream": False,
+        "response_format": "wav",
+        "task_type": "Base",
+        "voice": "clone",
+        "ref_audio": REF_AUDIO_URL,
+        "ref_text": REF_TEXT,
+    }
+    openai_client.send_audio_speech_request(request_config, request_num=get_max_batch_size("few"))


-class TestQwen3TTSBaseVoiceClone:
-    """Regression tests for Base voice-clone (fix #1663)."""
-
-    @pytest.mark.core_model
-    @pytest.mark.omni
-    @hardware_test(res={"cuda": "L4"}, num_cards=4)
-    def test_base_voice_clone_not_silence(self, omni_server) -> None:
-        """PCM output must contain real audio, not all-silence."""
-        response = make_base_speech_request(
-            host=omni_server.host,
-            port=omni_server.port,
-            response_format="pcm",
-        )
-
-        assert response.status_code == 200, f"Request failed: {response.text}"
-        assert len(response.content) > MIN_AUDIO_BYTES, f"Audio too small: {len(response.content)} bytes"
-        assert_not_silence(response.content)
-
-    @pytest.mark.core_model
-    @pytest.mark.omni
-    @hardware_test(res={"cuda": "L4"}, num_cards=4)
-    def test_base_voice_clone_whisper_transcription(self, omni_server) -> None:
-        """Whisper must transcribe the output as intelligible speech."""
-        response = make_base_speech_request(
-            host=omni_server.host,
-            port=omni_server.port,
-            response_format="wav",
-        )
-
-        assert response.status_code == 200, f"Request failed: {response.text}"
-        assert verify_wav_audio(response.content), "Response is not valid WAV"
-        assert len(response.content) > MIN_AUDIO_BYTES
-
-        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
-            f.write(response.content)
-            wav_path = f.name
-
-        try:
-            transcript = convert_audio_file_to_text(wav_path)
-            print(f"Whisper transcript: {transcript}")
-            assert len(transcript.strip()) > 0, "Whisper returned empty transcript — audio is likely silence"
-            similarity = cosine_similarity_text(transcript.lower(), SYN_TEXT.lower())
-            print(f"Cosine similarity: {similarity:.3f}")
-            assert similarity > 0.9, (
-                f"Transcript doesn't match input: similarity={similarity:.2f}, transcript='{transcript}'"
-            )
-        finally:
-            os.unlink(wav_path)
-
-    @pytest.mark.core_model
-    @pytest.mark.omni
-    @hardware_test(res={"cuda": "L4"}, num_cards=4)
-    def test_base_voice_clone_wav_format(self, omni_server) -> None:
-        """WAV response must have valid headers and sufficient size."""
-        response = make_base_speech_request(
-            host=omni_server.host,
-            port=omni_server.port,
-            response_format="wav",
-        )
-
-        assert response.status_code == 200, f"Request failed: {response.text}"
-        assert response.headers.get("content-type") == "audio/wav"
-        assert verify_wav_audio(response.content), "Response is not valid WAV"
-        assert len(response.content) > MIN_AUDIO_BYTES
-
-    @pytest.mark.core_model
-    @pytest.mark.omni
-    @hardware_test(res={"cuda": "L4"}, num_cards=4)
-    def test_base_voice_clone_no_distortion(self, omni_server) -> None:
-        """Audio must not be distorted (regression for ref_code context loss).
-
-        When the decoder loses ref_code speaker context on later streaming
-        chunks, HNR drops below 1.0 dB.  Clean voice clone should be > 1.2 dB.
-        See https://github.com/vllm-project/vllm-omni/issues/1944
-        """
-        response = make_base_speech_request(
-            host=omni_server.host,
-            port=omni_server.port,
-            response_format="pcm",
-        )
-
-        assert response.status_code == 200, f"Request failed: {response.text}"
-        assert len(response.content) > MIN_AUDIO_BYTES
-
-        pcm_samples = np.frombuffer(response.content, dtype=np.int16).astype(np.float32) / 32768.0
-        hnr = compute_hnr_db(pcm_samples)
-        print(f"Voice clone HNR: {hnr:.2f} dB (threshold: {MIN_HNR_DB} dB)")
-        assert hnr >= MIN_HNR_DB, (
-            f"Audio distortion detected: HNR={hnr:.2f} dB < {MIN_HNR_DB} dB. "
-            "Voice clone decoder may be losing ref_code speaker context on later chunks."
-        )
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
+def test_text_to_audio_002(omni_server, openai_client) -> None:
+    """
+    Test text input processing and audio output via OpenAI API.
+    Deploy Setting: default yaml
+    Input Modal: text
+    Output Modal: audio
+    Input Setting: stream=True
+    Datasets: single request
+    """
+    request_config = {
+        "model": omni_server.model,
+        "input": get_prompt(),
+        "stream": True,
+        "response_format": "wav",
+        "task_type": "Base",
+        "voice": "clone",
+        "ref_audio": REF_AUDIO_URL,
+        "ref_text": REF_TEXT,
+    }
+    openai_client.send_audio_speech_request(request_config)
@@ -0,0 +1,119 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+E2E Online tests for Qwen3-TTS model with text input and audio output.
+
+These tests verify the /v1/audio/speech endpoint works correctly with
+actual model inference, not mocks.
+"""
+
+import os
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
+
+from pathlib import Path
+
+import pytest
+
+from tests.conftest import OmniServerParams
+from tests.utils import hardware_test
+
+MODEL = "Qwen/Qwen3-TTS-12Hz-0.6B-Base"
+
+REF_AUDIO_URL = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-TTS-Repo/clone_2.wav"
+REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."
+
+
+def get_stage_config(name: str = "qwen3_tts.yaml"):
+    """Get the stage config path from vllm_omni model_executor stage_configs."""
+    return str(Path(__file__).parent.parent.parent.parent / "vllm_omni" / "model_executor" / "stage_configs" / name)
+
+
+def get_prompt(prompt_type="text"):
+    """Text prompt for text-to-audio tests (same as test_qwen3_omni - beijing test case)."""
+    prompts = {
+        "text": "The weather is nice today, perfect for a walk in the park.",
+    }
+    return prompts.get(prompt_type, prompts["text"])
+
+
+def get_max_batch_size(size_type="few"):
+    """Batch size for concurrent requests (same as test_qwen3_omni)."""
+    batch_sizes = {"few": 5, "medium": 100, "large": 256}
+    return batch_sizes.get(size_type, 5)
+
+
+tts_server_params = [
+    pytest.param(
+        OmniServerParams(
+            model=MODEL,
+            stage_config_path=get_stage_config("qwen3_tts.yaml"),
+            server_args=["--trust-remote-code", "--disable-log-stats"],
+        ),
+        id="async_chunk",
+    ),
+    pytest.param(
+        OmniServerParams(
+            model=MODEL,
+            stage_config_path=get_stage_config("qwen3_tts_no_async_chunk.yaml"),
+            server_args=["--trust-remote-code", "--disable-log-stats"],
+        ),
+        id="no_async_chunk",
+    ),
+]
+
+
+@pytest.mark.advanced_model
+@pytest.mark.core_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
+def test_voice_clone_streaming_001(omni_server, openai_client) -> None:
+    """
+    Test text input processing and audio output via OpenAI API.
+    Deploy Setting: default yaml
+    Input Modal: text
+    Output Modal: audio
+    Input Setting: stream=True
+    Datasets: few requests
+    """
+
+    request_config = {
+        "model": omni_server.model,
+        "input": get_prompt(),
+        "stream": True,
+        "response_format": "wav",
+        "task_type": "Base",
+        "voice": "clone",
+        "ref_audio": REF_AUDIO_URL,
+        "ref_text": REF_TEXT,
+    }
+    openai_client.send_audio_speech_request(request_config, request_num=get_max_batch_size("few"))
+
+
+@pytest.mark.advanced_model
+@pytest.mark.core_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
+def test_response_format_001(omni_server, openai_client) -> None:
+    """
+    Test text input processing and audio output via OpenAI API.
+    Deploy Setting: default yaml
+    Input Modal: text
+    Output Modal: audio
+    Input Setting: non-stream PCM
+    Datasets: few requests
+    """
+
+    request_config = {
+        "model": omni_server.model,
+        "input": get_prompt(),
+        "response_format": "pcm",
+        "task_type": "Base",
+        "voice": "clone",
+        "ref_audio": REF_AUDIO_URL,
+        "ref_text": REF_TEXT,
+    }
+    openai_client.send_audio_speech_request(request_config)
@@ -0,0 +1,105 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+E2E Online tests for Qwen3-TTS model with text input and audio output.
+
+These tests verify the /v1/audio/speech endpoint works correctly with
+actual model inference, not mocks.
+"""
+
+import os
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
+
+from pathlib import Path
+
+import pytest
+
+from tests.conftest import OmniServerParams
+from tests.utils import hardware_test
+
+MODEL = "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice"
+
+
+def get_stage_config(name: str = "qwen3_tts.yaml"):
+    """Get the stage config path from vllm_omni model_executor stage_configs."""
+    return str(Path(__file__).parent.parent.parent.parent / "vllm_omni" / "model_executor" / "stage_configs" / name)
+
+
+def get_prompt(prompt_type="text"):
+    """Text prompt for text-to-audio tests (same as test_qwen3_omni - beijing test case)."""
+    prompts = {
+        "text": "Beijing, China's capital, blends ancient wonders like the Great Wall with modern marvels. This vibrant metropolis offers rich culture, delicious Peking duck, and endless exploration opportunities.",
+    }
+    return prompts.get(prompt_type, prompts["text"])
+
+
+def get_max_batch_size(size_type="few"):
+    """Batch size for concurrent requests (same as test_qwen3_omni)."""
+    batch_sizes = {"few": 5, "medium": 100, "large": 256}
+    return batch_sizes.get(size_type, 5)
+
+
+tts_server_params = [
+    pytest.param(
+        OmniServerParams(
+            model=MODEL,
+            stage_config_path=get_stage_config("qwen3_tts.yaml"),
+            server_args=["--trust-remote-code", "--disable-log-stats"],
+        ),
+        id="async_chunk",
+    )
+]
+
+
+@pytest.mark.core_model
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
+def test_text_to_audio_001(omni_server, openai_client) -> None:
+    """
+    Test text input processing and audio output via OpenAI API.
+    Deploy Setting: default yaml
+    Input Modal: text
+    Output Modal: audio
+    Input Setting: stream=False
+    Datasets: few requests
+    """
+    request_config = {
+        "model": omni_server.model,
+        "input": get_prompt(),
+        "stream": False,
+        "response_format": "wav",
+        "task_type": "CustomVoice",
+        "voice": "vivian",
+    }
+
+    openai_client.send_audio_speech_request(request_config, request_num=get_max_batch_size())
+
+
+@pytest.mark.core_model
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
+def test_text_to_audio_002(omni_server, openai_client) -> None:
+    """
+    Test text input processing and audio output via OpenAI API.
+    Deploy Setting: default yaml
+    Input Modal: text
+    Output Modal: audio
+    Input Setting: stream=True
+    Datasets: single request
+    """
+    request_config = {
+        "model": omni_server.model,
+        "input": get_prompt(),
+        "stream": True,
+        "response_format": "wav",
+        "task_type": "CustomVoice",
+        "voice": "vivian",
+    }
+
+    openai_client.send_audio_speech_request(request_config)
@@ -0,0 +1,145 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+E2E Online tests for Qwen3-TTS model with text input and audio output.
+
+These tests verify the /v1/audio/speech endpoint works correctly with
+actual model inference, not mocks.
+"""
+
+import os
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
+
+from pathlib import Path
+
+import pytest
+
+from tests.conftest import OmniServerParams
+from tests.utils import hardware_test
+
+MODEL = "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice"
+
+
+def get_stage_config(name: str = "qwen3_tts.yaml"):
+    """Get the stage config path from vllm_omni model_executor stage_configs."""
+    return str(Path(__file__).parent.parent.parent.parent / "vllm_omni" / "model_executor" / "stage_configs" / name)
+
+
+def get_prompt(prompt_type="english"):
+    """Text prompt for text-to-audio tests (same as test_qwen3_omni - beijing test case)."""
+    prompts = {
+        "english": "Paris is a city in France. It has a tall tower by the river. The streets are old and small. You can walk and see art and food. The sky is light and soft. It is a calm place to rest.",
+        "chinese": "北京，中国的首都，是一座融合了长城等历史地点与现代建筑的国际化大都市，充满了独特的文化与活力",
+    }
+    return prompts.get(prompt_type, prompts["english"])
+
+
+def get_max_batch_size(size_type="few"):
+    """Batch size for concurrent requests (same as test_qwen3_omni)."""
+    batch_sizes = {"few": 3, "medium": 100, "large": 256}
+    return batch_sizes.get(size_type)
+
+
+tts_server_params = [
+    pytest.param(
+        OmniServerParams(
+            model=MODEL,
+            stage_config_path=get_stage_config("qwen3_tts.yaml"),
+            server_args=["--trust-remote-code", "--disable-log-stats"],
+        ),
+        id="async_chunk",
+    ),
+    pytest.param(
+        OmniServerParams(
+            model=MODEL,
+            stage_config_path=get_stage_config("qwen3_tts_no_async_chunk.yaml"),
+            server_args=["--trust-remote-code", "--disable-log-stats"],
+        ),
+        id="no_async_chunk",
+    ),
+]
+
+
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
+def test_voice_001(omni_server, openai_client) -> None:
+    """
+    Test text input processing and audio output via OpenAI API.
+    Deploy Setting: default yaml
+    Input Modal: text
+    Output Modal: audio
+    Input Setting: stream=False, voice=eric
+    Datasets: few requests
+    """
+    request_config = {
+        "model": omni_server.model,
+        "input": get_prompt(),
+        "stream": True,
+        "response_format": "wav",
+        "task_type": "CustomVoice",
+        "voice": "uncle_fu",
+    }
+
+    # Retry once on assertion failures from transcript similarity / gender checks (flaky ASR or estimators).
+    for attempt in range(2):
+        try:
+            openai_client.send_audio_speech_request(request_config, request_num=get_max_batch_size())
+            break
+        except AssertionError:
+            if attempt == 0:
+                continue
+            raise
+
+
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
+def test_voice_002(omni_server, openai_client) -> None:
+    """
+    Test text input processing and audio output via OpenAI API.
+    Deploy Setting: default yaml
+    Input Modal: text
+    Output Modal: audio
+    Input Setting: stream=False, language=chinese
+    Datasets: few requests
+    """
+    request_config = {
+        "model": omni_server.model,
+        "input": get_prompt(),
+        "stream": False,
+        "response_format": "wav",
+        "task_type": "CustomVoice",
+        "voice": "Serena",
+    }
+
+    openai_client.send_audio_speech_request(request_config)
+
+
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
+def test_language_001(omni_server, openai_client) -> None:
+    """
+    Test text input processing and audio output via OpenAI API.
+    Deploy Setting: default yaml
+    Input Modal: text
+    Output Modal: audio
+    Input Setting: stream=False, language=chinese
+    Datasets: few requests
+    """
+    request_config = {
+        "model": omni_server.model,
+        "input": get_prompt("chinese"),
+        "stream": False,
+        "response_format": "wav",
+        "task_type": "CustomVoice",
+        "voice": "vivian",
+    }
+
+    openai_client.send_audio_speech_request(request_config)
@@ -7,7 +7,7 @@ and are supported by Qwen-Image-Layered model.
 Kept cases (maximum feature coverage, no cpu-offload):
  sp_001              : cache_dit + Ulysses-SP 2          (2×H100)
  cfg_parallel_002    : cache_dit + CFG-Parallel 2        (2×H100)
-  layers_guard_001    : layerwise offload + layers=2      (1×H100, issue #1969 guard)
+  layers_guard_001    : layerwise offload + layers=3      (1×H100, issue #1969 guard)

 Total distinct features covered: cache_dit, Ulysses-SP, CFG-Parallel, layerwise-offload.
 """
@@ -102,8 +102,8 @@ LAYERS_GUARD_CASES = [
            model=MODEL,
            server_args=["--enable-layerwise-offload"],
        ),
-        2,
-        id="layers_guard_001_layers2",
+        3,
+        id="layers_guard_001_layers3",
        marks=SINGLE_CARD_FEATURE_MARKS,
    ),
 ]
@@ -33,7 +33,7 @@ stage_args:
    temperature: 0.4
    top_p: 0.9
    top_k: 1
-    max_tokens: 100
+    max_tokens: 150
    seed: 42
    ignore_eos: False
    detokenize: True
@@ -74,7 +74,7 @@ stage_args:
    devices: "1"
  engine_args:
    model_stage: code2wav
-    max_num_seqs: 1
+    max_num_seqs: 5
    model_arch: Qwen3OmniMoeForConditionalGeneration
    worker_type: generation
    scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
@@ -115,7 +115,7 @@ class MockGenerationResult:
 class FakeAsyncOmni:
    """Fake AsyncOmni that yields a single diffusion output."""

-    def __init__(self):
+    def __init__(self, images=None):
        self.stage_configs = [
            SimpleNamespace(stage_type="llm"),
            SimpleNamespace(stage_type="diffusion"),
@@ -123,11 +123,12 @@ class FakeAsyncOmni:
        self.default_sampling_params_list = [SamplingParams(temperature=0.1), OmniDiffusionSamplingParams()]
        self.captured_sampling_params_list = None
        self.captured_prompt = None
+        self._images = images or [Image.new("RGB", (64, 64), color="green")]

    async def generate(self, prompt, request_id, sampling_params_list):
        self.captured_sampling_params_list = sampling_params_list
        self.captured_prompt = prompt
-        images = [Image.new("RGB", (64, 64), color="green")]
+        images = [img.copy() for img in self._images]
        yield MockGenerationResult(images)


@@ -168,7 +169,13 @@ def test_client(mock_async_diffusion):
    app.state.engine_client = mock_async_diffusion
    app.state.diffusion_engine = mock_async_diffusion  # Also set for health endpoint
    app.state.stage_configs = [SimpleNamespace(stage_type="diffusion")]
-    app.state.diffusion_model_name = "Qwen/Qwen-Image"  # For models endpoint
+    from vllm.entrypoints.openai.models.protocol import BaseModelPath
+
+    from vllm_omni.entrypoints.openai.api_server import _DiffusionServingModels
+
+    app.state.openai_serving_models = _DiffusionServingModels(
+        [BaseModelPath(name="Qwen/Qwen-Image", model_path="Qwen/Qwen-Image")]
+    )
    app.state.args = Namespace(
        default_sampling_params='{"0": {"num_inference_steps":4, "guidance_scale":7.5}}',
        max_generated_image_size=4096,  # 64*64
@@ -199,6 +206,28 @@ def async_omni_test_client():
    return TestClient(app)


+@pytest.fixture
+def async_omni_rgba_test_client():
+    """Create test client with mocked AsyncOmni engine returning RGBA output."""
+    from fastapi import FastAPI
+
+    from vllm_omni.entrypoints.openai.api_server import router
+
+    app = FastAPI()
+    app.include_router(router)
+
+    app.state.engine_client = FakeAsyncOmni(images=[Image.new("RGBA", (64, 64), color=(0, 255, 0, 128))])
+    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}}',
+        max_generated_image_size=1048576,
+    )
+    return TestClient(app)
+
+
@pytest.fixture
 def async_omni_stage_configs_only_client():
    """Create test client with refactored AsyncOmni compatibility surface only."""
@@ -592,6 +621,13 @@ def test_parameter_validation():
    with pytest.raises(ValueError):
        ImageGenerationRequest(prompt="test", guidance_scale=21.0)

+    # Invalid layers for layered models (must stay within the backend-supported range)
+    with pytest.raises(ValueError):
+        ImageGenerationRequest(prompt="test", layers=2)
+
+    with pytest.raises(ValueError):
+        ImageGenerationRequest(prompt="test", layers=11)
+

 # Pass-Through Tests

@@ -711,6 +747,29 @@ def test_image_edit_images_processing(async_omni_test_client):
    assert processed_images[1].size == (24, 24)


+def test_image_edit_rejects_multiple_images_when_model_does_not_support_them(async_omni_test_client):
+    img_bytes_1 = make_test_image_bytes((16, 16))
+    img_bytes_2 = make_test_image_bytes((32, 32))
+
+    engine = async_omni_test_client.app.state.engine_client
+    engine.get_diffusion_od_config = lambda: SimpleNamespace(supports_multimodal_inputs=False)
+
+    response = async_omni_test_client.post(
+        "/v1/images/edits",
+        files=[
+            ("image", img_bytes_1),
+            ("image", img_bytes_2),
+        ],
+        data={"prompt": "hello world."},
+    )
+
+    assert response.status_code == 400
+    assert (
+        response.json()["detail"] == "Received multiple input images. Only a single image is supported by this model."
+    )
+    assert engine.captured_prompt is None
+
+
 def test_image_edit_parameter_pass(async_omni_test_client):
    img_bytes_1 = make_test_image_bytes((16, 16))

@@ -816,35 +875,36 @@ def test_image_edit_invalid_resolution(async_omni_test_client):


 def test_image_edit_invalid_layers(async_omni_test_client):
-    """Test that invalid layers values (< 1) are rejected with 400."""
+    """Test that layered image edits reject out-of-range layers with 400."""
    img_bytes = make_test_image_bytes((16, 16))

-    # Test layers = 0
+    # Test layers below the supported range
    response = async_omni_test_client.post(
        "/v1/images/edits",
        files=[("image", img_bytes)],
        data={
            "prompt": "test",
-            "layers": 0,
+            "layers": 2,
        },
    )
    assert response.status_code == 400
    detail = response.json()["detail"]
    assert "Invalid layers" in detail
-    assert "layers must be >= 1" in detail
+    assert "layers must be between 3 and 10 inclusive" in detail

-    # Test layers = -1
+    # Test layers above the supported range
    response = async_omni_test_client.post(
        "/v1/images/edits",
        files=[("image", img_bytes)],
        data={
            "prompt": "test",
-            "layers": -1,
+            "layers": 11,
        },
    )
    assert response.status_code == 400
    detail = response.json()["detail"]
    assert "Invalid layers" in detail
+    assert "layers must be between 3 and 10 inclusive" in detail


 def test_image_edit_resolution_and_size_conflict(async_omni_test_client):
@@ -977,6 +1037,27 @@ def test_image_edit_compression_jpeg(test_client):
    assert len(img_bytes_50) < len(img_bytes_100)


+def test_image_edit_rgba_output_converts_to_jpeg(async_omni_rgba_test_client):
+    img_bytes_1 = make_test_image_bytes((16, 16))
+
+    response = async_omni_rgba_test_client.post(
+        "/v1/images/edits",
+        files=[("image", img_bytes_1)],
+        data={
+            "prompt": "hello world.",
+            "output_format": "jpeg",
+        },
+    )
+    assert response.status_code == 200
+
+    data = response.json()
+    img_bytes = base64.b64decode(data["data"][0]["b64_json"])
+    img = Image.open(io.BytesIO(img_bytes))
+    assert img.format.lower() == "jpeg"
+    assert img.mode == "RGB"
+    assert data["output_format"] == "jpeg"
+
+
 def test_image_edit_compression_png(async_omni_test_client):
    img_bytes_1 = make_test_image_bytes((16, 16))
    # uploadfile with image key
@@ -0,0 +1,15 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+
+from vllm_omni.entrypoints.async_omni_diffusion import AsyncOmniDiffusion
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_get_diffusion_od_config_returns_direct_config():
+    diffusion = object.__new__(AsyncOmniDiffusion)
+    diffusion.od_config = object()
+
+    assert diffusion.get_diffusion_od_config() is diffusion.od_config
@@ -379,6 +379,19 @@ def test_openai_serving_models_can_consume_async_omni_compat_attrs():
    assert serving_models.input_processor is input_processor


+def test_get_diffusion_od_config_returns_diffusion_stage_config():
+    diffusion_od_config = object()
+    omni = object.__new__(AsyncOmni)
+    omni.engine = SimpleNamespace(
+        stage_clients=[
+            SimpleNamespace(stage_type="llm"),
+            SimpleNamespace(stage_type="diffusion", _engine=SimpleNamespace(od_config=diffusion_od_config)),
+        ]
+    )
+
+    assert omni.get_diffusion_od_config() is diffusion_od_config
+
+
@pytest.mark.asyncio
 async def test_async_omni_yields_only_final_stage_outputs(monkeypatch: pytest.MonkeyPatch):
    engine = FakeAsyncOmniEngine(
@@ -16,6 +16,8 @@ from vllm_omni.model_executor.stage_input_processors.qwen3_tts import (
    talker2code2wav_async_chunk,
 )

+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
 _FRAME = [1, 2, 3, 4]
 _Q = len(_FRAME)

@@ -90,7 +92,7 @@ def test_flush_on_finish():
        is_finished=True,
    )
    assert p is not None
-    assert p["finished"].item() is True
+    assert p["finished"] is True
    assert len(p["code_predictor_codes"]) == _Q * 24


@@ -158,24 +160,29 @@ def test_dynamic_ic_adapts_to_load():


 def test_ic_load_change_mid_request():
-    """IC stateless: load spike mid-request shifts initial_coverage."""
+    """IC is cached per request; a load spike only affects new requests."""
    tm = _tm(chunk_frames=25, left_context=25, max_num_seqs=8)

-    # Low load -> IC=2 -> emit at frame 2
+    # Low load -> IC=2 (cached for "r"), emit at frame 2
    p1 = _call(tm, "r", n_frames=2)
    assert p1 is not None

-    # Spike load: 6 others -> IC=16 -> initial_coverage=16
+    # Spike load: 6 others running
    for i in range(6):
        tm.code_prompt_token_ids[f"other-{i}"] = [[0]] * 10

-    # adjusted=25-16=9, 9%25!=0 -> hold
+    # IC for "r" is still cached as 2.
+    # initial_coverage = ((25-1)//2)*2 = 24, first normal emit at 24+25=49
    assert _call(tm, "r", n_frames=25) is None
-
-    # First normal emit at 16+25=41
-    p3 = _call(tm, "r", n_frames=41)
+    assert _call(tm, "r", n_frames=27) is None
+    p3 = _call(tm, "r", n_frames=49)
    assert p3 is not None
-    assert p3["left_context_size"] == 16
+
+    # A *new* request under high load gets IC=16 (not IC=2).
+    # Frame 2 would emit under IC=2 but must hold under IC=16.
+    assert _call(tm, "new_req", n_frames=2) is None
+    p4 = _call(tm, "new_req", n_frames=16)
+    assert p4 is not None


@pytest.mark.parametrize(
@@ -295,8 +302,13 @@ def test_non_async_processor_prepends_ref_code_and_sets_trim_context():
        ],
        dtype=torch.long,
    )
-    output = SimpleNamespace(multimodal_output={"audio_codes": audio_codes, "ref_code": ref_code})
-    stage = SimpleNamespace(engine_outputs=[SimpleNamespace(outputs=[output])])
+    output = SimpleNamespace(
+        multimodal_output={"audio_codes": audio_codes, "ref_code": ref_code},
+        token_ids=list(range(3)),
+    )
+    stage = SimpleNamespace(
+        engine_outputs=[SimpleNamespace(outputs=[output], finished=True)],
+    )

    prompts = talker2code2wav(stage_list=[stage], engine_input_source=[0])

@@ -335,8 +347,13 @@ def test_non_async_processor_filters_out_of_range_codec_values():
        ],
        dtype=torch.long,
    )
-    output = SimpleNamespace(multimodal_output={"audio_codes": audio_codes, "ref_code": ref_code})
-    stage = SimpleNamespace(engine_outputs=[SimpleNamespace(outputs=[output])])
+    output = SimpleNamespace(
+        multimodal_output={"audio_codes": audio_codes, "ref_code": ref_code},
+        token_ids=list(range(4)),
+    )
+    stage = SimpleNamespace(
+        engine_outputs=[SimpleNamespace(outputs=[output], finished=True)],
+    )

    prompts = talker2code2wav(stage_list=[stage], engine_input_source=[0])

@@ -20,6 +20,24 @@ from vllm_omni.diffusion.data import DiffusionCacheConfig
 logger = init_logger(__name__)


+def enable_hunyuan_image3_teacache(pipeline: Any, config: DiffusionCacheConfig) -> None:
+    """
+    Enable TeaCache for HunyuanImage3 model.
+
+    HunyuanImage3 uses a GPT-based architecture with KV cache, which is incompatible
+    with the standard hook-based TeaCache approach. Instead, we store the TeaCacheConfig
+    on the pipeline so the denoising loop can implement caching directly.
+    """
+    teacache_config = TeaCacheConfig(
+        transformer_type="HunyuanImage3Pipeline",
+        rel_l1_thresh=config.rel_l1_thresh,
+        coefficients=config.coefficients,
+    )
+    pipeline._tea_cache_config = teacache_config
+
+    logger.info(f"TeaCache enabled for HunyuanImage3 with rel_l1_thresh={teacache_config.rel_l1_thresh}")
+
+
 def enable_bagel_teacache(pipeline: Any, config: DiffusionCacheConfig) -> None:
    """
    Enable TeaCache for Bagel model.
@@ -70,6 +88,7 @@ def enable_flux2_klein_teacache(pipeline: Any, config: DiffusionCacheConfig) ->
 CUSTOM_TEACACHE_ENABLERS = {
    "BagelPipeline": enable_bagel_teacache,
    "Flux2KleinPipeline": enable_flux2_klein_teacache,
+    "HunyuanImage3Pipeline": enable_hunyuan_image3_teacache,
 }


@@ -157,6 +176,17 @@ class TeaCacheBackend(CacheBackend):
                                Currently not used by TeaCache but accepted for interface consistency.
            verbose: Whether to log refresh operations (default: True)
        """
+        # HunyuanImage3: tea cache state is managed inside the denoising loop,
+        # so refresh is a no-op (state is re-initialized every __call__).
+        if (
+            hasattr(pipeline, "_tea_cache_config")
+            and isinstance(pipeline._tea_cache_config, TeaCacheConfig)
+            and pipeline.__class__.__name__ == "HunyuanImage3Pipeline"
+        ):
+            if verbose:
+                logger.debug(f"TeaCache state refreshed for HunyuanImage3 (num_inference_steps={num_inference_steps})")
+            return
+
        # Extract transformer from pipeline
        transformer = pipeline.transformer

@@ -54,6 +54,16 @@ _MODEL_COEFFICIENTS = {
        -12.281286412689623,
        1.0733905006198015,
    ],
+    # HunyuanImage3 pipeline coefficients
+    # Calibrated via polyfit on 3920 data points (80 prompts × 49 steps).
+    # Maps time_embed rel_l1 (range ~0.12-0.54) to output rel_l1 (range ~0.01-0.53).
+    "HunyuanImage3Pipeline": [
+        1.04117826e02,
+        -1.26848482e02,
+        5.68168652e01,
+        -1.04182570e01,
+        6.78098549e-01,
+    ],
 }


@@ -17,12 +17,38 @@ from vllm_omni.diffusion.distributed.parallel_state import (
 )


+def _wrap(pred: torch.Tensor | tuple[torch.Tensor, ...]) -> tuple[torch.Tensor, ...]:
+    """Normalize prediction to tuple form."""
+    return pred if isinstance(pred, tuple) else (pred,)
+
+
+def _unwrap(pred: tuple[torch.Tensor, ...]) -> torch.Tensor | tuple[torch.Tensor, ...]:
+    """Unwrap single-element tuple to plain tensor; keep multi-element as tuple."""
+    return pred[0] if len(pred) == 1 else pred
+
+
+def _slice_pred(pred: tuple[torch.Tensor, ...], output_slice: int) -> tuple[torch.Tensor, ...]:
+    """Slice each element along dim 1."""
+    return tuple(p[:, :output_slice] for p in pred)
+
+
 class CFGParallelMixin(metaclass=ABCMeta):
    """
    Base Mixin class for Diffusion pipelines providing shared CFG methods.

    All pipelines should inherit from this class to reuse
    classifier-free guidance logic.
+
+    CFG Parallel Architecture:
+        When cfg_world_size > 1, each rank computes one branch (positive or
+        negative), then all_gather exchanges results. All ranks then compute
+        the CFG combine and scheduler step locally — no broadcast needed
+        because the operations are deterministic.
+
+    Multi-output models:
+        Models that return tuple from predict_noise() (e.g., video + audio)
+        should override combine_cfg_noise() to define per-element combine logic,
+        and set self.scheduler to a composite scheduler that handles tuples.
    """

    def predict_noise_maybe_with_cfg(
@@ -33,7 +59,7 @@ class CFGParallelMixin(metaclass=ABCMeta):
        negative_kwargs: dict[str, Any] | None,
        cfg_normalize: bool = True,
        output_slice: int | None = None,
-    ) -> torch.Tensor | None:
+    ) -> torch.Tensor | tuple[torch.Tensor, ...]:
        """
        Predict noise with optional classifier-free guidance.

@@ -43,57 +69,64 @@ class CFGParallelMixin(metaclass=ABCMeta):
            positive_kwargs: Kwargs for positive/conditional prediction
            negative_kwargs: Kwargs for negative/unconditional prediction
            cfg_normalize: Whether to normalize CFG output (default: True)
-            output_slice: If set, slice output to [:, :output_slice] for image editing
+            output_slice: If set, slice each output to [:, :output_slice] for image editing

        Returns:
-            Predicted noise tensor (only valid on rank 0 in CFG parallel mode)
+            Predicted noise tensor or tuple of tensors.
+            In CFG parallel mode, result is valid on ALL ranks (not just rank 0).
+
+        Note:
+            For multi-output models (e.g., video + audio where predict_noise
+            returns a tuple), override combine_cfg_noise() for per-element CFG
+            logic and set self.scheduler to a composite scheduler.
        """
        if do_true_cfg:
            # Automatically detect CFG parallel configuration
            cfg_parallel_ready = get_classifier_free_guidance_world_size() > 1

            if cfg_parallel_ready:
-                # Enable CFG-parallel: rank0 computes positive, rank1 computes negative.
                cfg_group = get_cfg_group()
                cfg_rank = get_classifier_free_guidance_rank()

-                if cfg_rank == 0:
-                    local_pred = self.predict_noise(**positive_kwargs)
-                else:
-                    local_pred = self.predict_noise(**negative_kwargs)
+                # Each rank computes one branch
+                kwargs = positive_kwargs if cfg_rank == 0 else negative_kwargs
+                local_pred = _wrap(self.predict_noise(**kwargs))

-                # Slice output for image editing pipelines (remove condition latents)
                if output_slice is not None:
-                    local_pred = local_pred[:, :output_slice]
+                    local_pred = _slice_pred(local_pred, output_slice)

-                gathered = cfg_group.all_gather(local_pred, separate_tensors=True)
+                # All-gather each element, reconstruct positive/negative tuples
+                gathered = [cfg_group.all_gather(p, separate_tensors=True) for p in local_pred]
+                positive_noise_pred = tuple(g[0] for g in gathered)
+                negative_noise_pred = tuple(g[1] for g in gathered)

-                if cfg_rank == 0:
-                    noise_pred = gathered[0]
-                    neg_noise_pred = gathered[1]
-                    noise_pred = self.combine_cfg_noise(noise_pred, neg_noise_pred, true_cfg_scale, cfg_normalize)
-                    return noise_pred
-                else:
-                    return None
+                # All ranks compute combine (deterministic, same result)
+                return self.combine_cfg_noise(
+                    positive_noise_pred,
+                    negative_noise_pred,
+                    true_cfg_scale,
+                    cfg_normalize,
+                )
            else:
                # Sequential CFG: compute both positive and negative
-                positive_noise_pred = self.predict_noise(**positive_kwargs)
-                negative_noise_pred = self.predict_noise(**negative_kwargs)
+                positive_noise_pred = _wrap(self.predict_noise(**positive_kwargs))
+                negative_noise_pred = _wrap(self.predict_noise(**negative_kwargs))

-                # Slice output for image editing pipelines
                if output_slice is not None:
-                    positive_noise_pred = positive_noise_pred[:, :output_slice]
-                    negative_noise_pred = negative_noise_pred[:, :output_slice]
+                    positive_noise_pred = _slice_pred(positive_noise_pred, output_slice)
+                    negative_noise_pred = _slice_pred(negative_noise_pred, output_slice)

-                noise_pred = self.combine_cfg_noise(
-                    positive_noise_pred, negative_noise_pred, true_cfg_scale, cfg_normalize
+                return self.combine_cfg_noise(
+                    positive_noise_pred,
+                    negative_noise_pred,
+                    true_cfg_scale,
+                    cfg_normalize,
                )
-                return noise_pred
        else:
            # No CFG: only compute positive/conditional prediction
            pred = self.predict_noise(**positive_kwargs)
            if output_slice is not None:
-                pred = pred[:, :output_slice]
+                pred = _unwrap(_slice_pred(_wrap(pred), output_slice))
            return pred

    def cfg_normalize_function(self, noise_pred: torch.Tensor, comb_pred: torch.Tensor) -> torch.Tensor:
@@ -113,35 +146,61 @@ class CFGParallelMixin(metaclass=ABCMeta):
        return noise_pred

    def combine_cfg_noise(
-        self, noise_pred: torch.Tensor, neg_noise_pred: torch.Tensor, true_cfg_scale: float, cfg_normalize: bool = False
-    ) -> torch.Tensor:
+        self,
+        positive_noise_pred: torch.Tensor | tuple[torch.Tensor, ...],
+        negative_noise_pred: torch.Tensor | tuple[torch.Tensor, ...],
+        true_cfg_scale: float,
+        cfg_normalize: bool = False,
+    ) -> torch.Tensor | tuple[torch.Tensor, ...]:
        """
        Combine conditional and unconditional noise predictions with CFG.

+        Accepts both plain tensors (backward-compatible, used by LTX2 etc.)
+        and tuples (multi-output models). Default implementation applies the
+        standard CFG formula to every element.
+
+        Multi-output models can override this to apply different logic per element.
+
+        Example override for a model returning (video_pred, audio_pred)::
+
+            def combine_cfg_noise(self, positive_noise_pred, negative_noise_pred, scale, normalize):
+                (video_pos, audio_pos) = positive_noise_pred
+                (video_neg, audio_neg) = negative_noise_pred
+                video_combined = super().combine_cfg_noise(video_pos, video_neg, scale, normalize)
+                return (video_combined, audio_pos)  # audio: positive only, no CFG
+
        Args:
-            noise_pred: Conditional noise prediction
-            neg_noise_pred: Unconditional noise prediction
+            positive_noise_pred: Positive/conditional prediction(s) — Tensor or tuple
+            negative_noise_pred: Negative/unconditional prediction(s) — Tensor or tuple
            true_cfg_scale: CFG scale factor
            cfg_normalize: Whether to normalize the combined prediction (default: False)

        Returns:
-            Combined noise prediction tensor
+            Combined noise prediction(s) — same type as inputs
        """
-        comb_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred)
+        pos_t = _wrap(positive_noise_pred)
+        neg_t = _wrap(negative_noise_pred)

-        if cfg_normalize:
-            noise_pred = self.cfg_normalize_function(noise_pred, comb_pred)
-        else:
-            noise_pred = comb_pred
+        results = []
+        for p, n in zip(pos_t, neg_t):
+            comb = n + true_cfg_scale * (p - n)
+            if cfg_normalize:
+                comb = self.cfg_normalize_function(p, comb)
+            results.append(comb)
+        return _unwrap(tuple(results))

-        return noise_pred
-
-    def predict_noise(self, *args: Any, **kwargs: Any) -> torch.Tensor:
+    def predict_noise(self, *args: Any, **kwargs: Any) -> torch.Tensor | tuple[torch.Tensor, ...]:
        """
        Forward pass through transformer to predict noise.

        Subclasses should override this if they need custom behavior,
        but the default implementation calls self.transformer.
+
+        Returns:
+            Single Tensor for standard models, or tuple of Tensors for
+            multi-output models (e.g., video + audio). Multi-output models
+            must also override combine_cfg_noise() and set self.scheduler
+            to a composite scheduler that handles tuples.
        """
        return self.transformer(*args, **kwargs)[0]

@@ -156,52 +215,89 @@ class CFGParallelMixin(metaclass=ABCMeta):
        Subclasses MUST implement this method to define the complete
        diffusion/denoising loop for their specific model.

-        Typical implementation pattern:
+        Typical implementation pattern (single output):
        ```python
        def diffuse(self, latents, timesteps, prompt_embeds, negative_embeds, ...):
            for t in timesteps:
-                # Prepare kwargs for positive and negative predictions
                positive_kwargs = {...}
                negative_kwargs = {...}

-                # Predict noise with automatic CFG handling
                noise_pred = self.predict_noise_maybe_with_cfg(
-                    do_true_cfg=True,
+                    do_true_cfg=do_true_cfg,
                    true_cfg_scale=self.guidance_scale,
                    positive_kwargs=positive_kwargs,
                    negative_kwargs=negative_kwargs,
                )

-                # Step scheduler with automatic CFG sync
                latents = self.scheduler_step_maybe_with_cfg(
-                    noise_pred, t, latents, do_true_cfg=True
+                    noise_pred, t, latents, do_true_cfg=do_true_cfg
                )

            return latents
        ```
+
+        Multi-output models (e.g., video + audio) should:
+        1. Override ``predict_noise()`` to return a tuple
+        2. Override ``combine_cfg_noise()`` for per-element CFG logic
+        3. Set ``self.scheduler`` to a composite scheduler that handles tuples
+
+        ```python
+        def diffuse(self, video_latents, audio_latents, timesteps_video, timesteps_audio, ...):
+            for t_v, t_a in zip(timesteps_video, timesteps_audio):
+                positive_kwargs = {...}
+                negative_kwargs = {...}
+
+                # Returns tuple: (video_pred, audio_pred)
+                video_pred, audio_pred = self.predict_noise_maybe_with_cfg(
+                    do_true_cfg=do_true_cfg,
+                    true_cfg_scale=self.guidance_scale,
+                    positive_kwargs=positive_kwargs,
+                    negative_kwargs=negative_kwargs,
+                )
+
+                # self.scheduler = VideoAudioScheduler(video_sched, audio_sched)
+                # which accepts and returns tuples
+                video_latents, audio_latents = self.scheduler_step_maybe_with_cfg(
+                    (video_pred, audio_pred),
+                    (t_v, t_a),
+                    (video_latents, audio_latents),
+                    do_true_cfg=do_true_cfg,
+                )
+
+            return video_latents, audio_latents
+        ```
        """
        raise NotImplementedError("Subclasses must implement diffuse")

    def scheduler_step(
        self,
-        noise_pred: torch.Tensor,
-        t: torch.Tensor,
-        latents: torch.Tensor,
+        noise_pred: torch.Tensor | tuple[torch.Tensor, ...],
+        t: torch.Tensor | tuple[torch.Tensor, ...],
+        latents: torch.Tensor | tuple[torch.Tensor, ...],
        per_request_scheduler: Any | None = None,
-    ) -> torch.Tensor:
+        generator: torch.Generator | None = None,
+    ) -> torch.Tensor | tuple[torch.Tensor, ...]:
        """
        Step the scheduler.

+        Default implementation passes inputs directly to ``self.scheduler.step()``.
+        For multi-output models, set ``self.scheduler`` to a composite scheduler
+        that handles tuples (e.g., ``VideoAudioScheduler``).
+
        Args:
-            noise_pred: Predicted noise
-            t: Current timestep
-            latents: Current latents
+            noise_pred: Predicted noise (Tensor or tuple for multi-output)
+            t: Current timestep (Tensor or tuple when schedulers differ per output)
+            latents: Current latents (Tensor or tuple for multi-output)
            per_request_scheduler: Optional request-scoped scheduler that
                overrides ``self.scheduler`` for this call. This is
                primarily used by step-wise execution, where each request
                may keep scheduler state in its own runner-managed state
                object. Request-level execution should usually leave this
                as ``None`` and continue using ``self.scheduler``.
+            generator: Optional torch Generator for reproducible sampling.
+                When using CFG parallel, both ranks should receive generators
+                initialized with the same seed so that non-deterministic
+                schedulers (e.g., DDPM) produce identical results.

        Returns:
            Updated latents after scheduler step
@@ -211,62 +307,49 @@ class CFGParallelMixin(metaclass=ABCMeta):
            raise ValueError("No scheduler is available. Set self.scheduler or pass per_request_scheduler.")
        if not callable(getattr(sched, "step", None)):
            raise TypeError("per_request_scheduler must provide a callable step(...) method.")
-        return sched.step(noise_pred, t, latents, return_dict=False)[0]
+        step_kwargs = dict(return_dict=False)
+        if generator is not None:
+            step_kwargs["generator"] = generator
+        return sched.step(noise_pred, t, latents, **step_kwargs)[0]

    def scheduler_step_maybe_with_cfg(
        self,
-        noise_pred: torch.Tensor,
-        t: torch.Tensor,
-        latents: torch.Tensor,
+        noise_pred: torch.Tensor | tuple[torch.Tensor, ...],
+        t: torch.Tensor | tuple[torch.Tensor, ...],
+        latents: torch.Tensor | tuple[torch.Tensor, ...],
        do_true_cfg: bool,
        per_request_scheduler: Any | None = None,
-    ) -> torch.Tensor:
+        generator: torch.Generator | None = None,
+    ) -> torch.Tensor | tuple[torch.Tensor, ...]:
        """
-        Step the scheduler with (maybe) automatic CFG parallel synchronization.
+        Step the scheduler with automatic CFG parallel handling.

-        In CFG parallel mode, only rank 0 computes the scheduler step,
-        then broadcasts the result to other ranks.
+        All ranks compute the scheduler step locally — no broadcast needed
+        because predict_noise_maybe_with_cfg already ensures all ranks have
+        identical noise_pred after all_gather + local combine.

        Args:
-            noise_pred: Predicted noise (only valid on rank 0 in CFG parallel)
-            t: Current timestep
-            latents: Current latents
+            noise_pred: Predicted noise (Tensor or tuple, valid on all ranks)
+            t: Current timestep (Tensor or tuple when schedulers differ per output)
+            latents: Current latents (Tensor or tuple)
            do_true_cfg: Whether CFG is enabled
            per_request_scheduler: Optional request-scoped scheduler that
                overrides ``self.scheduler`` for this call. This is mainly
                needed by step-wise execution, where scheduler state may be
                stored per request. Request-level execution should normally
                leave this as ``None``.
+            generator: Optional torch Generator for reproducible sampling.
+                When using CFG parallel, both ranks should receive generators
+                initialized with the same seed so that non-deterministic
+                schedulers (e.g., DDPM) produce identical results.

        Returns:
-            Updated latents (synchronized across all CFG ranks)
+            Updated latents (identical across all CFG ranks)
        """
-        # Automatically detect CFG parallel configuration
-        cfg_parallel_ready = do_true_cfg and get_classifier_free_guidance_world_size() > 1
-
-        if cfg_parallel_ready:
-            cfg_group = get_cfg_group()
-            cfg_rank = get_classifier_free_guidance_rank()
-
-            # Only rank 0 computes the scheduler step
-            if cfg_rank == 0:
-                latents = self.scheduler_step(
-                    noise_pred,
-                    t,
-                    latents,
-                    per_request_scheduler=per_request_scheduler,
-                )
-
-            # Broadcast the updated latents to all ranks
-            latents = latents.contiguous()
-            cfg_group.broadcast(latents, src=0)
-        else:
-            # No CFG parallel: directly compute scheduler step
-            latents = self.scheduler_step(
-                noise_pred,
-                t,
-                latents,
-                per_request_scheduler=per_request_scheduler,
-            )
-
-        return latents
+        return self.scheduler_step(
+            noise_pred,
+            t,
+            latents,
+            per_request_scheduler=per_request_scheduler,
+            generator=generator,
+        )
@@ -196,6 +196,8 @@ class FluxAttention(torch.nn.Module):
        attention_mask: torch.Tensor | None = None,
        **kwargs,
    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        # Ensure contiguous for FP8 quantized linear layers
+        hidden_states = hidden_states.contiguous()
        qkv, _ = self.to_qkv(hidden_states)
        q_size = self.to_qkv.num_heads * self.head_dim
        kv_size = self.to_qkv.num_kv_heads * self.head_dim
@@ -209,6 +211,7 @@ class FluxAttention(torch.nn.Module):
        key = self.norm_k(key)

        if self.added_kv_proj_dim is not None:
+            encoder_hidden_states = encoder_hidden_states.contiguous()
            encoder_qkv, _ = self.add_kv_proj(encoder_hidden_states)
            add_q_size = self.add_kv_proj.num_heads * self.head_dim
            add_kv_size = self.add_kv_proj.num_kv_heads * self.head_dim
@@ -720,6 +723,7 @@ class FluxKontextTransformer2DModel(FluxTransformer2DModel):
        guidance_embeds: bool = True,
        axes_dims_rope: tuple[int, int, int] = (16, 56, 56),
        theta: float = 10000.0,
+        quant_config: "QuantizationConfig | None" = None,
    ):
        super().__init__(
            od_config=od_config,
@@ -735,6 +739,7 @@ class FluxKontextTransformer2DModel(FluxTransformer2DModel):
            guidance_embeds=guidance_embeds,
            axes_dims_rope=axes_dims_rope,
            theta=theta,
+            quant_config=quant_config,
        )

    def forward(
@@ -130,6 +130,8 @@ class FluxKontextPipeline(nn.Module, FluxPipelineMixin, SupportImageInput):
        ).to(self._execution_device)

        transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, FluxKontextTransformer2DModel)
+        transformer_kwargs["od_config"] = od_config
+        transformer_kwargs["quant_config"] = od_config.quantization_config
        self.transformer = FluxKontextTransformer2DModel(**transformer_kwargs)

        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8
@@ -5,7 +5,7 @@
 import math
 from collections.abc import Iterable
 from functools import lru_cache
-from typing import Any
+from typing import TYPE_CHECKING, Any

 import torch
 import torch.nn as nn
@@ -29,6 +29,11 @@ from vllm_omni.diffusion.distributed.sp_plan import (
    SequenceParallelOutput,
 )

+if TYPE_CHECKING:
+    from vllm.model_executor.layers.quantization.base_config import (
+        QuantizationConfig,
+    )
+
 logger = init_logger(__name__)


@@ -96,7 +101,15 @@ class DistributedRMSNorm(nn.Module):
 class ColumnParallelGELU(nn.Module):
    """Column parallel linear with GELU activation."""

-    def __init__(self, dim_in: int, dim_out: int, *, approximate: str = "tanh", bias: bool = True):
+    def __init__(
+        self,
+        dim_in: int,
+        dim_out: int,
+        *,
+        approximate: str = "tanh",
+        bias: bool = True,
+        quant_config: "QuantizationConfig | None" = None,
+    ):
        super().__init__()
        self.proj = ColumnParallelLinear(
            dim_in,
@@ -104,6 +117,7 @@ class ColumnParallelGELU(nn.Module):
            bias=bias,
            gather_output=False,
            return_bias=False,
+            quant_config=quant_config,
        )
        self.approximate = approximate

@@ -121,18 +135,15 @@ class HeliosFeedForward(nn.Module):
        inner_dim: int,
        dim_out: int | None = None,
        bias: bool = True,
+        quant_config: "QuantizationConfig | None" = None,
    ) -> None:
        super().__init__()
        dim_out = dim_out or dim

-        self.net_0 = ColumnParallelGELU(dim, inner_dim, approximate="tanh", bias=bias)
+        self.net_0 = ColumnParallelGELU(dim, inner_dim, approximate="tanh", bias=bias, quant_config=quant_config)
        self.net_1 = nn.Identity()
        self.net_2 = RowParallelLinear(
-            inner_dim,
-            dim_out,
-            bias=bias,
-            input_is_parallel=True,
-            return_bias=False,
+            inner_dim, dim_out, bias=bias, input_is_parallel=True, return_bias=False, quant_config=quant_config
        )

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
@@ -262,6 +273,7 @@ class HeliosSelfAttention(nn.Module):
        dropout: float = 0.0,
        is_amplify_history: bool = False,
        history_scale_mode: str = "per_head",
+        quant_config: "QuantizationConfig | None" = None,
    ):
        super().__init__()

@@ -275,6 +287,7 @@ class HeliosSelfAttention(nn.Module):
            head_size=head_dim,
            total_num_heads=num_heads,
            bias=True,
+            quant_config=quant_config,
        )

        self.num_heads = self.to_qkv.num_heads
@@ -290,6 +303,7 @@ class HeliosSelfAttention(nn.Module):
            bias=True,
            input_is_parallel=True,
            return_bias=False,
+            quant_config=quant_config,
        )
        self.dropout = nn.Dropout(dropout)

@@ -368,6 +382,7 @@ class HeliosCrossAttention(nn.Module):
        head_dim: int,
        eps: float = 1e-5,
        dropout: float = 0.0,
+        quant_config: "QuantizationConfig | None" = None,
    ):
        super().__init__()

@@ -382,6 +397,7 @@ class HeliosCrossAttention(nn.Module):
            bias=True,
            gather_output=False,
            return_bias=False,
+            quant_config=quant_config,
        )
        self.to_k = ColumnParallelLinear(
            dim,
@@ -389,6 +405,7 @@ class HeliosCrossAttention(nn.Module):
            bias=True,
            gather_output=False,
            return_bias=False,
+            quant_config=quant_config,
        )
        self.to_v = ColumnParallelLinear(
            dim,
@@ -396,6 +413,7 @@ class HeliosCrossAttention(nn.Module):
            bias=True,
            gather_output=False,
            return_bias=False,
+            quant_config=quant_config,
        )

        tp_size = get_tensor_model_parallel_world_size()
@@ -411,6 +429,7 @@ class HeliosCrossAttention(nn.Module):
            bias=True,
            input_is_parallel=True,
            return_bias=False,
+            quant_config=quant_config,
        )
        self.dropout = nn.Dropout(dropout)

@@ -461,6 +480,7 @@ class HeliosTransformerBlock(nn.Module):
        guidance_cross_attn: bool = False,
        is_amplify_history: bool = False,
        history_scale_mode: str = "per_head",
+        quant_config: "QuantizationConfig | None" = None,
    ):
        super().__init__()

@@ -475,19 +495,17 @@ class HeliosTransformerBlock(nn.Module):
            eps=eps,
            is_amplify_history=is_amplify_history,
            history_scale_mode=history_scale_mode,
+            quant_config=quant_config,
        )

        # 2. Cross-attention
        self.attn2 = HeliosCrossAttention(
-            dim=dim,
-            num_heads=num_heads,
-            head_dim=head_dim,
-            eps=eps,
+            dim=dim, num_heads=num_heads, head_dim=head_dim, eps=eps, quant_config=quant_config
        )
        self.norm2 = FP32LayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity()

        # 3. Feed-forward
-        self.ffn = HeliosFeedForward(dim=dim, inner_dim=ffn_dim, dim_out=dim)
+        self.ffn = HeliosFeedForward(dim=dim, inner_dim=ffn_dim, dim_out=dim, quant_config=quant_config)
        self.norm3 = FP32LayerNorm(dim, eps, elementwise_affine=False)

        self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
@@ -602,6 +620,7 @@ class HeliosTransformer3DModel(nn.Module):
        has_multi_term_memory_patch: bool = True,
        is_amplify_history: bool = False,
        history_scale_mode: str = "per_head",
+        quant_config: "QuantizationConfig | None" = None,
    ) -> None:
        super().__init__()

@@ -689,6 +708,7 @@ class HeliosTransformer3DModel(nn.Module):
                    guidance_cross_attn=guidance_cross_attn,
                    is_amplify_history=is_amplify_history,
                    history_scale_mode=history_scale_mode,
+                    quant_config=quant_config,
                )
                for _ in range(num_layers)
            ]
@@ -9,7 +9,7 @@ import logging
 import math
 import os
 from collections.abc import Iterable
-from typing import Any
+from typing import TYPE_CHECKING, Any

 import numpy as np
 import torch
@@ -31,6 +31,11 @@ from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPi
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.platforms import current_omni_platform

+if TYPE_CHECKING:
+    from vllm.model_executor.layers.quantization.base_config import (
+        QuantizationConfig,
+    )
+
 logger = logging.getLogger(__name__)


@@ -80,7 +85,9 @@ def load_transformer_config(model_path: str, subfolder: str = "transformer", loc
    return load_json_config(model_path, subfolder, "config.json", local_files_only)


-def create_transformer_from_config(config: dict) -> HeliosTransformer3DModel:
+def create_transformer_from_config(
+    config: dict, quant_config: QuantizationConfig | None = None
+) -> HeliosTransformer3DModel:
    kwargs = {}

    key_map = [
@@ -112,7 +119,7 @@ def create_transformer_from_config(config: dict) -> HeliosTransformer3DModel:
                val = tuple(val)
            kwargs[key] = val

-    return HeliosTransformer3DModel(**kwargs)
+    return HeliosTransformer3DModel(quant_config=quant_config, **kwargs)


 def get_helios_post_process_func(
@@ -192,7 +199,9 @@ class HeliosPipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPip
        ).to(self.device)

        transformer_config = load_transformer_config(model, "transformer", local_files_only)
-        self.transformer = create_transformer_from_config(transformer_config)
+        self.transformer = create_transformer_from_config(
+            transformer_config, quant_config=od_config.quantization_config
+        )

        # Read scheduler config to determine scheduler type
        sched_cfg = load_json_config(model, "scheduler", "scheduler_config.json", local_files_only)
@@ -3,12 +3,28 @@

 from typing import Any

+import vllm.forward_context as _vllm_fc
 from vllm.model_executor.layers.fused_moe import SharedFusedMoE
 from vllm.utils.import_utils import resolve_obj_by_qualname

 from vllm_omni.platforms import current_omni_platform


+def _set_forward_context_num_tokens(num_tokens: int) -> None:
+    """Set num_tokens on the vLLM ForwardContext for MoE routing.
+
+    After the rebase to vLLM 0.18.0, SharedFusedMoE expects
+    ForwardContext.num_tokens to be set. Without it, MoE expert
+    routing may produce incorrect results (silent correctness bug).
+    """
+    if not _vllm_fc.is_forward_context_available():
+        return
+    forward_context = _vllm_fc.get_forward_context()
+    forward_context.num_tokens = num_tokens
+    if not hasattr(forward_context, "in_profile_run"):
+        forward_context.in_profile_run = False
+
+
 class HunyuanFusedMoEDefault(SharedFusedMoE):
    def __init__(self, *, prefix: str = "", **kwargs: Any) -> None:
        super().__init__(prefix=prefix, **kwargs)
@@ -21,6 +37,7 @@ class HunyuanFusedMoEDefault(SharedFusedMoE):
        self._init_hook_handle.remove()

    def forward(self, hidden_states: Any, router_logits: Any) -> Any:
+        _set_forward_context_num_tokens(hidden_states.shape[0])
        return super().forward(hidden_states, router_logits)


@@ -2667,6 +2667,15 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline):
        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
        self._num_timesteps = len(timesteps)

+        # ---- TeaCache initialization ----
+        tea_cache_config = getattr(self.model, "_tea_cache_config", None)
+        if tea_cache_config is not None:
+            tc_rescale = np.poly1d(tea_cache_config.coefficients)
+            tc_acc_dist = 0.0
+            tc_prev_mod = None
+            tc_prev_pred = None
+            tc_cnt = 0
+
        with self.progress_bar(total=num_inference_steps) as progress_bar:
            for i, t in enumerate(timesteps):
                # expand the latents if we are doing classifier free guidance
@@ -2675,17 +2684,42 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline):

                t_expand = t.repeat(latent_model_input.shape[0])

-                model_inputs = self.model.prepare_inputs_for_generation(
-                    input_ids,
-                    images=latent_model_input,
-                    timestep=t_expand,
-                    **model_kwargs,
-                )
+                # ---- TeaCache: decide whether to compute or reuse ----
+                should_compute = True
+                if tea_cache_config is not None:
+                    with torch.no_grad():
+                        # Use timestep embedding as the modulated input for cache decision
+                        cur_mod = self.model.time_embed(t.unsqueeze(0))
+                    if tc_cnt > 0 and tc_prev_mod is not None:
+                        rel_dist = (
+                            ((cur_mod - tc_prev_mod).abs().mean() / (tc_prev_mod.abs().mean() + 1e-8)).cpu().item()
+                        )
+                        rescaled = float(tc_rescale(rel_dist))
+                        tc_acc_dist += abs(rescaled)
+                        if tc_acc_dist < tea_cache_config.rel_l1_thresh:
+                            should_compute = False
+                        else:
+                            tc_acc_dist = 0.0
+                    tc_prev_mod = cur_mod.detach()

-                with torch.autocast(device_type=self.device.type, dtype=torch.bfloat16, enabled=True):
-                    model_output = self.model.forward_call(**model_inputs, first_step=(i == 0))
-                    pred = model_output["diffusion_prediction"]
-                pred = pred.to(dtype=torch.float32)
+                if should_compute:
+                    model_inputs = self.model.prepare_inputs_for_generation(
+                        input_ids,
+                        images=latent_model_input,
+                        timestep=t_expand,
+                        **model_kwargs,
+                    )
+
+                    with torch.autocast(device_type=self.device.type, dtype=torch.bfloat16, enabled=True):
+                        model_output = self.model.forward_call(**model_inputs, first_step=(i == 0))
+                        pred = model_output["diffusion_prediction"]
+                    pred = pred.to(dtype=torch.float32)
+
+                    if tea_cache_config is not None:
+                        tc_prev_pred = pred.clone()
+                else:
+                    # TeaCache fast path: reuse previous prediction
+                    pred = tc_prev_pred

                # perform guidance
                if self.do_classifier_free_guidance:
@@ -2695,7 +2729,7 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline):
                # compute the previous noisy sample x_t -> x_t-1
                latents = self.scheduler.step(pred, t, latents, **_scheduler_step_extra_kwargs, return_dict=False)[0]

-                if i != len(timesteps) - 1:
+                if i != len(timesteps) - 1 and should_compute:
                    model_kwargs = self.model._update_model_kwargs_for_generation(  # noqa
                        model_output,
                        model_kwargs,
@@ -2709,6 +2743,9 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline):
                    model_kwargs["query_lens"] = query_lens
                    model_kwargs["seq_lens"] = seq_lens

+                if tea_cache_config is not None:
+                    tc_cnt += 1
+
                if callback_on_step_end is not None:
                    callback_kwargs = {}
                    for k in callback_on_step_end_tensor_inputs:
@@ -118,6 +118,23 @@ class AsyncOmni(EngineClient, OmniBase):
        """Compatibility helper for call sites expecting async vllm config access."""
        return self.vllm_config

+    def get_diffusion_od_config(self) -> Any | None:
+        """Return the diffusion-stage config when the pipeline has one."""
+        for stage_client in self.engine.stage_clients:
+            if getattr(stage_client, "stage_type", None) != "diffusion":
+                continue
+
+            od_config = getattr(stage_client, "od_config", None)
+            if od_config is not None:
+                return od_config
+
+            inner_engine = getattr(stage_client, "_engine", None)
+            od_config = getattr(inner_engine, "od_config", None)
+            if od_config is not None:
+                return od_config
+
+        return None
+
    @property
    def model_config(self):
        """Return the model config for the comprehension stage when present."""
@@ -251,6 +251,10 @@ class AsyncOmniDiffusion:
            finished=True,
        )

+    def get_diffusion_od_config(self) -> OmniDiffusionConfig:
+        """Return the diffusion config used by this engine."""
+        return self.od_config
+
    # ------------------------------------------------------------------
    # Public generate API
    # ------------------------------------------------------------------
@@ -84,8 +84,10 @@ from vllm.utils.system_utils import decorate_logs
 from vllm_omni.entrypoints.async_omni import AsyncOmni
 from vllm_omni.entrypoints.openai.errors import InvalidInputReferenceError
 from vllm_omni.entrypoints.openai.image_api_utils import (
+    SUPPORTED_LAYERED_RESOLUTIONS,
    encode_image_base64,
    parse_size,
+    validate_layered_layers,
 )
 from vllm_omni.entrypoints.openai.protocol.audio import BatchSpeechRequest, OpenAICreateSpeechRequest
 from vllm_omni.entrypoints.openai.protocol.images import (
@@ -116,8 +118,6 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniSamplingParam
 logger = init_logger(__name__)
 router = APIRouter()

-# Supported resolution buckets for layered models (e.g., Qwen-Image-Layered)
-SUPPORTED_LAYERED_RESOLUTIONS = (640, 1024)
 profiler_router = APIRouter()


@@ -1207,38 +1207,16 @@ _remove_route_from_router(router, "/v1/models")

@router.get("/v1/models")
 async def show_available_models(raw_request: Request) -> JSONResponse:
-    """Show available models endpoint that works for both LLM and diffusion modes.
+    """Show available models for both LLM and diffusion modes.

-    Returns model information in OpenAI-compatible format.
+    Delegates to state.openai_serving_models which is set to either
+    OpenAIServingModels (LLM) or _DiffusionServingModels (pure diffusion).
    """
-    # Check if we're in diffusion mode
-    diffusion_model_name = getattr(raw_request.app.state, "diffusion_model_name", None)
-    if diffusion_model_name is not None:
-        # Diffusion mode - return the loaded model
-        return JSONResponse(
-            content={
-                "object": "list",
-                "data": [
-                    {
-                        "id": diffusion_model_name,
-                        "object": "model",
-                        "created": 0,
-                        "owned_by": "vllm-omni",
-                        "permission": [],
-                    }
-                ],
-            }
-        )
-
-    # LLM mode - delegate to openai_serving_models
-    openai_serving_models = getattr(raw_request.app.state, "openai_serving_models", None)
-    if openai_serving_models is not None:
-        models = await openai_serving_models.show_available_models()
+    handler = getattr(raw_request.app.state, "openai_serving_models", None)
+    if handler is not None:
+        models = await handler.show_available_models()
        return JSONResponse(content=models.model_dump())
-
-    return JSONResponse(
-        content={"object": "list", "data": []},
-    )
+    return JSONResponse(content={"object": "list", "data": []})


 # Image generation API endpoints
@@ -1314,6 +1292,7 @@ async def generate_images(request: ImageGenerationRequest, raw_request: Request)
            gen_params, "seed", request.seed if request.seed is not None else random.randint(0, 2**32 - 1)
        )
        _update_if_not_none(gen_params, "generator_device", request.generator_device)
+        _update_if_not_none(gen_params, "layers", request.layers)

        request_id = f"img_gen-{random_uuid()}"

@@ -1428,6 +1407,11 @@ async def edit_images(
        if not input_images_list:
            raise HTTPException(status_code=422, detail="Field 'image' or 'url' is required")
        pil_images = await _load_input_images(input_images_list)
+        if len(pil_images) > 1 and not _supports_multimodal_image_inputs(raw_request, engine_client):
+            raise HTTPException(
+                status_code=HTTPStatus.BAD_REQUEST.value,
+                detail="Received multiple input images. Only a single image is supported by this model.",
+            )
        prompt["multi_modal_data"] = {}
        prompt["multi_modal_data"]["image"] = pil_images

@@ -1462,11 +1446,13 @@ async def edit_images(
                detail=f"Invalid resolution {resolution}. Supported resolutions: {SUPPORTED_LAYERED_RESOLUTIONS}.",
            )
        # 3.2.1 Validate layers if provided
-        if layers is not None and layers < 1:
+        try:
+            layers = validate_layered_layers(layers)
+        except ValueError as e:
            raise HTTPException(
                status_code=HTTPStatus.BAD_REQUEST.value,
-                detail=f"Invalid layers value {layers}. layers must be >= 1.",
-            )
+                detail=str(e),
+            ) from e
        # 3.2.2 Check for conflicting size and resolution parameters
        if resolution is not None and size.lower() != "auto":
            raise HTTPException(
@@ -1606,6 +1592,20 @@ def _get_engine_and_model(raw_request: Request):
    return engine_client, model_name, normalized_stage_configs


+def _supports_multimodal_image_inputs(raw_request: Request, engine_client: Any) -> bool:
+    diffusion_engine = getattr(raw_request.app.state, "diffusion_engine", None) or engine_client
+    get_diffusion_od_config = getattr(diffusion_engine, "get_diffusion_od_config", None)
+    od_config = (
+        get_diffusion_od_config() if callable(get_diffusion_od_config) else getattr(diffusion_engine, "od_config", None)
+    )
+
+    if od_config is None:
+        # Preserve the existing compatibility behavior when the diffusion
+        # config is not exposed on the serving surface.
+        return True
+    return bool(getattr(od_config, "supports_multimodal_inputs", False))
+
+
 def _get_lora_from_json_str(lora_body):
    if lora_body is None:
        return None
@@ -1775,10 +1775,27 @@ def _choose_output_format(output_format: str | None, background: str | None) ->
    return "jpeg"


+def _prepare_image_for_output_format(image: Image.Image, format: str) -> Image.Image:
+    fmt = format.lower()
+    if fmt not in {"jpg", "jpeg"}:
+        return image
+
+    if image.mode == "RGB":
+        return image
+
+    if image.mode in {"RGBA", "LA"} or (image.mode == "P" and "transparency" in image.info):
+        alpha_image = image.convert("RGBA")
+        flattened = Image.new("RGB", alpha_image.size, (255, 255, 255))
+        flattened.paste(alpha_image, mask=alpha_image.getchannel("A"))
+        return flattened
+
+    return image.convert("RGB")
+
+
 def _encode_image_base64_with_compression(
    image: Image.Image, format: str = "png", output_compression: int = 100
 ) -> str:
-    """Encode PIL Image to base64 PNG string.
+    """Encode PIL Image to a base64 image string.

    Args:
        image: PIL Image object
@@ -1788,6 +1805,7 @@ def _encode_image_base64_with_compression(
        Base64-encoded image as string
    """
    buffer = io.BytesIO()
+    image = _prepare_image_for_output_format(image, format)
    save_kwargs = {}
    if format in ("jpg", "jpeg", "webp"):
        save_kwargs["quality"] = output_compression
@@ -13,6 +13,9 @@ import io

 import PIL.Image

+SUPPORTED_LAYERED_RESOLUTIONS = (640, 1024)
+SUPPORTED_LAYERED_LAYERS_RANGE = range(3, 11)
+

 def parse_size(size_str: str) -> tuple[int, int]:
    """Parse size string to width and height tuple.
@@ -63,3 +66,16 @@ def encode_image_base64(image: PIL.Image.Image) -> str:
    image.save(buffer, format="PNG")
    buffer.seek(0)
    return base64.b64encode(buffer.read()).decode("utf-8")
+
+
+def validate_layered_layers(layers: int | None) -> int | None:
+    """Validate the Qwen-Image-Layered ``layers`` parameter."""
+    if layers is None:
+        return None
+    if layers not in SUPPORTED_LAYERED_LAYERS_RANGE:
+        raise ValueError(
+            f"Invalid layers value {layers}. layers must be between "
+            f"{SUPPORTED_LAYERED_LAYERS_RANGE.start} and "
+            f"{SUPPORTED_LAYERED_LAYERS_RANGE.stop - 1} inclusive."
+        )
+    return layers
@@ -12,6 +12,8 @@ from typing import Any

 from pydantic import BaseModel, Field, field_validator

+from vllm_omni.entrypoints.openai.image_api_utils import validate_layered_layers
+

 class ResponseFormat(str, Enum):
    """Image response format"""
@@ -43,6 +45,10 @@ class ImageGenerationRequest(BaseModel):
    )
    response_format: ResponseFormat = Field(default=ResponseFormat.B64_JSON, description="Format of the returned image")
    user: str | None = Field(default=None, description="User identifier for tracking")
+    layers: int | None = Field(
+        default=None,
+        description="Number of output layers for layered image models. Supported range: 3-10.",
+    )

    @field_validator("size")
    @classmethod
@@ -67,6 +73,12 @@ class ImageGenerationRequest(BaseModel):
            raise ValueError(f"Only 'b64_json' response format is supported, got: {v}")
        return v

+    @field_validator("layers")
+    @classmethod
+    def validate_layers(cls, v):
+        """Validate the layers parameter for layered image models."""
+        return validate_layered_layers(v)
+
    # vllm-omni extensions for diffusion control
    negative_prompt: str | None = Field(default=None, description="Text describing what to avoid in the image")
    num_inference_steps: int | None = Field(
@@ -82,6 +82,7 @@ from vllm.tool_parsers.mistral_tool_parser import MistralToolCall
 from vllm.utils.collection_utils import as_list

 from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin
+from vllm_omni.entrypoints.openai.image_api_utils import validate_layered_layers
 from vllm_omni.entrypoints.openai.protocol import OmniChatCompletionStreamResponse
 from vllm_omni.entrypoints.openai.protocol.audio import AudioResponse, CreateAudio
 from vllm_omni.entrypoints.openai.utils import parse_lora_request
@@ -2094,6 +2095,11 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
            layers = extra_body.get("layers")
            resolution = extra_body.get("resolution")

+            try:
+                layers = validate_layered_layers(layers)
+            except ValueError as e:
+                return self._create_error_response(str(e), status_code=400)
+
            logger.info(
                "Diffusion chat request %s: prompt=%r, ref_images=%d, params=%s",
                request_id,
@@ -1,9 +1,10 @@
 """Qwen3-Omni Code Predictor -- optimized re-prefill, no KV cache.

 * SDPA attention (F.scaled_dot_product_attention) with native GQA support
+* HF-compatible numerics (float32 RMSNorm, float32 RoPE, separate linear layers)
 * Per-call embedding buffer to avoid cross-request aliasing
 * Pre-allocated position_ids (read-only, safe to persist)
-* torch.compile on inner transformer by default
+* torch.compile (epilogue_fusion=False) on inner transformer by default
 * Inline sampling (top-k + top-p) -- no custom op overhead
 """

@@ -12,15 +13,8 @@ import torch.nn as nn
 import torch.nn.functional as F
 from vllm.config import VllmConfig
 from vllm.logger import init_logger
-from vllm.model_executor.layers.layernorm import RMSNorm
-from vllm.model_executor.layers.linear import (
-    MergedColumnParallelLinear,
-    QKVParallelLinear,
-    RowParallelLinear,
-)
-from vllm.model_executor.layers.quantization import QuantizationConfig
-from vllm.model_executor.layers.rotary_embedding import get_rope
 from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader

 from vllm_omni.platforms import current_omni_platform

@@ -28,15 +22,84 @@ logger = init_logger(__name__)


 # ===================================================================
-#  Standalone Attention (SDPA, no KV cache, no HF backend fallback)
+# HF-numerics-compatible layers for code predictor
 # ===================================================================
+#
+# These use plain PyTorch ops (nn.Linear, manual RMSNorm in float32,
+# rotate_half RoPE) to produce outputs numerically identical to the
+# HuggingFace reference. vLLM's fused kernels (RMSNorm, QKVParallel,
+# get_rope) introduce small precision differences that compound across
+# the autoregressive steps of the code predictor, causing severe
+# audio quality degradation.
+#
+# See: https://github.com/vllm-project/vllm-omni/issues/2274
+
+
+class _RMSNorm(nn.Module):
+    """RMSNorm matching HuggingFace's implementation exactly.
+
+    Computes variance in float32 to avoid bfloat16 precision loss.
+    """
+
+    def __init__(self, hidden_size: int, eps: float = 1e-6) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
+
+
+def _rotate_half(x: torch.Tensor) -> torch.Tensor:
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+class _RotaryEmbedding(nn.Module):
+    """RoPE matching HuggingFace's implementation exactly.
+
+    Forces float32 computation for cos/sin, matching HF's torch.autocast(enabled=False).
+    """
+
+    def __init__(self, config) -> None:
+        super().__init__()
+        head_dim = getattr(
+            config,
+            "head_dim",
+            config.hidden_size // config.num_attention_heads,
+        )
+        rope_theta = getattr(config, "rope_theta", 10000.0)
+        inv_freq = 1.0 / (rope_theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+    def forward(self, x: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        # position_ids: [batch, seq_len]
+        inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
+        position_ids_expanded = position_ids[:, None, :].float()
+
+        # Force float32 (matching HF)
+        device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+        with torch.autocast(device_type=device_type, enabled=False):
+            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+            emb = torch.cat((freqs, freqs), dim=-1)
+            cos = emb.cos()
+            sin = emb.sin()
+
+        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)


 class Qwen3OmniCodePredictorAttention(nn.Module):
    """Multi-head self-attention for code predictor.

-    Uses ``F.scaled_dot_product_attention`` directly.  No KV cache -- the code
-    predictor always re-prefills the full (short) sequence each AR step.
+    Uses ``F.scaled_dot_product_attention`` with HF-compatible RoPE and RMSNorm.
+    No KV cache -- the code predictor always re-prefills the full (short)
+    sequence each AR step.

    Input : [B, seq_len, hidden_size]
    Output: [B, seq_len, hidden_size]
@@ -45,7 +108,6 @@ class Qwen3OmniCodePredictorAttention(nn.Module):
    def __init__(
        self,
        config,
-        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
    ):
        super().__init__()
@@ -58,58 +120,52 @@ class Qwen3OmniCodePredictorAttention(nn.Module):
            cp_cfg.hidden_size // cp_cfg.num_attention_heads,
        )
        self.hidden_size = cp_cfg.hidden_size
-        self.q_size = self.num_heads * self.head_dim
-        self.kv_size = self.num_kv_heads * self.head_dim
        self.scaling = self.head_dim**-0.5
        self._use_gqa = self.num_kv_heads != self.num_heads

-        self.qkv_proj = QKVParallelLinear(
-            hidden_size=self.hidden_size,
-            head_size=self.head_dim,
-            total_num_heads=self.num_heads,
-            total_num_kv_heads=self.num_kv_heads,
+        # Separate q/k/v projections matching HF (no fused packing)
+        self.q_proj = nn.Linear(
+            self.hidden_size,
+            self.num_heads * self.head_dim,
            bias=False,
-            quant_config=quant_config,
-            prefix=f"{prefix}.qkv_proj",
-            disable_tp=True,
        )
-        self.o_proj = RowParallelLinear(
-            input_size=self.num_heads * self.head_dim,
-            output_size=self.hidden_size,
+        self.k_proj = nn.Linear(
+            self.hidden_size,
+            self.num_kv_heads * self.head_dim,
            bias=False,
-            quant_config=quant_config,
-            prefix=f"{prefix}.o_proj",
-            disable_tp=True,
        )
-        self.rotary_emb = get_rope(
-            self.head_dim,
-            max_position=cp_cfg.max_position_embeddings,
-            rope_parameters=None,
-            dual_chunk_attention_config=None,
+        self.v_proj = nn.Linear(
+            self.hidden_size,
+            self.num_kv_heads * self.head_dim,
+            bias=False,
        )
-        self.q_norm = RMSNorm(self.head_dim, eps=cp_cfg.rms_norm_eps)
-        self.k_norm = RMSNorm(self.head_dim, eps=cp_cfg.rms_norm_eps)
+        self.o_proj = nn.Linear(
+            self.num_heads * self.head_dim,
+            self.hidden_size,
+            bias=False,
+        )
+        self.q_norm = _RMSNorm(self.head_dim, eps=cp_cfg.rms_norm_eps)
+        self.k_norm = _RMSNorm(self.head_dim, eps=cp_cfg.rms_norm_eps)

    def forward(
        self,
        hidden_states: torch.Tensor,
-        position_ids: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
    ) -> torch.Tensor:
        bsz, seq_len, _ = hidden_states.shape
+        hidden_shape_q = (bsz, seq_len, self.num_heads, self.head_dim)
+        hidden_shape_kv = (bsz, seq_len, self.num_kv_heads, self.head_dim)

-        # Flatten to 2-D so vLLM rotary_emb gets [num_tokens, size]
-        qkv, _ = self.qkv_proj(hidden_states.reshape(bsz * seq_len, -1))
-        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+        q = self.q_norm(self.q_proj(hidden_states).view(hidden_shape_q)).transpose(1, 2)
+        k = self.k_norm(self.k_proj(hidden_states).view(hidden_shape_kv)).transpose(1, 2)
+        v = self.v_proj(hidden_states).view(hidden_shape_kv).transpose(1, 2)

-        # QK-norm -> RoPE (both 2-D)
-        q = self.q_norm(q.view(-1, self.num_heads, self.head_dim)).view(q.shape)
-        k = self.k_norm(k.view(-1, self.num_kv_heads, self.head_dim)).view(k.shape)
-        q, k = self.rotary_emb(position_ids, q, k)
-
-        # [B, heads, seq, head_dim]
-        q = q.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
-        k = k.view(bsz, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
-        v = v.view(bsz, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        cos, sin = position_embeddings
+        # cos/sin are [batch, seq_len, head_dim], need unsqueeze at dim=1 for heads
+        cos = cos.unsqueeze(1)  # [batch, 1, seq_len, head_dim]
+        sin = sin.unsqueeze(1)
+        q = (q * cos) + (_rotate_half(q) * sin)
+        k = (k * cos) + (_rotate_half(k) * sin)

        attn_out = F.scaled_dot_product_attention(
            q,
@@ -120,9 +176,9 @@ class Qwen3OmniCodePredictorAttention(nn.Module):
            enable_gqa=self._use_gqa,
        )

-        attn_out = attn_out.transpose(1, 2).reshape(bsz * seq_len, -1)
-        output, _ = self.o_proj(attn_out)
-        return output.view(bsz, seq_len, -1)
+        attn_out = attn_out.transpose(1, 2).reshape(bsz, seq_len, -1)
+        output = self.o_proj(attn_out)
+        return output


 # ===================================================================
@@ -131,40 +187,23 @@ class Qwen3OmniCodePredictorAttention(nn.Module):


 class Qwen3OmniCodePredictorMLP(nn.Module):
-    """SiLU-gated MLP for code predictor."""
+    """SiLU-gated MLP for code predictor, matching HF's implementation."""

    def __init__(
        self,
        config,
-        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
    ):
        super().__init__()
        hidden_size = config.code_predictor_config.hidden_size
        intermediate_size = config.code_predictor_config.intermediate_size

-        self.gate_up_proj = MergedColumnParallelLinear(
-            input_size=hidden_size,
-            output_sizes=[intermediate_size, intermediate_size],
-            bias=False,
-            quant_config=quant_config,
-            prefix=f"{prefix}.gate_up_proj",
-            disable_tp=True,
-        )
-        self.down_proj = RowParallelLinear(
-            input_size=intermediate_size,
-            output_size=hidden_size,
-            bias=False,
-            quant_config=quant_config,
-            prefix=f"{prefix}.down_proj",
-            disable_tp=True,
-        )
+        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        gate_up, _ = self.gate_up_proj(hidden_states)
-        gate, up = gate_up.chunk(2, dim=-1)
-        down, _ = self.down_proj(F.silu(gate) * up)
-        return down
+        return self.down_proj(F.silu(self.gate_proj(hidden_states)) * self.up_proj(hidden_states))


 # ===================================================================
@@ -178,32 +217,29 @@ class Qwen3OmniCodePredictorDecoderLayer(nn.Module):
    def __init__(
        self,
        config,
-        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
    ) -> None:
        super().__init__()
        self.self_attn = Qwen3OmniCodePredictorAttention(
            config,
-            quant_config=quant_config,
            prefix=f"{prefix}.self_attn",
        )
        self.mlp = Qwen3OmniCodePredictorMLP(
            config,
-            quant_config=quant_config,
            prefix=f"{prefix}.mlp",
        )
        cp_cfg = config.code_predictor_config
-        self.input_layernorm = RMSNorm(cp_cfg.hidden_size, eps=cp_cfg.rms_norm_eps)
-        self.post_attention_layernorm = RMSNorm(cp_cfg.hidden_size, eps=cp_cfg.rms_norm_eps)
+        self.input_layernorm = _RMSNorm(cp_cfg.hidden_size, eps=cp_cfg.rms_norm_eps)
+        self.post_attention_layernorm = _RMSNorm(cp_cfg.hidden_size, eps=cp_cfg.rms_norm_eps)

    def forward(
        self,
        hidden_states: torch.Tensor,
-        position_ids: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
    ) -> torch.Tensor:
        residual = hidden_states
        hidden_states = self.input_layernorm(hidden_states)
-        hidden_states = self.self_attn(hidden_states, position_ids)
+        hidden_states = self.self_attn(hidden_states, position_embeddings)
        hidden_states = residual + hidden_states

        residual = hidden_states
@@ -238,13 +274,13 @@ class Qwen3OmniCodePredictorBaseModel(nn.Module):
            [
                Qwen3OmniCodePredictorDecoderLayer(
                    vllm_config.model_config.hf_config,
-                    quant_config=vllm_config.quant_config,
                    prefix=f"{prefix}.layers.{idx}",
                )
                for idx in range(config.num_hidden_layers)
            ]
        )
-        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.norm = _RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rotary_emb = _RotaryEmbedding(config)

    def forward(
        self,
@@ -252,8 +288,9 @@ class Qwen3OmniCodePredictorBaseModel(nn.Module):
        position_ids: torch.Tensor,
    ) -> torch.Tensor:
        hidden_states = inputs_embeds
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
        for layer in self.layers:
-            hidden_states = layer(hidden_states, position_ids)
+            hidden_states = layer(hidden_states, position_embeddings)
        hidden_states = self.norm(hidden_states)
        return hidden_states

@@ -332,7 +369,8 @@ class Qwen3OmniMoeTalkerCodePredictor(nn.Module):
        if self._pos_ids is not None and self._pos_ids.device == device:
            return
        max_seq = self.num_code_groups + 1
-        self._pos_ids = torch.arange(max_seq, dtype=torch.long, device=device)
+        # [1, max_seq] for HF-style RoPE (will be expanded to [bsz, seq_len] at use)
+        self._pos_ids = torch.arange(max_seq, dtype=torch.long, device=device).unsqueeze(0)

    def _ensure_cached_refs(self) -> None:
        if self._lm_heads is not None:
@@ -344,16 +382,18 @@ class Qwen3OmniMoeTalkerCodePredictor(nn.Module):
        if self._model_fwd is not None:
            return
        if current_omni_platform.supports_torch_inductor():
-            # torch.compile fuses the 5-layer transformer's small kernels,
-            # reducing BF16 intermediate round-trips and improving precision.
-            # mode="default" avoids Inductor's own CUDA graph capture so it
-            # doesn't conflict with the outer CUDAGraphWrapper.
+            # torch.compile fuses RMSNorm/RoPE in ways that lose float32
+            # precision, compounding across AR steps. Use epilogue_fusion=False
+            # to disable the problematic fusions while still getting kernel
+            # fusion benefits for the linear layers and SDPA.
            self._model_fwd = torch.compile(
                self.model.forward,
-                mode="default",
                dynamic=True,
+                options={
+                    "epilogue_fusion": False,
+                },
            )
-            logger.info("code_predictor: torch.compile enabled (mode=default)")
+            logger.info("code_predictor: torch.compile enabled (no epilogue fusion)")
        else:
            self._model_fwd = self.model.forward
            logger.info("code_predictor: using eager mode (no torch.compile)")
@@ -413,8 +453,8 @@ class Qwen3OmniMoeTalkerCodePredictor(nn.Module):
        for step in range(1, num_groups):
            seq_len = step + 1
            projected = proj_buf[:bsz, :seq_len, :]
-            step_pos_ids = pos_ids[:seq_len] if bsz == 1 else pos_ids[:seq_len].repeat(bsz)
-            assert step_pos_ids.shape[0] == bsz * seq_len
+            # position_ids: [batch, seq_len] for HF-style RoPE
+            step_pos_ids = pos_ids[:, :seq_len].expand(bsz, -1)

            hidden_out = model_fwd(projected, step_pos_ids)

@@ -446,21 +486,11 @@ class Qwen3OmniMoeTalkerCodePredictor(nn.Module):
    # ------------------------------------------------------------------

    def load_weights(self, weights: list[tuple[str, torch.Tensor]]) -> set[str]:
-        """Load weights with mapping for fused QKV and gate_up projections.
+        """Load weights directly (no fused projection remapping needed).

-        Maps original HF weights (q_proj, k_proj, v_proj, gate_proj, up_proj)
-        to fused vLLM weights (qkv_proj, gate_up_proj).
+        Since we use separate nn.Linear for q/k/v/o and gate/up/down,
+        weight names match the HF checkpoint directly.
        """
-        # Mapping for fused projections
-        stacked_params_mapping = [
-            # (param_name, shard_name, shard_id)
-            ("qkv_proj", "q_proj", "q"),
-            ("qkv_proj", "k_proj", "k"),
-            ("qkv_proj", "v_proj", "v"),
-            ("gate_up_proj", "gate_proj", 0),
-            ("gate_up_proj", "up_proj", 1),
-        ]
-
        params_dict = dict(self.named_parameters())
        loaded_params: set[str] = set()

@@ -469,36 +499,12 @@ class Qwen3OmniMoeTalkerCodePredictor(nn.Module):
            if "rotary_emb.inv_freq" in name:
                continue

-            # Handle stacked/fused parameters
-            for param_name, weight_name, shard_id in stacked_params_mapping:
-                if weight_name not in name:
-                    continue
-
-                name = name.replace(weight_name, param_name)
-                # Skip if parameter doesn't exist (e.g., bias)
-                if name.endswith(".bias") and name not in params_dict:
-                    continue
-                if name not in params_dict:
-                    continue
-
-                param = params_dict[name]
-                weight_loader = param.weight_loader
-                weight_loader(param, loaded_weight, shard_id)
-                break
-            else:
-                # Non-stacked parameters - use default loading
-                if name.endswith(".bias") and name not in params_dict:
-                    continue
-                if name not in params_dict:
-                    continue
-
-                param = params_dict[name]
-                weight_loader = getattr(param, "weight_loader", None)
-                if weight_loader is not None:
-                    weight_loader(param, loaded_weight)
-                else:
-                    param.data.copy_(loaded_weight)
+            param = params_dict.get(name)
+            if param is None:
+                continue

+            weight_loader = getattr(param, "weight_loader", default_weight_loader)
+            weight_loader(param, loaded_weight)
            loaded_params.add(name)

        return loaded_params
Author	SHA1	Message	Date
hsliuustc0106	8affabd527	update release Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>	2026-03-31 20:39:22 +08:00
noobHappylife	7cb7b3e8ad	Fix Qwen3-TTS gradio demo (#2372 ) Signed-off-by: noobhappylife <aratar1991@hotmail.com>	2026-03-31 11:42:09 +00:00
Reid	de1ac5ee88	[Bugfix] Fix case-sensitivity in Qwen3 TTS speaker name lookup (#2358 ) Signed-off-by: reidliu41 <reid201711@gmail.com>	2026-03-31 09:51:35 +00:00
WeiQing Chen	e68ca4aa5c	[CI]test: add wan22 i2v video similarity e2e (#2262 ) Signed-off-by: David Chen <530634352@qq.com>	2026-03-31 16:15:47 +08:00
WeiQing Chen	0cbd1cbeb8	[CI] fix test: use minimum supported layered output count (#2350 ) Signed-off-by: David Chen <530634352@qq.com>	2026-03-31 14:59:00 +08:00
wangyuandAlicia	c1a978ae8d	[Test] Add qwen3-omni tests for audio_in_video and one word prompt (#2097 ) Signed-off-by: yenuo26 <410167048@qq.com> Signed-off-by: wangyu <410167048@qq.com> Signed-off-by: wangyu <53896905+yenuo26@users.noreply.github.com> Co-authored-by: Alicia <115451386+congw729@users.noreply.github.com>	2026-03-31 14:30:12 +08:00
Yueqian Lin	1734c1ec23	[Test] Fix 4 broken Qwen3-TTS async chunk unit tests (#2351 ) Signed-off-by: linyueqian <linyueqian@outlook.com>	2026-03-31 00:45:13 -04:00
LLandCanlin Guo	12a8f57934	Fix Qwen3-TTS Base on NPU running failed (#2353 ) Signed-off-by: LL <rimchew@foxmail.com> Signed-off-by: Canlin Guo <canlinguosdu@gmail.com> Co-authored-by: Canlin Guo <canlinguosdu@gmail.com>	2026-03-31 12:42:21 +08:00
Yangshen DengandClaude Opus 4.6	837679dbaa	[Diffusion] Refactor CFG parallel for extensibility and performance (#2063 ) Signed-off-by: Yangshen Deng <yangshen.d@outlook.com> Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>	2026-03-30 16:48:58 +00:00
Zeyu Huang \| 黃澤宇andCopilot	0cf99144d6	[skip ci][Docs] reorganize multiple L4 test guidelines (#2119 ) Signed-off-by: Huang, Zeyu <11222265+fhfuih@users.noreply.github.com> Signed-off-by: Zeyu Huang \| 黃澤宇 <11222265+fhfuih@users.noreply.github.com> Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>	2026-03-31 00:12:14 +08:00
WeiQing Chen	1ca942999a	[Bugfix] fix: validate layered image layers range (#2334 ) Signed-off-by: David Chen <530634352@qq.com>	2026-03-30 21:27:09 +08:00
WeiQing Chen	1eee8435e1	[Bugfix] fix: return 400 for unsupported multi-image edits such as Qwen-Image-Layered (#2298 ) Signed-off-by: David Chen <530634352@qq.com>	2026-03-30 20:24:57 +08:00
Didan Deng	911be0d426	[skip ci][Doc] Refine the Diffusion Features User Guide (#1928 ) Signed-off-by: Didan Deng <33117903+wtomin@users.noreply.github.com>	2026-03-30 14:28:04 +08:00
wangyu yenuo26 R2-Y Yueqian Lin	c9370f493e	[Test] Add Qwen-tts test cases and unify the style of existing test cases (#2195 ) Signed-off-by: yenuo26 <410167048@qq.com> Signed-off-by: wangyu <53896905+yenuo26@users.noreply.github.com> Signed-off-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com> Signed-off-by: wangyu <410167048@qq.com> Co-authored-by: yenuo26 <410167048@example.com> Co-authored-by: R2-Y <ruiruyang2@gmail.com> Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>	2026-03-30 01:02:28 -04:00
Lancer	32af3af262	[Bugfix] Fix Qwen3TTSConfig init order to be compatible with newer Tansformers(5.x) (#2306 ) Signed-off-by: Lancer <maruixiang6688@gmail.com>	2026-03-30 09:50:16 +08:00
chickeyton	fa0792a0aa	[Doc] Add transformers version requirement in GLM-Image example doc (#2265 ) Signed-off-by: chickeyton <ngton2014@gmail.com>	2026-03-29 15:56:51 +00:00
WeiQing ChenandCanlin Guo	ecfee25c3b	[Bugfix] fix: handle Qwen-Image-Layered layered RGBA output for jpeg edits (#2297 ) Signed-off-by: David Chen <530634352@qq.com> Co-authored-by: Canlin Guo <canlinguosdu@gmail.com>	2026-03-29 11:33:37 +00:00
rongfu.lengandHongsheng Liu	9d9e6abae6	[Feat] add helios fp8 quantization (#1916 ) Signed-off-by: rongfu.leng <lenronfu@gmail.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-03-29 06:35:56 +00:00
TJian	f5c7ecb6cc	[ROCm] [CI] Migrate to use amd docker hub for ci (#2303 ) Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>	2026-03-29 19:07:03 +13:00
Lidang Jiang	168bea2440	[Misc] Make gradio an optional dependency and upgrade to >=6.7.0 (#2221 ) Signed-off-by: Lidang-Jiang <lidangjiang@gmail.com> Signed-off-by: Lidang Jiang <lidangjiang@gmail.com>	2026-03-28 12:33:55 -04:00
Ding Zuhaoand汪志鹏	283a62fd76	[Feature] [HunyuanImage3] Add TeaCache support for inference acceleration (#1927 ) Signed-off-by: Ding Zuhao <e1583181@u.nus.edu> Signed-off-by: 汪志鹏 <wangzhipeng628@gmail.com> Co-authored-by: 汪志鹏 <wangzhipeng628@gmail.com>	2026-03-28 15:24:21 +00:00
Junhong Liu	a3e1322ad9	Qwen3-Omni][Bugfix] Replace vLLM fused layers with HF-compatible numerics in code predictor (#2291 ) Signed-off-by: Junhong Liu <98734602+LJH-LBJ@users.noreply.github.com>	2026-03-28 23:02:44 +08:00
SYLAR	ef17011086	[skip ci] [Docs]: add CI Failures troubleshooting guide for contributors (#1259 ) Signed-off-by: lishunyang <lishunyang12@163.com>	2026-03-28 22:36:28 +08:00
lif	1414195fbb	Fix: Enable /v1/models endpoint for pure diffusion mode (#805 ) Signed-off-by: majiayu000 <1835304752@qq.com>	2026-03-28 22:17:38 +08:00
amy-why-3459	eced989b40	[Docs] Modify the documentation description for streaming output (#2300 ) Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com>	2026-03-28 11:24:15 +00:00
amy-why-3459	93bb988422	[Tests][Qwen3-Omni] Add performance test cases (#2011 ) Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com>	2026-03-28 19:01:44 +08:00
Lancer	64629cc174	[Feat] support quantization for Flux Kontext (#2184 ) Signed-off-by: Lancer <maruixiang6688@gmail.com>	2026-03-28 10:28:10 +00:00
Ding Zuhao	5cf6870277	[Bugfix][HunyuanImage3.0] Fix default guidance_scale from 1.0 to 4.0 and port GPU MoE ForwardContext fix from NPU (#2142 ) Signed-off-by: Ding Zuhao <e1583181@u.nus.edu>	2026-03-28 13:49:03 +08:00