[Rebase] Rebase to vllm v0.16.0 (#1357 )

Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk> Signed-off-by: tzhouam <tzhouam@connect.ust.hk> Signed-off-by: princepride <wangzhipeng628@gmail.com> Co-authored-by: 汪志鹏 <wangzhipeng628@gmail.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com> Co-authored-by: Isotr0py <Isotr0py@outlook.com> Co-authored-by: ZJY0516 <zhu.jiangyun@foxmail.com> Co-authored-by: youkaichao <youkaichao@gmail.com>
[BugFix]: Revert utils change (#1369 )
2026-02-13 18:05:07 +08:00 · 2026-02-13 17:54:09 +08:00 · 2026-02-12 23:11:57 +08:00 · 2026-02-12 21:33:06 +08:00 · 2026-02-12 20:49:34 +08:00 · 2026-02-12 12:17:54 +00:00
542 changed files with 78221 additions and 13446 deletions
@@ -9,27 +9,70 @@ steps:
    agents:
      queue: "cpu_queue_premerge"

-  # - label: "Test on NPU"
-  #   depends_on: ~
-  #   key: npu-test
-  #   commands:
-  #     - ".buildkite/scripts/hardware_ci/run_npu_test.sh"
-  #   agents:
-  #     queue: "ascend"
-
-  - label: "Simple Unit Test"
-    depends_on: ~
+  - label: "Upload Nightly Pipeline"
+    depends_on: image-build
+    if: build.env("NIGHTLY") == "1"
    commands:
-      - ".buildkite/scripts/simple_test.sh"
+      - buildkite-agent pipeline upload .buildkite/test-nightly.yaml
    agents:
      queue: "cpu_queue_premerge"

+  - label: "Simple Unit Test"
+    depends_on: image-build
+    commands:
+      - "pytest -v -s -m 'core_model and cpu' --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
+    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
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
  - label: "Diffusion Model Test"
    timeout_in_minutes: 20
    depends_on: image-build
    commands:
-      - pytest -s -v tests/e2e/offline_inference/test_t2i_model.py
-      - pytest -s -v tests/e2e/offline_inference/test_ovis_image.py
+      - pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "core_model and diffusion" --run-level "core_model"
+    agents:
+      queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Diffusion Model CPU offloading Test"
+    timeout_in_minutes: 20
+    depends_on: image-build
+    commands:
+      - pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
+      - pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
+    agents:
+      queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Audio Generation Model Test"
+    timeout_in_minutes: 20
+    depends_on: image-build
+    commands:
+      - pytest -s -v tests/e2e/offline_inference/test_stable_audio_model.py
    agents:
      queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
    plugins:
@@ -46,7 +89,7 @@ steps:
    timeout_in_minutes: 15
    depends_on: image-build
    commands:
-      - pytest -s -v tests/e2e/offline_inference/test_cache_dit.py tests/e2e/offline_inference/test_teacache.py
+      - pytest -s -v -m 'core_model and cache and diffusion and not distributed_cuda and L4'
    agents:
      queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
    plugins:
@@ -59,11 +102,11 @@ steps:
          volumes:
            - "/fsx/hf_cache:/fsx/hf_cache"

-  - label: "Diffusion Parallelism Test"
-    timeout_in_minutes: 15
+  - label: "Diffusion Sequence Parallelism Test"
+    timeout_in_minutes: 20
    depends_on: image-build
    commands:
-      - pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py
+      - pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py -m core_model
    agents:
      queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
    plugins:
@@ -77,13 +120,11 @@ steps:
          volumes:
            - "/fsx/hf_cache:/fsx/hf_cache"

-  - label: "Omni Model Test"
-    timeout_in_minutes: 15
+  - label: "Diffusion GPU Worker Test"
+    timeout_in_minutes: 20
    depends_on: image-build
    commands:
-      - export VLLM_LOGGING_LEVEL=DEBUG
-      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
+      - pytest -s -v tests/diffusion/test_diffusion_worker.py
    agents:
      queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
    plugins:
@@ -91,17 +132,20 @@ steps:
          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
+
+  - label: "Benchmark&Engine Test"
+    timeout_in_minutes: 15
    depends_on: image-build
    commands:
      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py tests/e2e/online_serving/test_async_omni.py
+      - pytest -s -v tests/benchmarks/test_serve_cli.py
+      - pytest -s -v tests/engine/test_async_omni_engine_abort.py
    agents:
      queue: "mithril-h100-pool"
    plugins:
@@ -131,12 +175,88 @@ steps:
                  path: /mnt/hf-cache
                  type: DirectoryOrCreate

+  - label: "Omni Model Test"
+    timeout_in_minutes: 15
+    depends_on: image-build
+    commands:
+      - export VLLM_LOGGING_LEVEL=DEBUG
+      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
+    agents:
+      queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Omni Model Test with H100"
+    timeout_in_minutes: 20
+    depends_on: image-build
+    commands:
+       - 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: public.ecr.aws/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: "Qwen3-TTS E2E Test"
+    timeout_in_minutes: 20
+    depends_on: image-build
+    commands:
+      - export VLLM_LOGGING_LEVEL=DEBUG
+      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py
+    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"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
  - label: "Diffusion Image Edit Test with H100 (1 GPU)"
    timeout_in_minutes: 20
    depends_on: image-build
    commands:
      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - pytest -s -v tests/e2e/online_serving/test_i2i_multi_image_input.py
+      - pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py
    agents:
      queue: "mithril-h100-pool"
    plugins:
@@ -165,3 +285,38 @@ steps:
                hostPath:
                  path: /mnt/hf-cache
                  type: DirectoryOrCreate
+
+  # - label: "Bagel Text2Img Model Test with H100"
+  #   timeout_in_minutes: 30
+  #   depends_on: image-build
+  #   commands:
+  #     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  #     - pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py
+  #   agents:
+  #     queue: "mithril-h100-pool"
+  #   plugins:
+  #     - kubernetes:
+  #         podSpec:
+  #           containers:
+  #             - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+  #               resources:
+  #                 limits:
+  #                   nvidia.com/gpu: 1
+  #               volumeMounts:
+  #                 - name: devshm
+  #                   mountPath: /dev/shm
+  #                 - name: hf-cache
+  #                   mountPath: /root/.cache/huggingface
+  #               env:
+  #                 - name: HF_HOME
+  #                   value: /root/.cache/huggingface
+  #           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
@@ -0,0 +1,73 @@
+#!/bin/bash
+# Helper function to safely login to ECR Public with per-job config isolation
+# Uses DOCKER_CONFIG environment variable to prevent race conditions
+#
+# This script prevents the "device or resource busy" error by giving each
+# Buildkite job its own isolated Docker config directory.
+#
+# Usage:
+#   source docker_login_ecr_public.sh && safe_docker_login_ecr_public
+
+set -euo pipefail
+
+# Configuration
+ECR_REGISTRY="public.ecr.aws"
+
+setup_isolated_docker_config() {
+    # Use BUILDKITE_JOB_ID for job-specific isolation
+    # Fallback to PID if running outside Buildkite
+    local job_id="${BUILDKITE_JOB_ID:-$$}"
+
+    # Set Docker config to job-specific directory
+    export DOCKER_CONFIG="/tmp/docker-config-${job_id}"
+
+    # Create directory if it doesn't exist
+    mkdir -p "$DOCKER_CONFIG"
+
+    echo "[docker-config] Using isolated Docker config: $DOCKER_CONFIG"
+}
+
+check_docker_auth() {
+    # Check if already authenticated to the given registry
+    # Returns 0 if authenticated, 1 if not
+    local registry="$1"
+
+    # Check if credentials exist in the isolated config
+    if [[ -f "$DOCKER_CONFIG/config.json" ]]; then
+        # Check if registry is present in config
+        if grep -q "$registry" "$DOCKER_CONFIG/config.json" 2>/dev/null; then
+            return 0
+        fi
+    fi
+
+    return 1
+}
+
+safe_docker_login_ecr_public() {
+    # Setup isolated config first
+    setup_isolated_docker_config
+
+    local registry="$ECR_REGISTRY"
+
+    # Check if already authenticated (within this job)
+    if check_docker_auth "$registry"; then
+        echo "[docker-login] Already authenticated to $registry in this job"
+        return 0
+    fi
+
+    # Perform login to isolated config directory
+    echo "[docker-login] Logging in to $ECR_REGISTRY (isolated config)..."
+    if aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$ECR_REGISTRY"; then
+        echo "[docker-login] Login successful (config: $DOCKER_CONFIG)"
+        return 0
+    else
+        local exit_code=$?
+        echo "[docker-login] ERROR: Login failed with exit code $exit_code" >&2
+        return $exit_code
+    fi
+}
+
+# Execute if run as script (not sourced)
+if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
+    safe_docker_login_ecr_public
+fi
@@ -66,6 +66,14 @@ done
 echo "--- Pulling container"
 image_name="public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:${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
+
 docker pull "${image_name}"

 remove_docker_container() {
@@ -108,6 +116,9 @@ if [[ $commands == *"--shard-id="* ]]; then
        --shm-size=16gb \
        --group-add "$render_gid" \
        --rm \
+        -e MIOPEN_DEBUG_CONV_DIRECT=0 \
+        -e MIOPEN_DEBUG_CONV_GEMM=0 \
+        -e VLLM_ROCM_USE_AITER=1 \
        -e HIP_VISIBLE_DEVICES="${GPU}" \
        -e HF_TOKEN \
        -e AWS_ACCESS_KEY_ID \
@@ -140,6 +151,9 @@ else
          --shm-size=16gb \
          --group-add "$render_gid" \
          --rm \
+          -e MIOPEN_DEBUG_CONV_DIRECT=0 \
+          -e MIOPEN_DEBUG_CONV_GEMM=0 \
+          -e VLLM_ROCM_USE_AITER=1 \
          -e HF_TOKEN \
          -e AWS_ACCESS_KEY_ID \
          -e AWS_SECRET_ACCESS_KEY \
@@ -1,54 +0,0 @@
-#!/usr/bin/env bash
-
-set -euo pipefail
-
-SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
-REPO_ROOT="$(cd -- "${SCRIPT_DIR}/../.." && pwd)"
-UV_BIN="uv"
-PYTHON_BOOTSTRAP="${PYTHON:-python3}"
-VENV_DIR="${VENV_DIR:-${REPO_ROOT}/.venv-simple-test}"
-TARGET_PY_VERSION="3.12"
-
-if ! command -v "${PYTHON_BOOTSTRAP}" >/dev/null 2>&1; then
-  PYTHON_BOOTSTRAP="python"
-fi
-
-cd "${REPO_ROOT}"
-
-# Ensure pip-installed scripts (including uv) are on PATH
-PYTHON_SCRIPTS="$(${PYTHON_BOOTSTRAP} -c 'import sysconfig; print(sysconfig.get_path("scripts"))')"
-USER_BASE="$(${PYTHON_BOOTSTRAP} -m site --user-base 2>/dev/null || true)"
-PATH="${PYTHON_SCRIPTS}:${PATH}"
-if [[ -n "${USER_BASE}" ]]; then
-  PATH="${USER_BASE}/bin:${PATH}"
-fi
-
-if ! command -v "${UV_BIN}" >/dev/null 2>&1; then
-  "${PYTHON_BOOTSTRAP}" -m pip install --upgrade pip
-  "${PYTHON_BOOTSTRAP}" -m pip install uv
-  hash -r
-fi
-
-TARGET_PY_PATH="$(${UV_BIN} python find "${TARGET_PY_VERSION}" 2>/dev/null | head -n1 || true)"
-if [[ -z "${TARGET_PY_PATH}" ]]; then
-  "${UV_BIN}" python install "${TARGET_PY_VERSION}"
-  TARGET_PY_PATH="$(${UV_BIN} python find "${TARGET_PY_VERSION}" 2>/dev/null | head -n1 || true)"
-fi
-
-if [[ -z "${TARGET_PY_PATH}" ]]; then
-  echo "Failed to locate python ${TARGET_PY_VERSION}" >&2
-  exit 1
-fi
-
-if [[ ! -d "${VENV_DIR}" ]]; then
-  "${UV_BIN}" venv --python "${TARGET_PY_PATH}" "${VENV_DIR}"
-fi
-
-VENV_PYTHON="${VENV_DIR}/bin/python"
-[[ -x "${VENV_PYTHON}" ]] || { echo "Python not found in ${VENV_DIR}"; exit 1; }
-
-"${UV_BIN}" pip install --python "${VENV_PYTHON}" vllm==0.12.0
-"${UV_BIN}" pip install --python "${VENV_PYTHON}" -e ".[dev]"
-"${VENV_PYTHON}" -m pytest -v -s tests/entrypoints/
-"${VENV_PYTHON}" -m pytest -v -s tests/diffusion/cache/
-"${VENV_PYTHON}" -m pytest -v -s tests/model_executor/models/qwen2_5_omni/test_audio_length.py
@@ -1,53 +1,116 @@
 steps:

 - label: "Diffusion Model Test"
-  timeout_in_minutes: 15
+  timeout_in_minutes: 20
  agent_pool: mi325_2
  depends_on: amd-build
-  mirror_hardwares: [amdexperimental, amdproduction, amdtentative]
+  mirror_hardwares: [amdproduction]
  grade: Blocking
  commands:
    - export GPU_ARCHS=gfx942
-    - export MIOPEN_DEBUG_CONV_DIRECT=0
-    - export MIOPEN_DEBUG_CONV_GEMM=0
-    - export VLLM_ROCM_USE_AITER=1
-    - export VLLM_ROCM_USE_AITER_MHA=1
-    - export VLLM_ROCM_USE_AITER_LINEAR=0
-    - export VLLM_ROCM_USE_AITER_RMSNORM=0
    - pytest -s -v tests/e2e/offline_inference/test_t2i_model.py

+- label: "Diffusion Images API LoRA E2E"
+  timeout_in_minutes: 20
+  agent_pool: mi325_1
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export GPU_ARCHS=gfx942
+    - export VLLM_LOGGING_LEVEL=DEBUG
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - pytest -s -v tests/e2e/online_serving/test_images_generations_lora.py
+
+- label: "Diffusion Model CPU offloading Test"
+  timeout_in_minutes: 20
+  agent_pool: mi325_1
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export GPU_ARCHS=gfx942
+    - export VLLM_LOGGING_LEVEL=DEBUG
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
+
 - label: "Diffusion Cache Backend Test"
  timeout_in_minutes: 15
  agent_pool: mi325_1
  depends_on: amd-build
-  mirror_hardwares: [amdexperimental, amdproduction, amdtentative]
+  mirror_hardwares: [amdproduction]
  grade: Blocking
  commands:
    - export GPU_ARCHS=gfx942
    - export VLLM_LOGGING_LEVEL=DEBUG
    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - export MIOPEN_DEBUG_CONV_DIRECT=0
-    - export MIOPEN_DEBUG_CONV_GEMM=0
-    - export VLLM_ROCM_USE_AITER=1
-    - export VLLM_ROCM_USE_AITER_MHA=1
-    - export VLLM_ROCM_USE_AITER_LINEAR=0
-    - export VLLM_ROCM_USE_AITER_RMSNORM=0
-    - pytest -s -v tests/e2e/offline_inference/test_cache_dit.py tests/e2e/offline_inference/test_teacache.py
+    - pytest -s -v -m 'core_model and cache and diffusion and not distributed_rocm and MI325'
+
+- label: "Diffusion Sequence Parallelism Test"
+  timeout_in_minutes: 20
+  agent_pool: mi325_2
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export GPU_ARCHS=gfx942
+    - export VLLM_LOGGING_LEVEL=DEBUG
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py -m core_model
+
+- label: "Diffusion Tensor Parallelism Test"
+  timeout_in_minutes: 20
+  agent_pool: mi325_2
+  depends_on: amd-build
+  commands:
+    - export GPU_ARCHS=gfx942
+    - export VLLM_LOGGING_LEVEL=DEBUG
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - pytest -s -v tests/e2e/offline_inference/test_zimage_parallelism.py
+
+- label: "Diffusion GPU Worker Test"
+  timeout_in_minutes: 20
+  agent_pool: mi325_2
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - pytest -s -v tests/diffusion/test_diffusion_worker.py

 - label: "Omni Model Test Qwen2-5-Omni"
  timeout_in_minutes: 15
  agent_pool: mi325_2
  depends_on: amd-build
-  mirror_hardwares: [amdexperimental, amdproduction, amdtentative]
+  mirror_hardwares: [amdproduction]
  grade: Blocking
  commands:
    - export GPU_ARCHS=gfx942
    - export VLLM_LOGGING_LEVEL=DEBUG
    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - export MIOPEN_DEBUG_CONV_DIRECT=0
-    - export MIOPEN_DEBUG_CONV_GEMM=0
-    - export VLLM_ROCM_USE_AITER=1
-    - export VLLM_ROCM_USE_AITER_MHA=1
-    - export VLLM_ROCM_USE_AITER_LINEAR=0
-    - export VLLM_ROCM_USE_AITER_RMSNORM=0
    - pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
+
+- label: "Omni Model Test Qwen3-Omni"
+  timeout_in_minutes: 15
+  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
+    - 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
+    - pytest -s -v tests/engine/test_async_omni_engine_abort.py
+
+- label: "Diffusion Image Edit Test"
+  timeout_in_minutes: 15
+  agent_pool: mi325_1
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export GPU_ARCHS=gfx942
+    - export VLLM_LOGGING_LEVEL=DEBUG
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py
@@ -0,0 +1,329 @@
+steps:
+  - label: ":docker: Build image"
+    key: image-build
+    commands:
+      - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
+      - "docker build --file docker/Dockerfile.ci -t vllm-omni-ci ."
+      - "docker tag vllm-omni-ci public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
+      - "docker push public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
+    agents:
+      queue: "cpu_queue_premerge"
+
+  - label: "Simple Unit Test"
+    depends_on: image-build
+    commands:
+      - "pytest -v -s -m 'core_model and cpu' --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
+    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
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Diffusion Model Test"
+    timeout_in_minutes: 20
+    depends_on: image-build
+    commands:
+      - pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "advanced_model and diffusion" --run-level "advanced_model"
+    agents:
+      queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Diffusion Images API LoRA E2E"
+    timeout_in_minutes: 20
+    depends_on: image-build
+    commands:
+      - pytest -s -v tests/e2e/online_serving/test_images_generations_lora.py
+    agents:
+      queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Diffusion Model CPU offloading Test"
+    timeout_in_minutes: 20
+    depends_on: image-build
+    commands:
+      - pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
+      - pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
+    agents:
+      queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Audio Generation Model Test"
+    timeout_in_minutes: 20
+    depends_on: image-build
+    commands:
+      - pytest -s -v tests/e2e/offline_inference/test_stable_audio_model.py
+    agents:
+      queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Diffusion Cache Backend Test"
+    timeout_in_minutes: 15
+    depends_on: image-build
+    commands:
+      - pytest -s -v tests/e2e/offline_inference/test_cache_dit.py tests/e2e/offline_inference/test_teacache.py
+    agents:
+      queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Diffusion Sequence Parallelism Test"
+    timeout_in_minutes: 20
+    depends_on: image-build
+    commands:
+      - pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py
+    agents:
+      queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
+    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: "Diffusion Tensor Parallelism Test"
+    timeout_in_minutes: 20
+    depends_on: image-build
+    commands:
+      - pytest -s -v tests/e2e/offline_inference/test_zimage_tensor_parallel.py
+    agents:
+      queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
+    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: "Diffusion GPU Worker Test"
+    timeout_in_minutes: 20
+    depends_on: image-build
+    commands:
+      - pytest -s -v tests/diffusion/test_diffusion_worker.py
+    agents:
+      queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
+    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: "Benchmark&Engine Test"
+    timeout_in_minutes: 15
+    depends_on: image-build
+    commands:
+      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - pytest -s -v tests/benchmarks/test_serve_cli.py
+      - pytest -s -v tests/engine/test_async_omni_engine_abort.py
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: public.ecr.aws/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"
+    timeout_in_minutes: 15
+    depends_on: image-build
+    commands:
+    - export VLLM_LOGGING_LEVEL=DEBUG
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
+    agents:
+      queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
+    plugins:
+    - docker#v5.2.0:
+        image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+        always-pull: true
+        propagate-environment: true
+        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: image-build
+    commands:
+       - 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 "advanced_model" --run-level "advanced_model"
+    agents:
+       queue: "mithril-h100-pool"
+    plugins:
+       - kubernetes:
+           podSpec:
+             containers:
+               - image: public.ecr.aws/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: "Diffusion Image Edit Test with H100 (1 GPU)"
+    timeout_in_minutes: 20
+    depends_on: image-build
+    commands:
+      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                resources:
+                  limits:
+                    nvidia.com/gpu: 1
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: hf-cache
+                    mountPath: /root/.cache/huggingface
+                env:
+                  - name: HF_HOME
+                    value: /root/.cache/huggingface
+            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: "Bagel Text2Img Model Test with H100"
+  #   timeout_in_minutes: 30
+  #   depends_on: image-build
+  #   commands:
+  #     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  #     - pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py
+  #   agents:
+  #     queue: "mithril-h100-pool"
+  #   plugins:
+  #     - kubernetes:
+  #         podSpec:
+  #           containers:
+  #             - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+  #               resources:
+  #                 limits:
+  #                   nvidia.com/gpu: 1
+  #               volumeMounts:
+  #                 - name: devshm
+  #                   mountPath: /dev/shm
+  #                 - name: hf-cache
+  #                   mountPath: /root/.cache/huggingface
+  #               env:
+  #                 - name: HF_HOME
+  #                   value: /root/.cache/huggingface
+  #           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
@@ -0,0 +1,94 @@
+steps:
+  - label: "Omni Model Test with H100"
+    timeout_in_minutes: 60
+    depends_on: image-build
+    if: build.env("NIGHTLY") == "1"
+    commands:
+      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - pytest -s -v tests/e2e/online_serving/test_qwen3_omni_expansion.py
+      - pytest -s -v tests/examples/online_serving/test_qwen3_omni.py
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: public.ecr.aws/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"
+    timeout_in_minutes: 60
+    depends_on: image-build
+    if: build.env("NIGHTLY") == "1"
+    commands:
+      - export VLLM_LOGGING_LEVEL=DEBUG
+      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - pytest -s -v tests/examples/online_serving/test_qwen2_5_omni.py
+    agents:
+      queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+
+  - label: "Omni Model Perf Test"
+    timeout_in_minutes: 120
+    depends_on: image-build
+    if: build.env("NIGHTLY") == "1"
+    commands:
+      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - pytest -s -v tests/perf/scripts/run_benchmark.py
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: public.ecr.aws/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
@@ -13,8 +13,8 @@
        depends_on: ~
        soft_fail: false
        commands:
-          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-          - "docker build -f docker/Dockerfile.rocm -t {{ docker_image_amd }} --progress plain ."
+          - "source .buildkite/scripts/docker_login_ecr_public.sh && safe_docker_login_ecr_public"
+          - "docker build -f docker/Dockerfile.rocm -t {{ docker_image_amd }} --target final --progress plain ."
          - "docker push {{ docker_image_amd }}"
        key: "amd-build"
        env:
@@ -40,6 +40,31 @@ body:
      </details>
  validations:
    required: true
+
+- type: textarea
+  attributes:
+    label: Your code version
+    description: |
+      Please provide a description of which commit id or version the code is.
+
+    value: |
+      <details>
+      <summary>The commit id or version of vllm</summary>
+
+      ```text
+
+      ```
+      </details>
+      <details>
+      <summary>The commit id or version of vllm-omni</summary>
+
+      ```text
+
+      ```
+      </details>
+  validations:
+    required: true
+
 - type: textarea
  attributes:
    label: 🐛 Describe the bug
@@ -20,6 +20,8 @@ body:
    label: Proposed Change.
    description: >
      The proposed change of the RFC.
+    value: |
+      Please provide the detailed design document of the RFC using the [template](https://docs.google.com/document/d/12YxSsVeD1jvL-InClkeAEnZyWFDndz_65JmXvsamuV4/edit?tab=t.0#heading=h.4ef4szrsgspp).
  validations:
    required: true
 - type: textarea
@@ -12,9 +12,9 @@ PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTT
 <summary> Essential Elements of an Effective PR Description Checklist </summary>

 - [ ] The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
- [ ] The test plan, such as providing test command.
- [ ] The test results, such as pasting the results comparison before and after, or e2e results
- [ ] (Optional) The necessary documentation update, such as updating `supported_models.md` and `examples` for a new model.
+- [ ] The test plan. Please providing the test scripts & test commands. Please state the reasons if your codes don't require additional test scripts. For test file guidelines, please check the [test style doc](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/tests_style/)
+- [ ] The test results. Please pasting the results comparison before and after, or e2e results.
+- [ ] (Optional) The necessary documentation update, such as updating `supported_models.md` and `examples` for a new model. **Please run `mkdocs serve` to sync the documentation editions to `./docs`.**
 - [ ] (Optional) Release notes update. If your change is user facing, please update the release notes draft.
 </details>

@@ -7,6 +7,7 @@ on:
    paths-ignore:
      - 'docs/**'
      - 'examples/**'
+      - 'apps/**'
      - '.github/**'
      - '.buildkite/**'
      - '**.md'
@@ -20,6 +21,7 @@ on:
    paths-ignore:
      - 'docs/**'
      - 'examples/**'
+      - 'apps/**'
      - '.github/**'
      - '.buildkite/**'
      - '*.yml'
@@ -0,0 +1,27 @@
+name: ComfyUI Integration -- Publish to Comfy registry
+on:
+  workflow_dispatch:
+  push:
+    tags:
+      - "*"
+    paths:
+      - "apps/ComfyUI-vLLM-Omni/**"
+
+permissions:
+  issues: write
+
+defaults:
+  run:
+    working-directory: apps/ComfyUI-vLLM-Omni
+
+jobs:
+  publish-node:
+    name: Publish Custom Node to registry
+    runs-on: ubuntu-latest
+    steps:
+      - name: ♻️ Check out code
+        uses: actions/checkout@v4
+      - name: 📦 Publish Custom Node
+        uses: Comfy-Org/publish-node-action@main
+        with:
+          personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }}
@@ -0,0 +1,40 @@
+name: ComfyUI Integration -- Validate backwards compatibility
+
+on:
+  pull_request:
+    branches:
+      - master
+      - main
+    paths:
+      - "apps/ComfyUI-vLLM-Omni/**"
+
+defaults:
+  run:
+    working-directory: apps/ComfyUI-vLLM-Omni
+
+jobs:
+  check-base-path:
+    runs-on: ubuntu-latest
+    outputs:
+      exists: ${{ steps.check.outputs.exists }}
+    defaults:
+      run:
+        working-directory: .
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          ref: ${{ github.event.pull_request.base.sha }}
+          fetch-depth: 1
+      - id: check
+        run: |
+          if [ -d "apps/ComfyUI-vLLM-Omni" ]; then
+            echo "exists=true" >> "$GITHUB_OUTPUT"
+          else
+            echo "exists=false" >> "$GITHUB_OUTPUT"
+          fi
+  validate:
+    needs: check-base-path
+    if: needs.check-base-path.outputs.exists == 'true'
+    runs-on: ubuntu-latest
+    steps:
+      - uses: comfy-org/node-diff@main
@@ -153,6 +153,14 @@ dmypy.json
 # Cython debug symbols
 cython_debug/

+# Claude
+CLAUDE.md
+.claude/
+
+# Codex
+AGENTS.md
+.codex/
+
 # cursor
 .cursor/

@@ -214,6 +222,7 @@ data/
 datasets/
 *.csv
 *.json
+!apps/ComfyUI-vLLM-Omni/example_workflows/*.json
 *.jsonl
 *.parquet

@@ -233,3 +242,6 @@ configs/development.yaml
 Dockerfile.dev
 discussion
 tmp_test
+
+# output files
+*.wav
@@ -29,7 +29,10 @@ repos:
        # only for staged files

  - repo: https://github.com/rhysd/actionlint
-    rev: v1.7.9
+    # v1.7.8+ sets `go 1.24.0` in go.mod, which older Go toolchains (and most
+    # current CI images) cannot parse. Pin to v1.7.7 until actionlint fixes the
+    # go.mod directive.
+    rev: v1.7.7
    hooks:
      - id: actionlint
        files: ^\.github/workflows/.*\.ya?ml$
@@ -9,13 +9,16 @@ Easy, fast, and cheap omni-modality model serving for everyone
 </h3>

 <p align="center">
-| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> |
+| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> | <a href="docs/assets/WeChat.jpg"><b>WeChat</b></a> | <a href="https://arxiv.org/abs/2602.02204"><b>Paper</b></a> | <a href="https://docs.google.com/presentation/d/1qv4qMW1rKAqDREMXiUDLIgqqHQe7TDPj/edit?usp=sharing&ouid=110473603432222024453&rtpof=true&sd=true"><b>Slides</b></a> |
 </p>

+
 ---

 *Latest News* 🔥

+- [2026/02] We released [0.14.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.14.0) - This is the first **stable release** of vLLM-Omni that expands Omni’s diffusion / image-video generation and audio / TTS stack, improves distributed execution and memory efficiency, and broadens platform/backend coverage (GPU/ROCm/NPU/XPU). It also brings meaningful upgrades to serving APIs, profiling & benchmarking, and overall stability. Please check our latest [paper](https://arxiv.org/abs/2602.02204) for architecture design and performance results.
+- [2026/01] We released [0.12.0rc1](https://github.com/vllm-project/vllm-omni/releases/tag/v0.12.0rc1) - a major RC milestone focused on maturing the diffusion stack, strengthening OpenAI-compatible serving, expanding omni-model coverage, and improving stability across platforms (GPU/NPU/ROCm).
 - [2025/11] vLLM community officially released [vllm-project/vllm-omni](https://github.com/vllm-project/vllm-omni) in order to support omni-modality models serving.

 ---
@@ -66,9 +69,26 @@ Visit our [documentation](https://vllm-omni.readthedocs.io/en/latest/) to learn
 We welcome and value any contributions and collaborations.
 Please check out [Contributing to vLLM-Omni](https://vllm-omni.readthedocs.io/en/latest/contributing/) for how to get involved.

+## Citation
+
+If you use vLLM-Omni for your research, please cite our [paper](https://arxiv.org/abs/2602.02204):
+
+```bibtex
+@article{yin2026vllmomni,
+  title={vLLM-Omni: Fully Disaggregated Serving for Any-to-Any Multimodal Models},
+  author={Peiqi Yin, Jiangyun Zhu, Han Gao, Chenguang Zheng, Yongxiang Huang, Taichang Zhou, Ruirui Yang, Weizhi Liu, Weiqing Chen, Canlin Guo, Didan Deng, Zifeng Mo, Cong Wang, James Cheng, Roger Wang, Hongsheng Liu},
+  journal={arXiv preprint arXiv:2602.02204},
+  year={2026}
+}
+```
+
 ## Join the Community
 Feel free to ask questions, provide feedbacks and discuss with fellow users of vLLM-Omni in `#sig-omni` slack channel at [slack.vllm.ai](https://slack.vllm.ai) or vLLM user forum at [discuss.vllm.ai](https://discuss.vllm.ai).

+## Star History
+
+[![Star History Chart](https://api.star-history.com/svg?repos=vllm-project/vllm-omni&type=date&legend=top-left)](https://www.star-history.com/#vllm-project/vllm-omni&type=date&legend=top-left)
+
 ## License

 Apache License 2.0, as found in the [LICENSE](./LICENSE) file.
@@ -0,0 +1,115 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# OSX useful to ignore
+*.DS_Store
+.AppleDouble
+.LSOverride
+
+# Thumbnails
+._*
+
+# Files that might appear in the root of a volume
+.DocumentRevisions-V100
+.fseventsd
+.Spotlight-V100
+.TemporaryItems
+.Trashes
+.VolumeIcon.icns
+.com.apple.timemachine.donotpresent
+
+# Directories potentially created on remote AFP share
+.AppleDB
+.AppleDesktop
+Network Trash Folder
+Temporary Items
+.apdisk
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+env/
+venv/
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*,cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+
+# Sphinx documentation
+docs/_build/
+
+# IntelliJ Idea
+.idea
+*.iml
+*.ipr
+*.iws
+
+# PyBuilder
+target/
+
+# Cookiecutter
+output/
+python_boilerplate/
+cookiecutter-pypackage-env/
+
+# vscode settings
+.history/
+*.code-workspace
+
+# Frontend extension
+node_modules/
+.env
+.env.local
+.env.development.local
+.env.test.local
+.env.production.local
+npm-debug.log*
+yarn-debug.log*
+yarn-error.log*
+node.zip
+.vscode/
+.claude/
+.codemate/
@@ -0,0 +1,15 @@
+Apache Software License 2.0
+
+Copyright (c) 2026, Zeyu Huang
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
@@ -0,0 +1,170 @@
+# vLLM-Omni
+
+vLLM-Omni offers a ComfyUI integration on top of its online serving API.
+It can send model inference requests to either a locally running vLLM-Omni service or a remote one.
+
+## Requirement
+
+- Python 3.12 or above
+- [ComfyUI installed](https://docs.comfy.org/installation/system_requirements)
+- [vLLM-Omni installed](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/) on either the same device or another device discoverable via the internet.
+- No need to install additional packages apart from those already required by ComfyUI.
+
+> [!TIP]
+> If you run both ComfyUI and vLLM-Omni on the same device, you can create separate virtual environments and use different Python versions for them.
+
+
+## Installation
+
+Copy this folder to the `custom_nodes` subfolder of your ComfyUI installation. Your directory should look like `ComfyUI/custom_nodes/ComfyUI-vLLM-Omni`.
+
+If you are running ComfyUI during copying, you should restart ComfyUI to load this extension.
+
+> [!TIP]
+> You can use utility websites such as https://download-directory.github.io/ to download a subdirectory of a repo. Also checkout community discussions (e.g., https://stackoverflow.com/questions/7106012/download-a-single-folder-or-directory-from-a-github-repository) for more info.
+
+On the device and virtual environment you run ComfyUI, launch ComfyUI with
+```bash
+cd ComfyUI
+
+# The regular way
+python main.py
+
+# If you are mainly using this node, launch it faster with
+python main.py --cpu
+```
+
+On the device and virtual environment you run vLLM-Omni, start a model service with
+```bash
+vllm serve The_Model_ID_to_Serve --omni --port 8000
+```
+
+Check **ComfyUI's sidebar -> Node Library**. There should be a new folder named **vLLM-Omni**.
+If no, check your shell running the ComfyUI process. There may be some error messages before the line `Import times for custom nodes:` and the line `To see the GUI go to: http://127.0.0.1:8188`.
+
+## Quickstart
+
+This extension offers the following nodes based on the output modalities (at **ComfyUI sidebar -> Node Library**):
+
+- **Generate Image** for text-to-image and image-to-image tasks
+- **Multimodality Comprehension** for multimodality-to-text and multimodality-to-audio tasks
+- **TTS** and **TTS Voice Clone** for TTS tasks
+
+This extension also offers example workflows (at **ComfyUI sidebar -> Templates -> vLLM-Omni**)
+
+> [!INFO]
+> The node UI and feature designs are intended to match vLLM-Omni online serving interfaces. It cannot offer more than what the interfaces support.
+
+To build a simple workflow yourself,
+
+- Drag a generation node onto the canvas.
+- Depending on your need, grab built-in multimedia file loader nodes, such as **image->Load Image**, **image->video->Load Video**, **audio->Load Audio**
+- Depending on your need, grab built-in multimedia file preview nodes, such as **image->Preview Image**, **image->video->Save Video**, **audio->Preview Audio**. For text output, you can install [ComfyUI-Custom-Scripts plugin](https://github.com/pythongosssss/ComfyUI-Custom-Scripts/) and grab its **utils->Show Text 🐍** node.
+- If you want to tune sampling parameters, grab corresponding nodes from **vLLM-Omni-> Sampling Params**.
+    - For multi-stage models, you can connect multiple **AR Sampling Params** and **Diffusion Sampling Params** nodes to a **Multi-Stage Sampling Params List** node, and connect this node to the generation node.
+    - For some multi-stage models like BAGEL, [only one stage's sampling parameters are exposed and tunable via vLLM-Omni's online serving API](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/examples/online_serving/bagel/). Thus, these models are treated as single-stage ones. Please check the vLLM-Omni documentation on how to correctly set each model's sampling parameters.
+    - For multi-stage models where all stages are either autoregression or diffusion, you can also connect only a single Sampling Params node, indicating that this set of sampling parameters will be used for all stages.
+
+**The following features are tested**:
+
+- Single-node workflows for
+    - Multimodal Comprehension (e.g., Qwen Omni, BAGEL)
+    - Text-to-Image Generation (e.g., Qwen-Image)
+    - Image-to-Image Generation (e.g., Qwen-Image-Edit)
+    - TTS (e.g., Qwen TTS, including VoiceDesign, VoiceClone, CustomVoice)
+
+**The following features are not currently tested**. They may work or break. You are welcomed to test it out and offer comments.
+
+- Multi-node workflow that connects multiple model services together.
+
+## Screenshots and Examples
+
+### Multimodal comprehension (e.g., Qwen Omni series, BAGEL)
+
+(Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni Annotated Example**)
+
+<p align="center">
+  <picture>
+    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-comprehension.jpg">
+    <img alt="vLLM-Omni Main Architecture" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-comprehension.jpg" width=55%>
+  </picture>
+</p>
+
+> [!TIP]
+> Although this node enables all-modality input, you should check whether the specific model you host and request for supports the modalities you connect to the node.
+
+You can configure per-stage sampling parameters for multi-stage models.
+
+<p align="center">
+  <picture>
+    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-multi-stage.jpg">
+    <img alt="vLLM-Omni Main Architecture" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-multi-stage.jpg" width=55%>
+  </picture>
+</p>
+
+### Text-to-image and image-to-image generation (e.g., Z-Image-Turbo, Qwen-Image-Edit, BAGEL)
+
+(Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni Image Generation**)
+
+<p align="center">
+  <picture>
+    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-image-generation.jpg">
+    <img alt="vLLM-Omni Main Architecture" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-image-generation.jpg" width=55%>
+  </picture>
+</p>
+
+> [!TIP]
+> The node automatically choose text-to-image or image-to-image API endpoints depending on whether you connect an image input or not.
+
+### TTS (e.g., Qwen TTS series)
+
+(Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni TTS**)
+
+<p align="center">
+  <picture>
+    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-tts.jpg">
+    <img alt="vLLM-Omni Main Architecture" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-tts.jpg" width=55%>
+  </picture>
+</p>
+
+> [!TIP]
+> There is a dedicated node for VoiceClone tasks with reference audio input. Other simple text-to-speech tasks should use the regular TTS node.
+
+## Develop
+
+Follow the [development convention and rules of vLLM-Omni](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/).
+
+## Limitation and Non-Goals
+
+- Single server mode only. No automatic load balancing or failover.
+- Features set is bounded to vLLM-Omni's online service capability, including
+    - The types of models supported in online mode,
+    - The types of sampling parameters supported in the online mode,
+    - The ways to send files (primarily through full-length base64 in JSON payload),
+    - Figuring out errors in the payload (such as unsupported fields by a specific model) if the endpoint does not explicitly return an error,
+    - (The lack of) Authentication
+    - (The lack of) Progress indicator
+
+## Support
+
+If you are new to ComfyUI, please check out [its documentation](https://docs.comfy.org/) for usage instructions.
+
+If you are new to vLLM-Omni, please also check out [its documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/) for usage instructions.
+
+Whenever you find an issue or problem, please
+
+- First find out if this is an upstream limitation of vLLM-Omni's online serving mode, by [checking their documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/examples/).
+- [Open an issue](https://github.com/vllm-project/vllm-omni/issues) that clearly describes this ComfyUI or online service problem.
+
+## Acknowledgements
+
+Features
+
+- https://github.com/dougbtv/comfyui-vllm-omni/ The official reference implementation for ComfyUI integration with vLLM-Omni's DALL-E compatible image generation API.
+- https://github.com/Comfy-Org/ComfyUI/tree/master/comfy_extras ComfyUI's built-in node implementations.
+
+UI/UX design references
+
+- https://github.com/sgl-project/sglang/pull/15271 SGLang Diffusion's official ComfyUI integration for image and video generation.
+- https://github.com/SXQBW/ComfyUI-Qwen-Omni A third party ComfyUI integration for Qwen Omni series.
+- https://github.com/flybirdxx/ComfyUI-Qwen-TTS https://github.com/DarioFT/ComfyUI-Qwen3-TTS Third  party ComfyUI integrations for Qwen TTS series.
@@ -0,0 +1,52 @@
+"""Top-level package for comfyui_vllm_omni."""  # noqa: N999  # This is not a python library intended to be imported
+
+__all__ = [
+    "NODE_CLASS_MAPPINGS",
+    "NODE_DISPLAY_NAME_MAPPINGS",
+    "WEB_DIRECTORY",
+]
+
+__author__ = """Zeyu Huang"""
+__email__ = "11222265+fhfuih@users.noreply.github.com"
+__version__ = "0.0.1"
+
+from .comfyui_vllm_omni.nodes import (
+    VLLMOmniARSampling,
+    VLLMOmniComprehension,
+    VLLMOmniDiffusionSampling,
+    VLLMOmniGenerateImage,
+    VLLMOmniQwenTTSParams,
+    VLLMOmniSamplingParamsList,
+    VLLMOmniTTS,
+    VLLMOmniVoiceClone,
+)
+
+# A dictionary that contains all nodes you want to export with their names
+NODE_CLASS_MAPPINGS = {
+    # === Generation ===
+    "VLLMOmniGenerateImage": VLLMOmniGenerateImage,
+    "VLLMOmniComprehension": VLLMOmniComprehension,
+    "VLLMOmniTTS": VLLMOmniTTS,
+    "VLLMOmniVoiceClone": VLLMOmniVoiceClone,
+    # === Params ===
+    "VLLMOmniARSampling": VLLMOmniARSampling,
+    "VLLMOmniDiffusionSampling": VLLMOmniDiffusionSampling,
+    "VLLMOmniSamplingParamsList": VLLMOmniSamplingParamsList,
+    "VLLMOmniQwenTTSParams": VLLMOmniQwenTTSParams,
+}
+
+# A dictionary that contains the friendly/humanly readable titles for the nodes
+NODE_DISPLAY_NAME_MAPPINGS = {
+    # === Generation ===
+    "VLLMOmniGenerateImage": "Generate Image",
+    "VLLMOmniComprehension": "Multimodality Comprehension",
+    "VLLMOmniTTS": "TTS (Text to Speech)",
+    "VLLMOmniVoiceClone": "TTS Voice Cloning",
+    # === Params ===
+    "VLLMOmniARSampling": "AR Sampling Params",
+    "VLLMOmniDiffusionSampling": "Diffusion Sampling Params",
+    "VLLMOmniSamplingParamsList": "Multi-Stage Sampling Params List",
+    "VLLMOmniQwenTTSParams": "Qwen TTS Params",
+}
+
+WEB_DIRECTORY = "./web"
@@ -0,0 +1 @@
+# noqa: N999   # This is not a python library intended to be imported
@@ -0,0 +1,569 @@
+from typing import Literal, cast
+
+import torch
+from comfy_api.input import AudioInput, VideoInput
+
+from .utils.api_client import VLLMOmniClient
+from .utils.logger import get_logger
+from .utils.models import lookup_model_spec
+from .utils.types import AudioFormat
+from .utils.validators import (
+    add_sampling_parameters_to_stage,
+    validate_model_and_sampling_params_types,
+)
+
+logger = get_logger(__name__)
+
+
+class _VLLMOmniGenerateBase:
+    """Base class for vLLM-Omni generation nodes with shared functionality."""
+
+    CATEGORY = "vLLM-Omni"
+
+    @classmethod
+    def VALIDATE_INPUTS(cls, url, model) -> str | Literal[True]:
+        """
+        Can only validate this model's own input. Cannot check inputs from other nodes.
+        See: https://docs.comfy.org/custom-nodes/backend/server_overview#validate_inputs
+        """
+        if not url:
+            return "URL must not be empty"
+        if not model:
+            return "Model must not be empty"
+        return True
+
+
+class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "url": ("STRING", {"default": "http://localhost:8000/v1"}),
+                "model": ("STRING", {"default": "Tongyi-MAI/Z-Image-Turbo"}),
+                "prompt": ("STRING", {"multiline": True}),
+                "negative_prompt": ("STRING", {"multiline": True, "default": ""}),
+                "width": ("INT", {"default": 512, "min": 64, "max": 2048}),
+                "height": ("INT", {"default": 512, "min": 64, "max": 2048}),
+            },
+            "optional": {
+                "image": ("IMAGE",),
+                "mask": ("MASK",),
+                # "video": ("VIDEO",),
+                # "audio": ("AUDIO",),
+                "sampling_params": ("SAMPLING_PARAMS",),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("image",)
+    FUNCTION = "generate"
+
+    async def generate(
+        self,
+        url: str,
+        model: str,
+        prompt: str,
+        width: int,
+        height: int,
+        negative_prompt: str | None = None,
+        image: torch.Tensor | None = None,
+        mask: torch.Tensor | None = None,
+        audio: AudioInput | None = None,  # Hidden & unused
+        video: VideoInput | None = None,  # Hidden & unused
+        sampling_params: dict | list[dict] | None = None,
+        **kwargs,
+    ):
+        logger.info("Uncaught kwargs: %s", kwargs)
+        logger.debug("Got sampling params: %s", sampling_params)
+        validate_model_and_sampling_params_types(model, sampling_params)
+        if image is None and mask is not None:
+            raise ValueError("Mask input provided without an image input.")
+
+        client = VLLMOmniClient(url)
+
+        spec, pattern = lookup_model_spec(model)
+        is_bagel = pattern is not None and "bagel" in pattern.lower()
+
+        # Prefer DALL-E compatible API for simple (one-stage) diffusion models
+        if (spec is None or spec["stages"] == ["diffusion"]) and not is_bagel:
+            sampling_params = cast(dict | None, sampling_params)
+            if audio is None and image is None and video is None:
+                # No multimodal input --- use DALL-E image generation
+                logger.info("Using DALL-E image generation endpoint")
+                output = await client.generate_image(
+                    model=model,
+                    prompt=prompt,
+                    width=width,
+                    height=height,
+                    negative_prompt=negative_prompt,
+                    sampling_params=sampling_params,
+                )
+                return (output,)
+            elif image is not None and audio is None and video is None:
+                # Image and text input --- use DALL-E image edit
+                logger.info("Using DALL-E image edit endpoint")
+                output = await client.edit_image(
+                    model=model,
+                    prompt=prompt,
+                    image=image,
+                    width=width,
+                    height=height,
+                    negative_prompt=negative_prompt,
+                    mask=mask,
+                    sampling_params=sampling_params,
+                )
+                return (output,)
+
+        logger.info("Using chat completion endpoint")
+        sampling_params = add_sampling_parameters_to_stage(
+            model, sampling_params, "diffusion", width=width, height=height
+        )
+        logger.debug("Edited sampling params: %s", sampling_params)
+
+        output = await client.generate_image_chat_completion(
+            model=model,
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            image=image,
+            audio=audio,
+            video=video,
+            sampling_params=sampling_params,
+        )
+
+        return (output,)
+
+
+class VLLMOmniComprehension(_VLLMOmniGenerateBase):
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "url": ("STRING", {"default": "http://localhost:8000/v1"}),
+                "model": ("STRING", {"default": "Qwen/Qwen2.5-Omni-7B"}),
+                "prompt": ("STRING", {"multiline": True}),
+                "output_text": ("BOOLEAN", {"default": True}),
+                "output_audio": ("BOOLEAN", {"default": True}),
+                "use_audio_in_video": ("BOOLEAN", {"default": True}),
+            },
+            "optional": {
+                "image": ("IMAGE",),
+                "video": ("VIDEO",),
+                "audio": ("AUDIO",),
+                "sampling_params": ("SAMPLING_PARAMS",),
+            },
+        }
+
+    RETURN_TYPES = ("STRING", "AUDIO")
+    RETURN_NAMES = ("text_response", "audio_response")
+    FUNCTION = "generate"
+
+    @classmethod
+    def VALIDATE_INPUTS(cls, url, model, output_text, output_audio) -> str | Literal[True]:  # type: ignore[reportIncompatibleMethodOverride]
+        super_validation = super().VALIDATE_INPUTS(url, model)
+        if isinstance(super_validation, str):
+            return super_validation
+        if not output_text and not output_audio:
+            return "At least one of output_text or output_audio must be True."
+        return True
+
+    async def generate(
+        self,
+        url: str,
+        model: str,
+        prompt: str,
+        image: torch.Tensor | None = None,
+        audio: AudioInput | None = None,
+        video: VideoInput | None = None,
+        sampling_params: dict | list[dict] | None = None,
+        output_text: bool = True,
+        output_audio: bool = True,
+        use_audio_in_video: bool = True,
+        **kwargs,
+    ) -> tuple[str, AudioInput]:
+        logger.info("Uncaught kwargs: %s", kwargs)
+        logger.debug("Got sampling params: %s", sampling_params)
+        validate_model_and_sampling_params_types(model, sampling_params)
+
+        client = VLLMOmniClient(url)
+        spec, pattern = lookup_model_spec(model)
+        is_bagel = pattern is not None and "bagel" in pattern.lower()
+
+        if is_bagel:
+            # A lot of special handlings here...
+            if output_audio:
+                raise ValueError("BAGEL models do not support audio output.")
+            if audio is not None or video is not None:
+                raise ValueError("BAGEL models do not support audio or video input.")
+            (
+                text_response,
+                _,
+            ) = await client.generate_comprehension_chat_completion(
+                model=model,
+                prompt=prompt,
+                image=image,
+                audio=None,
+                video=None,
+                sampling_params=sampling_params,
+                modalities=["text"],
+            )
+        else:
+            modalities = []
+            if output_text:
+                modalities.append("text")
+            if output_audio:
+                modalities.append("audio")
+
+            if use_audio_in_video and video is not None:
+                use_audio_in_video = True
+            else:
+                use_audio_in_video = False
+
+            (
+                text_response,
+                audio,
+            ) = await client.generate_comprehension_chat_completion(
+                model=model,
+                prompt=prompt,
+                image=image,
+                audio=audio,
+                video=video,
+                sampling_params=sampling_params,
+                modalities=modalities,
+                # == extra kwargs ==
+                mm_processor_kwargs={"use_audio_in_video": use_audio_in_video},
+            )
+
+        if text_response is None:
+            text_response = ""
+        if audio is None:
+            channels = 1
+            duration = 1
+            sample_rate = 44100
+            num_samples = int(round(duration * sample_rate))
+            waveform = torch.zeros((1, channels, num_samples), dtype=torch.float32)
+            audio = {"waveform": waveform, "sample_rate": sample_rate}
+
+        return (text_response, audio)
+
+
+class VLLMOmniTTS(_VLLMOmniGenerateBase):
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "url": ("STRING", {"default": "http://localhost:8000/v1"}),
+                "model": (
+                    "STRING",
+                    {"default": "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice"},
+                ),
+                "input": ("STRING", {"multiline": True}),
+                "voice": ("STRING", {"default": "Vivian"}),
+                "response_format": (["mp3", "opus", "aac", "flac", "wav", "pcm"],),
+                "speed": (
+                    "FLOAT",
+                    {"default": 1.0, "min": 0.25, "max": 4.0, "step": 0.01},
+                ),
+            },
+            "optional": {
+                "model_specific_params": ("TTS_PARAMS",),
+            },
+        }
+
+    RETURN_TYPES = ("AUDIO",)
+    RETURN_NAMES = ("audio",)
+    FUNCTION = "generate"
+
+    async def generate(
+        self,
+        url: str,
+        model: str,
+        input: str,
+        voice: str,
+        response_format: AudioFormat,
+        speed: float,
+        model_specific_params: dict | None,
+        **kwargs,
+    ) -> tuple[AudioInput]:
+        logger.info("Got extra kwargs in TTS: %s", kwargs)
+
+        is_qwen_tts = "qwen3-tts" in model.lower()
+        extra_params_type = None if model_specific_params is None else model_specific_params["type"]
+        if not is_qwen_tts and extra_params_type == "qwen-tts":
+            raise ValueError(
+                "You have provided Qwen-specific TTS params."
+                "However, the model appears to not be a Qwen TTS model (no 'Qwen3-TTS' in model name)."
+            )
+
+        combined_params = {**kwargs, **(model_specific_params or {})}
+        combined_params.pop("type", None)  # Internal fields in model_specific_params
+
+        client = VLLMOmniClient(url)
+
+        audio = await client.generate_speech(
+            model=model,
+            input=input,
+            voice=voice,
+            response_format=response_format,
+            speed=speed,
+            **combined_params,
+        )
+        return (audio,)
+
+
+class VLLMOmniVoiceClone(_VLLMOmniGenerateBase):
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "url": ("STRING", {"default": "http://localhost:8000/v1"}),
+                "model": ("STRING", {"default": "Qwen/Qwen3-TTS-12Hz-1.7B-Base"}),
+                "input": ("STRING", {"multiline": True}),
+                "voice": ("STRING", {"default": "Vivian"}),
+                "response_format": (["mp3", "opus", "aac", "flac", "wav", "pcm"],),
+                "speed": (
+                    "FLOAT",
+                    {"default": 1.0, "min": 0.25, "max": 4.0, "step": 0.01},
+                ),
+                "ref_audio": ("AUDIO",),
+                "ref_text": ("STRING", {"multiline": True}),
+                "x_vector_only_mode": ("BOOLEAN", {"default": False}),
+            },
+            "optional": {
+                "model_specific_params": ("TTS_PARAMS",),
+            },
+        }
+
+    RETURN_TYPES = ("AUDIO",)
+    RETURN_NAMES = ("audio",)
+    FUNCTION = "generate"
+
+    async def generate(
+        self,
+        url: str,
+        model: str,
+        input: str,
+        voice: str,
+        response_format: AudioFormat,
+        speed: float,
+        ref_audio: AudioInput,
+        ref_text: str,
+        x_vector_only_mode: bool,
+        model_specific_params: dict | None,
+        **kwargs,
+    ):
+        is_qwen_tts = "qwen3-tts" in model.lower()
+        extra_params_type = None if model_specific_params is None else model_specific_params["type"]
+        if not is_qwen_tts and extra_params_type == "qwen-tts":
+            raise ValueError(
+                "You have provided Qwen-specific TTS params."
+                "However, the model appears to not be a Qwen TTS model (no 'Qwen3-TTS' in model name)."
+            )
+
+        combined_params = {
+            "ref_audio": ref_audio,
+            "ref_text": ref_text,
+            "x_vector_only_mode": x_vector_only_mode,
+            **kwargs,
+            **(model_specific_params or {}),
+        }
+        combined_params.pop("type", None)  # Internal fields in model_specific_params
+
+        client = VLLMOmniClient(url)
+
+        audio = await client.generate_speech(
+            model=model,
+            input=input,
+            voice=voice,
+            response_format=response_format,
+            speed=speed,
+            **combined_params,
+        )
+        return (audio,)
+
+
+class VLLMOmniARSampling:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "max_tokens": ("INT", {"default": 100, "min": 1, "max": 10000}),
+                "temperature": (
+                    "FLOAT",
+                    {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.01},
+                ),
+                "top_p": (
+                    "FLOAT",
+                    {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01},
+                ),
+                "repetition_penalty": (
+                    "FLOAT",
+                    {"default": 1.0, "min": 0.0, "max": 5.0, "step": 0.01},
+                ),
+                # === Put seed at last. ===
+                # Whenever a field named "seed" is present, ComfyUI adds another field called "control after generate"
+                "seed": (
+                    "INT",
+                    {
+                        "default": -1,
+                        "min": -1,
+                        "step": 1,
+                        "tooltip": "-1 means to not provide a seed.",
+                    },
+                ),
+            }
+        }
+
+    RETURN_TYPES = ("SAMPLING_PARAMS",)
+    RETURN_NAMES = ("AR sampling params",)
+    FUNCTION = "get_params"
+    CATEGORY = "vLLM-Omni/Sampling Params"
+
+    def get_params(self, seed, **kwargs):
+        params = {
+            "type": "autoregression",  # for internal use, removed before sending the request
+            **kwargs,
+        }
+        if seed >= 0:
+            params["seed"] = seed
+        return (params,)
+
+
+class VLLMOmniDiffusionSampling:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "n": (
+                    "INT",
+                    {
+                        "default": 1,
+                        "min": 0,
+                        "max": 10,
+                        "step": 1,
+                        "tooltip": "Number of images to generate",
+                    },
+                ),
+                "num_inference_steps": (
+                    "INT",
+                    {
+                        "default": 50,
+                        "min": 1,
+                        "max": 1000,
+                        "tooltip": "Number of denoising steps (higher = better quality, slower).",
+                    },
+                ),
+                "guidance_scale": (
+                    "FLOAT",
+                    {
+                        "default": 7.5,
+                        "min": 0.0,
+                        "max": 20.0,
+                        "step": 0.1,
+                        "tooltip": "Classifier-free guidance scale (higher = more prompt adherence).",
+                    },
+                ),
+                "true_cfg_scale": (
+                    "FLOAT",
+                    {
+                        "default": 1.0,
+                        "min": 0.0,
+                        "max": 20.0,
+                        "step": 0.5,
+                        "tooltip": "True CFG scale for advanced control (model-specific).",
+                    },
+                ),
+                "vae_use_slicing": (
+                    "BOOLEAN",
+                    {
+                        "default": False,
+                        "tooltip": "Enable VAE slicing for reduced memory usage (slight quality trade-off)",
+                    },
+                ),
+                # === Put seed at last. ===
+                # Whenever a field named "seed" is present, ComfyUI adds another field called "control after generate"
+                "seed": (
+                    "INT",
+                    {
+                        "default": -1,
+                        "min": -1,
+                        "step": 1,
+                        "tooltip": "-1 means to not provide a seed.",
+                    },
+                ),
+            }
+        }
+
+    RETURN_TYPES = ("SAMPLING_PARAMS",)
+    RETURN_NAMES = ("diffusion sampling params",)
+    FUNCTION = "get_params"
+    CATEGORY = "vLLM-Omni/Sampling Params"
+
+    def get_params(self, seed, **kwargs):
+        params = {
+            "type": "diffusion",  # for internal use, removed before sending the request
+            **kwargs,
+        }
+        if seed >= 0:
+            params["seed"] = seed
+        return (params,)
+
+
+class VLLMOmniSamplingParamsList:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "param1": ("SAMPLING_PARAMS",),
+            },
+            "optional": {
+                "param2": ("SAMPLING_PARAMS",),
+                "param3": ("SAMPLING_PARAMS",),
+            },
+        }
+
+    RETURN_TYPES = ("SAMPLING_PARAMS",)
+    RETURN_NAMES = ("param list",)
+    FUNCTION = "aggregate"
+    CATEGORY = "vLLM-Omni/Sampling Params"
+
+    def aggregate(self, param1: dict, param2: dict | None = None, param3: dict | None = None):
+        for i, p in enumerate((param1, param2, param3)):
+            if isinstance(p, list):
+                raise ValueError(
+                    f"Input {i} is a Multi-Stage Sampling Params List. "
+                    f"Expected a single sampling parameters node (either AR or Diffusion)."
+                )
+
+        params = [param1]
+        if param2 is not None:
+            params.append(param2)
+        if param3 is not None:
+            params.append(param3)
+        return (params,)
+
+
+class VLLMOmniQwenTTSParams:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "task_type": (
+                    ["CustomVoice", "VoiceDesign", "Base"],
+                    {"default": "CustomVoice"},
+                ),
+                "language": (
+                    ["Auto", "Chinese", "English", "Japanese", "Korean"],
+                    {"default": "Auto"},
+                ),
+                "instructions": ("STRING", {"multiline": True}),
+                "max_new_tokens": ("INT", {"default": 2048, "min": 1}),
+            }
+        }
+
+    RETURN_TYPES = ("TTS_PARAMS",)
+    RETURN_NAMES = ("Qwen TTS params",)
+    FUNCTION = "get_params"
+    CATEGORY = "vLLM-Omni/TTS Params"
+
+    def get_params(self, **kwargs):
+        return ({"type": "qwen-tts", **kwargs},)
@@ -0,0 +1,473 @@
+"""
+An high-level API client adapter that forwards ComfyUI inputs to vLLM-Omni's REST API,
+and transforms the API responses back to ComfyUI formats.
+
+The image generation part is derived from dougbtv/comfyui-vllm-omni by Doug (@dougbtv).
+Original source at https://github.com/dougbtv/comfyui-vllm-omni, distributed under the MIT License.
+"""
+
+from typing import Any
+
+import aiohttp
+import av.error
+import torch
+from comfy_api.input import AudioInput, VideoInput
+
+from .format import (
+    audio_to_base64,
+    base64_to_audio,
+    base64_to_image_tensor,
+    bytes_to_audio,
+    image_tensor_to_base64,
+    image_tensor_to_png_bytes,
+    video_to_base64,
+)
+from .logger import get_logger, pretty_printer
+from .models import lookup_model_spec
+from .types import AudioFormat
+
+logger = get_logger(__name__)
+
+
+class VLLMOmniClient:
+    def __init__(self, base_url: str, timeout: float = 300.0):
+        self.base_url = base_url
+        self.timeout = aiohttp.ClientTimeout(total=timeout)
+
+    async def generate_image(
+        self,
+        *,
+        model: str,
+        prompt: str,
+        width: int,
+        height: int,
+        negative_prompt: str | None = None,
+        sampling_params: dict | None = None,
+    ) -> torch.Tensor:
+        """Run text-to-image generation via DALLE API"""
+        await self._check_model_exist(model)
+
+        size = f"{width}x{height}"
+        payload = {
+            "model": model,
+            "prompt": prompt,
+            "size": size,
+            "response_format": "b64_json",
+        }
+        if negative_prompt:
+            payload["negative_prompt"] = negative_prompt
+        if sampling_params is not None:
+            # Only select specific sampling params
+            for k in (
+                "n",
+                "num_inference_steps",
+                "guidance_scale",
+                "true_cfg_scale",
+                "vae_use_slicing",
+            ):
+                if k in sampling_params and sampling_params[k] is not None:
+                    payload[k] = sampling_params[k]
+        logger.debug("img gen payload: %s", payload)
+
+        url = self.base_url + "/images/generations"
+        async with aiohttp.ClientSession(timeout=self.timeout) as session:
+            try:
+                async with session.post(
+                    url,
+                    json=payload,
+                    headers={"Content-Type": "application/json"},
+                ) as response:
+                    if not response.ok:
+                        error_text = await response.text()
+                        raise (ValueError if response.status < 500 else RuntimeError)(
+                            f"vLLM-Omni API returned status {response.status}: {error_text}"
+                        )
+
+                    try:
+                        data = await response.json()
+                    except aiohttp.ContentTypeError as e:
+                        raise RuntimeError(f"Invalid JSON response from vLLM-Omni: {e}")
+                    if "data" not in data:
+                        raise RuntimeError("API response missing 'data' field - expected OpenAI DALL-E format")
+                    if not data["data"]:
+                        raise RuntimeError("API returned empty data array")
+
+                    image_tensors = []
+                    for idx, img in enumerate(data["data"]):
+                        if "b64_json" not in img:
+                            raise RuntimeError(f"API returned image #{idx} without 'b64_json' field")
+                        base64_str = img["b64_json"]
+                        tensor = base64_to_image_tensor(base64_str)
+                        image_tensors.append(tensor)
+                        logger.debug("Image #%d has shape %s", idx, tensor.shape)
+
+                    batch_tensor = torch.stack(image_tensors, dim=0)
+                    logger.debug("batch_tensor output has shape: %s", batch_tensor.shape)
+                    return batch_tensor
+
+            except aiohttp.ClientError as e:
+                raise RuntimeError(f"Network error connecting to vLLM-Omni at {url}: {e}")
+
+    async def edit_image(
+        self,
+        *,
+        model: str,
+        prompt: str,
+        image: torch.Tensor,
+        width: int,
+        height: int,
+        negative_prompt: str | None = None,
+        mask: torch.Tensor | None = None,
+        sampling_params: dict | None = None,
+    ) -> torch.Tensor:
+        """Run image editing via DALLE API"""
+        await self._check_model_exist(model)
+
+        size = f"{width}x{height}"
+        image_filename = "image.png"  # Required for multipart form
+        form = aiohttp.FormData()
+        form.add_field("model", model)
+        form.add_field(
+            "image",
+            image_tensor_to_png_bytes(image, image_filename),
+            filename=image_filename,
+            content_type="image/png",
+        )
+        form.add_field("prompt", prompt)
+        form.add_field("size", size)
+        if negative_prompt:
+            form.add_field("negative_prompt", negative_prompt)
+        if sampling_params is not None:
+            # Only select specific sampling params
+            for k in ("n", "num_inference_steps", "guidance_scale", "true_cfg_scale"):
+                if k in sampling_params and sampling_params[k] is not None:
+                    form.add_field(k, str(sampling_params[k]))
+        if mask is not None:
+            mask_filename = "mask.png"
+            form.add_field(
+                "mask",
+                image_tensor_to_png_bytes(mask, mask_filename),
+                filename=mask_filename,
+                content_type="image/png",
+            )
+
+        url = self.base_url + "/images/edits"
+        async with aiohttp.ClientSession(timeout=self.timeout) as session:
+            try:
+                async with session.post(url, data=form) as response:
+                    if not response.ok:
+                        error_text = await response.text()
+                        raise (ValueError if response.status < 500 else RuntimeError)(
+                            f"vLLM-Omni API returned status {response.status}: {error_text}"
+                        )
+
+                    try:
+                        data = await response.json()
+                    except aiohttp.ContentTypeError as e:
+                        raise RuntimeError(f"Invalid JSON response from vLLM-Omni: {e}")
+
+                    if "data" not in data:
+                        raise RuntimeError("API response missing 'data' field - expected OpenAI DALL-E format")
+                    if not data["data"]:
+                        raise RuntimeError("API returned empty data array")
+
+                    image_tensors = []
+                    for idx, img in enumerate(data["data"]):
+                        if "b64_json" not in img:
+                            raise RuntimeError(f"API returned image #{idx} without 'b64_json' field")
+                        base64_str = img["b64_json"]
+                        tensor = base64_to_image_tensor(base64_str)
+                        image_tensors.append(tensor)
+
+                    return torch.stack(image_tensors, dim=0)
+
+            except aiohttp.ClientError as e:
+                raise RuntimeError(f"Network error connecting to vLLM-Omni at {url}: {e}")
+
+    async def generate_image_chat_completion(
+        self,
+        *,
+        model: str,
+        prompt: str,
+        negative_prompt: str | None = None,
+        image: torch.Tensor | None = None,
+        audio: AudioInput | None = None,
+        video: VideoInput | None = None,
+        sampling_params: dict | list[dict] | None = None,
+    ) -> torch.Tensor:
+        payload = VLLMOmniClient._prepare_chat_completion_messages(
+            model=model,
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            image=image,
+            audio=audio,
+            video=video,
+            sampling_params=sampling_params,
+            modalities=["image"],
+        )
+        choices = await self._generate_base_chat_completion(model, payload)
+
+        image_tensors = []
+        for idx, img_content in enumerate(choices[0]["message"]["content"]):
+            base64_str = img_content.get("image_url", {}).get("url", "")
+            if not base64_str:
+                raise RuntimeError(f"API returned image #{idx} without image url")
+            tensor = base64_to_image_tensor(base64_str)
+            image_tensors.append(tensor)
+
+        return torch.stack(image_tensors, dim=0)
+
+    async def generate_comprehension_chat_completion(
+        self,
+        *,
+        model: str,
+        prompt: str,
+        image: torch.Tensor | None = None,
+        audio: AudioInput | None = None,
+        video: VideoInput | None = None,
+        sampling_params: dict | list[dict] | None = None,
+        modalities: list[str] = ["text", "audio"],
+        **extra_body,
+    ) -> tuple[str | None, AudioInput | None]:
+        # Response may contain two choices: one with text, one with audio
+        payload = VLLMOmniClient._prepare_chat_completion_messages(
+            model=model,
+            prompt=prompt,
+            negative_prompt=None,
+            image=image,
+            audio=audio,
+            video=video,
+            sampling_params=sampling_params,
+            modalities=modalities,
+            **extra_body,
+        )
+
+        choices = await self._generate_base_chat_completion(model, payload)
+        text_response = None
+        audio_base64 = None
+        for choice in choices:
+            try:
+                text_response = choice["message"]["content"]
+            except (KeyError, TypeError):
+                # Either this case (text response) or the audio response case will be hit. Checking None's later.
+                pass
+            try:
+                audio_base64 = choice["message"]["audio"]["data"]
+            except (KeyError, TypeError):
+                # Either this case (text response) or the audio response case will be hit. Checking None's later.
+                pass
+        if audio_base64 is None and text_response is None:
+            raise RuntimeError(
+                "API response missing both '.message.audio' and 'message.content' fields."
+                f"The choices object is {choices}"
+            )
+        if audio_base64 is not None:
+            audio = base64_to_audio(audio_base64)
+            logger.debug(
+                "audio sample rate %d, audio shape %s, duration in second %f",
+                audio["sample_rate"],
+                audio["waveform"].shape,
+                audio["waveform"].shape[2] / audio["sample_rate"],
+            )
+        else:
+            audio = None
+        return text_response, audio
+
+    async def generate_speech(
+        self,
+        *,
+        model: str,
+        input: str,
+        voice: str,
+        response_format: AudioFormat,
+        speed: float,
+        **extra_params,
+    ) -> AudioInput:
+        await self._check_model_exist(model)
+
+        ref_audio: AudioInput | None = extra_params.pop("ref_audio", None)
+
+        payload = {
+            "model": model,
+            "input": input,
+            "voice": voice,
+            "response_format": response_format,
+            "speed": speed,
+            **extra_params,
+        }
+
+        if ref_audio is not None:
+            audio_base64 = audio_to_base64(ref_audio)
+            payload["ref_audio"] = audio_base64
+
+        logger.debug("Omni TTS payload: %s", pretty_printer.pformat(payload))
+
+        url = self.base_url + "/audio/speech"
+        async with aiohttp.ClientSession(timeout=self.timeout) as session:
+            try:
+                async with session.post(
+                    url,
+                    json=payload,
+                    headers={"Content-Type": "application/json"},
+                ) as response:
+                    if not response.ok:
+                        error_text = await response.text()
+                        raise (ValueError if response.status < 500 else RuntimeError)(
+                            f"vLLM-Omni API returned status {response.status}: {error_text}"
+                        )
+
+                    try:
+                        audio_bytes = await response.read()
+                    except aiohttp.ContentTypeError as e:
+                        raise RuntimeError(f"Invalid JSON response from vLLM-Omni: {e}")
+
+                    try:
+                        audio = bytes_to_audio(audio_bytes)
+                    except av.error.InvalidDataError as e:
+                        raise ValueError(
+                            f"Invalid audio data received from vLLM-Omni: {e}"
+                            "Check if you have input unsupported arguments (such as 'voice')"
+                        )
+                    return audio
+
+            except aiohttp.ClientError as e:
+                raise RuntimeError(f"Network error connecting to vLLM-Omni at {url}: {e}")
+
+    async def _generate_base_chat_completion(self, model: str, payload: dict[str, Any]) -> list[dict[str, Any]]:
+        logger.debug("Omni payload: %s", pretty_printer.pformat(payload))
+        await self._check_model_exist(model)
+
+        url = self.base_url + "/chat/completions"
+        async with aiohttp.ClientSession(timeout=self.timeout) as session:
+            try:
+                async with session.post(
+                    url,
+                    json=payload,
+                    headers={"Content-Type": "application/json"},
+                ) as response:
+                    if not response.ok:
+                        error_text = await response.text()
+                        raise (ValueError if response.status < 500 else RuntimeError)(
+                            f"vLLM-Omni API returned status {response.status}: {error_text}"
+                        )
+
+                    try:
+                        data = await response.json()
+                    except aiohttp.ContentTypeError as e:
+                        raise RuntimeError(f"Invalid JSON response from vLLM-Omni: {e}")
+
+                    logger.debug(
+                        "chat completion response: %s",
+                        pretty_printer.pformat(data),
+                    )
+
+                    try:
+                        return data["choices"]
+                    except (KeyError, TypeError):
+                        raise RuntimeError("Invalid JSON response from vLLM-Omni: missing 'choices' field")
+
+            except aiohttp.ClientError as e:
+                raise RuntimeError(f"Network error connecting to vLLM-Omni at {self.base_url}: {e}")
+
+    async def _check_model_exist(self, model: str):
+        url = self.base_url + "/models"
+        async with aiohttp.ClientSession(timeout=self.timeout) as session:
+            try:
+                async with session.get(
+                    url,
+                    headers={"Content-Type": "application/json"},
+                ) as response:
+                    if not response.ok:
+                        error_text = await response.text()
+                        raise (ValueError if response.status < 500 else RuntimeError)(
+                            f"vLLM-Omni API returned status {response.status} "
+                            f"when getting hosted model list: {error_text}"
+                        )
+
+                    try:
+                        data = await response.json()
+                    except aiohttp.ContentTypeError as e:
+                        raise RuntimeError(f"Invalid JSON response when getting hosted model list from vLLM-Omni: {e}")
+
+            except aiohttp.ClientError as e:
+                raise RuntimeError(f"Network error connecting to vLLM-Omni at {self.base_url}: {e}")
+        try:
+            model_list = data["data"]
+            model_found = next((True for m in model_list if m["id"] == model), False)
+        except (KeyError, TypeError):
+            raise RuntimeError(f"Invalid JSON response of the hosted model list: {data}")
+
+        if not model_found:
+            raise ValueError(f"Model {model} not served at {self.base_url}.")
+
+    @staticmethod
+    def _prepare_chat_completion_messages(
+        *,
+        model: str,
+        prompt: str,
+        negative_prompt: str | None,
+        image: torch.Tensor | None = None,
+        audio: AudioInput | None = None,
+        video: VideoInput | None = None,
+        sampling_params: dict | list[dict] | None = None,
+        modalities: list[str] | None = None,  # diffusion don't have this field
+        **extra_body,
+    ):
+        message_content: list[dict] = [{"type": "text", "text": prompt}]
+        if image is not None:
+            message_content.append(
+                {
+                    "type": "image_url",
+                    "image_url": {"url": image_tensor_to_base64(image)},
+                }
+            )
+        if audio is not None:
+            message_content.append({"type": "audio_url", "audio_url": {"url": audio_to_base64(audio)}})
+        if video is not None:
+            message_content.append({"type": "video_url", "video_url": {"url": video_to_base64(video)}})
+        messages = [{"role": "user", "content": message_content}]
+
+        combined_extra_body: dict[str, Any] = {}
+        if sampling_params is not None:
+            spec, _ = lookup_model_spec(model)
+            is_single_sampling_param = isinstance(sampling_params, dict) or len(sampling_params) == 1
+
+            # Exclude internal key
+            if isinstance(sampling_params, dict):
+                sampling_params = {k: v for k, v in sampling_params.items() if k != "type"}
+            else:
+                sampling_params = [{k: v for k, v in sp.items() if k != "type"} for sp in sampling_params]
+
+            if (spec is None and is_single_sampling_param) or (spec is not None and spec["stages"] == ["diffusion"]):
+                # Diffusion format: extra_body directly contains sampling params.
+                # Validation should have taken care of matching sampling params' types.
+                # * Use this mode if the model is a simple one-stage diffusion model.
+                # * Fallback to this mode if model is not registered and a single sampling param is provided.
+                sampling_params = sampling_params if isinstance(sampling_params, dict) else sampling_params[0]
+                combined_extra_body: dict[str, Any] = {**sampling_params}
+                if "n" in combined_extra_body:
+                    combined_extra_body["num_outputs_per_prompt"] = combined_extra_body["n"]
+                    del combined_extra_body["n"]
+            else:
+                # Use AR style payload, extra_body has a sampling_params_list field
+                combined_extra_body: dict[str, Any] = {"sampling_params_list": sampling_params}
+
+        if negative_prompt:
+            combined_extra_body["negative_prompt"] = negative_prompt
+
+        if extra_body:
+            combined_extra_body.update(extra_body)
+
+        payload: dict[str, Any] = {"messages": messages, "model": model}
+        if combined_extra_body:
+            payload["extra_body"] = combined_extra_body
+        if modalities:
+            payload["modalities"] = modalities
+
+        spec, _ = lookup_model_spec(model)
+        if spec:
+            preprocessor = spec.get("payload_preprocessor", None)
+            if preprocessor is not None:
+                payload = preprocessor(payload)
+
+        return payload
@@ -0,0 +1,230 @@
+"""Image/tensor format helpers.
+
+The image generation part is derived from dougbtv/comfyui-vllm-omni by Doug (@dougbtv).
+Original source at https://github.com/dougbtv/comfyui-vllm-omni, distributed under the MIT License.
+"""
+
+import base64
+import mimetypes
+from io import BytesIO
+
+import av
+import numpy as np
+import torch
+from comfy_api.input import AudioInput, VideoInput
+from comfy_extras import nodes_audio
+from PIL import Image
+
+from .logger import get_logger
+
+logger = get_logger(__name__)
+
+
+def base64_to_image_tensor(base64_str: str, mode: str = "RGB") -> torch.Tensor:
+    """
+    Convert base64-encoded image to ComfyUI image tensor.
+
+    Args:
+        base64_str: Base64-encoded image string
+        mode: PIL image mode (default RGB for transparency support)
+
+    Returns:
+        torch.Tensor with shape (1, H, W, C) in float32 [0, 1] range
+
+    Raises:
+        ValueError: If base64 string is invalid or image cannot be decoded
+    """
+    if base64_str.startswith("data:image"):
+        _, base64_str = base64_str.split(",", 1)
+
+    try:
+        # Decode base64 to bytes
+        image_bytes = base64.b64decode(base64_str)
+    except Exception as e:
+        raise ValueError(f"Invalid base64 string: {e}")
+
+    # Create BytesIO object for PIL
+    image_bytesio = BytesIO(image_bytes)
+
+    # Open with PIL and convert to desired mode
+    try:
+        pil_image = Image.open(image_bytesio)
+        pil_image = pil_image.convert(mode)
+    except Exception as e:
+        raise RuntimeError(f"Failed to open image: {e}")
+
+    image_array = np.asarray(pil_image).astype(np.float32) / 255.0
+    image_tensor = torch.from_numpy(image_array)
+    return image_tensor
+
+
+def image_tensor_to_png_bytes(tensor: torch.Tensor, filename: str = "image.png") -> BytesIO:
+    """
+    Convert ComfyUI image tensor to PNG BytesIO for multipart upload.
+
+    This function converts a ComfyUI IMAGE tensor to a PNG-encoded BytesIO object
+    suitable for multipart/form-data upload. The BytesIO object has its .name
+    attribute set, which is required by aiohttp for file uploads.
+
+    Args:
+        tensor: ComfyUI IMAGE tensor with shape (B, H, W, C), dtype float32, range [0, 1]
+        filename: Name attribute to set on BytesIO (default: "image.png")
+
+    Returns:
+        BytesIO object containing PNG-encoded image with .name attribute set
+
+    Raises:
+        ValueError: If tensor format is invalid (not 4D, wrong dtype, etc.)
+    """
+    if tensor.ndim != 4:
+        raise ValueError(f"Expected 4D tensor with shape (B, H, W, C), got {tensor.ndim}D tensor")
+
+    image_tensor = tensor[0]  # Shape: (H, W, C)
+    image_np = (image_tensor.cpu().numpy() * 255.0).astype(np.uint8)
+    pil_image = Image.fromarray(image_np)
+
+    # Save to BytesIO as image file
+    img_bytes = BytesIO()
+    # Set name attribute (required for multipart upload and mimetype detection)
+    img_bytes.name = filename
+    try:
+        pil_image.save(img_bytes)
+    except Exception as e:
+        raise RuntimeError(f"Failed to save image as file: {e}")
+
+    # Reset position to beginning
+    img_bytes.seek(0)
+
+    return img_bytes
+
+
+def image_tensor_to_base64(tensor: torch.Tensor, filename: str = "image.png") -> str:
+    """
+    Convert ComfyUI image tensor to base64-encoded image string.
+
+    Args:
+        tensor: ComfyUI IMAGE tensor with shape (B, H, W, C), dtype float32, range [0, 1]
+        filename: Name attribute to set on BytesIO (default: "image.png")
+        format: File format of the output image file buffer (default: "PNG")
+
+    Returns:
+        Base64-encoded image string
+
+    Raises:
+        ValueError: If tensor format is invalid (not 4D, wrong dtype, etc.)
+    """
+    img_bytes = image_tensor_to_png_bytes(tensor, filename)
+    img_bytes.seek(0)
+    byte_data = img_bytes.read()
+    base64_str = base64.b64encode(byte_data).decode("utf-8")
+    mime_type = mimetypes.guess_type(filename)[0] or "application/octet-stream"
+    return f"data:{mime_type};base64,{base64_str}"
+
+
+def video_to_bytes(video: VideoInput, filename: str = "video.mp4") -> BytesIO:
+    output_buffer = BytesIO()
+    output_buffer.name = filename
+    video.save_to(output_buffer)
+    output_buffer.seek(0)
+    return output_buffer
+
+
+def video_to_base64(video: VideoInput, filename: str = "video.mp4") -> str:
+    video_buffer = video_to_bytes(video, filename)
+    video_buffer.seek(0)
+    byte_data = video_buffer.read()
+    base64_str = base64.b64encode(byte_data).decode("utf-8")
+    mime_type = mimetypes.guess_type(filename)[0] or "application/octet-stream"
+    return f"data:{mime_type};base64,{base64_str}"
+
+
+def audio_to_bytes(audio: AudioInput, filename: str = "audio.mp3", quality: str = "128k") -> BytesIO:
+    waveform = audio["waveform"][0]  # Shape: (C, T)
+    sample_rate = audio["sample_rate"]
+    format = filename.rsplit(".", maxsplit=1)[1]
+    layout = "mono" if waveform.shape[0] == 1 else "stereo"
+
+    output_buffer = BytesIO()
+    output_buffer.name = filename
+    output_container = av.open(output_buffer, mode="w", format=format)
+    if format == "opus":
+        out_stream = output_container.add_stream("libopus", rate=sample_rate, layout=layout)
+        if quality == "64k":
+            out_stream.bit_rate = 64000  # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py
+        elif quality == "96k":
+            out_stream.bit_rate = 96000  # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py
+        elif quality == "128k":
+            out_stream.bit_rate = 128000  # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py
+        elif quality == "192k":
+            out_stream.bit_rate = 192000  # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py
+        elif quality == "320k":
+            out_stream.bit_rate = 320000  # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py
+    elif format == "mp3":
+        out_stream = output_container.add_stream("libmp3lame", rate=sample_rate, layout=layout)
+        if quality == "V0":
+            out_stream.codec_context.qscale = 1  # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py
+        elif quality == "128k":
+            out_stream.bit_rate = 128000  # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py
+        elif quality == "320k":
+            out_stream.bit_rate = 320000  # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py
+    else:  # format == "flac":
+        out_stream = output_container.add_stream("flac", rate=sample_rate, layout=layout)
+
+    frame = av.AudioFrame.from_ndarray(
+        waveform.movedim(0, 1).reshape(1, -1).float().numpy(),
+        format="flt",
+        layout=layout,
+    )
+    frame.sample_rate = sample_rate
+    frame.pts = 0
+    output_container.mux(out_stream.encode(frame))  # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py
+    # Flush encoder
+    output_container.mux(out_stream.encode(None))  # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py
+    output_container.close()
+    output_buffer.seek(0)
+
+    return output_buffer
+
+
+def audio_to_base64(audio: AudioInput, filename: str = "audio.mp3", quality: str = "128k") -> str:
+    audio_buffer = audio_to_bytes(audio, filename, quality)
+    audio_buffer.seek(0)
+    byte_data = audio_buffer.read()
+    base64_str = base64.b64encode(byte_data).decode("utf-8")
+    mime_type = mimetypes.guess_type(filename)[0] or "application/octet-stream"
+    return f"data:{mime_type};base64,{base64_str}"
+
+
+def bytes_to_audio(audio_bytes: bytes) -> AudioInput:
+    """
+    Convert audio bytes to ComfyUI audio tensor.
+
+    Args:
+        audio_bytes: Audio file bytes
+    Returns:
+        torch.Tensor with shape (B, C, T) in float32 range [-1, 1]
+    """
+    audio_buffer = BytesIO(audio_bytes)
+    waveform, sample_rate = nodes_audio.load(audio_buffer)  # type: ignore # Although expect string argument, it calls av.open underneath, which supports BytesIO (file-like)
+    return {"waveform": waveform.unsqueeze(0), "sample_rate": sample_rate}
+
+
+def base64_to_audio(base64_str: str) -> AudioInput:
+    """
+    Convert base64-encoded audio to ComfyUI audio tensor.
+
+    Args:
+        base64_str: Base64-encoded audio string
+    Returns:
+        torch.Tensor with shape (B, C, T) in float32 range [-1, 1]
+    """
+    if base64_str.startswith("data:audio"):
+        _, base64_str = base64_str.split(",", 1)
+
+    try:
+        # Decode base64 to bytes
+        audio_bytes = base64.b64decode(base64_str)
+    except Exception as e:
+        raise ValueError(f"Invalid base64 string: {e}")
+
+    return bytes_to_audio(audio_bytes)
@@ -0,0 +1,125 @@
+"""Centralized logger configuration for vLLM-Omni ComfyUI."""
+
+import logging
+import pprint
+import sys
+from typing import Any
+
+
+def get_logger(name: str) -> logging.Logger:
+    """
+    Get or create a logger with proper formatting.
+
+    Args:
+        name: Logger name (typically __name__ of the calling module)
+
+    Returns:
+        Configured logger instance
+    """
+    logger = logging.getLogger(name)
+
+    # Only configure if not already configured
+    if not logger.handlers:
+        logger.setLevel(logging.DEBUG)
+
+        # Create console handler
+        handler = logging.StreamHandler(sys.stdout)
+        handler.setLevel(logging.INFO)
+
+        # Create formatter
+        formatter = logging.Formatter(
+            fmt="(ComfyUI-vLLM-Omni) [%(levelname)s] %(asctime)s [%(filename)s:%(lineno)s] %(message)s",
+            datefmt="%Y-%m-%d %H:%M:%S",
+        )
+        handler.setFormatter(formatter)
+
+        # Add handler to logger
+        logger.addHandler(handler)
+
+        # Prevent propagation to root logger
+        logger.propagate = False
+
+    return logger
+
+
+class OmitBase64PrettyPrinter(pprint.PrettyPrinter):
+    """
+    A PrettyPrinter that redacts specific field names with '...'
+    wherever they appear in nested structures.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def _format(self, obj: Any, stream, indent: int, allowance: int, context, level: int) -> None:
+        # Check if this is a dict with redacted keys
+        if isinstance(obj, dict):
+            # Create a copy with redacted values
+            display_obj = {}
+            for key, value in obj.items():
+                if key == "data" or key == "url":
+                    if value.startswith("data:"):
+                        base64_header = value.split(",", 1)[0]
+                        display_obj[key] = f"{base64_header},***"
+                    elif value.startswith("http://") or value.startswith("https://"):
+                        display_obj[key] = value
+                    elif len(value) > 10:
+                        display_obj[key] = f"{value[:10]}***"
+                    else:
+                        display_obj[key] = value
+                else:
+                    display_obj[key] = value
+            obj = display_obj
+
+        # Handle list/tuple/set containing dicts that might have redacted keys
+        # (pprint will recursively call _format on nested items, so this
+        # handles arbitrary nesting automatically)
+
+        super()._format(obj, stream, indent, allowance, context, level)
+
+
+pretty_printer = OmitBase64PrettyPrinter()
+
+
+# ========== PPrint EXAMPLES ==========
+
+if __name__ == "__main__":
+    data = {
+        "messages": [
+            {
+                "role": "system",
+                "content": (
+                    "You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group,"
+                    "capable of perceiving auditory and visual inputs, as well as generating text and speech."
+                ),
+            },
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "text",
+                        "text": "What sound is it, and what is the drawing about?",
+                    },
+                    {
+                        "type": "text",
+                        "text": "What sound is it, and what is the drawing about?",
+                    },
+                    {
+                        "type": "image_url",
+                        "image_url": {"url": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAwAAAA"},
+                    },
+                    {
+                        "type": "audio_url",
+                        "audio_url": {"url": "data:audio/mpeg;base64,SUQzBAAAAAAAIlRTU0UAAAAOAAADT="},
+                    },
+                ],
+            },
+        ],
+        "extra_body": {"mm_processor_kwargs": {"use_audio_in_video": False}},
+        "modalities": ["text"],
+    }
+
+    # Create printer that redacts 'password' and 'token' fields
+
+    print("\nRedactingPrettyPrinter (hides secrets):")
+    pretty_printer.pprint(data)
@@ -0,0 +1,119 @@
+import re
+
+from .types import Modality, ModelMode, Spec
+
+
+def _bagel_payload_preprocessor(payload: dict) -> dict:
+    try:
+        for message in payload["messages"]:
+            for content in message["content"]:
+                if content["type"] == "text":
+                    content["text"] = "<|im_start|>" + content["text"] + "<|im_end|>"
+    except (KeyError, TypeError):
+        raise RuntimeError("Internal Error: malformatted BAGEL payload")
+    extra_body = payload.pop("extra_body", {})
+    return {**payload, **extra_body}
+
+
+def _qwen25_payload_preprocessor(payload: dict) -> dict:
+    if payload["messages"][0]["role"] != "system":
+        payload["messages"] = [
+            {
+                "role": "system",
+                "content": (
+                    "You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group,"
+                    "capable of perceiving auditory and visual inputs, as well as generating text and speech."
+                ),
+            },
+            *payload["messages"],
+        ]
+    return payload
+
+
+_MODEL_PIPELINE_SPECS: dict[str, Spec] = {
+    r"BAGEL-7B-MoT": {
+        "stages": [
+            "diffusion"  # The vLLM-Omni interface treats it as a single-stage diffusion model
+        ],
+        "modes": [
+            {
+                "mode": ModelMode.COMPREHENSION,
+                "input_modalities": [Modality.TEXT, Modality.IMAGE],
+            }
+        ],
+        "payload_preprocessor": _bagel_payload_preprocessor,
+    },
+    r"Qwen2.5-Omni*": {
+        "stages": ["autoregression", "autoregression", "autoregression"],
+        "payload_preprocessor": _qwen25_payload_preprocessor,
+        "modes": [
+            {
+                "mode": ModelMode.COMPREHENSION,
+                "input_modalities": [
+                    Modality.TEXT,
+                    Modality.IMAGE,
+                    Modality.VIDEO,
+                    Modality.AUDIO,
+                ],
+            }
+        ],
+    },
+    r"Qwen3-Omni*": {
+        "stages": ["autoregression", "autoregression", "autoregression"],
+        "modes": [
+            {
+                "mode": ModelMode.COMPREHENSION,
+                "input_modalities": [
+                    Modality.TEXT,
+                    Modality.IMAGE,
+                    Modality.VIDEO,
+                    Modality.AUDIO,
+                ],
+            }
+        ],
+    },
+}
+# Convert dict keys to regex patterns
+MODEL_PIPELINE_SPECS: dict[re.Pattern, Spec] = {}
+for k, v in _MODEL_PIPELINE_SPECS.items():
+    MODEL_PIPELINE_SPECS[re.compile(k)] = v
+del _MODEL_PIPELINE_SPECS
+
+
+def lookup_model_spec(model: str) -> tuple[Spec | None, str | None]:
+    try:
+        last_component = model.rstrip("/").rsplit("/", 1)[-1]
+    except IndexError:
+        last_component = model
+    for pattern, spec in MODEL_PIPELINE_SPECS.items():
+        if pattern.search(last_component):
+            return spec, pattern.pattern
+    return None, None
+
+
+# ============== DEMONSTRATION ==============
+
+if __name__ == "__main__":
+    test_paths = [
+        "Qwen/Qwen2.5-Omni-7B",
+        "MyModels/Qwen2.5-Omni-3B",
+        "/root/home/Qwen2.5-Omni-7B",
+        "Qwen/Qwen3-Omni",
+        "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+        "Custom/Path/UnknownModel-Instruct",
+        "Not/Matching/Anything",
+    ]
+
+    test_payload = {"messages": [{"role": "user", "content": "prompt"}]}
+
+    print("Testing registry lookups:\n")
+    for path in test_paths:
+        spec, _ = lookup_model_spec(path)
+        if spec:
+            if preprocessor := spec.get("payload_preprocessor"):
+                result = preprocessor(test_payload)
+                print(f"✓ {path:<40} → {result}")
+            else:
+                print(f"✓ {path:<40} → No preprocessor")
+        else:
+            print(f"✗ {path:<40} → No match")
@@ -0,0 +1,39 @@
+from collections.abc import Callable
+from enum import Enum, auto
+from typing import (
+    Any,
+    Literal,
+    NotRequired,
+    TypeAlias,
+    TypedDict,
+)
+
+AudioFormat: TypeAlias = Literal["mp3", "opus", "aac", "flac", "wav", "pcm"]
+
+
+class ModelMode(Enum):
+    IMAGE_GENERATION = auto()
+    VIDEO_GENERATION = auto()
+    AUDIO_GENERATION = auto()
+    COMPREHENSION = auto()
+
+
+class Modality(Enum):
+    TEXT = auto()  # maybe not useful. Prompt is always required
+    IMAGE = auto()
+    VIDEO = auto()
+    AUDIO = auto()
+
+
+class ModelModeSpec(TypedDict):
+    mode: ModelMode
+    input_modalities: list[Modality]
+
+
+PayloadPreprocessor: TypeAlias = Callable[[dict[str, Any]], dict[str, Any]]
+
+
+class Spec(TypedDict):
+    stages: list[Literal["diffusion", "autoregression"]]
+    modes: list[ModelModeSpec]
+    payload_preprocessor: NotRequired[PayloadPreprocessor]
@@ -0,0 +1,84 @@
+from .logger import get_logger
+from .models import lookup_model_spec
+
+logger = get_logger(__name__)
+
+
+def validate_model_and_sampling_params_types(
+    model_name: str,
+    sampling_param_list: dict | list[dict] | None = None,
+):
+    # Check if model name exists
+    if not model_name:
+        raise ValueError("Model name must not be empty.")
+
+    # Skip if no spec or no sampling params
+    pipeline_spec, _ = lookup_model_spec(model_name)
+    if pipeline_spec is None:
+        logger.info("skipping sampling params check because spec is not found")
+        return
+    if sampling_param_list is None:
+        return
+
+    # Check the number of stages and their data types
+    stages = pipeline_spec["stages"]
+    if isinstance(sampling_param_list, list):
+        # Check that the lengths match
+        if len(stages) != len(sampling_param_list):
+            raise ValueError(
+                f"Sampling parameter list length {len(sampling_param_list)} does not match "
+                f"number of stages {len(stages)} for model {model_name}."
+            )
+        # Check that each stage's type match
+        for i, sp in enumerate(sampling_param_list):
+            if "type" not in sp:
+                raise RuntimeError("Internal error: unknown sampling parameter type")
+            if sp["type"] != stages[i]:
+                raise ValueError(
+                    f"Sampling parameter type ({sp['type']}) does not match "
+                    f"stage type ({stages[i]}) at index {i} for model {model_name}."
+                )
+    elif isinstance(sampling_param_list, dict):
+        if "type" not in sampling_param_list:
+            raise RuntimeError("Internal error: unknown sampling parameter type")
+        # Check that the provided single sampling param matches all stages
+        elif any(stage != sampling_param_list["type"] for stage in stages):
+            raise ValueError(
+                f"When passing a single sampling parameter node, all stages of the model must match "
+                f"the provided sampling parameter's type. "
+                f"However, the stages of model {model_name} are: {stages}. "
+                f"The provided sampling parameter is {sampling_param_list['type']}"
+            )
+
+
+def add_sampling_parameters_to_stage(
+    model_name: str,
+    sampling_param_list: dict | list[dict] | None,
+    stage_type: str,
+    /,
+    **params_to_add,
+) -> dict | list[dict]:
+    """
+    Given a model's name and the sampling parameter list to query this model,
+    add arbitrary additional parameters to the sampling parameters of all stages of the given type.
+    """
+    pipeline_spec, _ = lookup_model_spec(model_name)
+    if not pipeline_spec:
+        logger.warning(
+            f"Since the model {model_name} is not in our list, we cannot ensure if "
+            f"the fields ({tuple(params_to_add.keys())}) are added to the correct stage's sampling params. "
+            f"We will do it heuristically."
+        )
+        pipeline_spec = {"stages": ["diffusion"]}
+
+    stages = pipeline_spec["stages"]
+    if isinstance(sampling_param_list, dict):
+        sampling_param_list.update(params_to_add)
+    elif sampling_param_list is None:
+        sampling_param_list = params_to_add.copy()
+    else:
+        for i, stage in enumerate(stages):
+            if stage == stage_type:
+                sampling_param_list[i].update(params_to_add)
+
+    return sampling_param_list
@@ -0,0 +1 @@
+{"id":"1c99f525-0a37-45ba-a28a-7df7c3af66b4","revision":0,"last_node_id":12,"last_link_id":14,"nodes":[{"id":1,"type":"VLLMOmniComprehension","pos":[1191.2177053682556,144.66829928181377],"size":[400,268],"flags":{},"order":8,"mode":0,"inputs":[{"localized_name":"image","name":"image","shape":7,"type":"IMAGE","link":1},{"localized_name":"video","name":"video","shape":7,"type":"VIDEO","link":2},{"localized_name":"audio","name":"audio","shape":7,"type":"AUDIO","link":3},{"localized_name":"sampling_params","name":"sampling_params","shape":7,"type":"SAMPLING_PARAMS","link":14},{"localized_name":"url","name":"url","type":"STRING","widget":{"name":"url"},"link":null},{"localized_name":"model","name":"model","type":"STRING","widget":{"name":"model"},"link":null},{"localized_name":"prompt","name":"prompt","type":"STRING","widget":{"name":"prompt"},"link":null},{"localized_name":"output_text","name":"output_text","type":"BOOLEAN","widget":{"name":"output_text"},"link":null},{"localized_name":"output_audio","name":"output_audio","type":"BOOLEAN","widget":{"name":"output_audio"},"link":null},{"localized_name":"use_audio_in_video","name":"use_audio_in_video","type":"BOOLEAN","widget":{"name":"use_audio_in_video"},"link":null}],"outputs":[{"localized_name":"text_response","name":"text_response","type":"STRING","links":[8]},{"localized_name":"audio_response","name":"audio_response","type":"AUDIO","links":[9]}],"properties":{"Node name for S&R":"VLLMOmniComprehension"},"widgets_values":["http://localhost:8000/v1","Qwen/Qwen2.5-Omni-7B","",true,true,true]},{"id":3,"type":"LoadVideo","pos":[729.5984141255855,-198.631920454299],"size":[282.798828125,233.0743408203125],"flags":{},"order":0,"mode":0,"inputs":[{"localized_name":"file","name":"file","type":"COMBO","widget":{"name":"file"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"IMAGEUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"VIDEO","name":"VIDEO","type":"VIDEO","links":[2]}],"properties":{"Node name for S&R":"LoadVideo"},"widgets_values":["draw.mp4","image"]},{"id":4,"type":"LoadAudio","pos":[729.8037086965753,99.86963519703949],"size":[282.798828125,136],"flags":{},"order":1,"mode":0,"inputs":[{"localized_name":"audio","name":"audio","type":"COMBO","widget":{"name":"audio"},"link":null},{"localized_name":"audioUI","name":"audioUI","type":"AUDIO_UI","widget":{"name":"audioUI"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"AUDIOUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"AUDIO","name":"AUDIO","type":"AUDIO","links":[3]}],"properties":{"Node name for S&R":"LoadAudio"},"widgets_values":["Megan-Fox.mp3",null,null]},{"id":5,"type":"VLLMOmniARSampling","pos":[510.3517536642828,658.073751009259],"size":[270,178],"flags":{},"order":2,"mode":0,"inputs":[{"localized_name":"max_tokens","name":"max_tokens","type":"INT","widget":{"name":"max_tokens"},"link":null},{"localized_name":"temperature","name":"temperature","type":"FLOAT","widget":{"name":"temperature"},"link":null},{"localized_name":"top_p","name":"top_p","type":"FLOAT","widget":{"name":"top_p"},"link":null},{"localized_name":"repetition_penalty","name":"repetition_penalty","type":"FLOAT","widget":{"name":"repetition_penalty"},"link":null},{"localized_name":"seed","name":"seed","type":"INT","widget":{"name":"seed"},"link":null}],"outputs":[{"localized_name":"AR sampling params","name":"AR sampling params","type":"SAMPLING_PARAMS","links":[12]}],"properties":{"Node name for S&R":"VLLMOmniARSampling"},"widgets_values":[100,1,1,1,-1,"randomize"]},{"id":7,"type":"VLLMOmniARSampling","pos":[503.33235181647115,419.34158016181806],"size":[270,178],"flags":{},"order":3,"mode":0,"inputs":[{"localized_name":"max_tokens","name":"max_tokens","type":"INT","widget":{"name":"max_tokens"},"link":null},{"localized_name":"temperature","name":"temperature","type":"FLOAT","widget":{"name":"temperature"},"link":null},{"localized_name":"top_p","name":"top_p","type":"FLOAT","widget":{"name":"top_p"},"link":null},{"localized_name":"repetition_penalty","name":"repetition_penalty","type":"FLOAT","widget":{"name":"repetition_penalty"},"link":null},{"localized_name":"seed","name":"seed","type":"INT","widget":{"name":"seed"},"link":null}],"outputs":[{"localized_name":"AR sampling params","name":"AR sampling params","type":"SAMPLING_PARAMS","links":[5,13]}],"properties":{"Node name for S&R":"VLLMOmniARSampling"},"widgets_values":[100,1,1,1,-1,"randomize"]},{"id":8,"type":"VLLMOmniSamplingParamsList","pos":[820.6056617389042,426.38372037182273],"size":[263.066015625,66],"flags":{},"order":7,"mode":0,"inputs":[{"localized_name":"param1","name":"param1","type":"SAMPLING_PARAMS","link":5},{"localized_name":"param2","name":"param2","shape":7,"type":"SAMPLING_PARAMS","link":12},{"localized_name":"param3","name":"param3","shape":7,"type":"SAMPLING_PARAMS","link":13}],"outputs":[{"localized_name":"param list","name":"param list"Line truncated
@@ -0,0 +1 @@
+{"id":"91f75acc-8040-40f6-865a-2e8a7cfd6672","revision":0,"last_node_id":11,"last_link_id":23,"nodes":[{"id":3,"type":"PreviewImage","pos":[1281.8455167767304,-69.02638461454333],"size":[305.6628685610174,307.0336884172169],"flags":{},"order":4,"mode":0,"inputs":[{"localized_name":"images","name":"images","type":"IMAGE","link":22}],"outputs":[],"properties":{"Node name for S&R":"PreviewImage"},"widgets_values":[]},{"id":11,"type":"VLLMOmniGenerateImage","pos":[816.0962515873642,-137.51112584854445],"size":[400,278],"flags":{},"order":3,"mode":0,"inputs":[{"localized_name":"image","name":"image","shape":7,"type":"IMAGE","link":21},{"localized_name":"mask","name":"mask","shape":7,"type":"MASK","link":null},{"localized_name":"sampling_params","name":"sampling_params","shape":7,"type":"SAMPLING_PARAMS","link":23},{"localized_name":"url","name":"url","type":"STRING","widget":{"name":"url"},"link":null},{"localized_name":"model","name":"model","type":"STRING","widget":{"name":"model"},"link":null},{"localized_name":"prompt","name":"prompt","type":"STRING","widget":{"name":"prompt"},"link":null},{"localized_name":"negative_prompt","name":"negative_prompt","type":"STRING","widget":{"name":"negative_prompt"},"link":null},{"localized_name":"width","name":"width","type":"INT","widget":{"name":"width"},"link":null},{"localized_name":"height","name":"height","type":"INT","widget":{"name":"height"},"link":null}],"outputs":[{"localized_name":"image","name":"image","type":"IMAGE","links":[22]}],"properties":{"Node name for S&R":"VLLMOmniGenerateImage"},"widgets_values":["http://localhost:8000/v1","Qwen/Qwen-Image-Edit","Put this figure in a realistic mountain view","",512,512]},{"id":8,"type":"VLLMOmniDiffusionSampling","pos":[478.59266934006774,183.67711984955648],"size":[284.205078125,202],"flags":{},"order":0,"mode":0,"inputs":[{"localized_name":"n","name":"n","type":"INT","widget":{"name":"n"},"link":null},{"localized_name":"num_inference_steps","name":"num_inference_steps","type":"INT","widget":{"name":"num_inference_steps"},"link":null},{"localized_name":"guidance_scale","name":"guidance_scale","type":"FLOAT","widget":{"name":"guidance_scale"},"link":null},{"localized_name":"true_cfg_scale","name":"true_cfg_scale","type":"FLOAT","widget":{"name":"true_cfg_scale"},"link":null},{"localized_name":"vae_use_slicing","name":"vae_use_slicing","type":"BOOLEAN","widget":{"name":"vae_use_slicing"},"link":null},{"localized_name":"seed","name":"seed","type":"INT","widget":{"name":"seed"},"link":null}],"outputs":[{"localized_name":"diffusion sampling params","name":"diffusion sampling params","type":"SAMPLING_PARAMS","links":[23]}],"properties":{"Node name for S&R":"VLLMOmniDiffusionSampling"},"widgets_values":[4,50,1,1,false,42,"fixed"]},{"id":10,"type":"MarkdownNote","pos":[227.99306819575278,-231.24143306069843],"size":[240.3326478922474,136.3505820791881],"flags":{},"order":2,"mode":0,"inputs":[],"outputs":[],"title":"Note: Task and Input","properties":{},"widgets_values":["vLLM-Omni nodes are categorized based on the output modality. The \"Generate Image\" node supports both text-to-image generation or image-to-image generation (a.k.a. image editing). The node will route to the correct endpoint depending on whether an input image is present or not."],"color":"#432","bgcolor":"#000"},{"id":4,"type":"LoadImage","pos":[496.31859627609606,-229.71277089860084],"size":[270,314],"flags":{},"order":1,"mode":0,"inputs":[{"localized_name":"image","name":"image","type":"COMBO","widget":{"name":"image"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"IMAGEUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"IMAGE","name":"IMAGE","type":"IMAGE","links":[21]},{"localized_name":"MASK","name":"MASK","type":"MASK","links":null}],"properties":{"Node name for S&R":"LoadImage"},"widgets_values":["example.png","image"]}],"links":[[21,4,0,11,0,"IMAGE"],[22,11,0,3,0,"IMAGE"],[23,8,0,11,2,"SAMPLING_PARAMS"]],"groups":[{"id":1,"title":"Input","bounding":[213.1706010087147,-313.20095750667554,560.8246471437027,407.4246532472182],"color":"#3f789e","font_size":24,"flags":{}}],"config":{},"extra":{"workflowRendererVersion":"LG","ds":{"scale":1.1469075819486894,"offset":[276.8427259980847,549.4885226647582]}},"version":0.4}
@@ -0,0 +1 @@
+{"id":"c0a11241-f0e6-4f5c-8488-b3d344a9df3d","revision":0,"last_node_id":7,"last_link_id":3,"nodes":[{"id":4,"type":"PreviewAudio","pos":[690.1759063763212,226.76444957981545],"size":[270,88],"flags":{},"order":5,"mode":0,"inputs":[{"localized_name":"audio","name":"audio","type":"AUDIO","link":2},{"localized_name":"audioUI","name":"audioUI","type":"AUDIO_UI","widget":{"name":"audioUI"},"link":null}],"outputs":[],"properties":{"Node name for S&R":"PreviewAudio"},"widgets_values":[]},{"id":5,"type":"LoadAudio","pos":[-83.61093919682321,103.80510898465587],"size":[282.798828125,136],"flags":{},"order":0,"mode":0,"inputs":[{"localized_name":"audio","name":"audio","type":"COMBO","widget":{"name":"audio"},"link":null},{"localized_name":"audioUI","name":"audioUI","type":"AUDIO_UI","widget":{"name":"audioUI"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"AUDIOUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"AUDIO","name":"AUDIO","type":"AUDIO","links":[3]}],"properties":{"Node name for S&R":"LoadAudio"},"widgets_values":["Megan-Fox.mp3",null,null]},{"id":3,"type":"VLLMOmniVoiceClone","pos":[248.63466280064262,168.13112893706375],"size":[400,306],"flags":{},"order":4,"mode":0,"inputs":[{"localized_name":"ref_audio","name":"ref_audio","type":"AUDIO","link":3},{"localized_name":"model_specific_params","name":"model_specific_params","shape":7,"type":"TTS_PARAMS","link":1},{"localized_name":"url","name":"url","type":"STRING","widget":{"name":"url"},"link":null},{"localized_name":"model","name":"model","type":"STRING","widget":{"name":"model"},"link":null},{"localized_name":"input","name":"input","type":"STRING","widget":{"name":"input"},"link":null},{"localized_name":"voice","name":"voice","type":"STRING","widget":{"name":"voice"},"link":null},{"localized_name":"response_format","name":"response_format","type":"COMBO","widget":{"name":"response_format"},"link":null},{"localized_name":"speed","name":"speed","type":"FLOAT","widget":{"name":"speed"},"link":null},{"localized_name":"ref_text","name":"ref_text","type":"STRING","widget":{"name":"ref_text"},"link":null},{"localized_name":"x_vector_only_mode","name":"x_vector_only_mode","type":"BOOLEAN","widget":{"name":"x_vector_only_mode"},"link":null}],"outputs":[{"localized_name":"audio","name":"audio","type":"AUDIO","links":[2]}],"properties":{"Node name for S&R":"VLLMOmniVoiceClone"},"widgets_values":["http://localhost:8000/v1","Qwen/Qwen3-TTS-12Hz-1.7B-Base","Someone just spilled a cup of coffee on my jacket this morning in the subway. Now I have to wear this stained jacket for a whole day!","","mp3",1,"",true]},{"id":6,"type":"MarkdownNote","pos":[255.52667408769662,27.26842775454861],"size":[381.58204108637847,88],"flags":{},"order":2,"mode":0,"inputs":[],"outputs":[],"title":"Note: TTS Nodes","properties":{},"widgets_values":["Apart from the Voice Cloning node, there is also another TTS node for simple text-to-speech tasks (without extra inputs for reference audio)."],"color":"#432","bgcolor":"#000"},{"id":7,"type":"MarkdownNote","pos":[-196.99453003423977,598.3693693727904],"size":[382.4539509209102,88],"flags":{},"order":3,"mode":0,"inputs":[],"outputs":[],"title":"Note: Model-Specific Parameters","properties":{},"widgets_values":["TTS models often require some tailor-made parameters. If you need to customize these parameters, grab one that matches the model you are requesting from \"TTS Params\" subfolder."],"color":"#432","bgcolor":"#000"},{"id":2,"type":"VLLMOmniQwenTTSParams","pos":[-207.8925310560497,342.7942867650348],"size":[400,200],"flags":{},"order":1,"mode":0,"inputs":[{"localized_name":"task_type","name":"task_type","type":"COMBO","widget":{"name":"task_type"},"link":null},{"localized_name":"language","name":"language","type":"COMBO","widget":{"name":"language"},"link":null},{"localized_name":"instructions","name":"instructions","type":"STRING","widget":{"name":"instructions"},"link":null},{"localized_name":"max_new_tokens","name":"max_new_tokens","type":"INT","widget":{"name":"max_new_tokens"},"link":null}],"outputs":[{"localized_name":"Qwen TTS params","name":"Qwen TTS params","type":"TTS_PARAMS","links":[1]}],"properties":{"Node name for S&R":"VLLMOmniQwenTTSParams"},"widgets_values":["Base","Auto","Super angry",2048]}],"links":[[1,2,0,3,1,"TTS_PARAMS"],[2,3,0,4,0,"AUDIO"],[3,5,0,3,0,"AUDIO"]],"groups":[{"id":1,"title":"Model-Specific Parameters","bounding":[-228.38328407738007,264.42790274716276,443.9887293164586,439.6853171248897],"color":"#3f789e","font_size":24,"flags":{}}],"config":{},"extra":{"workflowRendererVersion":"LG","ds":{"scale":1.1469075819486911,"offset":[898.8819468322333,258.7179979718389]}},"version":0.4}
@@ -0,0 +1,21 @@
+/**
+ * @file This file is intended to add dynamic fields to vLLM-Omni nodes
+ * based on widget (in-node form fields) and input (connection link) values and changes.
+ * However, this functionality is currently disabled/commented out.
+ * Because it introduces too much complexity,
+ * and it may even conflict with the current backend (Python) validation for unknown reasons (pending ComfyUI upstream fixes).
+ */
+
+import { app } from "../../scripts/app.js";
+app.registerExtension({
+    name: "vllm.vllm_omni",
+    async beforeRegisterNodeDef(nodeType, nodeData, app) {
+        if (!nodeData.name.startsWith("VLLMOmni")) {
+            return
+        }
+        // Stub frontend plugin for now
+    },
+    async setup() {
+        console.info("vLLM-Omni Setup complete!")
+    },
+})
@@ -56,7 +56,7 @@ python extract_tts_prompts.py -i meta.lst -o top100.txt -n 100
 **Options:**
 - `-i, --input`: Input metadata file (default: `meta.lst`)
 - `-o, --output`: Output prompts file (default: `prompts.txt`)
- `-n, --num_lines`: Number of prompts to extract (required)
+- `-n, --num-lines`: Number of prompts to extract (required)

 ### 7. Clean Up (Optional)

@@ -56,7 +56,7 @@ def main():
        "-o", "--output", type=str, default="prompts.txt", help="Output txt file path (default: prompts.txt)"
    )
    parser.add_argument(
-        "-n", "--num_lines", type=int, required=True, help="Number of lines to extract from the beginning"
+        "-n", "--num-lines", type=int, required=True, help="Number of lines to extract from the beginning"
    )

    args = parser.parse_args()
@@ -0,0 +1,117 @@
+
+# Diffusion Serving Benchmark (Image/Video)
+
+This folder contains an online-serving benchmark script for diffusion models.
+It sends requests to a vLLM OpenAI-compatible endpoint and reports throughput,
+latency percentiles, and optional SLO attainment.
+
+The main entrypoint is:
+
+- `benchmarks/diffusion/diffusion_benchmark_serving.py`
+
+## 1. Quick Start
+
+1. Start the server:
+
+```bash
+vllm serve Qwen/Qwen-Image --omni --port 8099
+```
+
+2. Run a minimal benchmark:
+
+```bash
+python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
+	--base-url http://localhost:8099 \
+	--model Qwen/Qwen-Image \
+	--task t2i \
+	--dataset vbench \
+	--num-prompts 5
+```
+
+**Notes**
+
+- The benchmark talks to `http://<host>:<port>/v1/chat/completions`.
+- If you run the server on another host or port, pass `--base-url` accordingly.
+
+## 2. Supported Datasets
+
+The benchmark supports three dataset modes via `--dataset`:
+
+- `vbench`: Built-in prompt/data loader.
+- `trace`: Heterogeneous request traces (each request can have different resolution/frames/steps).
+- `random`: Synthetic prompts for quick smoke tests.
+
+### VBench dataset
+
+If you use i2v/i2i bench datasets and need auto-download support, you may need:
+
+```bash
+uv pip install gdown
+```
+
+### Trace dataset
+
+Use `--dataset trace` to replay a trace file. The trace can specify per-request fields such as:
+
+- `width`, `height`
+- `num_frames` (video)
+- `num_inference_steps`
+- `seed`, `fps`
+- optional `slo_ms` (per-request SLO target)
+
+By default (when `--dataset-path` is not provided), the script downloads a default trace from
+the HuggingFace dataset repo `asukaqaqzz/Dit_Trace`. The default filename can depend on `--task`
+(e.g., `t2v` uses a video trace).
+
+Current defaults:
+
+- `--task t2i` -> `sd3_trace.txt`
+- `--task t2v` -> `cogvideox_trace.txt`
+
+You can point to your own trace using `--dataset-path`.
+
+## 3. Benchmark Parameters
+
+### Basic flags
+
+- `--base-url`: Server address (the script calls `.../v1/chat/completions`).
+- `--model`: The OpenAI-compatible `model` field.
+- `--task`: Task type (e.g., `t2i`, `t2v`, `i2i`, `i2v`).
+- `--dataset`: Dataset mode (`vbench` / `trace` / `random`).
+- `--num-prompts`: Number of requests to send.
+
+Common optional flags:
+
+- `--output-file`: Write metrics to a JSON file.
+- `--disable-tqdm`: Disable the progress bar.
+
+### Resolution / frames / steps: CLI defaults vs dataset fields
+
+Related flags: `--width`, `--height`, `--num-frames`, `--fps`, `--num-inference-steps`.
+
+- For `vbench` / `random`: these CLI flags act as global defaults for all generated requests.
+- For `trace`: each request can carry its own fields (e.g., `width/height/num_frames/num_inference_steps`).
+
+Precedence rules for `trace` (i.e., what actually gets sent):
+
+- `width/height`: if either `--width` or `--height` is explicitly set, it overrides per-request values from the trace; otherwise per-request values are used when present.
+- `num_frames`: per-request `num_frames` takes precedence; otherwise fall back to `--num-frames`.
+- `num_inference_steps`: per-request `num_inference_steps` takes precedence; otherwise fall back to `--num-inference-steps`.
+
+### SLO, warmup, and max concurrency
+
+Enable SLO evaluation with `--slo`.
+
+- If a request in the trace already has `slo_ms`, that value is used.
+- Otherwise, the script runs warmup requests to infer a base unit time, estimates `expected_ms` by linearly scaling with area/frames/steps, and then sets `slo_ms = expected_ms * --slo-scale`.
+
+Warmup flags:
+
+- `--warmup-requests`: Number of warmup requests.
+- `--warmup-num-inference-steps`: Steps used during warmup.
+- For `--task t2v`: warmup requests are forced to use `num_frames=1` to make warmup faster and less noisy.
+
+Traffic / concurrency flags:
+
+- `--request-rate`: Target request rate (requests/second). If set to `inf`, the script sends all requests immediately.
+- `--max-concurrency`: Max number of in-flight requests (default: `1`). This can hard-cap the achieved QPS: if it is too small, requests will queue behind the semaphore, and both achieved throughput and observed SLO attainment can be skewed.
@@ -9,28 +9,29 @@ If you want to use i2v, i2i dataset, you should `uv pip install gdown` first
 Usage:
    # Video
    t2v:
-    python3 benchmarks/diffusion_benchmark_serving.py \
+    python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
        --dataset vbench --task t2v --num-prompts 10 \
        --height 480 --width 640 --fps 16 --num-frames 80

    i2v:
-    python3 benchmarks/diffusion_benchmark_serving.py \
+    python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
        --dataset vbench --task i2v --num-prompts 10


    # Image
    t2i:
-    python3 benchmarks/diffusion_benchmark_serving.py \
+    python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
        --dataset vbench --task t2i --num-prompts 10 \
        --height 1024 --width 1024

    i2i:
-    python3 benchmarks/diffusion_benchmark_serving.py \
+    python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
        --dataset vbench --task i2i --num-prompts 10

 """

 import argparse
+import ast
 import asyncio
 import base64
 import glob
@@ -40,6 +41,7 @@ import os
 import time
 import uuid
 from abc import ABC, abstractmethod
+from collections.abc import AsyncGenerator
 from dataclasses import dataclass, field, replace
 from typing import Any

@@ -60,6 +62,8 @@ class RequestFuncInput:
    num_inference_steps: int | None = None
    seed: int | None = None
    fps: int | None = None
+    timestamp: float | None = None
+    slo_ms: float | None = None
    extra_body: dict[str, Any] = field(default_factory=dict)
    image_paths: list[str] | None = None
    request_id: str = field(default_factory=lambda: str(uuid.uuid4()))
@@ -73,6 +77,7 @@ class RequestFuncOutput:
    start_time: float = 0.0
    response_body: dict[str, Any] = field(default_factory=dict)
    peak_memory_mb: float = 0.0
+    slo_achieved: bool | None = None


 class BaseDataset(ABC):
@@ -310,6 +315,226 @@ class VBenchDataset(BaseDataset):
        return [self[i] for i in range(len(self))]


+class TraceDataset(BaseDataset):
+    """Trace-based dataset loader for heterogeneous diffusion requests."""
+
+    DEFAULT_REPO_ID = "asukaqaqzz/Dit_Trace"
+    DEFAULT_FILENAME = "sd3_trace.txt"
+    DEFAULT_FILENAME_BY_TASK: dict[str, str] = {
+        # Text-to-image traces (e.g., SD3)
+        "t2i": "sd3_trace.txt",
+        # Text-to-video traces (e.g., CogVideoX)
+        "t2v": "cogvideox_trace.txt",
+    }
+
+    def __init__(self, args, api_url: str, model: str):
+        super().__init__(args, api_url, model)
+        self.cache_dir = os.path.join(os.path.expanduser("~"), ".cache", "vllm-omni", "trace")
+        self.default_filename = self.DEFAULT_FILENAME_BY_TASK.get(getattr(args, "task", ""), self.DEFAULT_FILENAME)
+        dataset_root = args.dataset_path
+        if not dataset_root:
+            dataset_root = self._download_default_trace()
+        self.items = self._load_items(dataset_root)
+
+    @staticmethod
+    def _coerce_int(x: Any) -> int | None:
+        if x is None:
+            return None
+        if isinstance(x, bool):
+            return None
+        if isinstance(x, int):
+            return x
+        try:
+            s = str(x).strip()
+            if not s:
+                return None
+            return int(float(s))
+        except Exception:
+            return None
+
+    @staticmethod
+    def _coerce_float(x: Any) -> float | None:
+        if x is None:
+            return None
+        if isinstance(x, float):
+            return x
+        if isinstance(x, int):
+            return float(x)
+        try:
+            s = str(x).strip()
+            if not s:
+                return None
+            return float(s)
+        except Exception:
+            return None
+
+    def _download_default_trace(self) -> str:
+        """Download default trace file from HuggingFace Hub if not provided."""
+
+        try:
+            from huggingface_hub import hf_hub_download
+        except ImportError as exc:
+            raise ImportError(
+                "huggingface_hub is required to download the default trace dataset. "
+                "Install via `pip install huggingface_hub`."
+            ) from exc
+
+        os.makedirs(self.cache_dir, exist_ok=True)
+        return hf_hub_download(
+            repo_id=self.DEFAULT_REPO_ID,
+            filename=self.default_filename,
+            repo_type="dataset",
+            local_dir=self.cache_dir,
+            local_dir_use_symlinks=False,
+        )
+
+    def _expand_paths(self, dataset_path: str | None) -> list[str]:
+        if not dataset_path:
+            return []
+
+        parts = [p.strip() for p in str(dataset_path).split(",") if p.strip()]
+        paths: list[str] = []
+        for p in parts:
+            if any(ch in p for ch in ["*", "?", "["]):
+                paths.extend(sorted(glob.glob(p)))
+            elif os.path.isdir(p):
+                paths.extend(sorted(glob.glob(os.path.join(p, "**", "*.txt"), recursive=True)))
+            else:
+                paths.append(p)
+
+        seen = set()
+        unique_paths = []
+        for p in paths:
+            if p not in seen:
+                seen.add(p)
+                unique_paths.append(p)
+        return unique_paths
+
+    def _parse_trace_file(self, path: str) -> list[dict[str, Any]]:
+        rows: list[dict[str, Any]] = []
+
+        def parse_request_repr_line(line: str) -> dict[str, Any] | None:
+            text = line.strip()
+            if not text:
+                return None
+            if not (text.startswith("Request(") and text.endswith(")")):
+                return None
+            inner = text[len("Request(") : -1]
+            try:
+                expr = ast.parse(f"f({inner})", mode="eval")
+                if not isinstance(expr.body, ast.Call):
+                    return None
+                call = expr.body
+                out: dict[str, Any] = {}
+                for kw in call.keywords:
+                    if kw.arg is None:
+                        continue
+                    out[kw.arg] = ast.literal_eval(kw.value)
+                return out
+            except Exception:
+                return None
+
+        # detect first non-empty line to pick parser
+        first_non_empty = None
+        with open(path, encoding="utf-8") as f:
+            for _ in range(50):
+                pos = f.tell()
+                line = f.readline()
+                if not line:
+                    break
+                if line.strip():
+                    first_non_empty = line.strip()
+                    f.seek(pos)
+                    break
+
+        if first_non_empty is None:
+            return rows
+
+        if first_non_empty.startswith("Request("):
+            with open(path, encoding="utf-8") as f:
+                for line in f:
+                    parsed = parse_request_repr_line(line)
+                    if isinstance(parsed, dict):
+                        rows.append(parsed)
+            return rows
+
+        # txt fallback: parse Request(...) lines only
+        with open(path, encoding="utf-8") as f:
+            for line in f:
+                parsed = parse_request_repr_line(line)
+                if isinstance(parsed, dict):
+                    rows.append(parsed)
+        return rows
+
+    def _load_items(self, dataset_root: str) -> list[dict[str, Any]]:
+        paths = self._expand_paths(dataset_root)
+        if not paths:
+            raise ValueError("No trace files found. Provide --dataset-path or rely on default HuggingFace download.")
+
+        items: list[dict[str, Any]] = []
+        for p in paths:
+            if not os.path.exists(p):
+                continue
+            for row in self._parse_trace_file(p):
+                if isinstance(row, dict):
+                    row = dict(row)
+                    row.setdefault("_source", p)
+                    items.append(row)
+
+        if not items:
+            raise ValueError("Trace dataset is empty after parsing provided paths.")
+
+        if self.args.num_prompts is not None:
+            items = items[: self.args.num_prompts]
+
+        return items
+
+    def __len__(self) -> int:
+        return len(self.items)
+
+    def __getitem__(self, idx: int) -> RequestFuncInput:
+        row = self.items[idx]
+        prompt = row.get("prompt") or row.get("text") or ""
+
+        row_height = self._coerce_int(row.get("height"))
+        row_width = self._coerce_int(row.get("width"))
+        num_frames = self._coerce_int(row.get("num_frames"))
+        num_steps = self._coerce_int(row.get("num_inference_steps"))
+        seed = self._coerce_int(row.get("seed"))
+        fps = self._coerce_int(row.get("fps"))
+        timestamp = self._coerce_float(row.get("timestamp"))
+        slo_ms = self._coerce_float(row.get("slo_ms"))
+        image_paths = row.get("image_paths")
+
+        override_w = self.args.width
+        override_h = self.args.height
+        if override_w is not None or override_h is not None:
+            width = override_w
+            height = override_h
+        else:
+            width = row_width
+            height = row_height
+
+        return RequestFuncInput(
+            prompt=str(prompt),
+            api_url=self.api_url,
+            model=self.model,
+            width=width,
+            height=height,
+            num_frames=num_frames if num_frames is not None else self.args.num_frames,
+            num_inference_steps=num_steps if num_steps is not None else self.args.num_inference_steps,
+            seed=seed if seed is not None else self.args.seed,
+            fps=fps if fps is not None else self.args.fps,
+            timestamp=timestamp,
+            slo_ms=slo_ms,
+            image_paths=image_paths,
+            request_id=str(row.get("request_id")) if row.get("request_id") is not None else str(uuid.uuid4()),
+        )
+
+    def get_requests(self) -> list[RequestFuncInput]:
+        return [self[i] for i in range(len(self))]
+
+
 class RandomDataset(BaseDataset):
    def __init__(self, args, api_url: str, model: str):
        self.args = args
@@ -337,6 +562,128 @@ class RandomDataset(BaseDataset):
        return [self[i] for i in range(len(self))]


+def _compute_expected_latency_ms_from_base(req: RequestFuncInput, args, base_time_ms: float | None) -> float | None:
+    """Compute expected execution time (ms) based on a base per-step-per-frame unit time.
+
+    Assumes linear scaling with pixel area, frame count, and num_inference_steps.
+    The base unit represents latency for a 16x16 resolution, single frame, single step.
+    """
+
+    if base_time_ms is None:
+        return None
+
+    width = req.width if req.width is not None else args.width
+    height = req.height if req.height is not None else args.height
+    if width is None or height is None:
+        return None
+
+    frames = req.num_frames if req.num_frames is not None else args.num_frames
+    steps = req.num_inference_steps if req.num_inference_steps is not None else args.num_inference_steps
+
+    frame_scale = frames if isinstance(frames, int) and frames > 0 else 1
+    step_scale = steps if isinstance(steps, int) and steps > 0 else 1
+
+    area_units = max((float(width) * float(height)) / float(16 * 16), 1.0)
+    return float(base_time_ms) * area_units * frame_scale * step_scale
+
+
+def _infer_slo_base_time_ms_from_warmups(
+    warmup_pairs: list[tuple[RequestFuncInput, RequestFuncOutput]],
+    args,
+) -> float | None:
+    """Infer base SLO unit time from warmup requests.
+
+    Returns the median base latency (ms) for a 16x16 resolution, single-frame,
+    single-step request. Only uses warmups that succeeded and have resolvable
+    width/height.
+    """
+
+    candidates_ms: list[float] = []
+    for req, out in warmup_pairs:
+        if not out.success or out.latency <= 0:
+            continue
+
+        width = req.width if req.width is not None else args.width
+        height = req.height if req.height is not None else args.height
+        if width is None or height is None:
+            continue
+
+        frames = req.num_frames if req.num_frames is not None else args.num_frames
+        steps = req.num_inference_steps if req.num_inference_steps is not None else args.num_inference_steps
+
+        frame_scale = int(frames) if isinstance(frames, int) and frames > 0 else 1
+        step_scale = int(steps) if isinstance(steps, int) and steps > 0 else 1
+
+        area_units = max((float(width) * float(height)) / float(16 * 16), 1.0)
+        denom = area_units * float(frame_scale) * float(step_scale)
+        if denom <= 0:
+            continue
+
+        candidates_ms.append((out.latency * 1000.0) / denom)
+
+    if not candidates_ms:
+        return None
+    return float(np.median(candidates_ms))
+
+
+def _populate_slo_ms_from_warmups(
+    requests_list: list[RequestFuncInput],
+    warmup_pairs: list[tuple[RequestFuncInput, RequestFuncOutput]],
+    args,
+) -> list[RequestFuncInput]:
+    """Populate missing RequestFuncInput.slo_ms using warmup outputs.
+
+    - If a request already has slo_ms (e.g., trace-provided), it is kept as-is.
+    - If any request has slo_ms is None and we can infer base time from warmups,
+      we estimate each missing request's expected execution time and set:
+        req.slo_ms = expected_latency_ms * args.slo_scale
+
+    Returns updated requests_list.
+    """
+
+    if not any(req.slo_ms is None for req in requests_list):
+        return requests_list
+
+    base_time_ms = _infer_slo_base_time_ms_from_warmups(warmup_pairs, args)
+    if base_time_ms is None:
+        return requests_list
+
+    slo_scale = float(getattr(args, "slo_scale", 3.0))
+    if slo_scale <= 0:
+        raise ValueError(f"slo_scale must be positive, got {slo_scale}.")
+
+    updated: list[RequestFuncInput] = []
+    for req in requests_list:
+        if req.slo_ms is not None:
+            updated.append(req)
+            continue
+        expected_ms = _compute_expected_latency_ms_from_base(req, args, base_time_ms)
+        updated.append(replace(req, slo_ms=(expected_ms * slo_scale) if expected_ms is not None else None))
+
+    return updated
+
+
+async def iter_requests(
+    requests_list: list[RequestFuncInput],
+    request_rate: float,
+) -> AsyncGenerator[RequestFuncInput, None]:
+    """Yield requests using a fixed interval if request_rate is set.
+
+    - If request_rate is inf, all requests are yielded immediately (no sleep).
+    - Otherwise, requests are emitted at a fixed cadence of 1 / request_rate seconds.
+    """
+
+    if request_rate != float("inf"):
+        if request_rate <= 0:
+            raise ValueError(f"request_rate must be positive or inf, got {request_rate}.")
+        interval_s = 1.0 / float(request_rate)
+
+    for i, req in enumerate(requests_list):
+        if request_rate != float("inf") and i > 0:
+            await asyncio.sleep(interval_s)
+        yield req
+
+
 def _guess_mime_type(path: str) -> str:
    mime, _ = mimetypes.guess_type(path)
    return mime or "application/octet-stream"
@@ -415,12 +762,21 @@ async def async_request_chat_completions(

    output.latency = time.perf_counter() - output.start_time

+    if output.success and input.slo_ms is not None:
+        output.slo_achieved = (output.latency * 1000.0) <= float(input.slo_ms)
+
    if pbar:
        pbar.update(1)
    return output


-def calculate_metrics(outputs: list[RequestFuncOutput], total_duration: float):
+def calculate_metrics(
+    outputs: list[RequestFuncOutput],
+    total_duration: float,
+    requests_list: list[RequestFuncInput],
+    args,
+    slo_enabled: bool,
+):
    success_outputs = [o for o in outputs if o.success]
    error_outputs = [o for o in outputs if not o.success]

@@ -442,6 +798,29 @@ def calculate_metrics(outputs: list[RequestFuncOutput], total_duration: float):
        "peak_memory_mb_median": np.median(peak_memories) if peak_memories else 0,
    }

+    if slo_enabled:
+        slo_defined_total = 0
+        slo_met_success = 0
+
+        for req, out in zip(requests_list, outputs):
+            if req.slo_ms is None:
+                continue
+            slo_defined_total += 1
+            if out.slo_achieved is None:
+                continue
+            if out.slo_achieved:
+                slo_met_success += 1
+
+        slo_attain_all = (slo_met_success / slo_defined_total) if slo_defined_total > 0 else 0.0
+
+        metrics.update(
+            {
+                "slo_attainment_rate": slo_attain_all,
+                "slo_met_success": slo_met_success,
+                "slo_scale": getattr(args, "slo_scale", 3.0),
+            }
+        )
+
    return metrics


@@ -475,6 +854,8 @@ async def benchmark(args):

    if args.dataset == "vbench":
        dataset = VBenchDataset(args, api_url, args.model)
+    elif args.dataset == "trace":
+        dataset = TraceDataset(args, api_url, args.model)
    elif args.dataset == "random":
        dataset = RandomDataset(args, api_url, args.model)
    else:
@@ -501,6 +882,7 @@ async def benchmark(args):
    pbar = tqdm(total=len(requests_list), disable=args.disable_tqdm)

    async with aiohttp.ClientSession() as session:
+        warmup_pairs: list[tuple[RequestFuncInput, RequestFuncOutput]] = []
        if args.warmup_requests and requests_list:
            print(
                f"Running {args.warmup_requests} warmup request(s) \
@@ -513,16 +895,20 @@ async def benchmark(args):
                        warm_req,
                        num_inference_steps=args.warmup_num_inference_steps,
                    )
-                await limited_request_func(warm_req, session, None)
+                warm_out = await limited_request_func(warm_req, session, None)
+                warmup_pairs.append((warm_req, warm_out))
+
+        if args.slo:
+            # Prefer trace-provided per-request slo_ms. Only populate when missing.
+            requests_list = _populate_slo_ms_from_warmups(
+                requests_list=requests_list,
+                warmup_pairs=warmup_pairs,
+                args=args,
+            )

        start_time = time.perf_counter()
        tasks = []
-        for req in requests_list:
-            if args.request_rate != float("inf"):
-                # Poisson process: inter-arrival times follow exponential distribution
-                interval = np.random.exponential(1.0 / args.request_rate)
-                await asyncio.sleep(interval)
-
+        async for req in iter_requests(requests_list=requests_list, request_rate=args.request_rate):
            task = asyncio.create_task(limited_request_func(req, session, pbar))
            tasks.append(task)

@@ -532,7 +918,7 @@ async def benchmark(args):
    pbar.close()

    # Calculate metrics
-    metrics = calculate_metrics(outputs, total_duration)
+    metrics = calculate_metrics(outputs, total_duration, requests_list, args, args.slo)

    print("\n{s:{c}^{n}}".format(s=" Serving Benchmark Result ", n=60, c="="))

@@ -561,6 +947,12 @@ async def benchmark(args):
    print("{:<40} {:<15.4f}".format("Latency Median (s):", metrics["latency_median"]))
    print("{:<40} {:<15.4f}".format("Latency P99 (s):", metrics["latency_p99"]))

+    if args.slo:
+        print(f"{'-' * 50}")
+        print("{:<40} {:<15.2%}".format("SLO Attainment Rate (all):", metrics.get("slo_attainment_rate", 0.0)))
+        print("{:<40} {:<15}".format("SLO Met (success count):", str(metrics.get("slo_met_success", 0))))
+        print("{:<40} {:<15}".format("SLO Scale:", str(metrics.get("slo_scale", 3.0))))
+
    if metrics["peak_memory_mb_max"] > 0:
        print(f"{'-' * 50}")
        print("{:<40} {:<15.2f}".format("Peak Memory Max (MB):", metrics["peak_memory_mb_max"]))
@@ -581,16 +973,16 @@ if __name__ == "__main__":
        "--base-url",
        type=str,
        default=None,
-        help="Base URL of the server (e.g., http://localhost:30000). Overrides host/port.",
+        help="Base URL of the server (e.g., http://localhost:8091). Overrides host/port.",
    )
    parser.add_argument("--host", type=str, default="localhost", help="Server host.")
-    parser.add_argument("--port", type=int, default=30000, help="Server port.")
+    parser.add_argument("--port", type=int, default=8091, help="Server port.")
    parser.add_argument("--model", type=str, default="default", help="Model name.")
    parser.add_argument(
        "--dataset",
        type=str,
        default="vbench",
-        choices=["vbench", "random"],
+        choices=["vbench", "trace", "random"],
        help="Dataset to use.",
    )
    parser.add_argument(
@@ -645,7 +1037,7 @@ if __name__ == "__main__":
    parser.add_argument(
        "--num-inference-steps",
        type=int,
-        default=None,
+        default=50,
        help="Number of inference steps (for diffusion models).",
    )
    parser.add_argument(
@@ -656,6 +1048,21 @@ if __name__ == "__main__":
    )
    parser.add_argument("--fps", type=int, default=None, help="FPS (for video).")
    parser.add_argument("--output-file", type=str, default=None, help="Output JSON file for metrics.")
+    parser.add_argument(
+        "--slo",
+        action="store_true",
+        help=(
+            "Enable SLO calculation and reporting. If trace provides per-request slo_ms, it is used. "
+            "Otherwise, warmup request(s) are used to infer expected execution time assuming linear "
+            "scaling by resolution, frames, and steps, then slo_ms = expected_time * --slo-scale."
+        ),
+    )
+    parser.add_argument(
+        "--slo-scale",
+        type=float,
+        default=3.0,
+        help="SLO target multiplier: slo_ms = estimated_exec_time_ms * slo_scale (default: 3).",
+    )
    parser.add_argument("--disable-tqdm", action="store_true", help="Disable progress bar.")

    args = parser.parse_args()
@@ -36,7 +36,7 @@ bash benchmarks/qwen3-omni/transformers/eval_qwen3_moe_omni_transformers.sh
 ```

 What it does:
- Runs `qwen3_omni_moe_transformers.py` over `top100.txt` with `--num_prompts 100`.
+- Runs `qwen3_omni_moe_transformers.py` over `top100.txt` with `--num-prompts 100`.
 - Outputs to `benchmarks/qwen3-omni/transformers/benchmark_results/`:
  - `perf_stats.json` — aggregated & per-prompt TPS/latency (thinker/talker/code2wav/overall).
  - `results.json` — per-prompt outputs and audio paths.
@@ -53,7 +53,7 @@ bash benchmarks/qwen3-omni/vllm_omni/eval_qwen3_moe_omni.sh
 ```

 What it does:
- Runs `examples/offline_inference/qwen3_omni/end2end.py` with `--enable-stats`.
+- Runs `examples/offline_inference/qwen3_omni/end2end.py` with `--log-stats`.
 - Uses `benchmarks/build_dataset/top100.txt` and writes to:
  - Logs: `benchmarks/qwen3-omni/vllm_omni/logs/`
    - `omni_llm_pipeline_text.orchestrator.stats.jsonl` — per-stage latency stats.
@@ -16,7 +16,7 @@ if [[ ! -f "benchmarks/qwen3-omni/transformers/qwen3_omni_moe_transformers.py" ]
 else
    cd benchmarks/qwen3-omni/transformers

-    python qwen3_omni_moe_transformers.py --prompts_file ../../build_dataset/top100.txt --num_prompts 100
+    python qwen3_omni_moe_transformers.py --prompts-file ../../build_dataset/top100.txt --num-prompts 100

    echo "Logs and outputs are saved to $(pwd)/benchmark_results:"
    echo "  - perf_stats.json    Aggregated/per-prompt TPS and latency (thinker/talker/code2wav/overall)"
@@ -227,17 +227,17 @@ def print_stats(stats: dict):
 def main():
    parser = argparse.ArgumentParser(description="Qwen3-Omni Benchmark Script")
    parser.add_argument(
-        "--prompts_file",
+        "--prompts-file",
        type=str,
        default="benchmark/build_dataset/top100.txt",
        help="Path to the prompts file (one prompt per line)",
    )
    parser.add_argument(
-        "--output_dir", type=str, default="benchmark_results", help="Directory to save benchmark results"
+        "--output-dir", type=str, default="benchmark_results", help="Directory to save benchmark results"
    )
-    parser.add_argument("--model_path", type=str, default=MODEL_PATH, help="Path to the model")
+    parser.add_argument("--model-path", type=str, default=MODEL_PATH, help="Path to the model")
    parser.add_argument("--speaker", type=str, default="Ethan", help="Speaker voice for audio output")
-    parser.add_argument("--num_prompts", type=int, default=None, help="Number of prompts to process (default: all)")
+    parser.add_argument("--num-prompts", type=int, default=None, help="Number of prompts to process (default: all)")
    args = parser.parse_args()

    # Load model and processor
@@ -23,7 +23,7 @@ else
    python $end2end_script_path --output-wav $outputs_dir \
                      --query-type text \
                      --txt-prompts $build_dataset_path \
-                      --enable-stats \
+                      --log-stats \
                      --log-dir $log_dir
    echo "Logs and outputs are saved in ${log_dir} and ${outputs_dir} respectively:"
    echo "  - omni_llm_pipeline_text                       run dir/base name"
@@ -57,6 +57,7 @@ SystemEnv = namedtuple(
        "cpu_info",
        "rocm_version",  # vllm specific field
        "vllm_version",  # vllm specific field
+        "vllm_omni_version",  # vllm-omni specific field
        "vllm_build_flags",  # vllm specific field
        "gpu_topo",  # vllm specific field
        "env_vars",
@@ -289,6 +290,31 @@ def get_vllm_version():
    return __version__


+def get_vllm_omni_version(run_lambda):
+    try:
+        import vllm_omni
+        from vllm_omni import __version__, __version_tuple__
+
+        version_str = __version_tuple__[-1]
+        if isinstance(version_str, str) and version_str.startswith("g"):
+            if "." in version_str:
+                git_sha = version_str.split(".")[0][1:]
+                date = version_str.split(".")[-1][1:]
+                return f"{__version__} (git sha: {git_sha}, date: {date})"
+            else:
+                git_sha = version_str[1:]
+                return f"{__version__} (git sha: {git_sha})"
+
+        package_dir = os.path.dirname(os.path.abspath(vllm_omni.__file__))
+        git_sha = run_and_read_all(run_lambda, f"git -C {package_dir} rev-parse --short HEAD")
+        if git_sha:
+            return f"{__version__} (git sha: {git_sha})"
+
+        return __version__
+    except ImportError:
+        return "N/A (vllm_omni not installed)"
+
+
 def summarize_vllm_build_flags():
    # This could be a static method if the flags are constant, or dynamic if you need to check environment variables, etc.
    return "CUDA Archs: {}; ROCm: {}".format(
@@ -524,6 +550,7 @@ def get_env_info():

    rocm_version = get_rocm_version(run_lambda)
    vllm_version = get_vllm_version()
+    vllm_omni_version = get_vllm_omni_version(run_lambda)
    vllm_build_flags = summarize_vllm_build_flags()
    gpu_topo = get_gpu_topo(run_lambda)

@@ -555,6 +582,7 @@ def get_env_info():
        cpu_info=get_cpu_info(run_lambda),
        rocm_version=rocm_version,
        vllm_version=vllm_version,
+        vllm_omni_version=vllm_omni_version,
        vllm_build_flags=vllm_build_flags,
        gpu_topo=gpu_topo,
        env_vars=get_env_vars(),
@@ -621,6 +649,7 @@ env_info_fmt += """
 ==============================
 ROCM Version                 : {rocm_version}
 vLLM Version                 : {vllm_version}
+vLLM-Omni Version            : {vllm_omni_version}
 vLLM Build Flags:
  {vllm_build_flags}
 GPU Topology:
@@ -1,14 +1,45 @@
 ARG VLLM_BASE_IMAGE=vllm/vllm-openai
-ARG VLLM_BASE_TAG=v0.12.0
+ARG VLLM_BASE_TAG=v0.15.0
 FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG}
 ARG APP_DIR=/workspace/vllm-omni
 WORKDIR ${APP_DIR}

 COPY . .

-# Install vllm-omni into the same uv-managed Python environment used by the base image.
-RUN uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" --no-cache-dir ".[dev]"
+# Install system dependencies
+RUN apt-get update && \
+    apt-get install -y ffmpeg sox libsox-fmt-all && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/*

+# Install vLLM wheel for the specific commit
+# URL format: https://wheels.vllm.ai/<commit_hash>/
+# Current vLLM wheel commit: 2d5be1dd5ce2e44dfea53ea03ff61143da5137eb
+# Note: --prerelease=allow is required because nightly/commit wheels use
+# pre-release version numbers (e.g. 2.10.0.dev...) which uv ignores by default,
+# causing it to fall back to the older stable release from PyPI.
+RUN uv pip install --system --upgrade --force-reinstall --prerelease=allow vllm --extra-index-url https://wheels.vllm.ai/2d5be1dd5ce2e44dfea53ea03ff61143da5137eb
+
+
+RUN uv pip uninstall --system nvidia-cublas-cu12
+# Install vllm-omni into the same uv-managed Python environment used by the base image.
+# Use bash -c so that $(python3 -c ...) is expanded inside the container.
+RUN uv pip install --system --no-cache-dir ".[dev]"
+
+# pplx_kernels is optional; the base image may already ship a version.
+# Reinstall attempt may fail if no compatible wheel exists for the current
+# PyTorch/CUDA combination – allow the build to continue regardless.
+RUN pip install --upgrade --force-reinstall "pplx_kernels" || true
+
+RUN /bin/bash -c 'FLASHINFER_CUDA_TAG="$(python3 -c "import torch; print((torch.version.cuda or \"12.4\").replace(\".\", \"\"))")" && \
+    uv pip install --system --upgrade --force-reinstall \
+      "flashinfer-python==0.6.3" \
+      "flashinfer-cubin==0.6.3" \
+      "flashinfer-jit-cache==0.6.3" \
+      --extra-index-url "https://flashinfer.ai/whl/cu${FLASHINFER_CUDA_TAG}"'
+# Pin nvidia-cublas-cu12 after vllm-omni install so dependency resolution does not overwrite it
+RUN uv pip install --system --upgrade --force-reinstall "nvidia-cublas-cu12==12.9.1.4"
+RUN uv pip install --system --upgrade --force-reinstall "numpy==2.2.6"
 RUN ln -sf /usr/bin/python3 /usr/bin/python

 ENTRYPOINT []
@@ -1,5 +1,5 @@
-ARG VLLM_ASCEND_IMAGE=quay.nju.edu.cn/ascend/vllm-ascend
-ARG VLLM_ASCEND_TAG=v0.11.0rc2
+ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
+ARG VLLM_ASCEND_TAG=v0.14.0rc1
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}

 ARG APP_DIR=/vllm-workspace/vllm-omni
@@ -8,7 +8,7 @@ WORKDIR ${APP_DIR}
 COPY . .

 # Install vllm-omni with dev dependencies
-RUN pip install --no-cache-dir -e ".[dev]"
+RUN pip install --no-cache-dir -e . --no-build-isolation

 ENV VLLM_WORKER_MULTIPROC_METHOD=spawn

@@ -0,0 +1,15 @@
+ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
+ARG VLLM_ASCEND_TAG=v0.14.0rc1-a3
+FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
+
+ARG APP_DIR=/vllm-workspace/vllm-omni
+WORKDIR ${APP_DIR}
+
+COPY . .
+
+# Install vllm-omni with dev dependencies
+RUN pip install --no-cache-dir -e . --no-build-isolation
+
+ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
+
+ENTRYPOINT []
@@ -1,9 +1,7 @@
-ARG BASE_IMAGE=rocm/vllm-dev:nightly_main_20251205
-FROM ${BASE_IMAGE}
+ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.15.0
+FROM ${BASE_IMAGE} AS final

 ARG COMMON_WORKDIR=/app
-ARG VLLM_VERSION=v0.12.0
-ARG PYTORCH_ROCM_ARCH="gfx942;gfx950"

 WORKDIR ${COMMON_WORKDIR}

@@ -13,30 +11,18 @@ RUN apt-get update && \
    apt-get clean && \
    rm -rf /var/lib/apt/lists/*

-# Step 2: Reinstall vllm from source
-RUN python3 -m pip uninstall -y vllm && rm -rf vllm &&\
-    git clone https://github.com/vllm-project/vllm.git && \
-    cd vllm && \
-    git checkout ${VLLM_VERSION} && \
-    python3 -m pip install -r requirements/rocm.txt && \
-    python3 setup.py clean --all && \
-    PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH} python3 setup.py develop && \
-    cd ../ && \
-    rm -rf vllm/.git
-
 RUN mkdir -p ${COMMON_WORKDIR}/vllm-omni

-# Step 3: Copy vllm-omni code and install without uv
+# Step 2: Copy vllm-omni code and install without uv
 COPY . ${COMMON_WORKDIR}/vllm-omni
-RUN cd ${COMMON_WORKDIR}/vllm-omni && python3 -m pip install --no-cache-dir ".[dev]"
+RUN cd ${COMMON_WORKDIR}/vllm-omni && uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" --no-cache-dir ".[dev]" --no-build-isolation

-# Create python symlink
-# `GPU_ARCHS` is an environment variable that is used to set the GPU archs for the AITER.
-# This is needed to prevent the AITER automatic GPU arch detection from failing on MI325X.
-# The AITER version used in this dockerfile has issues with handling
-# the GPU archs of MI325X (CI machine) correctly. So we manually set the GPU archs here.
-# We reuse AITER_ROCM_ARCH from the base image to avoid duplication.
-ENV GPU_ARCHS=${AITER_ROCM_ARCH}
 RUN ln -sf /usr/bin/python3 /usr/bin/python

+CMD ["/bin/bash"]
+
 ENTRYPOINT []
+
+#Set entrypoint for vllm-openai official images
+FROM final AS vllm-openai
+ENTRYPOINT ["vllm", "serve", "--omni"]
@@ -3,36 +3,52 @@ nav:
 - User Guide:
  - Getting Started:
    - getting_started/quickstart.md
-    - getting_started/installation
+    - getting_started/installation/*
  - Serving:
    - OpenAI-Compatible API:
      - Image Generation: serving/image_generation_api.md
+      - Image Edit: serving/image_edit_api.md
+      - Text to Speech: serving/speech_api.md
  - Examples:
    - examples/README.md
    - Offline Inference:
+      - BAGEL-7B-MoT: user_guide/examples/offline_inference/bagel.md
      - Image-To-Image: user_guide/examples/offline_inference/image_to_image.md
+      - Image-To-Video: user_guide/examples/offline_inference/image_to_video.md
+      - LoRA Inference(Diffusion): user_guide/examples/offline_inference/lora_inference.md
      - Qwen2.5-Omni: user_guide/examples/offline_inference/qwen2_5_omni.md
      - Qwen3-Omni: user_guide/examples/offline_inference/qwen3_omni.md
+      - Qwen3-TTS: user_guide/examples/offline_inference/qwen3_tts.md
+      - Text-To-Audio: user_guide/examples/offline_inference/text_to_audio.md
      - Text-To-Image: user_guide/examples/offline_inference/text_to_image.md
      - Text-To-Video: user_guide/examples/offline_inference/text_to_video.md
    - Online Serving:
+      - BAGEL-7B-MoT: user_guide/examples/online_serving/bagel.md
      - Image-To-Image: user_guide/examples/online_serving/image_to_image.md
+      - LoRA Inference(Diffusion): user_guide/examples/online_serving/lora_inference.md
      - Qwen2.5-Omni: user_guide/examples/online_serving/qwen2_5_omni.md
      - Qwen3-Omni: user_guide/examples/online_serving/qwen3_omni.md
+      - Qwen3-TTS: user_guide/examples/online_serving/qwen3_tts.md
      - Text-To-Image: user_guide/examples/online_serving/text_to_image.md
  - General:
    - usage/*
  - Configuration:
    - configuration/README.md
    - configuration/*
-  - Diffusion Acceleration:
-    - Overview: user_guide/diffusion_acceleration.md
-    - Acceleration Methods:
-      - TeaCache: user_guide/acceleration/teacache.md
-      - Cache-DiT: user_guide/acceleration/cache_dit_acceleration.md
-      - Parallelism Acceleration: user_guide/acceleration/parallelism_acceleration.md
  - Models:
    - models/supported_models.md
+  - Features:
+    - Sleep Mode: features/sleep_mode.md
+    - Diffusion Features:
+      - Overview: user_guide/diffusion_acceleration.md
+      - TeaCache: user_guide/diffusion/teacache.md
+      - Cache-DiT: user_guide/diffusion/cache_dit_acceleration.md
+      - Quantization:
+        - Overview: user_guide/diffusion/quantization/overview.md
+        - FP8: user_guide/diffusion/quantization/fp8.md
+      - Parallelism Acceleration: user_guide/diffusion/parallelism_acceleration.md
+      - CPU Offloading: user_guide/diffusion/cpu_offload_diffusion.md
+    - ComfyUI: features/comfyui.md
 - Developer Guide:
  - General:
    - contributing/README.md
@@ -42,15 +58,20 @@ nav:
    - contributing/model/README.md
    - contributing/model/adding_omni_model.md
    - contributing/model/adding_diffusion_model.md
+    - Advanced Features:
+      - contributing/features/cfg_parallel.md
+      - contributing/features/sequence_parallel.md
+      - contributing/features/tensor_parallel.md
+      - contributing/features/cache_dit.md
+      - contributing/features/teacache.md
  - CI: contributing/ci
-  - Tests: contributing/tests
  - Design Documents:
    - design/index.md
    - design/architecture_overview.md
    - Feature Design:
      - design/feature/disaggregated_inference.md
-      - design/feature/multi_request_streaming.md
      - design/feature/ray_based_execution.md
+      - design/feature/omni_connectors/
    - Module Design:
      - design/module/ar_module.md
      - design/module/dit_module.md
@@ -60,5 +60,5 @@ vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, in

 For more information, checkout the following:

- [vllm-omni architecture design and recent roadmaps](https://docs.google.com/presentation/d/1Y7t2Zm3BIISPN-_X_sDpxZMNCWir2rAx/edit?usp=drive_link&ouid=110473603432222024453&rtpof=true&sd=true)
+- [vllm-omni architecture design and recent roadmaps](https://docs.google.com/presentation/d/1qv4qMW1rKAqDREMXiUDLIgqqHQe7TDPj/edit?usp=sharing&ouid=110473603432222024453&rtpof=true&sd=true)
 - [vllm-omni announcement blogpost](https://blog.vllm.ai/2025/11/30/vllm-omni.html)
@@ -7,14 +7,12 @@ Main entry points for vLLM-Omni inference and serving.
 - [vllm_omni.entrypoints.async_omni.AsyncOmni][]
 - [vllm_omni.entrypoints.async_omni_diffusion.AsyncOmniDiffusion][]
 - [vllm_omni.entrypoints.async_omni_llm.AsyncOmniLLM][]
- [vllm_omni.entrypoints.chat_utils.OmniAsyncMultiModalContentParser][]
- [vllm_omni.entrypoints.chat_utils.OmniAsyncMultiModalItemTracker][]
- [vllm_omni.entrypoints.chat_utils.parse_chat_messages_futures][]
+- [vllm_omni.entrypoints.chat_utils.extract_audio_from_video_async][]
+- [vllm_omni.entrypoints.cli.benchmark.base.OmniBenchmarkSubcommandBase][]
+- [vllm_omni.entrypoints.cli.benchmark.main.OmniBenchmarkSubcommand][]
+- [vllm_omni.entrypoints.cli.benchmark.serve.OmniBenchmarkServingSubcommand][]
 - [vllm_omni.entrypoints.cli.serve.OmniServeCommand][]
 - [vllm_omni.entrypoints.client_request_state.ClientRequestState][]
- [vllm_omni.entrypoints.log_utils.OrchestratorMetrics][]
- [vllm_omni.entrypoints.log_utils.StageRequestMetrics][]
- [vllm_omni.entrypoints.log_utils.StageStats][]
 - [vllm_omni.entrypoints.omni.Omni][]
 - [vllm_omni.entrypoints.omni.OmniBase][]
 - [vllm_omni.entrypoints.omni_diffusion.OmniDiffusion][]
@@ -26,17 +24,17 @@ Main entry points for vLLM-Omni inference and serving.

 Input data structures for multi-modal inputs.

+- [vllm_omni.inputs.data.OmniDiffusionSamplingParams][]
 - [vllm_omni.inputs.data.OmniEmbedsPrompt][]
+- [vllm_omni.inputs.data.OmniTextPrompt][]
 - [vllm_omni.inputs.data.OmniTokenInputs][]
 - [vllm_omni.inputs.data.OmniTokensPrompt][]
- [vllm_omni.inputs.parse.parse_singleton_prompt_omni][]
 - [vllm_omni.inputs.preprocess.OmniInputPreprocessor][]

 ## Engine

 Engine classes for offline and online inference.

- [vllm_omni.diffusion.diffusion_engine.BackgroundResources][]
 - [vllm_omni.diffusion.diffusion_engine.DiffusionEngine][]
 - [vllm_omni.engine.AdditionalInformationEntry][]
 - [vllm_omni.engine.AdditionalInformationPayload][]
@@ -54,40 +52,17 @@ Engine classes for offline and online inference.

 Core scheduling and caching components.

+- [vllm_omni.core.sched.omni_ar_scheduler.KVCacheTransferData][]
 - [vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler][]
 - [vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler][]
+- [vllm_omni.core.sched.output.OmniCachedRequestData][]
 - [vllm_omni.core.sched.output.OmniNewRequestData][]
-
-## Model Executor
-
-Model execution components.
-
- [vllm_omni.model_executor.custom_process_mixin.CustomProcessMixin][]
- [vllm_omni.model_executor.models.output_templates.OmniOutput][]
- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni.Qwen2_5OmniForConditionalGeneration][]
- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_talker.Qwen2_5OmniTalkerForConditionalGeneration][]
- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker.Qwen2_5OmniConditionalGenerationMixin][]
- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker.Qwen2_5OmniThinkerForConditionalGeneration][]
- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_token2wav.Qwen2_5OmniToken2WavBigVGANModel][]
- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_token2wav.Qwen2_5OmniToken2WavDiTModel][]
- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_token2wav.Qwen2_5OmniToken2WavForConditionalGenerationVLLM][]
- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_token2wav.Qwen2_5OmniToken2WavModel][]
- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_old.Qwen2EmbeddingModel][]
- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_old.Qwen2ForCausalLM][]
- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_old.Qwen2Model][]
- [vllm_omni.model_executor.models.qwen3_omni.qwen3_moe.Qwen3MoeForCausalLM][]
- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni.Qwen3OmniMoeForConditionalGeneration][]
- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_code2wav.Qwen3OmniMoeCode2Wav][]
- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_code_predictor_mtp.Qwen3OmniCodePredictorBaseModel][]
- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_code_predictor_mtp.Qwen3OmniMoeTalkerCodePredictor][]
- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_talker.Qwen3OmniMoeModel][]
- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_talker.Qwen3OmniMoeTalkerForConditionalGeneration][]
- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3MoeLLMForCausalLM][]
- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3MoeLLMModel][]
- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3OmniMoeConditionalGenerationMixin][]
- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3OmniMoeThinkerForConditionalGeneration][]
- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3OmniMoeThinkerMultiModalProcessor][]
- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3OmniMoeThinkerProcessingInfo][]
+- [vllm_omni.core.sched.output.OmniSchedulerOutput][]
+- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.vq.core_vq.DistributedGroupResidualVectorQuantization][]
+- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.vq.core_vq.DistributedResidualVectorQuantization][]
+- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.vq.core_vq.EuclideanCodebook][]
+- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.vq.core_vq.VectorQuantization][]
+- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.vq.core_vq.preprocess][]

 ## Configuration

@@ -97,24 +72,41 @@ Configuration classes.
 - [vllm_omni.diffusion.cache.teacache.config.TeaCacheConfig][]
 - [vllm_omni.distributed.omni_connectors.utils.config.ConnectorSpec][]
 - [vllm_omni.distributed.omni_connectors.utils.config.OmniTransferConfig][]
+- [vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts.Qwen3TTSConfig][]
+- [vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts.Qwen3TTSSpeakerEncoderConfig][]
+- [vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts.Qwen3TTSTalkerCodePredictorConfig][]
+- [vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts.Qwen3TTSTalkerConfig][]
+- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2.Qwen3TTSTokenizerV2Config][]
+- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2.Qwen3TTSTokenizerV2DecoderConfig][]
+- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1Config][]
+- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1DecoderBigVGANConfig][]
+- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1DecoderConfig][]
+- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1DecoderDiTConfig][]
+- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1EncoderConfig][]

 ## Workers

 Worker classes and model runners for distributed inference.

- [vllm_omni.diffusion.worker.gpu_worker.GPUWorker][]
- [vllm_omni.diffusion.worker.gpu_worker.WorkerProc][]
- [vllm_omni.diffusion.worker.npu.npu_worker.NPUWorker][]
- [vllm_omni.diffusion.worker.npu.npu_worker.NPUWorkerProc][]
+- [vllm_omni.diffusion.worker.diffusion_model_runner.DiffusionModelRunner][]
+- [vllm_omni.diffusion.worker.diffusion_worker.DiffusionWorker][]
+- [vllm_omni.diffusion.worker.diffusion_worker.WorkerProc][]
+- [vllm_omni.platforms.npu.worker.npu_ar_model_runner.ExecuteModelState][]
+- [vllm_omni.platforms.npu.worker.npu_ar_model_runner.NPUARModelRunner][]
+- [vllm_omni.platforms.npu.worker.npu_ar_worker.NPUARWorker][]
+- [vllm_omni.platforms.npu.worker.npu_generation_model_runner.NPUGenerationModelRunner][]
+- [vllm_omni.platforms.npu.worker.npu_generation_worker.NPUGenerationWorker][]
+- [vllm_omni.platforms.npu.worker.npu_model_runner.OmniNPUModelRunner][]
+- [vllm_omni.platforms.xpu.worker.xpu_ar_model_runner.XPUARModelRunner][]
+- [vllm_omni.platforms.xpu.worker.xpu_ar_worker.XPUARWorker][]
+- [vllm_omni.platforms.xpu.worker.xpu_generation_model_runner.XPUGenerationModelRunner][]
+- [vllm_omni.platforms.xpu.worker.xpu_generation_worker.XPUGenerationWorker][]
+- [vllm_omni.worker.base.OmniGPUWorkerBase][]
 - [vllm_omni.worker.gpu_ar_model_runner.ExecuteModelState][]
 - [vllm_omni.worker.gpu_ar_model_runner.GPUARModelRunner][]
 - [vllm_omni.worker.gpu_ar_worker.GPUARWorker][]
 - [vllm_omni.worker.gpu_generation_model_runner.GPUGenerationModelRunner][]
 - [vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker][]
+- [vllm_omni.worker.gpu_memory_utils.parse_cuda_visible_devices][]
 - [vllm_omni.worker.gpu_model_runner.OmniGPUModelRunner][]
- [vllm_omni.worker.npu.npu_ar_model_runner.ExecuteModelState][]
- [vllm_omni.worker.npu.npu_ar_model_runner.NPUARModelRunner][]
- [vllm_omni.worker.npu.npu_ar_worker.NPUARWorker][]
- [vllm_omni.worker.npu.npu_generation_model_runner.NPUGenerationModelRunner][]
- [vllm_omni.worker.npu.npu_generation_worker.NPUGenerationWorker][]
- [vllm_omni.worker.npu.npu_model_runner.OmniNPUModelRunner][]
+- [vllm_omni.worker.mixins.OmniWorkerMixin][]
@@ -22,3 +22,21 @@ If you have custom stage configs file, launch the server with command below
 ```bash
 vllm serve Qwen/Qwen2.5-Omni-7B --omni --stage-configs-path /path/to/stage_configs_file
 ```
+
+
+## bench
+
+Run benchmark tests for online serving throughput.
+Available Commands:
+
+```bash
+vllm bench serve --omni \
+    --model Qwen/Qwen2.5-Omni-7B \
+    --host server-host \
+    --port server-port \
+    --random-input-len 32 \
+    --random-output-len 4  \
+    --num-prompts  5
+```
+
+See [vllm bench serve](./bench/serve.md) for the full reference of all available arguments.
@@ -0,0 +1,359 @@
+# vLLM-Omni Benchmark CLI Guide
+The vllm bench command launches the vLLM-Omni benchmark to evaluate the performance of multimodal models.
+
+## Notes
+We currently only support using the "openai-chat-omni" backend.
+
+## Basic Parameter Description
+You can use `vllm bench serve --omni --help=all` to get descriptions of all parameters. The commonly used parameters are described below:
+- `--omni`  
+  Enable Omni (multimodal) mode, supporting multimodal inputs and outputs such as images, videos, and audio.
+
+- `--backend`  
+  Specify the backend adapter as openai-chat-omni, using OpenAI Chat compatible API behavior as the protocol. Currently only openai-chat-omni is supported.
+
+- `--model`  
+  The model identifier to load, filled according to the models supported by vLLM-Omni.
+
+- `--endpoint`  
+  The API endpoint exposed externally, to which clients send their requests.
+
+- `--dataset-name`  
+  The name of the dataset used; random-mm indicates generating random multimodal inputs (images, videos, audio).
+
+- `--num-prompts`  
+  The total number of requests to send, an integer.
+
+- `--max-concurrency`  
+  "Maximum number of concurrent requests. This can be used "
+        "to help simulate an environment where a higher level component "
+        "is enforcing a maximum number of concurrent requests. While the "
+        "--request-rate argument controls the rate at which requests are "
+        "initiated, this argument will control how many are actually allowed "
+        "to execute at a time. This means that when used in combination, the "
+        "actual request rate may be lower than specified with --request-rate, "
+        "if the server is not processing requests fast enough to keep up."
+
+- `--request-rate`  
+  "Number of requests per second. If this is inf, "
+        "then all the requests are sent at time 0. "
+        "Otherwise, we use Poisson process or gamma distribution "
+        "to synthesize the request arrival times."
+
+- `--ignore-eos`  
+  "Set ignore_eos flag when sending the benchmark request."
+
+- `--metric-percentiles`  
+  Comma-separated list of percentiles for selected metrics. "
+        "To report 25-th, 50-th, and 75-th percentiles, use \"25,50,75\". "
+        "Default value is \"99\"."
+        "Use \"--percentile-metrics\" to select metrics.
+
+- `--percentile-metrics`  
+        "Comma-separated list of selected metrics to report percentiles."
+                    "This argument specifies the metrics to report percentiles."
+                    'Allowed metric names are "ttft", "tpot", "itl", "e2el", "audio_ttfp", "audio_rtf", "audio_duration". '
+
+- `--save-result`  
+Specify to save benchmark results to a json file
+
+- `--save-detailed`  
+"When saving the results, whether to include per request "
+        "information such as response, error, ttfs, tpots, etc."
+
+- `--result-dir`  
+ "Specify directory to save benchmark json results."
+        "If not specified, results are saved in the current directory."
+
+- `--result-filename`  
+"Specify the filename to save benchmark json results."
+        "If not specified, results will be saved in "
+        "{label}-{args.request_rate}qps-{base_model_id}-{current_dt}.json"
+
+- `--random-prefix-len`  
+  Number of fixed prefix tokens before the random context in a request.
+  The total input length is the sum of random-prefix-len and a random
+  context length sampled from [input_len * (1 - range_ratio),
+  input_len * (1 + range_ratio)].Only the random and random-mm modes
+  support this parameter.
+
+- `--random-input-len`  
+  Number of input tokens per request.Only the random and random-mm modes support this parameter.
+
+- `--random-output-len`  
+  Number of output tokens per request.Only the random and random-mm modes support this parameter.
+
+- `--random-range-ratio`  
+  Range ratio for sampling input/output length,
+  used only for random sampling. Must be in the range [0, 1) to define
+  a symmetric sampling range
+  [length * (1 - range_ratio), length * (1 + range_ratio)].
+  Only the random and random-mm modes support this parameter.
+
+- `--random-mm-base-items-per-request`  
+  Base number of multimodal items per request for random-mm.
+  Actual per-request count is sampled around this base using
+  --random-mm-num-mm-items-range-ratio.
+  Only the random-mm mode supports this parameter.
+
+- `--random-mm-limit-mm-per-prompt`  
+  Per-modality hard caps for items attached per request, e.g.
+  '{"image": 3, "video": 1, "audio": 1}'. The sampled per-request item
+  count is clamped to the sum of these limits. When a modality
+  reaches its cap, its buckets are excluded and probabilities are
+  renormalized.
+  Only the random-mm mode supports this parameter.
+
+- `--random-mm-num-mm-items-range-ratio`  
+  Range ratio r in [0, 1] for sampling items per request.
+  We sample uniformly from the closed integer range
+  [floor(n*(1-r)), ceil(n*(1+r))]
+  where n is the base items per request.
+  r=0 keeps it fixed; r=1 allows 0 items. The maximum is clamped
+  to the sum of per-modality limits from
+  --random-mm-limit-mm-per-prompt.
+  An error is raised if the computed min exceeds the max.
+  Only the random-mm mode supports this parameter.
+
+- `--random-mm-bucket-config`  
+  The bucket config is a dictionary mapping a multimodal item
+  sampling configuration to a probability.
+  Currently allows for 3 modalities: audio, images and videos.
+  A bucket key is a tuple of (height, width, num_frames)
+  The value is the probability of sampling that specific item.
+  Example:
+  --random-mm-bucket-config
+  "{(256, 256, 1): 0.5, (720, 1280, 16): 0.4, (0, 1, 5): 0.10}"
+  First item: images with resolution 256x256 w.p. 0.5
+  Second item: videos with resolution 720x1280 and 16 frames
+  Third item: audios with 1s duration and 5 channels w.p. 0.1
+  OBS.: If the probabilities do not sum to 1, they are normalized.
+  Only the random-mm mode supports this parameter
+
+## Usage Examples
+
+### Online Benchmark
+<details class="admonition abstract" markdown="1">
+<summary>Show more</summary>
+
+First start serving your model:
+
+```bash
+vllm serve Qwen/Qwen2.5-Omni-7B --omni
+```
+
+Then run the benchmarking for sharegpt:
+
+```bash
+# download dataset
+# wget https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_V3_unfiltered_cleaned_split.json
+vllm bench serve \
+  --omni \
+  --port 43845 \
+  --model /home/models/Qwen/Qwen3-Omni-30B-A3B-Instruct \
+  --endpoint /v1/chat/completions \
+  --backend openai-chat-omni \
+  --num-prompts 2 \
+  --dataset-name sharegpt \
+  --dataset-path ShareGPT_V3_unfiltered_cleaned_split.json \
+  --percentile-metrics ttft,tpot,itl,e2el
+```
+If successful, you will see the following output:
+```text
+============ Serving Benchmark Result ============
+Successful requests:                     2
+Failed requests:                         0
+Benchmark duration (s):                  81.63
+Request throughput (req/s):              0.02
+Peak concurrent requests:                2.00
+----------------End-to-end Latency----------------
+Mean E2EL (ms):                          56966.13
+Median E2EL (ms):                        56966.13
+P99 E2EL (ms):                           81016.80
+================== Text Result ===================
+Total input tokens:                      36
+Total generated tokens:                  5926
+Output token throughput (tok/s):         72.60
+Peak output token throughput (tok/s):    103.00
+Peak concurrent requests:                2.00
+Total Token throughput (tok/s):          73.04
+---------------Time to First Token----------------
+Mean TTFT (ms):                          124.76
+Median TTFT (ms):                        124.76
+P99 TTFT (ms):                           156.10
+-----Time per Output Token (excl. 1st token)------
+Mean TPOT (ms):                          481.30
+Median TPOT (ms):                        481.30
+P99 TPOT (ms):                           947.55
+---------------Inter-token Latency----------------
+Mean ITL (ms):                           25.11
+Median ITL (ms):                         0.33
+P99 ITL (ms):                            25.17
+================== Audio Result ==================
+Total audio duration generated(s):       3.95
+Total audio frames generated:            94890
+Audio throughput(audio duration/s):      0.05
+==================================================
+```
+
+Or run the benchmarking for random:
+
+```bash
+vllm bench serve \
+  --omni \
+  --port 43845 \
+  --endpoint /v1/chat/completions \
+  --backend openai-chat-omni \
+  --model /home/models/Qwen/Qwen3-Omni-30B-A3B-Instruct \
+  --dataset-name random \
+  --num-prompts 2 \
+  --random-prefix-len 5 \
+  --random-input-len 10 \
+  --random-output-len 100 \
+  --percentile-metrics ttft,tpot,itl,e2el,audio_ttfp,audio_rtf \
+  --ignore-eos
+```
+
+If successful, you will see the following output:
+
+```text
+============ Serving Benchmark Result ============
+Successful requests:                     2
+Failed requests:                         0
+Benchmark duration (s):                  3.89
+Request throughput (req/s):              0.51
+Peak concurrent requests:                2.00
+----------------End-to-end Latency----------------
+Mean E2EL (ms):                          3824.76
+Median E2EL (ms):                        3824.76
+P99 E2EL (ms):                           3888.54
+================== Text Result ===================
+Total input tokens:                      30
+Total generated tokens:                  10101
+Output token throughput (tok/s):         2595.57
+Peak output token throughput (tok/s):    111.00
+Peak concurrent requests:                2.00
+Total Token throughput (tok/s):          2603.28
+---------------Time to First Token----------------
+Mean TTFT (ms):                          117.15
+Median TTFT (ms):                        117.15
+P99 TTFT (ms):                           142.69
+-----Time per Output Token (excl. 1st token)------
+Mean TPOT (ms):                          0.73
+Median TPOT (ms):                        0.73
+P99 TPOT (ms):                           0.74
+---------------Inter-token Latency----------------
+Mean ITL (ms):                           16.47
+Median ITL (ms):                         16.19
+P99 ITL (ms):                            52.55
+================== Audio Result ==================
+Total audio duration generated(s):       15.79
+Total audio frames generated:            379050
+Audio throughput(audio duration/s):      4.06
+---------------Time to First Packet---------------
+Mean AUDIO_TTFP (ms):                    3701.37
+Median AUDIO_TTFP (ms):                  3701.37
+P99 AUDIO_TTFP (ms):                     3762.25
+-----------------Real Time Factor-----------------
+Mean AUDIO_RTF:                          0.47
+Median AUDIO_RTF:                        0.47
+P99 AUDIO_RTF:                           0.48
+==================================================
+```
+Notes:
+We use audio generation time / audio duration to calculate RTF.
+
+</details>
+
+### Multi-Modal Benchmark
+
+<details class="admonition abstract" markdown="1">
+<summary>Show more</summary>
+
+Benchmark the performance of multi-modal requests in vLLM-Omni.
+
+Generate synthetic image、video、audio inputs alongside random text prompts to stress-test vision models without external datasets.
+
+Notes:
+
+- Works only with online benchmark via the OpenAI backend (`--backend openai-chat-omni`) and endpoint `/v1/chat/completions`.
+
+Start the server (example):
+
+```bash
+vllm serve Qwen/Qwen2.5-Omni-7B --omni
+```
+
+It is recommended to use the flag `--ignore-eos` to simulate real responses. You can set the size of the output via the arg `random-output-len`.
+
+Then run the benchmarking script:
+```bash
+vllm bench serve \
+    --omni \
+  --dataset-name random-mm \
+  --port 40849 \
+  --model /home/models/Qwen/Qwen3-Omni-30B-A3B-Instruct \
+  --endpoint /v1/chat/completions \
+  --backend openai-chat-omni \
+  --request-rate 1 \
+  --num-prompts 1 \
+  --random-input-len 10 \
+  --random-range-ratio 0.0 \
+  --random-mm-base-items-per-request 2 \
+  --random-mm-num-mm-items-range-ratio 0 \
+  --random-mm-limit-mm-per-prompt '{"image":1,"video":1, "audio": 1}' \
+  --random-mm-bucket-config '{"(32, 32, 1)": 0.5, "(0, 1, 1)": 0.1, "(32, 32, 2)":0.4}' \
+  --ignore-eos \
+  --percentile-metrics ttft,tpot,itl \
+  --random-output-len 2 \
+  --extra_body '{"modalities": ["text"]}'
+```
+
+If successful, you will see the following output:
+
+```text
+============ Serving Benchmark Result ============
+Successful requests:                     1
+Failed requests:                         0
+Request rate configured (RPS):           1.00
+Benchmark duration (s):                  1.21
+Request throughput (req/s):              0.83
+Peak concurrent requests:                1.00
+================== Text Result ===================
+Total input tokens:                      10
+Total generated tokens:                  3
+Output token throughput (tok/s):         2.49
+Peak output token throughput (tok/s):    3.00
+Peak concurrent requests:                1.00
+Total Token throughput (tok/s):          10.77
+---------------Time to First Token----------------
+Mean TTFT (ms):                          179.74
+Median TTFT (ms):                        179.74
+P99 TTFT (ms):                           179.74
+-----Time per Output Token (excl. 1st token)------
+Mean TPOT (ms):                          12.76
+Median TPOT (ms):                        12.76
+P99 TPOT (ms):                           12.76
+---------------Inter-token Latency----------------
+Mean ITL (ms):                           12.76
+Median ITL (ms):                         12.76
+P99 ITL (ms):                            25.24
+================== Audio Result ==================
+Total audio duration generated(s):       0.00
+Total audio frames generated:            0
+Audio throughput(audio duration/s):      0.00
+==================================================
+```
+
+Behavioral notes:
+
+- If the requested base item count cannot be satisfied under the provided per-prompt limits, the tool raises an error rather than silently clamping.
+
+How sampling works:
+
+- Determine per-request item count k by sampling uniformly from the integer range defined by `--random-mm-base-items-per-request` and `--random-mm-num-mm-items-range-ratio`, then clamp k to at most the sum of per-modality limits.
+- For each of the k items, sample a bucket (H, W, T) according to the normalized probabilities in `--random-mm-bucket-config`, while tracking how many items of each modality have been added.
+- If a modality (e.g., image) reaches its limit from `--random-mm-limit-mm-per-prompt`, all buckets of that modality are excluded and the remaining bucket probabilities are renormalized before continuing.
+This should be seen as an edge case, and if this behavior can be avoided by setting `--random-mm-limit-mm-per-prompt` to a large number. Note that this might result in errors due to engine config `--limit-mm-per-prompt`.
+- The resulting request contains synthetic image data in `multi_modal_data` (OpenAI Chat format). When `random-mm` is used with the OpenAI Chat backend, prompts remain text and MM content is attached via `multi_modal_data`.
+</details>
@@ -2,7 +2,7 @@

 This section lists the most common options for running vLLM-Omni.

-For options within a vLLM Engine. Please refer to [vLLM Configuration](https://docs.vllm.ai/en/v0.11.0/configuration/index.html)
+For options within a vLLM Engine. Please refer to [vLLM Configuration](https://docs.vllm.ai/en/v0.14.0/configuration/index.html)

 Currently, the main options are maintained by stage configs for each model.

@@ -16,6 +16,6 @@ For introduction, please check [Introduction for stage config](./stage_configs.m

 ## Optimization Features

- **[TeaCache Configuration](../user_guide/acceleration/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/acceleration/cache_dit_acceleration.md)** - Enable Cache-DiT as cache acceleration backends for DiT models
- **[Parallelism Configuration](../user_guide/acceleration/parallelism_acceleration.md)** - Enable parallelism (e.g., sequence parallelism) for for DiT models
+- **[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
@@ -40,7 +40,7 @@ stage_args:
    engine_args: # Engine arguments for a certain engine
      model_stage: thinker
      model_arch: Qwen2_5OmniForConditionalGeneration # The model implementation registered in model_executor/models/registry.py
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker # The specific worker used
+      worker_type: ar # The specific worker used
      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler # The specific scehduler used
      gpu_memory_utilization: 0.8 # The gpu memory allocation for the stage within a single chip
      enforce_eager: true  # Now we only support eager mode
@@ -66,7 +66,7 @@ stage_args:
    engine_args:
      model_stage: talker
      model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_type: ar
      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
      gpu_memory_utilization: 0.8
      enforce_eager: true
@@ -92,7 +92,7 @@ stage_args:
    engine_args:
      model_stage: code2wav
      model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
+      worker_type: generation
      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
      gpu_memory_utilization: 0.15
      enforce_eager: true
@@ -8,7 +8,7 @@ stage_args:
    engine_args:
      model_stage: thinker
      model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_type: ar
      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
      gpu_memory_utilization: 0.8
      enforce_eager: true  # Now we only support eager mode
@@ -34,7 +34,7 @@ stage_args:
    engine_args:
      model_stage: talker
      model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_type: ar
      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
      gpu_memory_utilization: 0.8
      enforce_eager: true
@@ -60,7 +60,7 @@ stage_args:
    engine_args:
      model_stage: code2wav
      model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
+      worker_type: generation
      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
      gpu_memory_utilization: 0.15
      enforce_eager: true
@@ -0,0 +1,678 @@
+# Multi-Level Automated Testing System Documentation
+
+## Document Overview
+
+This testing system aims to build a complete, efficient, and well-structured quality assurance framework for the development, integration, and release of model services. It draws on the concept of the test pyramid from modern software engineering, progressively expanding testing activities from basic code logic verification to complex end-to-end (E2E) functionality, performance, accuracy, and even long-term stability validation.
+
+Through five levels (L1-L5) and common (Common) specifications, the system clarifies the testing objectives, scope, execution frequency, and required resources for different development stages (e.g., each commit, PR merge, daily build, pre-release). This ensures that models meet high standards for functionality, performance, and reliability across various deployment scenarios (online serving and offline inference).
+
+<table>
+  <thead>
+    <tr>
+      <th>Level</th>
+      <th>Scope & Focus</th>
+      <th>Time Cost</th>
+      <th>Test Dir</th>
+      <th>Doc</th>
+      <th>Frequency</th>
+      <th>Hardware</th>
+    </tr>
+  </thead>
+  <tbody>
+    <tr>
+      <td rowspan="2"><strong>Common</strong></td>
+      <td>Contribution Guideline & PR checklist</td>
+      <td>/</td>
+      <td>/</td>
+      <td>.github/PULL_REQUEST_TEMPLATE.md <a href="../tests_style/"> Test Style (PR Checklist)</a></td>
+      <td>/</td>
+      <td>/</td>
+    </tr>
+    <tr>
+      <td>CI Failure Description</td>
+      <td>/</td>
+      <td>/</td>
+      <td><a href="../failures/"> CI Failures</a></td>
+      <td>/</td>
+      <td>/</td>
+    </tr>
+    <tr>
+      <td><strong>L1</strong><br>(Unit & Logic)</td>
+      <td>Unit tests for components like entrypoints, models</td>
+      <td rowspan="2">&lt;15min</td>
+      <td>/tests/{component_name}/test_xxx</td>
+      <td>
+        <a href="#chapter-1-l1-l2-level-testing-unit-testing-and-basic-end-to-end-verification">Chapter 1</a><br>
+        Section 1 L1&amp;L2: Purpose, Test Content, Directory Location, Example
+      </td>
+      <td>PR with ready label (also can run locally)</td>
+      <td>CPU</td>
+    </tr>
+    <tr>
+      <td><strong>L2</strong><br>(E2E across models & GPU-required UT)</td>
+      <td>Online & Offline (basic deployment scenarios):<br>dummy, normal inference function (output format, stream), some instance startup UT</td>
+      <td>
+        /tests/e2e/online_serving/test_{model_name}.py<br>
+        /tests/e2e/offline_inference/test_{model_name}.py
+      </td>
+      <td>
+        <a href="#chapter-1-l1-l2-level-testing-unit-testing-and-basic-end-to-end-verification">Chapter 1</a><br>
+        L1&amp;L2: Purpose, Test Content, Directory Location, Example
+      </td>
+      <td>PR with ready label</td>
+      <td>GPU</td>
+    </tr>
+    <tr>
+      <td><strong>L3</strong><br>(Important Perf & Integration & Accuracy)</td>
+      <td>Online & Offline (multiple deployment scenarios):<br>real model, normal inference function, normal accuracy</td>
+      <td>&lt;30min</td>
+      <td>
+        /tests/e2e/online_serving/test_{model_name}_expansion.py<br>
+        /tests/e2e/offline_inference/test_{model_name}_expansion.py
+      </td>
+      <td>
+        <a href="#chapter-2-l3-level-testing-core-integration-performance-and-accuracy-verification">Chapter 2</a><br>
+        L3: Purpose, Test Content, Directory Location, Example
+      </td>
+      <td>PR Merged (Also run L1&L2 Tests)</td>
+      <td>GPU</td>
+    </tr>
+    <tr>
+      <td><strong>L4</strong><br>(Perf & Integration & Accuracy)</td>
+      <td>Online & Offline: full functional scenarios + performance test + doc test</td>
+      <td>&lt;3 hour</td>
+      <td>
+        <strong>Full Function:</strong><br>
+        /tests/e2e/online_serving/test_{model_name}_expansion.py<br>
+        /tests/e2e/offline_inference/test_{model_name}_expansion.py<br>
+        <strong>Performance:</strong><br>
+        /tests/e2e/perf/nightly.json<br>
+        <strong>Doc Test:</strong><br>
+        tests/example/online_serving/test_{model_name}.py<br>
+        tests/example/offline_inference/test_{model_name}.py
+      </td>
+      <td>
+        <a href="#chapter-3-l4-level-testing-full-functionality-performance-and-documentation-testing">Chapter 3</a><br>
+        L4: Purpose, Test Content, Directory Location, Example
+      </td>
+      <td>Nightly</td>
+      <td>GPU</td>
+    </tr>
+    <tr>
+      <td><strong>L5</strong><br>(Stability & Reliability)</td>
+      <td>Online & Offline: long-term stability test + reliability test</td>
+      <td> Depends on reality </td>
+      <td>
+        <strong>Stability:</strong><br>
+        tests/e2e/stability/weekly.json<br>
+        <strong>Reliability:</strong><br>
+        tests/e2e/reliability/test_{model_name}.py
+      </td>
+      <td>
+        <a href="#chapter-4-l5-level-testing-stability-and-reliability-testing">Chapter 4</a><br>
+        L5: Purpose, Test Content, Directory Location, Example
+      </td>
+      <td>Weekly / Days before Release</td>
+      <td>GPU</td>
+    </tr>
+  </tbody>
+</table>
+
+
+---
+<details>
+<summary> The folder structure for tests file based on the 5 levels design</summary>
+Legend: `✅` = test exists, `⬜` = suggested to add.
+```
+vllm_omni/                                    tests/
+├── config/                             →     ├── config/
+│   ├── model.py                              │   └── test_model.py                    ⬜
+│   └── lora.py                               │   └── test_lora.py                      ⬜
+│
+├── core/                               →     ├── core/
+│   └── sched/                                 │   └── sched/
+│       ├── omni_ar_scheduler.py               │       ├── test_omni_ar_scheduler.py    ⬜
+│       ├── omni_generation_scheduler.py       │       ├── test_omni_generation_scheduler.py  ⬜
+│       └── output.py                          │       └── test_output.py               ✅ currently in entrypoints/test_omni_new_request_data.py (tests output.OmniNewRequestData)
+│
+├── diffusion/                          →     ├── diffusion/
+│   ├── diffusion_engine.py                    │   ├── test_diffusion_engine.py          ⬜
+│   ├── attention/                             │   ├── attention/
+│   │   ├── layer.py                            │   │   ├── test_attention_sp.py         ✅
+│   │   └── backends/                           │   │   └── test_flash_attn.py           ✅
+│   ├── distributed/                           │   ├── distributed/
+│   │   └── ...                                 │   │   ├── test_comm.py                 ✅
+│   │                                            │   │   ├── test_cfg_parallel.py        ✅
+│   │                                            │   │   └── test_sp_plan_hooks.py       ✅
+│   ├── lora/                                   │   ├── lora/
+│   │   └── ...                                 │   │   ├── test_base_linear.py          ✅
+│   │                                            │   │   └── test_lora_manager.py        ✅
+│   ├── models/                                 │   ├── models/
+│   │   ├── qwen_image/                         │   │   ├── qwen_image/                 (e2e coverage)
+│   │   ├── z_image/                            │   │   └── z_image/
+│   │   └── ...                                 │   │       └── test_zimage_tp_constraints.py  ✅
+│   └── worker/                                 │   └── worker/
+│       ├── diffusion_worker.py                 │       └── test_diffusion_worker.py   ✅ file at diffusion/test_diffusion_worker.py
+│       └── diffusion_model_runner.py            │
+│
+├── distributed/                         →     ├── distributed/
+│   └── omni_connectors/                         │   └── omni_connectors/
+│       ├── adapter.py                           │       ├── test_adapter_and_flow.py   ✅
+│       ├── kv_transfer_manager.py               │       ├── test_basic_connectors.py   ✅
+│       ├── connectors/                           │       ├── test_kv_flow.py             ✅
+│       └── utils/                               │       └── test_omni_connector_configs.py  ✅
+│
+├── engine/                             →     ├── engine/
+│   ├── input_processor.py                      │   ├── test_input_processor.py         ⬜  (no processor.py in source)
+│   ├── output_processor.py                     │   └── test_output_processor.py         ⬜
+│   └── arg_utils.py                            │   └── test_arg_utils.py               ⬜
+│
+├── entrypoints/                        →     ├── entrypoints/
+│   ├── omni.py                                 │   ├── test_omni.py                    ⬜  (E2E covered by e2e/offline, e2e/online)
+│   ├── omni_llm.py                             │   ├── test_omni_llm.py                ✅
+│   ├── omni_stage.py                            │   ├── test_omni_stage.py              ⬜  (partial in test_omni_stage_diffusion_config.py)
+│   ├── omni_diffusion.py                       │   ├── test_omni_diffusion.py          ✅
+│   ├── async_omni.py                            │   ├── test_async_omni.py              ✅ actually in e2e/online_serving/test_async_omni.py
+│   ├── async_omni_diffusion.py                 │   ├── test_async_omni_diffusion_config.py  ✅
+│   ├── stage_utils.py                          │   ├── test_stage_utils.py            ✅
+│   ├── cli/                                     │   ├── cli/                           (benchmarks/test_serve_cli.py covers CLI serve)
+│   │   └── ...                                  │   │   └── test_*.py                  ⬜
+│   └── openai/                                  │   └── openai_api/                    # maps to entrypoints/openai/
+│       ├── api_server.py                        │       ├── test_api_server.py         ⬜  (e2e indirect coverage)
+│       ├── serving_chat.py                       │       ├── test_serving_chat_sampling_params.py  ✅
+│       ├── serving_speech.py                     │       ├── test_serving_speech.py     ✅
+│       └── image_api_utils.py                   │       └── test_image_server.py      ✅
+│
+├── inputs/                             →     ├── inputs/
+│   ├── data.py                                 │   ├── test_data.py                   ⬜
+│   ├── parse.py                                │   ├── test_parse.py                 ⬜
+│   └── preprocess.py                            │   └── test_preprocess.py            ✅ currently in entrypoints/test_omni_input_preprocessor.py
+│
+├── model_executor/                     →     ├── model_executor/
+│   ├── layers/                                  │   ├── layers/
+│   │   └── mrope.py                             │   │   └── test_mrope.py              ⬜
+│   ├── model_loader/                            │   ├── model_loader/
+│   │   └── weight_utils.py                      │   │   └── test_weight_utils.py      ⬜
+│   ├── models/                                  │   ├── models/
+│   │   ├── qwen2_5_omni/                         │   │   ├── qwen2_5_omni/
+│   │   │   ├── qwen2_5_omni_thinker.py           │   │   │   ├── test_audio_length.py  ✅
+│   │   │   ├── qwen2_5_omni_talker.py            │   │   │   ├── test_qwen2_5_omni_thinker.py  ⬜
+│   │   │   └── qwen2_5_omni_token2wav.py         │   │   │   ├── test_qwen2_5_omni_talker.py  ⬜
+│   │   └── qwen3_omni/                          │   │   │   └── test_qwen2_5_omni_token2wav.py  ⬜
+│   │       └── ...                               │   │   └── qwen3_omni/
+│   ├── stage_configs/                           │   │       └── test_*.py              ⬜
+│   │   └── *.yaml                               │   └── stage_configs/                 (used by e2e, test_*.py can be added)  ⬜
+│   └── stage_input_processors/                  │   └── stage_input_processors/
+│       └── ...                                  │       └── test_*.py                 ⬜
+│
+├── sample/                             →     ├── sample/
+│   └── __init__.py                             │   └── test_*.py                      ⬜
+│
+├── utils/                              →     ├── utils/
+│   └── __init__.py                             │   └── test_*.py                       ⬜  (no platform_utils.py currently)
+│
+├── worker/                             →     ├── worker/
+│   ├── gpu_ar_model_runner.py                  │   ├── test_gpu_ar_model_runner.py    ⬜
+│   ├── gpu_ar_worker.py                        │   ├── test_gpu_ar_worker.py           ⬜
+│   ├── gpu_generation_model_runner.py          │   ├── test_gpu_generation_model_runner.py  ✅
+│   ├── gpu_generation_worker.py                │   ├── test_gpu_generation_worker.py  ⬜
+│   ├── gpu_model_runner.py                     │   ├── test_omni_gpu_model_runner.py   ✅
+│   └── mixins.py                               │   └── (npu under platforms/npu/worker/)  # not worker/npu/
+│
+├── platforms/                          →     (no tests/platforms/, e2e and stage_configs provide indirect coverage)
+│   ├── cuda/
+│   ├── npu/worker/                             # NPU worker here, not vllm_omni/worker/npu/
+│   ├── rocm/
+│   └── xpu/worker/
+│
+├── outputs.py                          →     test_outputs.py                         ✅ (at tests root)
+├── (logger, patch, request, version)    →     (no corresponding unit test)
+│
+└── e2e (tests side only)               →     ├── e2e/
+                                               ├── online_serving/                     ✅ non-empty
+                                               │   ├── test_async_omni.py
+                                               │   ├── test_qwen3_omni.py
+                                               │   ├── test_qwen3_omni_expansion.py
+                                               │   ├── test_image_gen_edit.py
+                                               │   ├── test_images_generations_lora.py
+                                               │   └── stage_configs/
+                                               └── offline_inference/                  ✅
+                                                   ├── test_qwen2_5_omni.py
+                                                   ├── test_qwen3_omni.py
+                                                   ├── test_bagel_text2img.py
+                                                   ├── test_t2i_model.py
+                                                   ├── test_t2v_model.py
+                                                   ├── test_ovis_image.py
+                                                   ├── test_zimage_tensor_parallel.py
+                                                   ├── test_cache_dit.py
+                                                   ├── test_teacache.py
+                                                   ├── test_stable_audio_model.py
+                                                   ├── test_diffusion_cpu_offload.py
+                                                   ├── test_diffusion_layerwise_offload.py
+                                                   ├── test_diffusion_lora.py
+                                                   ├── test_sequence_parallel.py
+                                                   └── stage_configs/
+                                                       ├── qwen2_5_omni_ci.yaml
+                                                       ├── qwen3_omni_ci.yaml
+                                                       ├── bagel_*.yaml
+                                                       └── npu/, rocm/, etc.
+```
+
+
+</details>
+
+
+## Common Specifications
+
+Before entering specific testing levels, the project establishes two common specifications aimed at standardizing the development process and quickly locating issues.
+
+1.  ***PR Checklist ([Tests Style](../ci/tests_style.md))***: This template defines the self-check items that must be completed before submitting a code review (Pull Request). It ensures that each code change meets basic requirements such as code style, dependency updates, and documentation synchronization before entering the automated testing pipeline, serving as the first manual line of defense for quality assurance.
+2.  ***CI Failure Explanation ([CI Failures](../ci/failures.md))***: This document archives and explains common failure patterns in the Continuous Integration (CI) pipeline, error log interpretation, and preliminary troubleshooting steps. It helps developers and testers quickly diagnose the causes of automated test failures, improving problem-solving efficiency.
+
+## Chapter 1: L1 & L2 Level Testing - Unit Testing and Basic End-to-End Verification
+
+### 1.1 Testing Purpose
+
+L1 and L2 level testing form the foundation of the quality assurance system. L1 level testing focuses on verifying the internal logic correctness of code units (e.g., functions, classes), ensuring each independent component behaves as designed.
+
+L2 level testing builds upon L1 by introducing GPU resources and verifying that the end-to-end (E2E) process of the model in basic deployment scenarios is smooth. For example, it uses dummy models to confirm that core interfaces like the inference pipeline, output format, and streaming response work properly. The common goal of these two levels is to provide developers with rapid feedback, discovering and fixing issues early in the development cycle  .
+
+
+
+### 1.2 Testing Content and Scope
+
+-   ***L1 (Unit & Logic Testing)***:
+-   -   ***Scope***: Tests internal functions and methods of core components such as `entrypoints`, `models`.
+    -   ***Focus***: Branch coverage, exception handling, algorithm logic correctness. Does not involve external dependencies or the complete service stack.
+    -   ***Time Cost***: Execution time is controlled within ***15 minutes*** to ensure fast feedback.
+-   ***L2 (Basic End-to-End Testing)***:
+-   -   ***Scope***: Covers two basic deployment scenarios: `online` (serving) and `offline` (inference).
+    -   ***Focus***: Uses `dummy` models or lightweight real models to verify that the entire chain from request input to result output works normally, including output data structure, streaming (stream) support, etc. Also includes some unit tests that require launching independent service instances.
+    -   ***Characteristic***: Requires ***GPU*** resources to perform model computations.
+
+### 1.3 Test Directory and Execution Files
+
+A clear directory structure is key to managing test cases efficiently.
+
+-   ***L1 Test Directory***: `/tests/{component_name}/test_xxx.py`
+-   -   Here, `{component_name}` corresponds to modules in the source code, such as `distributed`, `entrypoints`, etc., and `test_xxx.py` is the specific test file.
+-   ***L2 Test Directory***:
+-   -   Online Serving: `/tests/e2e/online_serving/test_{model_name}.py`
+    -   Offline Inference: `/tests/e2e/offline_inference/test_{model_name}.py`
+
+### 1.4 Execution Method and Example
+
+-   ***Trigger Timing***: **`PR with ready label`**. That is, when a developer adds a "ready for review" or similar label to a PR on platforms like GitHub, L1 and L2 tests are automatically triggered.
+-   ***Execution Environment***: L1 uses ***CPU*** environment; L2 requires ***GPU*** environment.
+-   ***Script Example***:
+
+<details>
+<summary> L1 Test Examples</summary>
+
+Examples from `tests/model_executor/models/qwen2_5_omni/test_audio_length.py`
+```python
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+def test_resolve_max_mel_frames_default():
+    from vllm_omni.model_executor.models.qwen2_5_omni.audio_length import resolve_max_mel_frames
+
+    assert resolve_max_mel_frames(None, default=30000) == 30000
+    assert resolve_max_mel_frames(None, default=6000) == 6000
+
+
+def test_resolve_max_mel_frames_explicit():
+    from vllm_omni.model_executor.models.qwen2_5_omni.audio_length import resolve_max_mel_frames
+
+    # Explicit argument always wins over default
+    assert resolve_max_mel_frames(123, default=30000) == 123
+    assert resolve_max_mel_frames(6000, default=30000) == 6000
+    assert resolve_max_mel_frames(0, default=30000) == 0
+
+
+@pytest.mark.parametrize("repeats", [2, 4])
+@pytest.mark.parametrize("code_len", [0, 1, 32768])
+@pytest.mark.parametrize("max_mel_frames", [None, -1, 0, 1, 6000, 30000])
+def test_cap_and_align_mel_length_no_mismatch(repeats, code_len, max_mel_frames):
+    """Guard that any max_mel_frames yields a mel length aligned to repeats, and
+    consistent with the truncated code length (prevents concat mismatch).
+    """
+    from vllm_omni.model_executor.models.qwen2_5_omni.audio_length import cap_and_align_mel_length
+
+    target_code_len, target_mel_len = cap_and_align_mel_length(
+        code_len=code_len,
+        repeats=repeats,
+        max_mel_frames=max_mel_frames,
+    )
+
+    assert isinstance(target_code_len, int)
+    assert isinstance(target_mel_len, int)
+
+    if code_len == 0:
+        assert target_code_len == 0
+        assert target_mel_len == 0
+        return
+
+    assert target_code_len >= 1
+    assert target_mel_len >= repeats
+    assert target_mel_len % repeats == 0
+    assert target_mel_len == target_code_len * repeats
+    assert target_code_len <= code_len
+
+    if max_mel_frames is not None and int(max_mel_frames) > 0 and int(max_mel_frames) >= repeats:
+        assert target_mel_len <= int(max_mel_frames)
+```
+</details>
+
+<details>
+<summary> L2 Test Examples</summary>
+You can refer to Test Examples in Chapter 2 to see example test cases that incorporate both L2 and L3 testing logic.
+</details>
+
+-   -   ***Run Command***:
+
+    `pytest -s -v /tests/e2e/online_serving/test_{model_name}.py`
+    `pytest -s -v -m 'core_model and cpu' --run-level=core_model`
+
+## Chapter 2: L3 Level Testing - Core Integration, Performance, and Accuracy Verification
+
+### 2.1 Testing Purpose
+
+L3 level testing executes after code is merged into the main branch. Its core purpose is to verify the integration effect, key performance indicators, and output accuracy of ***real models*** in ***multiple deployment scenarios***
+
+. It acts as the "quality gatekeeper" for the main branch, ensuring that no merge breaks the core capabilities of the model service. Testing needs to provide clear conclusions within a relatively short time (<30min), balancing test depth with feedback speed.
+
+
+
+### 2.2 Testing Content and Scope
+
+-   ***Deployment Scenarios***: Covers richer `online` and `offline` deployment configurations, which may include different hardware configurations, batch sizes, concurrency levels, etc.
+-   ***Core Verification***:
+-   1.  ***Inference Functionality***: Ensures real models can perform forward computation normally and return results.
+    2.  ***Accuracy Compliance***: Verifies that the model's evaluation metrics (e.g., accuracy) meet the expected baseline, preventing code changes from introducing accuracy issues.
+    3.  ***Important Performance***: Verifies whether performance (e.g., P99 latency, throughput) in core scenarios meets preset thresholds.
+
+### 2.3 Test Directory and Execution Files
+
+-   ***Functional Testing***:
+-   -   Online Serving: `/tests/e2e/online_serving/test_{model_name}_expansion.py`
+    -   Offline Inference: `/tests/e2e/offline_inference/test_{model_name}_expansion.py`
+    -   (Note: `_expansion.py` likely means it contains more comprehensive scenario cases compared to L2 tests).
+
+### 2.4 Execution Method and Example
+
+-   ***Trigger Timing***: **`PR Merged`**. Automatically triggered after code review is approved and merged into the main branch.
+-   ***Execution Environment***: ***GPU*** servers.
+-   ***Script Example***:
+
+???+ example "Test Examples"
+
+    **2.4.1 Mark Declaration Section**
+
+    ```python
+    @pytest.mark.advanced_model
+    @pytest.mark.core_model
+    @pytest.mark.parametrize("omni_server", test_params, indirect=True)
+    ```
+
+    **Explanation**:
+
+    @pytest.mark.advanced_model: Marks the test as L3 or L4 level, indicating that this test case performs deep validation, using real models for performance, integration, and accuracy testing. This forms a "basic-advanced" correspondence with the core_model mark at the L2 level.
+
+    @pytest.mark.core_model: Marks the test as L1 or L2 level, indicating that this test case validates the basic functionality of the core model. It uses mock weights and only checks if the relevant interface functions correctly.
+
+    @pytest.mark.parametrize: A parameterization decorator that allows abstracting test data into parameters, enabling reuse of the same test logic across different data configurations. indirect=True indicates that parameters will be passed to the fixture for processing.
+
+    **Notes**: If you believe the test case only needs to execute basic run logic at the PR-level CI, you can mark it only with @pytest.mark.core_model. If you believe it only needs to execute deep validation run logic at the merge or nightly level, you can mark it only with @pytest.mark.advanced_model. If you believe the test case needs to accommodate both basic run and deep validation test logic, you should mark it with both @pytest.mark.core_model and @pytest.mark.advanced_model.
+
+    **2.4.2 Test Function Definition and Documentation**
+
+    ```python
+    def test_mix_to_text_audio_001(omni_server, openai_client) -> None:
+        """
+        Test multi-modal input processing and text/audio output generation via OpenAI API.
+        Deploy Setting: default yaml
+        Input Modal: text + audio + video + image
+        Output Modal: text + audio
+        Input Setting: stream=True
+        Datasets: single request
+        """
+    ```
+
+    **Explanation**:
+
+    **Function Naming Convention**: Uses the test_ prefix, describes the test scenario mix_to_text_audio, and the number 001 indicates the first test case for this scenario.
+
+    **Parameter Explanation**:
+
+    omni_server: Omni server instance obtained via fixture, containing model information and configuration.
+
+    openai_client: Unified OpenAI client processing instance, encapsulating request sending and response validation logic.
+
+    Docstring: Describes the test purpose, deployment settings, input/output modalities, streaming settings, and dataset type in detail, providing clear context for test maintenance.
+
+    **2.4.3 Multimodal Data Preparation**
+
+    ```python
+    video_data_url = f"data:video/mp4;base64,{generate_synthetic_video(224, 224, 300)['base64']}"
+    image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(224, 224)['base64']}"
+    audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(5, 1)['base64']}"
+    ```
+
+    **Explanation**:
+
+    **Data Generation Functions**: Use the generate_synthetic_* series of functions to generate synthetic test data, avoiding reliance on external resources and ensuring test reproducibility and stability.
+
+    **Parameter Explanation**:
+
+    Video: width, height, duration_frames
+
+    Image: width, height
+
+    Audio: duration_seconds, channels
+
+    **2.4.4 Request Configuration and Keyword Validation**
+
+    ```python
+    request_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "stream": True,
+        "key_words": {
+            "audio": ["water", "cricket"],
+            "video": ["sphere", "globe", "circle", "round"],
+            "image": ["square", "quadrate"],
+            "text": ["beijing"]
+        },
+    }
+    ```
+
+    **Explanation**:
+
+    **Model Specification**: Uses omni_server.model to ensure the test aligns with the model configured on the server.
+
+    **Keyword Validation Mechanism**: This is an innovative design of the template to address the specific needs of multimodal testing:
+
+    Audio Keywords: Validate whether the generated text's description of audio content contains expected elements (e.g., "water" for water sounds, "cricket" for cricket sounds). If you provide multiple keywords, the validation is considered successful if at least one keyword is present.
+
+    **Video Keywords**: Validate whether the generated text's description of video content contains expected elements. If you provide multiple keywords, the validation is considered successful if at least one keyword is present.
+
+    Image Keywords: Validate whether the generated text's description of image content contains expected elements. If you provide multiple keywords, the validation is considered successful if at least one keyword is present.
+
+    Text Keywords: Validate whether the generated text contains expected elements. If you provide multiple keywords, the validation is considered successful if at least one keyword is present.
+
+    **2.4.5 Request Execution**
+
+    ```python
+    openai_client.send_request(request_config, request_num=1)
+    ```
+
+    **Explanation**:
+
+    **Unified Client**: Uses the OpenAIClientHandler instance to send requests. This client encapsulates error handling, retry mechanisms, and response validation logic.
+
+    **Single Request**: The comment clearly states this is a single-request completion test. For concurrent testing, it can be extended to multiple requests using request_num = n.
+
+    **Implicit Validation**: The send_request method internally includes validation logic dynamically selected based on the --run-level parameter: core_model performs basic validation, while advanced_model performs deep validation.
+
+-   ***Run Command***: `pytest -s -v /tests/e2e/online_serving/test_{model_name}.py -m advanced_model --run-level=advanced_model`
+
+## Chapter 3: L4 Level Testing - Full Functionality, Performance, and Documentation Testing
+
+### 3.1 Testing Purpose
+
+L4 level testing is a comprehensive quality audit before a version release. It expands upon L3, executing ***full*** functional scenarios, conducting systematic ***performance stress tests***, and simultaneously verifying the correctness of accompanying ***example documentation***. Its purpose is to perform deep validation of the system during off-peak nighttime hours, providing quality trend reports for daytime development and data support for release decisions.
+
+
+
+### 3.2 Testing Content and Scope
+
+-   ***Full Functionality Testing***: Executes all test cases defined in `test_{model_name}_expansion.py`, covering all implemented features, positive flows, boundary conditions, and exception handling.
+-   ***Performance Testing***: Uses the `/tests/e2e/perf/nightly.json` configuration file to drive performance testing tools for stress, load, and endurance tests, collecting metrics like throughput, response time, and resource utilization.
+-   ***Documentation Testing***: Verifies whether the example code provided to users is runnable and its results match the description.
+
+### 3.3 Test Directory and Execution Files
+
+-   ***Functional Testing***: Same directories as L3.
+-   ***Performance Test Configuration***: `/tests/e2e/perf/nightly.json`
+-   ***Documentation Example Tests***:
+-   -   `tests/example/online_serving/test_{model_name}.py`
+    -   `tests/example/offline_inference/test_{model_name}.py`
+
+### 3.4 Execution Method and Example
+
+-   ***Trigger Timing***: **`Nightly`**, automatically executed every night.
+-   ***Execution Environment***: ***GPU*** server clusters to meet the resource demands of performance testing.
+-   ***Script Example***:
+<details>
+<summary> Test Examples</summary>
+When you want to add L4-level performance test cases, you can refer to the following format for case addition in tests/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],
+            "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
+            }
+        }
+    ]
+}
+```
+
+#### 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.  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`).
+
+## Chapter 4: L5 Level Testing - Stability and Reliability Testing
+
+### 4.1 Testing Purpose
+
+L5 level testing focuses on the performance of model services under ***long-running*** and ***abnormal fault*** scenarios. It aims to uncover deep-seated issues that only manifest under sustained pressure or extreme conditions, such as memory leaks, resource contention, gradual performance degradation, and lack of fault tolerance mechanisms. This is the final, yet crucial, line of defense for ensuring service high availability and production environment robustness.
+
+
+
+### 4.2 Testing Content and Scope
+
+-   ***Long-term Stability (Stability) Testing***: Uses the `tests/e2e/stability/weekly.json` configuration to run the service under moderate load for an extended period (e.g., over 12 hours), monitoring whether metrics like memory/VRAM usage, response time, and throughput degrade over time, and whether the service process remains stable.
+-   ***Reliability Testing***: Uses `tests/e2e/reliability/test_{model_name}.py` to actively simulate various fault and abnormal scenarios, such as: dependent service interruption, abnormal input data, network flicker, hardware resource preemption, etc., to verify the system's fault tolerance, self-healing, and graceful degradation capabilities.
+
+### 4.3 Test Directory and Execution Files
+
+-   ***Stability Test Configuration***: `tests/e2e/stability/weekly.json`
+-   ***Reliability Test Suite***: `tests/e2e/reliability/test_{model_name}.py`
+
+### 4.4 Execution Method and Example
+
+-   ***Trigger Timing***: **`Weekly`** (weekly) or **`Days before Release`** (several days before a major release). Due to long execution times, the frequency is lower.
+-   ***Execution Environment***: ***GPU*** servers, requiring a stable and exclusive testing environment.
+-   ***Script Example***:
+<details>
+<summary> Test Examples</summary>
+```python
+# WIP
+```
+</details>
+
+-   -   ***Stability***: (Execution would be driven by the configuration in `weekly.json`).
+    -   ***Reliability***: `pytest -s -v tests/e2e/reliability/test_{model_name}.py`
+
+## Summary
+
+This multi-level testing system achieves continuous, progressive validation of model service quality by tightly integrating testing activities with the development workflow (commit, review, merge, release). From rapid unit testing to comprehensive end-to-end testing, and further to in-depth performance, stability, and reliability verification, each level has clear objectives, collectively building a robust quality protection net. By following this system, teams can deliver high-quality, highly reliable model services more efficiently.
@@ -1,4 +1,10 @@
 # CI Failures

-What should I do when a CI job fails on my PR, but I don't think my PR caused
-the failure?
+What should I do when a CI job fails on my PR, but I don't think my PR caused the failure?
+
+
+## Common Case 1:
+
+🚨 Error: The command was interrupted by a signal: signal: terminated
+
+Reason: The test is terminated due to exceed the time limits. Sometimes rebuild the test will pass.
@@ -0,0 +1,157 @@
+# Markers for Tests
+
+By adding markers before test functions, tests can later be executed uniformly by simply declaring the corresponding marker type.
+
+## Current Markers
+Defined in `pyproject.toml`:
+
+| Marker             | Description                                               |
+| ------------------ | --------------------------------------------------------- |
+| `core_model`       | L1&L2 tests (run in each PR)                              |
+| `advanced_model`   | L3&L4 level tests (run in each merge or nightly)          |
+| `diffusion`        | Diffusion model tests                                     |
+| `omni`             | Omni model tests                                          |
+| `cache`            | Cache backend tests                                       |
+| `parallel`         | Parallelism/distributed tests                             |
+| `cpu`              | Tests that run on CPU                                     |
+| `gpu`              | Tests that run on GPU *                                   |
+| `cuda`             | Tests that run on CUDA *                                  |
+| `rocm`             | Tests that run on AMD/ROCm *                              |
+| `npu`              | Tests that run on NPU/Ascend *                            |
+| `H100`             | Tests that require H100 GPU  *                            |
+| `L4`               | Tests that require L4 GPU *                               |
+| `MI325`            | Tests that require MI325 GPU (AMD/ROCm) *                 |
+| `A2`               | Tests that require A2 NPU *                               |
+| `A3`               | Tests that require A3 NPU *                               |
+| `distributed_cuda` | Tests that require multi cards on CUDA platform *         |
+| `distributed_rocm` | Tests that require multi cards on ROCm platform  *        |
+| `distributed_npu`  | Tests that require multi cards on NPU platform  *         |
+| `skipif_cuda`      | Skip if the num of CUDA cards is less than the required * |
+| `skipif_rocm`      | Skip if the num of ROCm cards is less than the required * |
+| `skipif_npu`       | Skip if the num of NPU cards is less than the required *  |
+| `slow`             | Slow tests (may skip in quick CI)                         |
+| `benchmark`        | Benchmark tests                                           |
+
+\* Means those markers are auto-added, and they will be added by the `@hardware_test` decorator.
+
+### Example usage for markers
+
+```python
+from tests.utils import hardware_test
+
+@pytest.mark.core_model
+@pytest.mark.omni
+@hardware_test(
+   res={"cuda": "L4", "rocm": "MI325", "npu": "A2"},
+   num_cards=2,
+)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_video_to_audio()
+    ...
+```
+### Decorator: `@hardware_test`
+
+This decorator is intended to make hardware-aware, cross-platform test authoring easier and more robust for CI/CD environments. The `hardware_test` decorator in `vllm-omni/tests/utils.py` performs the following actions:
+
+1. **Applies platform and resource markers**  
+   Adds the appropriate pytest markers for each specified hardware platform (e.g., `cuda`, `rocm`, `npu`) and resource type (e.g., `L4`, `H100`, `MI325`, `A2`, `A3`).
+   ```
+   @pytest.mark.cuda
+   @pytest.mark.L4
+   ```
+2. **Handles multi-card (distributed) scenarios**  
+   For tests requiring multiple cards, it automatically adds distributed markers such as `distributed_cuda`, `distributed_rocm`, or `distributed_npu`.
+   ```
+   @pytest.mark.distributed_cuda(num_cards=num_cards)
+   ```
+3. **Supports flexible card requirements**  
+   Accepts `num_cards` as either a single integer for all platforms or as a dictionary with per-platform values. If not specified, defaults to 1 card per platform.
+
+4. **Integrates resource validation**  
+   On CUDA, adds a skip marker (`skipif_cuda`) if the system does not have the required number of devices.
+   Support for `skipif_rocm` and `skipif_npu` will be implemented later.
+
+
+5. **Works with pytest filtering**  
+   Allows tests to be filtered and selected at runtime using standard pytest marker expressions (e.g., `-m "distributed_cuda and L4"`).
+
+#### Example usage for decorator
+- Single call for multiple platforms:
+    ```python
+    @hardware_test(
+        res={"cuda": "L4", "rocm": "MI325", "npu": "A2"},
+        num_cards={"cuda": 2, "rocm": 2, "npu": 2},
+    )
+    ```
+    or
+    ```python
+    @hardware_test(
+        res={"cuda": "L4", "rocm": "MI325", "npu": "A2"},
+        num_cards=2,
+    )
+    ```
+- `res` must be a dict; supported resources: CUDA (L4/H100), ROCm (MI325), NPU (A2/A3)
+- `num_cards` can be int (all platforms) or dict (per platform); defaults to 1 when missing
+- Distributed markers (`distributed_cuda`, `distributed_rocm`, `distributed_npu`) are auto-added for multi-card cases
+- Filtering examples:
+    - CUDA only: `pytest -m "distributed_cuda and L4"`
+    - ROCm only: `pytest -m "distributed_rocm and MI325"`
+    - NPU only: `pytest -m "distributed_npu"`
+
+## Add Support for a New Platform
+
+If you want to add support for a new platform (e.g., "tpu" for a new accelerator), follow these steps:
+
+1. **Extend the marker list in your pytest config** so that platform/resource markers are defined:
+   ```toml
+   # In pyproject.toml or pytest.ini
+   [tool.pytest.ini_options]
+   markers = [
+       # ... existing markers ...
+       "tpu: Tests that require TPU device",
+       "TPU_V3: Tests that require TPU v3 hardware",
+       "distributed_tpu: Tests that require multiple TPU devices",
+   ]
+   ```
+2. **Implement a marker construction function for your platform** in `vllm-omni/tests/utils.py`:
+   ```python
+   # In vllm-omni/tests/utils.py
+
+   def tpu_marks(*, res: str, num_cards: int):
+       test_platform = pytest.mark.tpu
+       if res == "TPU_V3":
+           test_resource = pytest.mark.TPU_V3
+       else:
+           raise ValueError(
+               f"Invalid TPU resource type: {res}. Supported: TPU_V3")
+
+       if num_cards == 1:
+           return [test_platform, test_resource]
+       else:
+           test_distributed = pytest.mark.distributed_tpu(num_cards=num_cards)
+           # Optionally: add skipif_tpu when implemented
+           return [test_platform, test_resource, test_distributed]
+   ```
+3. **Update `hardware_test` to recognize your new platform**:
+    In the relevant place (see the `hardware_test` implementation), add:
+    ```python
+    if platform == "tpu":
+        marks = tpu_marks(res=resource, num_cards=cards)
+    ```
+4. **(Recommended) Add a test using your new markers**:
+   ```python
+   @hardware_test(
+       res={"tpu": "TPU_V3"},
+       num_cards=2,
+   )
+   def test_my_tpu_feature():
+       ...
+   ```
+
+**Summary**:  
+- Add pytest markers for your new platform/resources  
+- Implement a marker function (`xxx_marks`)  
+- Plug into `hardware_test`  
+- You're done: tests decorated with `@hardware_test` using your platform now automatically get the correct markers, distribution, and isolation!
+
+See code in `vllm-omni/tests/utils.py` for existing examples (`cuda_marks`, `rocm_marks`, `npu_marks`).
@@ -0,0 +1,467 @@
+# Test File Structure and Style Guide
+
+To ensure project maintainability and sustainable development, we encourage contributors to submit test code (unit tests, system tests, or end-to-end tests) alongside their code changes. This document outlines the guidelines for organizing and naming test files.
+
+## Checklist before submitting your test files
+
+1. The file is saved in an appropriate place and the file name is clear.
+2. The coding style follows the requirements outlined below.
+3. All test functions have appropriate pytest markers.
+4. For tests that need run in CI, please ensure it labeled as ``@pytest.mark.core_model` the test is configured under the `./buildkite/` folder.
+
+
+## Test Types
+
+For more details about our [Five Levels Tests design](../ci/CI_5levels.md).
+
+### Unit Tests and System Tests
+For unit tests and system tests, we strongly recommend placing test files in the same directory structure as the source code being tested, using the naming convention `test_*.py`.
+
+### End-to-End (E2E) Tests for Models
+End-to-end tests verify the complete functionality of a system or component. For our project, the E2E tests for different omni models are organized into two subdirectories:
+
+- **`tests/e2e/offline_inference/`**: Tests for offline inference modes (e.g., Qwen3Omni offline inference)
+
+- **`tests/e2e/online_serving/`**: Tests for online serving scenarios (e.g., API server tests)
+
+**Example:** The test file for `vllm_omni/entrypoints/omni_llm.py` should be located at `tests/entrypoints/test_omni_llm.py`.
+
+## Test Directory Structure
+
+The ideal directory structure mirrors the source code organization. Legend: `✅` = test exists, `⬜` = suggested to add.
+
+```
+vllm_omni/                                    tests/
+├── config/                             →     ├── config/
+│   ├── model.py                              │   └── test_model.py                    ⬜
+│   └── lora.py                               │   └── test_lora.py                      ⬜
+│
+├── core/                               →     ├── core/
+│   └── sched/                                 │   └── sched/
+│       ├── omni_ar_scheduler.py               │       ├── test_omni_ar_scheduler.py    ⬜
+│       ├── omni_generation_scheduler.py       │       ├── test_omni_generation_scheduler.py  ⬜
+│       └── output.py                          │       └── test_output.py               ✅ currently in entrypoints/test_omni_new_request_data.py (tests output.OmniNewRequestData)
+│
+├── diffusion/                          →     ├── diffusion/
+│   ├── diffusion_engine.py                    │   ├── test_diffusion_engine.py          ⬜
+│   ├── attention/                             │   ├── attention/
+│   │   ├── layer.py                            │   │   ├── test_attention_sp.py         ✅
+│   │   └── backends/                           │   │   └── test_flash_attn.py           ✅
+│   ├── distributed/                           │   ├── distributed/
+│   │   └── ...                                 │   │   ├── test_comm.py                 ✅
+│   │                                            │   │   ├── test_cfg_parallel.py        ✅
+│   │                                            │   │   └── test_sp_plan_hooks.py       ✅
+│   ├── lora/                                   │   ├── lora/
+│   │   └── ...                                 │   │   ├── test_base_linear.py          ✅
+│   │                                            │   │   └── test_lora_manager.py        ✅
+│   ├── models/                                 │   ├── models/
+│   │   ├── qwen_image/                         │   │   ├── qwen_image/                 (e2e coverage)
+│   │   ├── z_image/                            │   │   └── z_image/
+│   │   └── ...                                 │   │       └── test_zimage_tp_constraints.py  ✅
+│   └── worker/                                 │   └── worker/
+│       ├── diffusion_worker.py                 │       └── test_diffusion_worker.py   ✅ file at diffusion/test_diffusion_worker.py
+│       └── diffusion_model_runner.py            │
+│
+├── distributed/                         →     ├── distributed/
+│   └── omni_connectors/                         │   └── omni_connectors/
+│       ├── adapter.py                           │       ├── test_adapter_and_flow.py   ✅
+│       ├── kv_transfer_manager.py               │       ├── test_basic_connectors.py   ✅
+│       ├── connectors/                           │       ├── test_kv_flow.py             ✅
+│       └── utils/                               │       └── test_omni_connector_configs.py  ✅
+│
+├── engine/                             →     ├── engine/
+│   ├── input_processor.py                      │   ├── test_input_processor.py         ⬜  (no processor.py in source)
+│   ├── output_processor.py                     │   └── test_output_processor.py         ⬜
+│   └── arg_utils.py                            │   └── test_arg_utils.py               ⬜
+│
+├── entrypoints/                        →     ├── entrypoints/
+│   ├── omni.py                                 │   ├── test_omni.py                    ⬜  (E2E covered by e2e/offline, e2e/online)
+│   ├── omni_llm.py                             │   ├── test_omni_llm.py                ✅
+│   ├── omni_stage.py                            │   ├── test_omni_stage.py              ⬜  (partial in test_omni_stage_diffusion_config.py)
+│   ├── omni_diffusion.py                       │   ├── test_omni_diffusion.py          ✅
+│   ├── async_omni.py                            │   ├── test_async_omni.py              ✅ actually in e2e/online_serving/test_async_omni.py
+│   ├── async_omni_diffusion.py                 │   ├── test_async_omni_diffusion_config.py  ✅
+│   ├── stage_utils.py                          │   ├── test_stage_utils.py            ✅
+│   ├── cli/                                     │   ├── cli/                           (benchmarks/test_serve_cli.py covers CLI serve)
+│   │   └── ...                                  │   │   └── test_*.py                  ⬜
+│   └── openai/                                  │   └── openai_api/                    # maps to entrypoints/openai/
+│       ├── api_server.py                        │       ├── test_api_server.py         ⬜  (e2e indirect coverage)
+│       ├── serving_chat.py                       │       ├── test_serving_chat_sampling_params.py  ✅
+│       ├── serving_speech.py                     │       ├── test_serving_speech.py     ✅
+│       └── image_api_utils.py                   │       └── test_image_server.py      ✅
+│
+├── inputs/                             →     ├── inputs/
+│   ├── data.py                                 │   ├── test_data.py                   ⬜
+│   ├── parse.py                                │   ├── test_parse.py                 ⬜
+│   └── preprocess.py                            │   └── test_preprocess.py            ✅ currently in entrypoints/test_omni_input_preprocessor.py
+│
+├── model_executor/                     →     ├── model_executor/
+│   ├── layers/                                  │   ├── layers/
+│   │   └── mrope.py                             │   │   └── test_mrope.py              ⬜
+│   ├── model_loader/                            │   ├── model_loader/
+│   │   └── weight_utils.py                      │   │   └── test_weight_utils.py      ⬜
+│   ├── models/                                  │   ├── models/
+│   │   ├── qwen2_5_omni/                         │   │   ├── qwen2_5_omni/
+│   │   │   ├── qwen2_5_omni_thinker.py           │   │   │   ├── test_audio_length.py  ✅
+│   │   │   ├── qwen2_5_omni_talker.py            │   │   │   ├── test_qwen2_5_omni_thinker.py  ⬜
+│   │   │   └── qwen2_5_omni_token2wav.py         │   │   │   ├── test_qwen2_5_omni_talker.py  ⬜
+│   │   └── qwen3_omni/                          │   │   │   └── test_qwen2_5_omni_token2wav.py  ⬜
+│   │       └── ...                               │   │   └── qwen3_omni/
+│   ├── stage_configs/                           │   │       └── test_*.py              ⬜
+│   │   └── *.yaml                               │   └── stage_configs/                 (used by e2e, test_*.py can be added)  ⬜
+│   └── stage_input_processors/                  │   └── stage_input_processors/
+│       └── ...                                  │       └── test_*.py                 ⬜
+│
+├── sample/                             →     ├── sample/
+│   └── __init__.py                             │   └── test_*.py                      ⬜
+│
+├── utils/                              →     ├── utils/
+│   └── __init__.py                             │   └── test_*.py                       ⬜  (no platform_utils.py currently)
+│
+├── worker/                             →     ├── worker/
+│   ├── gpu_ar_model_runner.py                  │   ├── test_gpu_ar_model_runner.py    ⬜
+│   ├── gpu_ar_worker.py                        │   ├── test_gpu_ar_worker.py           ⬜
+│   ├── gpu_generation_model_runner.py          │   ├── test_gpu_generation_model_runner.py  ✅
+│   ├── gpu_generation_worker.py                │   ├── test_gpu_generation_worker.py  ⬜
+│   ├── gpu_model_runner.py                     │   ├── test_omni_gpu_model_runner.py   ✅
+│   └── mixins.py                               │   └── (npu under platforms/npu/worker/)  # not worker/npu/
+│
+├── platforms/                          →     (no tests/platforms/, e2e and stage_configs provide indirect coverage)
+│   ├── cuda/
+│   ├── npu/worker/                             # NPU worker here, not vllm_omni/worker/npu/
+│   ├── rocm/
+│   └── xpu/worker/
+│
+├── outputs.py                          →     test_outputs.py                         ✅ (at tests root)
+├── (logger, patch, request, version)    →     (no corresponding unit test)
+│
+└── e2e (tests side only)               →     ├── e2e/
+                                               ├── online_serving/                     ✅ non-empty
+                                               │   ├── test_async_omni.py
+                                               │   ├── test_qwen3_omni.py
+                                               │   ├── test_qwen3_omni_expansion.py
+                                               │   ├── test_image_gen_edit.py
+                                               │   ├── test_images_generations_lora.py
+                                               │   └── stage_configs/
+                                               └── offline_inference/                  ✅
+                                                   ├── test_qwen2_5_omni.py
+                                                   ├── test_qwen3_omni.py
+                                                   ├── test_bagel_text2img.py
+                                                   ├── test_t2i_model.py
+                                                   ├── test_t2v_model.py
+                                                   ├── test_ovis_image.py
+                                                   ├── test_zimage_tensor_parallel.py
+                                                   ├── test_cache_dit.py
+                                                   ├── test_teacache.py
+                                                   ├── test_stable_audio_model.py
+                                                   ├── test_diffusion_cpu_offload.py
+                                                   ├── test_diffusion_layerwise_offload.py
+                                                   ├── test_diffusion_lora.py
+                                                   ├── test_sequence_parallel.py
+                                                   └── stage_configs/
+                                                       ├── qwen2_5_omni_ci.yaml
+                                                       ├── qwen3_omni_ci.yaml
+                                                       ├── bagel_*.yaml
+                                                       └── npu/, rocm/, etc.
+```
+
+
+
+### Naming Conventions
+
+- **Unit Tests**: Use `test_<module_name>.py` format. Example: `omni_llm.py` → `test_omni_llm.py`
+
+- **E2E Tests**: Place in `tests/e2e/offline_inference/` or `tests/e2e/online_serving/` with descriptive names. Example: `tests/e2e/offline_inference/test_qwen3_omni.py`, `tests/e2e/offline_inference/test_diffusion_model.py`
+
+- **Expansion Tests**
+
+### Best Practices
+
+1. **Mirror Source Structure**: Test directories should mirror the source code structure
+2. **Test Type Indicators**: Use comments to indicate test types (UT for unit tests, E2E for end-to-end tests)
+3. **Shared Resources**: Place shared test configurations (e.g., CI configs) in appropriate subdirectories
+4. **Consistent Naming**: Follow the `test_*.py` naming convention consistently across all test files
+
+
+## Test codes requirements
+
+### Coding style
+
+1. **File header**: Add SPDX license header to all test files
+2. **Imports**: Pls don't use manual `sys.path` modifications, use standard imports instead.
+3. **Test type differentiation**:
+
+      - Unit tests: Maintain mock style
+      - E2E tests for models: Consider using OmniRunner uniformly, avoid decorators
+
+4. **Documentation**: Add docstrings to all test functions
+5. **Environment variables**: Set uniformly in `conftest.py` or at the top of files
+6. **Type annotations**: Add type annotations to all test function parameters
+7. **Pytest Markers**: Add necessary markers like `@pytest.mark.core_model` and use `@hardware_test` to declare hardware requirements (check detailed in [Markers for Tests](../ci/tests_markers.md)).
+
+### Template
+#### E2E - Online serving
+
+E2E Online tests for Qwen3-Omni model with mix input and audio+text output. Based on `tests/e2e/online_serving/test_qwen3_omni.py`.
+
+```python
+"""
+E2E Online tests for Qwen3-Omni model with mix input and audio+text output.
+"""
+
+import os
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
+
+import threading
+from pathlib import Path
+
+import openai
+import pytest
+
+from tests.conftest import (
+    OmniServer,
+    convert_audio_to_text,
+    cosine_similarity_text,
+    dummy_messages_from_mix_data,
+    generate_synthetic_video,
+    merge_base64_and_convert_to_text,
+)
+from vllm_omni.platforms import current_omni_platform
+
+# Edit: model name and stage config path
+models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]
+
+#If you use the default configuration file, you can directly use the following address.
+def get_default_config():
+    return str(Path(__file__).parent.parent / "stage_configs" / "qwen3_omni_ci.yaml")
+
+#If you need to modify the configuration file, you can use modify_stage_config.
+def get_chunk_config():
+    path = modify_stage_config(
+        get_default_config(),
+        updates={
+            "async_chunk": True,
+            "stage_args": {
+                0: {
+                    "engine_args.custom_process_next_stage_input_func": "vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_async_chunk"
+                },
+                1: {
+                    "engine_args.custom_process_next_stage_input_func": "vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav_async_chunk"
+                },
+            },
+        },
+        deletes={"stage_args": {2: ["custom_process_input_func"]}},
+    )
+    return path
+
+stage_configs = [get_default_config(), CHUNK_CONFIG_PATH]
+
+test_params = [(model, stage_config) for model in models for stage_config in stage_configs]
+
+
+#Please use this method to launch the online instance.
+_omni_server_lock = threading.Lock()
+
+@pytest.fixture(scope="module")
+def omni_server(request):
+    """Start vLLM-Omni server as a subprocess. Use module scope for multi-stage init (10-20+ min)."""
+    with _omni_server_lock:
+        model, stage_config_path = request.param
+        with OmniServer(
+            model,
+            ["--stage-configs-path", stage_config_path, "--stage-init-timeout", "120"],
+        ) as server:
+            yield server
+
+
+@pytest.fixture
+def client(omni_server):
+    """OpenAI client for the running vLLM-Omni server."""
+    return openai.OpenAI(
+        base_url=f"http://{omni_server.host}:{omni_server.port}/v1",
+        api_key="EMPTY",
+    )
+
+#Please use function definitions above the test function to define the prompts and other parameters you need.
+def get_system_prompt():
+    return {
+        "role": "system",
+        "content": [
+            {
+                "type": "text",
+                "text": (
+                    "You are Qwen, a virtual human developed by the Qwen Team, "
+                    "Alibaba Group, capable of perceiving auditory and visual inputs, "
+                    "as well as generating text and speech."
+                ),
+            }
+        ],
+    }
+
+...
+
+#Please define test case tags according to the instructions in the marker documentation.
+@pytest.mark.core_model
+@pytest.mark.omni
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_mix_to_text_audio_001(client: openai.OpenAI, omni_server, request) -> None:
+    # PLEASE FOLLOW THESE TEMPLATE INSTRUCTIONS:
+    # ============================================================================
+    # TEMPLATE USAGE GUIDE:
+    # 1. Copy this entire function as a starting point for multi-modal tests
+    # 2. Update the test name to reflect your specific test scenario
+    # 3. Modify input/output modalities as needed (see OPTIONS section below)
+    # 4. Adjust assertions based on your expected outcomes
+    # 5. Add custom validation logic for your specific use case
+    # ============================================================================
+
+    #Please list the relevant test points.
+    """
+    Test multi-modal input processing and text/audio output generation via OpenAI API.
+    Deploy Setting: default yaml
+    Input Modal: text + audio + video + image
+    Output Modal: text + audio
+    Input Setting: stream=True
+    Datasets: single request
+    """
+    # SECTION 1: TEST SETUP AND INITIALIZATION
+    # =========================================
+    # INSTRUCTIONS: Initialize test variables and prepare test environment
+    # MODIFY: Add any additional test setup required for your scenario
+    e2e_list = list()
+    # SECTION 2: TEST DATA GENERATION
+    # ================================
+    # INSTRUCTIONS: Generate or load test data for each input modality
+    # MODIFY: Replace synthetic generators with your actual data sources
+    # VIDEO DATA - Generate synthetic video for testing
+    # FORMAT: data:video/mp4;base64,{base64_encoded_video}
+    # PARAMETERS: width, height, duration_frames
+    video_data_url = f"data:video/mp4;base64,{generate_synthetic_video(224, 224, 300)['base64']}"
+    # IMAGE DATA - Generate synthetic image for testing
+    # FORMAT: data:image/jpeg;base64,{base64_encoded_image}
+    # PARAMETERS: width, height
+    image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(224, 224)['base64']}"
+    # AUDIO DATA - Generate synthetic audio for testing
+    # FORMAT: data:audio/wav;base64,{base64_encoded_audio}
+    # PARAMETERS: duration_seconds, channels
+    audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(5, 1)['base64']}"
+
+    # SECTION 3: MESSAGE CONSTRUCTION
+    # ================================
+    # INSTRUCTIONS: Assemble the complete message payload for API request
+    # MODIFY: Add/remove modalities or change prompt structure as needed
+
+    # USAGE: Construct a message containing all input modalities
+    # IMPORTANT: Ensure the message structure matches OpenAI API expectations
+    # CUSTOMIZATION POINTS:
+    #   - system_prompt: Controls the assistant's behavior
+    #   - content_text: The user's text prompt/question
+    #   - *_data_url: URLs for media content (video/image/audio)
+    messages = dummy_messages_from_mix_data(
+        system_prompt=get_system_prompt(),
+        video_data_url=video_data_url,
+        image_data_url=image_data_url,
+        audio_data_url=audio_data_url,
+        content_text=get_prompt("mix"),
+    )
+
+    # SECTION 4: API REQUEST EXECUTION
+    # =================================
+    # INSTRUCTIONS: Make the API call and measure performance
+    # MODIFY: Add timeout, retry logic, or additional parameters
+    start_time = time.perf_counter()
+    chat_completion = client.chat.completions.create(model=omni_server.model, messages=messages, stream=True)
+
+    #Call using your preferred method and obtain the final audio and text outputs.
+    ...
+
+    # SECTION 5: OUTPUT VALIDATION
+    # =============================
+    # INSTRUCTIONS: Verify that outputs meet expected criteria
+    # MODIFY: Adjust validation logic for your specific requirements
+
+    # ASSERTION 1: E2E Validation
+    # PURPOSE: Verify that the E2E latency is less than the baseline.
+    current_e2e = time.perf_counter() - start_time
+    print(f"the request e2e is: {current_e2e}")
+    e2e_list.append(current_e2e)
+
+    print(f"the avg e2e is: {sum(e2e_list) / len(e2e_list)}")
+
+
+
+    # ASSERTION 2: Text Output Validation
+    # PURPOSE: Verify that text output was generated with keyword content
+    assert text_content is not None and len(text_content) >= 2, "No text output is generated"
+    assert any(
+        keyword in text_content.lower() for keyword in ["square", "quadrate", "sphere", "globe", "circle", "round"]
+    ), "The output does not contain any of the keywords."
+
+
+    # ASSERTION 3: Cross-Modal Consistency
+    # PURPOSE: Verify text and audio outputs convey the same information
+    # CUSTOMIZATION: Adjust similarity threshold (0.9) based on accuracy requirements
+    assert audio_data is not None, "No audio output is generated"
+    audio_content = merge_base64_and_convert_to_text(audio_data)
+    print(f"text content is: {text_content}")
+    print(f"audio content is: {audio_content}")
+    similarity = cosine_similarity_text(audio_content.lower(), text_content.lower())
+    print(f"similarity is: {similarity}")
+    assert similarity > 0.9, "The audio content is not same as the text"
+```
+
+
+#### E2E - Offline inference
+```python
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Offline E2E smoke test for an omni model (video → audio).
+"""
+
+import os
+from pathlib import Path
+
+import pytest
+from vllm.assets.video import VideoAsset
+
+from tests.utils import hardware_test
+from ..multi_stages.conftest import OmniRunner
+
+# Optional: set process start method for workers
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+models = ["{your model name}"] #Edit here to load your model
+stage_configs = [str(Path(__file__).parent / "stage_configs" / {your model yaml})] #Edit here to load your model yaml
+
+# Create parameter combinations for model and stage config
+test_params = [(model, stage_config) for model in models for stage_config in stage_configs]
+
+# function name: test_{input_modality}_to_{output_modality}
+# modality candidate: text, image, audio, video, mixed_modalities
+@pytest.mark.core_model
+@pytest.mark.omni
+@hardware_test(
+    res={"cuda": "L4", "rocm": "MI325", "npu": "A2"},
+    num_cards=2,
+)
+@pytest.mark.parametrize("test_config", test_params)
+def test_video_to_audio(omni_runner: type[OmniRunner], model: str) -> None:
+    """Offline inference: video input, audio output."""
+    model, stage_config_path = test_config
+    with omni_runner(model, seed=42, stage_configs_path=stage_config_path) as runner:
+        # Prepare inputs
+        video = VideoAsset(name="sample", num_frames=4).np_ndarrays
+
+        outputs = runner.generate_multimodal(
+            prompts="Describe this video briefly.",
+            videos=video,
+        )
+
+        # Minimal assertions: got outputs and at least one audio result
+        assert outputs
+        has_audio = any(o.final_output_type == "audio" for o in outputs)
+        assert has_audio
+```
@@ -0,0 +1,286 @@
+# Support Cache-DiT
+
+This section describes how to add cache-dit acceleration to a new diffusion pipeline. We use the Qwen-Image pipeline and LongCat-Image pipeline as reference implementations.
+
+---
+
+## Table of Contents
+
+- [Overview](#overview)
+- [Standard Models: Automatic Support](#standard-models-automatic-support)
+- [Custom Architectures: Writing Custom Implementation](#custom-architectures-writing-custom-implementation)
+- [Testing](#testing)
+- [Troubleshooting](#troubleshooting)
+- [Reference Implementations](#reference-implementations)
+- [Summary](#summary)
+
+---
+
+## Overview
+
+### What is Cache-DiT?
+
+Cache-DiT is an acceleration library for Diffusion Transformers (DiT) that caches intermediate computation results across denoising steps. The core insight is that adjacent denoising steps often produce similar intermediate features, so we can skip redundant computations by reusing cached results.
+
+The library supports three main caching strategies:
+
+- **DBCache:** Dynamic block-level caching that selectively computes or caches transformer blocks based on residual differences
+- **TaylorSeer:** Calibration-based prediction that estimates block outputs using Taylor expansion
+- **SCM (Step Computation Masking):** Dynamic step skipping based on configurable policies
+
+### Architecture
+
+vLLM-omni integrates cache-dit through the `CacheDiTBackend` class, which provides a unified interface for managing cache-dit acceleration on diffusion models.
+
+| Method/Class | Purpose | Behavior |
+|--------------|---------|----------|
+| [`CacheDiTBackend`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/#vllm_omni.diffusion.cache.CacheBackend) | Unified backend interface | Automatically handles enabler selection and cache refresh |
+| [`enable_cache_for_dit()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/cache_dit_backend/#vllm_omni.diffusion.cache.cache_dit_backend.enable_cache_for_dit) | Apply caching to transformer | Configures DBCache on transformer blocks |
+
+**Key APIs from Cache-DiT:**
+
+[Cache-DiT API Reference](https://cache-dit.readthedocs.io/en/latest/user_guide/CACHE_API/)
+
+| API | Description |
+|-----|-------------|
+| `BlockAdapter` | Core abstraction for applying cache-dit to transformers. Specifies transformer module(s), block list(s), and forward signature pattern(s). |
+| `ForwardPattern` | Defines block forward signature patterns: `Pattern_0`, `Pattern_1`, `Pattern_2` |
+| `ParamsModifier` | Per-transformer or per-block-list cache configuration customization |
+| `DBCacheConfig` | Configuration for DBCache parameters (warmup steps, cached steps, thresholds) |
+| `refresh_context()` | Update cache context | Called when `num_inference_steps` changes |
+
+---
+
+## Standard Models: Automatic Support
+
+Most DiT models follow this pattern:
+- Single transformer with one `ModuleList` of blocks
+- Standard forward signature
+- Compatible with cache-dit's automatic detection
+
+**Examples:** Qwen-Image, Z-Image
+
+For standard single-transformer models, **no code changes are needed**. The `CacheDiTBackend` automatically uses `enable_cache_for_dit()`:
+
+```python
+from vllm_omni import Omni
+
+# Works automatically for standard models
+omni = Omni(
+    model="Qwen/Qwen-Image",  # Standard single-transformer model
+    cache_backend="cache_dit",
+    cache_config={
+        "Fn_compute_blocks": 1,
+        "Bn_compute_blocks": 0,
+        "max_warmup_steps": 4,
+    }
+)
+```
+
+**What happens automatically:**
+
+```python
+def enable_cache_for_dit(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
+    """Default enabler for standard single-transformer DiT models."""
+
+    # Build cache configuration
+    db_cache_config = DBCacheConfig(
+        num_inference_steps=None,  # Will be set during first inference
+        Fn_compute_blocks=cache_config.Fn_compute_blocks,
+        Bn_compute_blocks=cache_config.Bn_compute_blocks,
+        max_warmup_steps=cache_config.max_warmup_steps,
+        max_cached_steps=cache_config.max_cached_steps,
+        max_continuous_cached_steps=cache_config.max_continuous_cached_steps,
+        residual_diff_threshold=cache_config.residual_diff_threshold,
+    )
+
+    # Enable cache-dit on transformer
+    cache_dit.enable_cache(
+        pipeline.transformer,
+        cache_config=db_cache_config,
+    )
+
+    # Return refresh function for dynamic num_inference_steps updates
+    def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True):
+        cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose)
+
+    return refresh_cache_context
+```
+
+---
+
+## Custom Architectures: Writing Custom Implementation
+
+Some models require custom handling:
+
+- **Dual-transformer:** Models with separate high-noise and low-noise transformers (e.g., Wan2.2)
+- **Multi-block-list:** Models with multiple block lists in one transformer (e.g., LongCatImage with `transformer_blocks` + `single_transformer_blocks`)
+- **Special forward patterns:** Models with non-standard block execution patterns
+
+### Example 1: Dual-Transformer Model (Wan2.2)
+
+Wan2.2 uses two transformers: one for high-noise steps and one for low-noise steps.
+
+**Key difference:** Use `BlockAdapter` to wrap multiple transformers with separate configurations.
+
+```python
+# Standard: cache_dit.enable_cache(pipeline.transformer, ...)
+# Custom: Use BlockAdapter to handle multiple transformers
+cache_dit.enable_cache(
+    BlockAdapter(
+        transformer=[pipeline.transformer, pipeline.transformer_2],  # Multiple transformers
+        blocks=[pipeline.transformer.blocks, pipeline.transformer_2.blocks],
+        forward_pattern=[ForwardPattern.Pattern_2, ForwardPattern.Pattern_2],
+        params_modifiers=[
+            ParamsModifier(...),  # Config for high-noise transformer
+            ParamsModifier(...),  # Config for low-noise transformer (different params)
+        ],
+    ),
+    cache_config=db_cache_config,
+)
+```
+
+**Key difference:** `refresh_context` must be called on each transformer separately.
+
+```python
+# Standard: cache_dit.refresh_context(pipeline.transformer, num_inference_steps=N)
+# Custom: Refresh each transformer with its own step count
+def refresh_cache_context(pipeline, num_inference_steps, verbose=True):
+    high_steps, low_steps = _split_inference_steps(num_inference_steps)
+    cache_dit.refresh_context(pipeline.transformer, num_inference_steps=high_steps, ...)
+    cache_dit.refresh_context(pipeline.transformer_2, num_inference_steps=low_steps, ...)
+```
+
+### Example 2: Multi-Block-List Model (LongCatImage)
+
+LongCatImage has a single transformer with two block lists: `transformer_blocks` and `single_transformer_blocks`.
+
+**Key difference:** Use `BlockAdapter` to specify multiple block lists within one transformer.
+
+```python
+# Standard: cache_dit.enable_cache(pipeline.transformer, ...)
+#   - Automatically detects single block list
+# Custom: Use BlockAdapter to specify multiple block lists
+cache_dit.enable_cache(
+    BlockAdapter(
+        transformer=pipeline.transformer,  # Single transformer
+        blocks=[
+            pipeline.transformer.transformer_blocks,        # Block list 1
+            pipeline.transformer.single_transformer_blocks, # Block list 2
+        ],
+        forward_pattern=[ForwardPattern.Pattern_1, ForwardPattern.Pattern_1],
+        params_modifiers=[modifier],
+    ),
+    cache_config=db_cache_config,
+)
+```
+
+> **Note:** For single transformer with multiple block lists, `refresh_context` works the same as standard models.
+
+### Registering Custom Implementations
+
+After writing your custom enabler, register it in `CUSTOM_DIT_ENABLERS` in `vllm_omni/diffusion/cache/cache_dit_backend.py`:
+
+```python
+CUSTOM_DIT_ENABLERS = {
+    "Wan22Pipeline": enable_cache_for_wan22,
+    "LongCatImagePipeline": enable_cache_for_longcat_image,
+    "YourCustomPipeline": enable_cache_for_your_model,  # Add here
+}
+```
+
+---
+
+## Testing
+
+After adding cache-dit support, test with:
+
+```python
+from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+# Test your custom model
+omni = Omni(
+    model="your-model-name",
+    cache_backend="cache_dit",
+    cache_config={
+        "Fn_compute_blocks": 1,
+        "Bn_compute_blocks": 0,
+        "max_warmup_steps": 4,
+        "residual_diff_threshold": 0.24,
+    }
+)
+
+images = omni.generate(
+    "a beautiful landscape",
+    OmniDiffusionSamplingParams(num_inference_steps=50),
+)
+```
+
+**Verify:**
+
+1. Cache is applied (check logs for "Cache-dit enabled successfully on xxx")
+2. Performance improvement (should be around 1.5x-2x faster)
+3. Image quality (compare with `cache_backend=None`)
+
+---
+
+## Troubleshooting
+
+### Issue: Cache not applied
+
+**Symptoms:** No speedup observed, no cache-related log messages.
+
+**Causes & Solutions:**
+
+- **Enabler not registered:**
+
+**Problem:** Pipeline name not in `CUSTOM_DIT_ENABLERS` registry.
+
+**Solution:** Verify `pipeline.__class__.__name__` matches the registry key and add your enabler to `CUSTOM_DIT_ENABLERS`.
+
+### Issue: Quality degradation
+
+**Symptoms:** Generated images have artifacts or lower quality compared to non-cached inference.
+
+**Causes & Solutions:**
+
+- **Cache parameters too aggressive:**
+
+**Solution:**
+```python
+cache_config={
+    "residual_diff_threshold": 0.12,  # Lower from 0.24 (try 0.12-0.18)
+    "max_warmup_steps": 6,            # Increase from 4 (try 6-8)
+    "max_continuous_cached_steps": 2, # Reduce if higher
+}
+```
+
+Check the [user guide for cache_dit](../../user_guide/diffusion/cache_dit_acceleration.md) for more adjustable parameters.
+
+---
+
+## Reference Implementations
+
+Complete examples in the codebase:
+
+| Model | Path | Pattern | Notes |
+|-------|------|---------|-------|
+| **Standard DiT** | [`cache_dit_backend.py::enable_cache_for_dit`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/cache_dit_backend/#vllm_omni.diffusion.cache.cache_dit_backend.enable_cache_for_dit) | Default enabler | Single transformer, automatic |
+| **Wan2.2** | [`cache_dit_backend.py::enable_cache_for_wan22`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/cache_dit_backend/#vllm_omni.diffusion.cache.cache_dit_backend.enable_cache_for_wan22) | Dual-transformer | Separate high/low noise transformers |
+| **LongCat** | [`cache_dit_backend.py::enable_cache_for_longcat_image`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/cache_dit_backend/#vllm_omni.diffusion.cache.cache_dit_backend.enable_cache_for_longcat_image) | Multi-block-list | Two block lists in one transformer |
+| **BAGEL** | [`cache_dit_backend.py::enable_cache_for_bagel`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/cache_dit_backend/#vllm_omni.diffusion.cache.cache_dit_backend.enable_cache_for_bagel) | Omni model | Complex architecture |
+
+---
+
+## Summary
+
+Adding cache-dit support:
+
+1. ✅ **Check model type** - Standard models work automatically, custom architectures need enablers
+2. ✅ **Write enabler** (if needed) - Use `BlockAdapter` for complex architectures
+3. ✅ **Register enabler** (if needed) - Add to `CUSTOM_DIT_ENABLERS` dictionary
+4. ✅ **Return refresh function** (if needed) - Handle `num_inference_steps` changes
+5. ✅ **Test** - Verify with `cache_backend="cache_dit"`
+
+For most models, the default enabler is sufficient. Only write custom enablers for complex architectures!
@@ -0,0 +1,249 @@
+# Support CFG-Parallel
+
+This section describes how to add CFG-Parallel (Classifier-Free Guidance Parallel) to a diffusion pipeline. We use the Qwen-Image pipeline as the reference implementation.
+
+---
+
+## Table of Contents
+
+- [Overview](#overview)
+- [Step-by-Step Implementation](#step-by-step-implementation)
+- [Customization](#customization)
+- [Testing](#testing)
+- [Troubleshooting](#troubleshooting)
+- [Reference Implementations](#reference-implementations)
+- [Summary](#summary)
+
+---
+
+## Overview
+
+### What is CFG-Parallel?
+
+In standard Classifier-Free Guidance, each diffusion step requires two forward passes through the transformer:
+
+1. **Positive/Conditional**: Guided by the text prompt
+2. **Negative/Unconditional**: Typically using empty or negative prompt
+
+CFG-Parallel eliminates this bottleneck by distributing the two forward passes across different GPU ranks, allowing them to execute simultaneously rather than sequentially.
+
+### Architecture
+
+vLLM-omni provides `CFGParallelMixin` that encapsulates all CFG parallel logic. Pipelines inherit from this mixin and implement a `diffuse()` method that orchestrates the denoising loop.
+
+| 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 |
+| [`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 |
+
+### How It Works
+
+`predict_noise_maybe_with_cfg()` automatically detects and switches between two execution modes:
+
+- **CFG-Parallel mode** (when `cfg_world_size > 1`):
+  - 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
+
+- **Sequential mode** (when `cfg_world_size == 1`):
+  - Single rank computes both positive and negative predictions
+  - Directly combines them with CFG formula
+
+`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
+
+---
+
+## Step-by-Step Implementation
+
+### Step 1: Inherit `CFGParallelMixin`
+
+Allow your pipeline to inherit from `CFGParallelMixin` and implements the `diffuse()` method for your specific model.
+
+**Example (Qwen-Image):**
+
+```python
+from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
+import torch.nn as nn
+class YourModelPipeline(nn.Module, CFGParallelMixin):
+    def diffuse(self, ...) -> torch.Tensor:
+        for i, t in enumerate(timesteps):
+            # Prepare positive_kwargs (conditional) and negative_kwargs (unconditional)
+            positive_kwargs = {...}  # hidden_states, encoder_hidden_states, etc.
+            negative_kwargs = {...} if do_true_cfg else None
+
+            # Key method 1: Predict noise with automatic CFG parallel handling
+            noise_pred = self.predict_noise_maybe_with_cfg(
+                do_true_cfg=do_true_cfg,
+                true_cfg_scale=true_cfg_scale,
+                positive_kwargs=positive_kwargs,
+                negative_kwargs=negative_kwargs,
+            )
+
+            # Key method 2: Step scheduler with automatic CFG synchronization
+            latents = self.scheduler_step_maybe_with_cfg(
+                noise_pred, t, latents, do_true_cfg
+            )
+
+        return latents
+```
+
+**Key Points:**
+
+- `positive_kwargs`: transformer arguments for conditional (text-guided) prediction
+- `negative_kwargs`: transformer arguments for unconditional prediction (set to `None` if CFG disabled)
+- For image editing pipelines, add `output_slice=image_seq_len` to extract the generative image portion
+
+### Step 2: Call `diffuse`
+
+Call `self.diffuse` in your pipeline's forward function:
+
+```python
+import torch.nn as nn
+class YourModelPipeline(nn.Module, CFGParallelMixin):
+    def forward(
+        self,
+        prompt: str,
+        negative_prompt: str | None = None,
+        guidance_scale: float = 3.5,
+        num_inference_steps: int = 50,
+        **kwargs,
+    ):
+        # Encode prompts, Initialize latents, Get timesteps
+        ...
+        # Run diffusion loop (calls the mixin's diffuse method)
+        latents = self.diffuse(
+            prompt_embeds=prompt_embeds,
+            prompt_embeds_mask=prompt_embeds_mask,
+            negative_prompt_embeds=negative_embeds,
+            negative_prompt_embeds_mask=negative_mask,
+            latents=latents,
+            timesteps=timesteps,
+            do_true_cfg=do_true_cfg,
+            true_cfg_scale=guidance_scale,
+            ...
+        )
+```
+
+---
+
+## Customization
+
+### Override `predict_noise()` for Custom Transformer Calls
+
+If your transformer requires custom prediction function, you can rewrite `predict_noise` function. Taking Wan2.2 as an example, which has two transformer models. The actual transformer to be called is determined by `self.transformer`.
+
+```python
+class Wan22Pipeline(nn.Module, CFGParallelMixin):
+    def predict_noise(self, current_model: nn.Module | None = None, **kwargs: Any) -> torch.Tensor:
+        if current_model is None:
+            current_model = self.transformer
+        return current_model(**kwargs)[0]
+```
+
+
+### Override `cfg_normalize_function()` for Custom Normalization
+
+Some models have their own normalization function. Taking LongCat Image model as an example:
+
+```python
+class LongCatImagePipeline(nn.Module, CFGParallelMixin):
+    def cfg_normalize_function(self, noise_pred, comb_pred, cfg_renorm_min=0.0):
+        """
+        Normalize the combined noise prediction.
+        """
+        cond_norm = torch.norm(noise_pred, dim=-1, keepdim=True)
+        noise_norm = torch.norm(comb_pred, dim=-1, keepdim=True)
+        scale = (cond_norm / (noise_norm + 1e-8)).clamp(min=cfg_renorm_min, max=1.0)
+        noise_pred = comb_pred * scale
+        return noise_pred
+
+        # The original cfg_normalize_function function in CFGParallelMixin
+        # cond_norm = torch.norm(noise_pred, dim=-1, keepdim=True)
+        # noise_norm = torch.norm(comb_pred, dim=-1, keepdim=True)
+        # noise_pred = comb_pred * (cond_norm / noise_norm)
+        # return noise_pred
+```
+
+---
+
+## Testing
+
+After adding CFG-Parallel support, test with:
+
+```bash
+cd examples/offline_inference/text_to_image
+python text_to_image.py \
+    --model Your-org/your-model \
+    --prompt "a cup of coffee on the table" \
+    --negative-prompt "ugly, unclear" \
+    --cfg-scale 4.0 \
+    --num-inference-steps 50 \
+    --output "cfg_enabled.png" \
+    --cfg-parallel-size 2
+```
+
+**Verify:**
+
+1. Check logs for CFG parallel being activated
+2. Record the `e2e_time_ms` in the log and compare with CFG-Parallel disabled
+3. Compare the generated result quality with baseline
+4. Record comparison results in your PR
+
+---
+
+## Troubleshooting
+
+### Issue: CFG parallel not activating
+
+**Symptoms:** Generation still slow, logs don't show CFG parallel being used.
+
+**Causes & Solutions:**
+
+- **CFG is not enabled:**
+
+**Problem:** Guidance scale too low or negative prompt not provided.
+
+**Solution:** Ensure `guidance_scale > 1.0` and negative prompt is provided:
+```python
+images = pipeline(
+    prompt="a cat",
+    negative_prompt="",  # Must provide (even if empty)
+    guidance_scale=3.5,   # Must be > 1.0
+)
+```
+
+---
+
+## Reference Implementations
+
+Complete examples in the codebase:
+
+| Model | Path | Pattern | Notes |
+|-------|------|---------|-------|
+| **Qwen-Image** | `vllm_omni/diffusion/models/qwen_image/cfg_parallel.py` | Mixin | Dual-stream transformer |
+| **Qwen-Image-Edit** | `vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py` | Mixin | Image editing with `output_slice` |
+| **Wan2.2** | `vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py` | Mixin | Dual-transformer architecture |
+| **CFGParallelMixin** | `vllm_omni/diffusion/distributed/cfg_parallel.py` | Base implementation | Core mixin class |
+
+---
+
+## Summary
+
+Adding CFG-Parallel support:
+
+1. ✅ **Create mixin** - Inherit from `CFGParallelMixin` and implement `diffuse()` method
+2. ✅ **(Optional) Customize** - Override `predict_noise()` or `cfg_normalize_function()` for custom behavior
+3. ✅ **Test** - Verify with `--cfg-parallel-size 2` and compare performance
@@ -0,0 +1,405 @@
+# Support Sequence Parallel
+
+This section describes how to add Sequence Parallel (SP) to a diffusion transformer model. We use the Qwen-Image transformer and Wan2.2 transformer as reference implementations.
+
+---
+
+## Table of Contents
+
+- [Overview](#overview)
+- [Approach 1: Non-Intrusive `_sp_plan` (Recommended)](#approach-1-non-intrusive-_sp_plan-recommended)
+- [Approach 2: Intrusive Modification (For Complex Cases)](#approach-2-intrusive-modification-for-complex-cases)
+- [Testing](#testing)
+- [Troubleshooting](#troubleshooting)
+- [Reference Implementations](#reference-implementations)
+- [Summary](#summary)
+
+---
+
+## Overview
+
+
+### What is Sequence Parallel?
+
+**Terminology Note:** Our "Sequence Parallelism" (SP) corresponds to "Context Parallelism" (CP) in the [diffusers library](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/_modeling_parallel.py). We use "Sequence Parallelism" to align with vLLM-Omni's terminology.
+
+Diffusion transformers process long sequences of image patches or video frames. For high-resolution generation, these sequences can become very large. Enabling SP allows each GPU to process only a portion of the sequence, with attention mechanisms (Ulysses/Ring) handling cross-GPU communication transparently.
+
+### Architecture
+
+The major APIs for Sequence Parallel:
+
+```python
+from vllm_omni.diffusion.distributed.sp_plan import (
+    SequenceParallelInput,   # For sharding (splitting) tensors
+    SequenceParallelOutput,  # For gathering tensors
+)
+from vllm_omni.diffusion.distributed.sp_sharding import sp_shard, sp_gather
+```
+
+| Method/Class | Purpose | Behavior |
+|--------------|---------|----------|
+| `SequenceParallelInput` | Declare input sharding in `_sp_plan` | Auto-shards tensors at module input |
+| `SequenceParallelOutput` | Declare output gathering in `_sp_plan` | Auto-gathers tensors at module output |
+| `sp_shard()` | Manual tensor sharding | Splits tensor across SP workers |
+| `sp_gather()` | Manual tensor gathering | Gathers sharded tensors from all workers |
+
+---
+
+## Approach 1: Non-Intrusive `_sp_plan` (Recommended)
+
+The `_sp_plan` mechanism allows SP **without modifying `forward()` logic**. The framework automatically registers hooks to shard inputs and gather outputs at module boundaries.
+
+**When to use:**
+- Standard transformer architectures
+- Tensor operations happen at `nn.Module` boundaries
+- Predictable sharding/gathering patterns
+
+**How it works:**
+1. Declare `_sp_plan` dict in your transformer class
+2. Framework automatically applies hooks when `sequence_parallel_size > 1`
+3. Hooks shard/gather tensors at specified module boundaries
+4. Attention layers handle cross-GPU communication internally
+
+```python
+class StandardTransformer(nn.Module):
+    _sp_plan = {
+        # Shard hidden_states at first transformer block input
+        "blocks.0": {
+            "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3),
+        },
+        # Gather at final output projection
+        "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3),
+    }
+```
+
+`StandardTransformer` has a transformer blocks list `self.blocks = nn.ModuleList([...])`, and a projection output layer `self.proj_out`. The `_sp_plan` above defines that when SP is enabled, sharding the input tensor to the first transformer block, and gathering the sharded tensor at the final output projection layer.
+
+**Requirements:**
+- Tensor operations that need sharding/gathering must happen at **`nn.Module` boundaries**
+- Inline Python operations (e.g., `torch.cat`, `pad_sequence`) **cannot be hooked**
+
+**Solution for inline operations:** Extract into a submodule (see Step 2 below).
+
+### Step 1: Understand Module Boundaries
+
+Identify where tensors need to be sharded or gathered in your model's forward pass:
+
+```python
+class MyTransformer(nn.Module):
+    def __init__(self):
+        self.patch_embed = PatchEmbed()      # ← Boundary 1
+        self.pos_embed = RoPE()              # ← Boundary 2
+        self.blocks = nn.ModuleList([...])   # ← Boundary 3
+        self.norm_out = LayerNorm()
+        self.proj_out = Linear()             # ← Boundary 4
+
+    def forward(self, x):
+        x = self.patch_embed(x)              # ← Shard before this?
+        pos = self.pos_embed(x)              # ← Shard RoPE outputs?
+        for block in self.blocks:
+            x = block(x, pos)                # ← Blocks process sharded x
+        x = self.norm_out(x)
+        output = self.proj_out(x)            # ← Gather after this?
+        return output
+```
+
+### Step 2: Handle Inline Operations
+
+If your `forward()` contains inline tensor operations, **extract them into submodules**.
+
+**Example: Z-Image concatenates image + text features inline**
+
+```python
+# ❌ BAD: Inline operation - hooks cannot intercept
+class ZImageTransformer(nn.Module):
+    def forward(self, x, cap_feats):
+        # This concatenation happens inline - _sp_plan can't shard it!
+        unified = torch.cat([x, cap_feats], dim=1)
+
+        for layer in self.layers:
+            unified = layer(unified)
+
+        return unified
+
+# ✅ GOOD: Extract into submodule
+class UnifiedPrepare(nn.Module):
+    """Submodule to concatenate image and text features."""
+    def forward(self, x, cap_feats):
+        return torch.cat([x, cap_feats], dim=1)
+
+class ZImageTransformer(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.unified_prepare = UnifiedPrepare()  # Now a module!
+        self.layers = nn.ModuleList([...])
+
+    def forward(self, x, cap_feats):
+        # Now _sp_plan can shard the output of unified_prepare!
+        unified = self.unified_prepare(x, cap_feats)
+
+        for layer in self.layers:
+            unified = layer(unified)
+
+        return unified
+```
+
+**Other common cases:**
+- `pad_sequence()` → `PadSequenceModule`
+- `torch.cat()` → `ConcatModule`
+- `tensor.reshape()` → `ReshapeModule`
+- Complex preprocessing → `PreprocessModule`
+
+### Step 3: Write `_sp_plan` for Your Model
+
+Create a class-level `_sp_plan` dictionary specifying where to shard/gather tensors.
+
+Typically, there are two patterns for diffusion models:
+
+**Pattern 1: Shard at first block, gather at output projection**
+
+Most common pattern for standard transformers:
+
+```python
+from vllm_omni.diffusion.distributed.sp_plan import (
+    SequenceParallelInput,   # For sharding (splitting) tensors
+    SequenceParallelOutput,  # For gathering tensors
+)
+class StandardTransformer(nn.Module):
+    _sp_plan = {
+        # Shard hidden_states at first transformer block input
+        "blocks.0": {
+            "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3),
+        },
+        # Gather at final output projection
+        "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3),
+    }
+```
+
+**Pattern 2: Shard RoPE embeddings separately**
+
+When RoPE is computed in a separate module:
+
+```python
+from vllm_omni.diffusion.distributed.sp_plan import (
+    SequenceParallelInput,   # For sharding (splitting) tensors
+    SequenceParallelOutput,  # For gathering tensors
+)
+class TransformerWithRoPE(nn.Module):
+    _sp_plan = {
+        # Shard RoPE module OUTPUTS (returns tuple of cos, sin)
+        "rope": {
+            0: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True),  # cos
+            1: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True),  # sin
+        },
+        # Shard transformer block INPUT
+        "blocks.0": {
+            "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3),
+        },
+        # Gather at output
+        "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3),
+    }
+```
+
+### API Reference
+
+**SequenceParallelInput Parameters:**
+
+| Parameter | Type | Description |
+|-----------|------|-------------|
+| `split_dim` | int | Dimension to split (usually `1` for sequence) |
+| `expected_dims` | int \| None | Expected tensor rank for validation (optional) |
+| `split_output` | bool | `False`: shard **input** params; `True`: shard **output** tensors |
+| `auto_pad` | bool | Auto-pad if sequence not divisible by world_size (default: `False`) |
+
+**SequenceParallelOutput Parameters:**
+
+| Parameter | Type | Description |
+|-----------|------|-------------|
+| `gather_dim` | int | Dimension to gather (usually `1` for sequence) |
+| `expected_dims` | int \| None | Expected tensor rank for validation (optional) |
+
+**Module Naming Conventions:**
+
+| Key | Meaning | Python equivalent |
+|-----|---------|-------------------|
+| `""` | Root model | `model` |
+| `"blocks.0"` | First element of ModuleList | `model.blocks[0]` |
+| `"blocks.*"` | All elements of ModuleList | `for b in model.blocks` |
+| `"rope"` | Named submodule | `model.rope` |
+| `"outputs.main"` | ModuleDict entry | `model.outputs["main"]` |
+
+**Dictionary Value Types:**
+
+| Key type | `split_output` | Description |
+|----------|----------------|-------------|
+| `"param_name"` (str) | `False` | Shard **input parameter** by name |
+| `0`, `1`, ... (int) | `True` | Shard **output tuple** by index |
+
+---
+
+## Approach 2: Intrusive Modification (For Complex Cases)
+
+For models with dynamic sharding logic that cannot be expressed via `_sp_plan`, manually insert shard/gather calls.
+
+**When to use:**
+- Dynamic/conditional sharding logic
+- Complex tensor manipulations that can't be encapsulated
+- Temporary workaround during development
+
+```python
+from vllm_omni.diffusion.distributed.sp_sharding import sp_shard, sp_gather
+
+def forward(self, hidden_states, ...):
+    if self.parallel_config.sequence_parallel_size > 1:
+        hidden_states = sp_shard(hidden_states, dim=1)
+
+    # ... computation ...
+
+    if self.parallel_config.sequence_parallel_size > 1:
+        output = sp_gather(output, dim=1)
+
+    return output
+```
+
+---
+
+## Testing
+
+After implementing Sequence Parallel support, thoroughly test your implementation to ensure correctness and performance across different configurations.
+
+**Test Different `sp_size`:**
+
+Test your model with various sequence parallel world sizes to verify correctness and identify optimal configurations:
+
+```bash
+cd examples/offline_inference/text_to_image
+python text_to_image.py \
+    --model Your-org/your-model \
+    --prompt "a cup of coffee on the table" \
+    --num-inference-steps 50 \
+    --ulysses-degree 2 \
+    --ring-degree 2 \
+    --output sp_test_image_ulysses=2_ring=2.png
+```
+
+**Verify:**
+
+1. **Correctness:** Output should be identical across all `sp_size` values
+2. **Speed:** Throughput should remain stable or improve (especially for large sequences)
+3. **Logs:** Check for any shape mismatch or communication errors
+
+**Test with Tensor Parallel:**
+
+Sequence Parallel can be combined with other parallelism strategies:
+
+```bash
+cd examples/offline_inference/text_to_image
+python text_to_image.py \
+    --model Your-org/your-model \
+    --prompt "a cup of coffee on the table" \
+    --num-inference-steps 50 \
+    --ulysses-degree 2 \
+    --tensor-parallel-size 2 \
+    --output sp_test_image_ulysses=2_tp=2.png
+```
+
+---
+
+## Troubleshooting
+
+### Issue: Shape mismatch errors
+
+**Symptoms:** `RuntimeError: shape mismatch` during forward pass.
+
+**Causes & Solutions:**
+
+- **RoPE dimension mismatch:**
+
+**Problem:** RoPE embeddings not sharded, but hidden_states is sharded.
+
+**Solution:** Shard RoPE outputs in `_sp_plan`:
+```python
+_sp_plan = {
+    "rope": {
+        0: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True),
+        1: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True),
+    },
+    ...
+}
+```
+
+- **Sequence Length not divisible by sp_size:**
+
+**Problem:** `SequenceParallelInput(auto_pad=False)` - auto_pad should be True to enable automatic sequence padding.
+
+**Solution:** In `SequenceParallelInput`, set `auto_pad=True`:
+```python
+"blocks.0": {
+    "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3, auto_pad=True)
+}
+```
+
+### Issue: Inline operations not sharded
+
+**Symptoms:** Some tensors remain full-sized, not sharded.
+
+**Causes & Solutions:**
+
+- **Operations happen inline in `forward()`, not at module boundaries:**
+
+**Problem:**
+```python
+def forward(self, x, cap):
+    unified = torch.cat([x, cap], dim=1)  # ← Inline operation!
+    # _sp_plan can't hook this
+```
+
+**Solution:** Extract into submodule:
+```python
+class ConcatModule(nn.Module):
+    def forward(self, x, cap):
+        return torch.cat([x, cap], dim=1)
+
+class MyModel(nn.Module):
+    _sp_plan = {
+        "concat": {
+            0: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True),
+            1: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True),
+        },
+        ...
+    }
+    def __init__(self):
+        self.concat = ConcatModule()  # Now hookable!
+
+    def forward(self, x, cap):
+        unified = self.concat(x, cap)  # ← Can be sharded via _sp_plan
+```
+
+---
+
+## Reference Implementations
+
+Complete examples in the codebase:
+
+| Model | Path | Pattern | Notes |
+|-------|------|---------|-------|
+| **Qwen-Image** | `vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py` | Dual-stream + preprocessing | auto_pad, separate RoPE |
+| **Wan2.2** | `vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py` | Dual-Transformer + RoPE | Video transformer |
+| **Z-Image** | `vllm_omni/diffusion/models/z_image/z_image_transformer.py` | Unified sequence | Concatenated input |
+| **SP Plan Types** | `vllm_omni/diffusion/distributed/sp_plan.py` | Type definitions | SequenceParallelInput/Output |
+| **Hook Implementation** | `vllm_omni/diffusion/hooks/sequence_parallel.py` | Hook mechanics | How hooks work |
+| **Tests** | `tests/diffusion/distributed/test_sp_plan_hooks.py` | Test examples | Validation patterns |
+
+---
+
+## Summary
+
+Adding Sequence Parallel support to a transformer:
+
+1. ✅ **Choose approach** - Use `_sp_plan` for standard cases, intrusive modification for complex cases
+2. ✅ **Identify sharding boundaries** - Where should tensors be split/gathered?
+3. ✅ **Extract inline operations** - Move `torch.cat`, `pad_sequence`, etc. to submodules
+4. ✅ **Define `_sp_plan`** - Declare shard/gather points as class attribute
+5. ✅ **Use `auto_pad` for variable lengths** - Support non-uniform sequences
+6. ✅ **Test** - Verify with different `ulysses_degree` and `ring_degree` combinations
@@ -0,0 +1,459 @@
+# Support TeaCache
+
+This section describes how to add TeaCache to a diffusion transformer model. We use the Qwen-Image transformer as the reference implementation.
+
+---
+
+## Table of Contents
+
+- [Overview](#overview)
+- [Step-by-Step Implementation](#step-by-step-implementation)
+- [Customization](#customization)
+- [Testing](#testing)
+- [Troubleshooting](#troubleshooting)
+- [Reference Implementations](#reference-implementations)
+- [Summary](#summary)
+
+---
+
+## Overview
+
+### What is TeaCache?
+
+TeaCache speeds up diffusion inference by caching transformer block computations when consecutive timesteps are similar. It provides **1.5x-2.0x speedup** with minimal quality loss.
+
+The core insight is that the modulated input (after normalization and timestep conditioning) changes gradually across timesteps. By measuring the L1 distance between consecutive modulated inputs and comparing it to a threshold, TeaCache decides whether to execute the full transformer blocks or reuse the cached residual from the previous step.
+
+vLLM-omni provides a **hook-based** TeaCache system that requires **zero changes to model code**. The hook completely intercepts the transformer's forward pass and implements adaptive caching transparently. This design allows easy integration with any transformer model by simply writing an extractor function.
+
+### Architecture
+
+The TeaCache system consists of three main components:
+
+| Component | Purpose | Location |
+|-----------|---------|----------|
+| [`CacheContext`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/#vllm_omni.diffusion.cache.CacheContext) | Dataclass containing model-specific information for caching | `vllm_omni/diffusion/cache/teacache/context.py` |
+| [`EXTRACTOR_REGISTRY`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/teacache/extractors/#vllm_omni.diffusion.cache.teacache.extractors.EXTRACTOR_REGISTRY) | Maps transformer class names to extractor functions | `vllm_omni/diffusion/cache/teacache/extractors.py` |
+| [`TeaCacheConfig`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/#vllm_omni.diffusion.cache.TeaCacheConfig) | Configuration including thresholds and polynomial coefficients | `vllm_omni/diffusion/cache/teacache/config.py` |
+
+The hook handles all caching logic automatically, including:
+
+- CFG-aware state management (separate states for positive/negative branches)
+- CFG-parallel compatibility
+- L1 distance computation with polynomial rescaling
+- Residual caching and reuse
+
+
+---
+
+## Step-by-Step Implementation
+
+To add TeaCache support for a new model, you need to:
+
+1. Write an **extractor function** that returns a `CacheContext` object
+2. Register the extractor in the `EXTRACTOR_REGISTRY`
+3. Add model-specific polynomial coefficients to `TeaCacheConfig`
+
+### Step 1: Model-Specific Preprocessing
+
+Extract and process model inputs. This typically involves:
+- Embedding image/latent inputs
+- Processing text encoder outputs (if dual-stream)
+- Creating timestep embeddings
+- Computing positional embeddings
+
+**Example (Qwen-Image):**
+
+```python
+def extract_qwen_context(
+    module: nn.Module,
+    hidden_states: torch.Tensor,
+    encoder_hidden_states: torch.Tensor,
+    encoder_hidden_states_mask: torch.Tensor,
+    timestep: torch.Tensor,
+    img_shapes: torch.Tensor,
+    txt_seq_lens: torch.Tensor,
+    guidance: torch.Tensor | None = None,
+    **kwargs: Any,
+) -> CacheContext:
+    # Validate model structure
+    if not hasattr(module, "transformer_blocks") or len(module.transformer_blocks) == 0:
+        raise ValueError("Module must have transformer_blocks")
+
+    # Preprocessing: embed inputs
+    hidden_states = module.img_in(hidden_states)
+    timestep = timestep.to(device=hidden_states.device, dtype=hidden_states.dtype)
+    encoder_hidden_states = module.txt_norm(encoder_hidden_states)
+    encoder_hidden_states = module.txt_in(encoder_hidden_states)
+
+    # Create timestep embedding
+    if guidance is not None:
+        guidance = guidance.to(hidden_states.dtype) * 1000
+    temb = (
+        module.time_text_embed(timestep, hidden_states)
+        if guidance is None
+        else module.time_text_embed(timestep, guidance, hidden_states)
+    )
+
+    # Compute position embeddings
+    image_rotary_emb = module.pos_embed(img_shapes, txt_seq_lens, device=hidden_states.device)
+```
+
+### Step 2: Extract Modulated Input
+
+The modulated input is used for cache decisions. Extract it from the **first transformer block** after normalization and modulation.
+
+**Example (Qwen-Image):**
+
+```python
+    # Extract modulated input from first transformer block
+    block = module.transformer_blocks[0]
+    img_mod_params = block.img_mod(temb)
+    img_mod1, _ = img_mod_params.chunk(2, dim=-1)
+    img_modulated, _ = block.img_norm1(hidden_states, img_mod1)
+```
+
+**Key Points:**
+
+- Use the **first block** to extract modulated input early
+- Apply the same normalization and modulation as the actual forward pass
+- The tensor should represent the processed features that will change across timesteps
+
+### Step 3: Define Transformer Execution
+
+Create a callable that executes all transformer blocks. This encapsulates the main computation loop.
+
+**Example (Qwen-Image dual-stream):**
+
+```python
+    def run_transformer_blocks():
+        """Execute all Qwen transformer blocks."""
+        h = hidden_states
+        e = encoder_hidden_states
+
+        for block in module.transformer_blocks:
+            e, h = block(
+                hidden_states=h,
+                encoder_hidden_states=e,
+                encoder_hidden_states_mask=encoder_hidden_states_mask,
+                temb=temb,
+                image_rotary_emb=image_rotary_emb,
+            )
+        return (h, e)  # Return both image and text hidden states
+```
+
+**Example (Single-stream model like Flux):**
+
+```python
+    def run_transformer_blocks():
+        """Execute all Flux transformer blocks."""
+        h = hidden_states
+
+        for block in module.transformer_blocks:
+            h = block(h, temb=temb)
+        return (h,)  # Return only image hidden states
+```
+
+**Key Points:**
+
+- Return format:
+- For single-stream models: return `(hidden_states,)`
+- For dual-stream models: return `(hidden_states, encoder_hidden_states)`
+
+### Step 4: Define Postprocessing
+
+Create a callable that applies final transformations to produce the model output.
+
+**Example (Qwen-Image):**
+
+```python
+    return_dict = kwargs.get("return_dict", True)
+
+    def postprocess(h):
+        """Apply Qwen-specific output postprocessing."""
+        h = module.norm_out(h, temb)
+        output = module.proj_out(h)
+        if not return_dict:
+            return (output,)
+        return Transformer2DModelOutput(sample=output)
+```
+
+### Step 5: Return CacheContext
+
+Package all information into a `CacheContext` object.
+
+```python
+    return CacheContext(
+        modulated_input=img_modulated,
+        hidden_states=hidden_states,
+        encoder_hidden_states=encoder_hidden_states,  # or None for single-stream
+        temb=temb,
+        run_transformer_blocks=run_transformer_blocks,
+        postprocess=postprocess,
+    )
+```
+
+**CacheContext Fields:**
+
+| Field | Type | Purpose |
+|-------|------|---------|
+| `modulated_input` | `torch.Tensor` | Tensor used for cache decision (similarity comparison) |
+| `hidden_states` | `torch.Tensor` | Current hidden states (will be modified by caching) |
+| `encoder_hidden_states` | `torch.Tensor | None` | Encoder states for dual-stream models, `None` for single-stream |
+| `temb` | `torch.Tensor` | Timestep embedding tensor |
+| `run_transformer_blocks` | `Callable[[], tuple]` | Executes transformer blocks, returns `(hidden_states, [encoder_hidden_states])` |
+| `postprocess` | `Callable[[torch.Tensor], Any]` | Applies final transformations to produce model output |
+| `extra_states` | `dict | None` | Optional dict for additional model-specific state |
+
+### Step 6: Register the Extractor
+
+Add your extractor to the `EXTRACTOR_REGISTRY` in `vllm_omni/diffusion/cache/teacache/extractors.py`:
+
+```python
+EXTRACTOR_REGISTRY: dict[str, Callable] = {
+    "QwenImageTransformer2DModel": extract_qwen_context,
+    "Bagel": extract_bagel_context,
+    "ZImageTransformer2DModel": extract_zimage_context,
+    "YourModelTransformer2DModel": extract_your_model_context,  # Add here
+}
+```
+
+**Key:** Use the transformer class name (`module.__class__.__name__`)
+
+### Step 7: Add Model Coefficients
+
+Add polynomial rescaling coefficients to `vllm_omni/diffusion/cache/teacache/config.py`:
+
+```python
+_MODEL_COEFFICIENTS = {
+    "QwenImageTransformer2DModel": [
+        -4.50000000e02,
+        2.80000000e02,
+        -4.50000000e01,
+        3.20000000e00,
+        -2.00000000e-02,
+    ],
+    "YourModelTransformer2DModel": [  # Add your model's coefficients
+        # 5 polynomial coefficients (can reuse similar model's coefficients initially)
+    ],
+}
+```
+
+
+**Initial approach:** Start with coefficients from a similar model architecture, then tune empirically following [Customization](#customization) section.
+
+---
+
+## Customization
+
+### Coefficient Estimation
+
+While you can start with coefficients from a similar model architecture, estimating custom coefficients for your specific model typically improves TeaCache performance.
+
+**Why Estimate Coefficients?**
+
+The polynomial coefficients rescale L1 distances between consecutive modulated inputs to better predict when cached residuals can be reused. Model-specific coefficients account for:
+
+- Architecture differences (layer count, hidden size, attention patterns)
+- Training data characteristics
+- Noise prediction behavior across timesteps
+
+| Approach | Performance | Effort |
+|----------|-------------|--------|
+| Using defaults from similar model | Within 5-10% of optimal | Low |
+| Estimating custom coefficients | Best performance | Medium |
+
+#### Implement Data Collection Adapter
+
+Add an adapter in `vllm_omni/diffusion/cache/teacache/coefficient_estimator.py`:
+
+```python
+class YourModelAdapter:
+    """Adapter for coefficient estimation on your model."""
+
+    @staticmethod
+    def load_pipeline(model_path: str, device: str, dtype: torch.dtype) -> Any:
+        """Load your diffusion pipeline."""
+        from your_model_package import YourModelPipeline
+
+        pipeline = YourModelPipeline.from_pretrained(
+            model_path,
+            torch_dtype=dtype,
+        )
+        pipeline = pipeline.to(device)
+        return pipeline
+
+    @staticmethod
+    def get_transformer(pipeline: Any) -> tuple[Any, str]:
+        """Extract transformer from pipeline."""
+        return pipeline.transformer, "YourTransformer2DModel"
+
+    @staticmethod
+    def install_hook(transformer: Any, hook: DataCollectionHook) -> None:
+        """Install data collection hook on transformer."""
+        from vllm_omni.diffusion.hooks import HookRegistry
+
+        registry = HookRegistry.get_or_create(transformer)
+        registry.register_hook(hook._HOOK_NAME, hook)
+
+
+# Register your adapter
+_MODEL_ADAPTERS["YourModel"] = YourModelAdapter
+```
+
+#### Collect Data and Estimate
+
+```python
+from vllm_omni.diffusion.cache.teacache.coefficient_estimator import (
+    TeaCacheCoefficientEstimator,
+)
+from datasets import load_dataset
+from tqdm import tqdm
+
+# Initialize estimator
+estimator = TeaCacheCoefficientEstimator(
+    model_path="/path/to/your/model",
+    model_type="YourModel",
+)
+
+# Load diverse prompts (paper recommends ~70 prompts)
+dataset = load_dataset("nateraw/parti-prompts", split="train")
+prompts = dataset["Prompt"][:70]
+
+# Collect data
+for prompt in tqdm(prompts, desc="Collecting data"):
+    estimator.collect_from_prompt(prompt=prompt, num_inference_steps=50)
+
+# Estimate coefficients
+coeffs = estimator.estimate(poly_order=4)
+print(f"Estimated coefficients: {coeffs.tolist()}")
+```
+
+**Data Statistics Guide:**
+
+| Metric | Good Range | Warning Signs |
+|--------|------------|---------------|
+| **Count** | 2000-5000+ | < 1000: too few prompts |
+| **Input Diffs (x)** | 0.01-0.10 | Very small (<0.001): model may not modulate properly |
+| **Output Diffs (y)** | Should correlate with x | No correlation: check extractor |
+| **Coefficient magnitude** | -1e6 to 1e6 | > 1e8: numerical instability |
+
+---
+
+## Testing
+
+After adding TeaCache support, test with:
+
+```python
+from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+omni = Omni(
+    model="your-model-name",
+    cache_backend="tea_cache",
+    cache_config={
+        "rel_l1_thresh": 0.2,
+        "coefficients": [1.33e6, -1.69e5, 7.95e3, -1.64e2, 1.26],  # Your coefficients
+    }
+)
+
+images = omni.generate(
+    "a beautiful landscape",
+    OmniDiffusionSamplingParams(num_inference_steps=50),
+)
+```
+
+**Verify:**
+
+1. **Check logs** - Look for TeaCache initialization messages
+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).
+
+---
+
+## Troubleshooting
+
+### Issue: "Unknown model type"
+
+**Symptoms:** Error message indicating the model type is not recognized when enabling TeaCache.
+
+**Causes & Solutions:**
+
+- **Extractor not registered:**
+
+**Problem:** The transformer class name doesn't exist in `EXTRACTOR_REGISTRY`.
+
+**Solution:** Check the class name and add to registry:
+```python
+# Check transformer class name
+print(pipeline.transformer.__class__.__name__)
+
+# Add to EXTRACTOR_REGISTRY
+EXTRACTOR_REGISTRY["YourTransformer2DModel"] = extract_your_context
+```
+
+- **Transformer class name mismatch:**
+
+**Solution:** Ensure the registry key matches exactly with `module.__class__.__name__`.
+
+### Issue: "Cannot find coefficients"
+
+**Symptoms:** Error when initializing TeaCache about missing model coefficients.
+
+**Causes & Solutions:**
+
+- **Missing coefficients in config:**
+
+**Solution:** Add coefficients to `_MODEL_COEFFICIENTS` in `config.py`, or pass custom coefficients:
+```python
+omni = Omni(
+    model="your-model",
+    cache_backend="tea_cache",
+    cache_config={"coefficients": [1.0, -0.5, 0.1, -0.01, 0.001]}
+)
+```
+
+### Issue: Quality Degradation
+
+**Symptoms:** Output images look noticeably different or have artifacts compared to baseline.
+
+**Causes & Solutions:**
+
+- **Threshold too high:**
+
+**Problem:** `rel_l1_thresh` is too aggressive, causing cache reuse when outputs differ significantly.
+
+**Solution:** Lower the threshold:
+```python
+cache_config={"rel_l1_thresh": 0.1}  # Try 0.1-0.2
+```
+
+- **Coefficients not tuned:**
+
+**Solution:** Estimate model-specific coefficients using the coefficient estimation process described above.
+
+---
+
+## Reference Implementations
+
+Complete examples in the codebase:
+
+| Model | Path | Pattern | Notes |
+|-------|------|---------|-------|
+| **Qwen-Image** | `vllm_omni/diffusion/cache/teacache/extractors.py` | Dual-stream | `extract_qwen_context` |
+| **Bagel** | `vllm_omni/diffusion/cache/teacache/extractors.py` | Omni model | `extract_bagel_context` |
+| **TeaCache Core** | `vllm_omni/diffusion/cache/teacache/` | Base implementation | Hook and config |
+| **Coefficient Estimator** | `vllm_omni/diffusion/cache/teacache/coefficient_estimator.py` | Estimation tool | Adapter pattern |
+
+---
+
+## Summary
+
+Adding TeaCache support:
+
+1. ✅ **Write extractor** - Create function returning `CacheContext` with model-specific preprocessing
+2. ✅ **Register extractor** - Add to `EXTRACTOR_REGISTRY` with transformer class name
+3. ✅ **Add coefficients** - Add polynomial coefficients to `_MODEL_COEFFICIENTS`
+4. ✅ **Test** - Verify with `cache_backend="tea_cache"`
@@ -0,0 +1,279 @@
+# Support Tensor Parallel
+
+This section describes how to add Tensor Parallel (TP) to a diffusion transformer model. We use the Z-Image transformer as the reference implementation.
+
+---
+
+## Table of Contents
+
+- [Overview](#overview)
+- [Step-by-Step Implementation](#step-by-step-implementation)
+- [Testing](#testing)
+- [Troubleshooting](#troubleshooting)
+- [Reference Implementations](#reference-implementations)
+- [Summary](#summary)
+
+---
+
+## Overview
+
+### What is Tensor Parallel?
+
+Tensor Parallel (TP) is a model parallelism technique that **shards model weights** across multiple GPUs. Each GPU holds only a portion of the model's parameters and computes only part of each layer's output.
+
+Diffusion transformers contain large attention and MLP layers. We can use Tensor Parallel to shard the model dimension across multiple GPUs, allowing larger models to fit in memory while achieving near-linear speedup.
+
+### Architecture
+
+The Tensor Parallel implementation relies vLLM's Parallel Layers:
+
+[vLLM Parallel Layers API Reference](https://docs.vllm.ai/en/latest/contributing/model/basic/?h=column#3-optional-implement-tensor-parallelism-and-quantization-support)
+
+**Parallel Layer Types:**
+
+| Layer Type | Purpose | Weight Partitioning |
+|------------|---------|---------------------|
+| `ColumnParallelLinear` | First FFN layer, separated QKV | Columns (output dimension) |
+| `RowParallelLinear` | Second FFN layer, attention output | Rows (input dimension) |
+| `QKVParallelLinear` | Multi-head/grouped-query attention QKV | Handles head replication automatically |
+| `ReplicatedLinear` | Layers that shouldn't be sharded | No partitioning (replicated) |
+
+---
+
+## Step-by-Step Implementation
+
+
+### Step 1: Identify Linear Layers
+
+Find all `nn.Linear` layers in your transformer that need to be sharded.
+
+**Key questions:**
+- Which layers should be column parallel (weight split by columns)?
+- Which layers should be row parallel (weight split by rows)?
+
+### Step 2: Replace Linear Layers with Parallel Equivalents
+
+Replace `nn.Linear` with parallel layers from `vllm.model_executor.layers.linear`.
+
+**Example (MLP Block - Up-Down Pattern):**
+
+```python
+class FeedForward(nn.Module):
+    def __init__(self, dim: int, hidden_dim: int):
+        super().__init__()
+        # Column parallel: weight split by columns [hidden_dim/N, dim]
+        self.w1 = ColumnParallelLinear(
+            dim,
+            hidden_dim,
+            bias=False,
+            return_bias=False,
+        )
+        self.act = nn.GELU()
+
+        self.w2 = RowParallelLinear(
+            hidden_dim,
+            dim,
+            bias=False,
+            input_is_parallel=True,  # Input already sharded from w1
+            return_bias=False,
+        )
+
+    def forward(self, x):
+        # x: [batch, seq, dim] (replicated on all GPUs)
+        # w1 outputs sharded [batch, seq, hidden_dim/N]
+        x = self.w1(x)
+        # act operates on sharded tensors (no communication)
+        x = self.act(x)
+        # w2 outputs full dim [batch, seq, dim] via all-reduce
+        x = self.w2(x)
+        return x
+```
+
+**Example (Attention - QKV-Out Pattern):**
+
+```python
+from vllm_omni.diffusion.attention.layer import Attention
+class YourModelAttention(nn.Module):
+    def __init__(self, dim: int, num_heads: int, num_kv_heads: int):
+        super().__init__()
+        self.head_dim = dim // num_heads
+
+        # Column parallel: QKV weight split by columns
+        # Each GPU gets num_heads/N heads
+        self.to_qkv = QKVParallelLinear(
+            hidden_size=dim,
+            head_size=self.head_dim,
+            total_num_heads=num_heads,
+            total_num_kv_heads=num_kv_heads,
+            bias=False,
+            return_bias=False,
+        )
+
+        # Row parallel: output weight split by rows
+        self.to_out = RowParallelLinear(
+            dim,
+            dim,
+            bias=False,
+            input_is_parallel=True,  # Input sharded from attention
+            return_bias=False,
+        )
+
+        self.attn = Attention(
+            num_heads=self.to_qkv.num_heads, # Each GPU gets num_heads/N heads
+            head_size=self.head_dim,
+            softmax_scale=1.0 / (self.head_dim**0.5),
+            causal=False,
+            num_kv_heads=self.to_qkv.num_kv_heads,
+        )
+
+    def forward(self, x):
+        # x: [batch, seq, dim] (replicated)
+        # to_qkv outputs sharded [batch, seq, (q+k+v) * head_dim/N]
+        qkv = self.to_qkv(x)
+        # Split into Q, K, V (each sharded on heads)
+        q, k, v = qkv.split([...], dim=-1)
+        # Attention computed independently on each GPU
+        out = self.attn(q, k, v)
+        # to_out all-reduces to full dim
+        out = self.to_out(out)
+        return out
+```
+
+**Key Points:**
+
+- `ColumnParallelLinear` → `RowParallelLinear` is the standard pairing
+- Set `input_is_parallel=True` on `RowParallelLinear` when input comes from `ColumnParallelLinear`
+- Use `QKVParallelLinear` for attention projections (handles head replication automatically)
+
+### Step 3: Validate TP Constraints
+
+For correct TP operation, these dimensions **must be divisible** by `tensor_parallel_size`:
+
+| Dimension | Reason | Example Error |
+|-----------|--------|---------------|
+| `num_heads` | Heads sharded by QKVParallelLinear | `num_heads=30, tp=4` ❌ (30 % 4 ≠ 0) |
+| `num_kv_heads` | KV heads sharded by QKVParallelLinear | `num_kv_heads=30, tp=4` ❌ (30 % 4 ≠ 0) |
+
+---
+
+## Testing
+
+After adding Tensor Parallel support, test with:
+
+```python
+from vllm_omni import Omni
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+parallel_config = DiffusionParallelConfig(tensor_parallel_size=2)
+omni = Omni(model="your-model-name", parallel_config=parallel_config)
+
+output = omni.generate(
+    "a cup of coffee on the table",
+    OmniDiffusionSamplingParams(num_inference_steps=50),
+)
+```
+
+**Or via command line:**
+
+```bash
+cd examples/offline_inference/text_to_image
+python text_to_image.py \
+    --model Your-org/your-model \
+    --prompt "a cup of coffee on the table" \
+    --negative-prompt "ugly, unclear" \
+    --cfg-scale 4.0 \
+    --num-inference-steps 50 \
+    --output "tp_enabled.png" \
+    --tensor-parallel-size 2
+```
+
+**Verify:**
+
+1. Check the `e2e_time_ms` in the log for speedup
+2. Compare generated image quality with TP disabled
+3. Verify memory usage is reduced proportionally
+4. Record comparison results in your PR
+
+---
+
+## Troubleshooting
+
+### Issue: TP not activating
+
+**Symptoms:** Model runs on single GPU, no memory savings or speedup.
+
+**Causes & Solutions:**
+
+- **Still using `nn.Linear`:**
+
+**Problem:** Linear layers not replaced with parallel equivalents.
+
+**Solution:** Replace with parallel layers:
+```python
+# ❌ BAD
+self.proj = nn.Linear(dim, dim)
+
+# ✅ GOOD
+self.proj = RowParallelLinear(dim, dim, input_is_parallel=True)
+```
+
+### Issue: Dimension mismatch errors
+
+**Symptoms:** `RuntimeError: shape mismatch` during forward pass.
+
+**Causes & Solutions:**
+
+- **Missing `input_is_parallel=True`:**
+
+**Problem:** RowParallelLinear expects sharded input but receives full tensor.
+
+**Solution:** Set `input_is_parallel=True` when input comes from ColumnParallelLinear:
+```python
+# ✅ GOOD: Correct pairing
+self.w1 = ColumnParallelLinear(dim, hidden_dim, return_bias=False,)
+self.w2 = RowParallelLinear(
+    hidden_dim,
+    dim,
+    input_is_parallel=True,  # Input sharded from w1
+    return_bias=False,
+)
+```
+
+- **Incorrect split dimensions:**
+
+**Problem:** QKV split sizes don't match sharded dimensions.
+
+**Solution:** Use `self.to_qkv.num_heads` (local heads per GPU):
+```python
+# ❌ BAD: Uses total heads
+q_size = self.total_num_heads * self.head_dim
+
+# ✅ GOOD: Uses local heads
+q_size = self.to_qkv.num_heads * self.head_dim
+```
+
+---
+
+## Reference Implementations
+
+Complete examples in the codebase:
+
+| Model | Path | Pattern | Notes |
+|-------|------|---------|-------|
+| **Z-Image** | `vllm_omni/diffusion/models/z_image/z_image_transformer.py` | Standard TP | Full implementation with validation |
+| **FLUX** | `vllm_omni/diffusion/models/flux/flux_transformer.py` | Dual-stream | Image + text streams |
+| **Qwen-Image** | `vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py` | Standard TP | With RoPE |
+| **TP Tests** | `tests/e2e/offline_inference/test_zimage_parallelism.py` | E2E testing | TP correctness and performance |
+| **Constraint Tests** | `tests/diffusion/models/z_image/test_zimage_tp_constraints.py` | Unit testing | Validation logic |
+
+---
+
+## Summary
+
+Adding Tensor Parallel support to a transformer:
+
+1. ✅ **Identify linear layers** - Which layers should be sharded?
+2. ✅ **Replace with parallel layers** - Use QKVParallelLinear, ColumnParallelLinear, RowParallelLinear
+3. ✅ **Validate TP constraints** - Ensure dimensions divisible by TP size
+4. ✅ **Test** - Verify with `tensor_parallel_size=N`, check memory, speed, and quality
@@ -0,0 +1,165 @@
+
+# Metrics vLLM-Omni:
+
+You can use these metrics in production to monitor the health and performance of the vLLM-omni system. Typical scenarios include:
+- **Performance Monitoring**: Track throughput (e.g., `e2e_avg_tokens_per_s`), latency (e.g., `e2e_total_ms`), and resource utilization to verify that the system meets expected standards.
+- **Debugging and Troubleshooting**: Use detailed per-request metrics to diagnose issues, such as high transfer times or unexpected token counts.
+
+## How to Enable and View Metrics
+
+### 1. Start the Service with Metrics Logging
+
+```bash
+vllm serve /workspace/models/Qwen3-Omni-30B-A3B-Instruct --omni --port 8014 --log-stats
+```
+
+### 2. Send a Request
+
+```bash
+python openai_chat_completion_client_for_multimodal_generation.py --query-type use_image
+```
+
+### 3. What You Will See
+
+With `--log-stats` enabled, the server will output detailed metrics logs after each request. Example output:
+
+
+#### Overall Summary
+
+| Field                       | Value        |
+|-----------------------------|--------------|
+| e2e_requests                | 1            |
+| e2e_wall_time_ms            | 41,299.190   |
+| e2e_total_tokens            | 5,202        |
+| e2e_avg_time_per_request_ms | 41,299.190   |
+| e2e_avg_tokens_per_s        | 125.959      |
+| e2e_stage_0_wall_time_ms    | 10,192.289   |
+| e2e_stage_1_wall_time_ms    | 30,541.409   |
+| e2e_stage_2_wall_time_ms    |    207.496   |
+
+#### RequestE2EStats
+
+| Field                   | Value      |
+|-------------------------|------------|
+| e2e_total_ms            | 41,299.133 |
+| e2e_total_tokens        | 5,202      |
+| transfers_total_time_ms | 245.895    |
+| transfers_total_kbytes  | 138,089.939|
+
+#### StageRequestStats
+
+| Field                  | 0      | 1      | 2      |
+|------------------------|--------|--------|--------|
+| audio_generated_frames | 0      | 0      | 525,525|
+| batch_id               | 38     | 274    | 0      |
+| batch_size             | 1      | 1      | 1      |
+| num_tokens_in          | 4,860  | 4,826  | 4,384  |
+| num_tokens_out         | 67     | 275    | 0      |
+| postprocess_time_ms    | 256.158| 0.491  | 0.000  |
+| stage_gen_time_ms      | 9,910.007|30,379.198|160.745|
+
+#### TransferEdgeStats
+
+| Field               | 0->1        | 1->2       |
+|---------------------|-------------|------------|
+| size_kbytes         | 109,277.349 | 28,812.591 |
+| tx_time_ms          | 78.701      | 18.790     |
+| rx_decode_time_ms   | 111.865     | 31.706     |
+| in_flight_time_ms   | 2.015       | 2.819      |
+
+
+These logs include:
+- **Overall summary**: total requests, wall time, average tokens/sec, etc.
+- **E2E table**: per-request latency and token counts.
+- **Stage table**: per-stage batch and timing details.
+- **Transfer table**: data transfer and timing for each edge.
+
+You can use these logs to monitor system health, debug performance, and analyze request-level metrics as described above.
+
+
+## Metrics Scope: Offline vs Online Inference
+
+For **offline inference** (batch mode), the summary includes both system-level metrics (aggregated across all requests) and per-request metrics. In this case, `e2e_requests` can be greater than 1, reflecting multiple completed requests in a batch.
+
+For **online inference** (serving mode), the summary is always per-request. `e2e_requests` is always 1, and only request-level metrics are reported for each completion.
+
+---
+
+## Parameter Details
+
+| Field                     | Meaning                                                                                       |
+|---------------------------|----------------------------------------------------------------------------------------------|
+| `e2e_requests`            | Number of completed requests.                                                                |
+| `e2e_wall_time_ms`        | Wall-clock time span from run start to last completion, in ms.                               |
+| `e2e_total_tokens`        | Total tokens counted across all completed requests (stage0 input + all stage outputs).       |
+| `e2e_avg_time_per_request_ms` | Average wall time per request: `e2e_wall_time_ms / e2e_requests`.                        |
+| `e2e_avg_tokens_per_s`    | Average token throughput over wall time: `e2e_total_tokens * 1000 / e2e_wall_time_ms`.      |
+| `e2e_stage_{i}_wall_time_ms` | Wall-clock time span for stage i, in ms. Each stage's wall time is reported as a separate field, e.g., `e2e_stage_0_wall_time_ms`, `e2e_stage_1_wall_time_ms`, etc. |
+
+---
+
+## E2E Table (per request)
+
+| Field                     | Meaning                                                               |
+|---------------------------|-----------------------------------------------------------------------|
+| `e2e_total_ms`            | End-to-end latency in ms.                                             |
+| `e2e_total_tokens`        | Total tokens for the request (stage0 input + all stage outputs).      |
+| `transfers_total_time_ms` | Sum of transfer edge `total_time_ms` for this request.                |
+| `transfers_total_kbytes`  | Sum of transfer kbytes for this request.                              |
+
+
+---
+
+## Stage Table (per stage event / request)
+
+| Field                     | Meaning                                                                                         |
+|---------------------------|-------------------------------------------------------------------------------------------------|
+| `batch_id`                | Batch index.                                                                                    |
+| `batch_size`              | Batch size.                                                                                     |
+| `num_tokens_in`           | Input tokens to the stage.                                                                      |
+| `num_tokens_out`          | Output tokens from the stage.                                                                   |
+| `stage_gen_time_ms`       | Stage compute time in ms, excluding postprocessing time (reported separately as `postprocess_time_ms`). |
+| `image_num`               | Number of images generated (for diffusion/image stages).                                        |
+| `resolution`              | Image resolution (for diffusion/image stages).                                                                  |
+| `postprocess_time_ms` | Diffusion/image: post-processing time in ms.                                                    |
+
+---
+
+## Transfer Table (per edge / request)
+
+| Field                | Meaning                                                                   |
+|----------------------|---------------------------------------------------------------------------|
+| `size_kbytes`        | Total kbytes transferred.                                                 |
+| `tx_time_ms`         | Sender transfer time in ms.                                               |
+| `rx_decode_time_ms`  | Receiver decode time in ms.                                               |
+| `in_flight_time_ms`  | In-flight time in ms.                                                     |
+
+
+## Expectation of the Numbers (Verification)
+
+**Formulas:**
+- `e2e_total_tokens = Stage0's num_tokens_in + sum(all stages' num_tokens_out)`
+- `transfers_total_time_ms = sum(tx_time_ms + rx_decode_time_ms + in_flight_time_ms)` for every edge
+
+**Using the example above:**
+
+### e2e_total_tokens
+- Stage0's `num_tokens_in`: **4,860**
+- Stage0's `num_tokens_out`: **67**
+- Stage1's `num_tokens_out`: **275**
+- Stage2's `num_tokens_out`: **0**
+
+So,
+```
+e2e_total_tokens = 4,860 + 67 + 275 + 0 = 5,202
+```
+This matches the table value: `e2e_total_tokens = 5,202`.
+
+### transfers_total_time_ms
+For each edge:
+- 0->1: tx_time_ms (**78.701**) + rx_decode_time_ms (**111.865**) + in_flight_time_ms (**2.015**) = **192.581**
+- 1->2: tx_time_ms (**18.790**) + rx_decode_time_ms (**31.706**) + in_flight_time_ms (**2.819**) = **53.315**
+
+Sum: 192.581 + 53.315 = **245.896**
+
+The table shows `transfers_total_time_ms = 245.895`, which matches the calculation (difference is due to rounding).
@@ -572,7 +572,7 @@ def talker2code2wav(

 ## Testing

-For comprehensive testing guidelines, please refer to the [Test File Structure and Style Guide](../tests/tests_style.md).
+For comprehensive testing guidelines, please refer to the [Test File Structure and Style Guide](../ci/tests_style.md).

 ## Adding a Model Recipe

@@ -1,3 +1,149 @@
 # Profiling vLLM-Omni

-Profiling is only intended for vLLM-Omni developers and maintainers to understand the proportion of time spent in different parts of the codebase. **vLLM-Omni end-users should never turn on profiling** as it will significantly slow down the inference.
+> **Warning:** Profiling incurs significant overhead. Use only for development and debugging, never in production.
+
+vLLM-Omni uses the PyTorch Profiler to analyze performance across both **multi-stage omni-modality models** and **diffusion models**.
+
+### 1. Set the Output Directory
+Before running any script, set this environment variable. The system detects this and automatically saves traces here.
+
+```bash
+export VLLM_TORCH_PROFILER_DIR=./profiles
+```
+
+### 2. Profiling Omni-Modality Models
+
+It is best to limit profiling to one iteration to keep trace files manageable.
+
+```bash
+export VLLM_PROFILER_MAX_ITERS=1
+```
+
+**Selective Stage Profiling**
+The profiler is default to function across all stages. But It is highly recommended to profile specific stages by passing the stages list, preventing from producing too large trace files:
+```python
+# Profile all stages
+omni_llm.start_profile()
+
+# Only profile Stage 1
+omni_llm.start_profile(stages=[1])
+```
+
+```python
+# Stage 0 (Thinker) and Stage 2 (Audio Decoder) for qwen omni
+omni_llm.start_profile(stages=[0, 2])
+```
+
+**Python Usage**: Wrap your generation logic with `start_profile()` and `stop_profile()`.
+
+```python
+from vllm_omni import omni_llm
+
+profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
+
+# 1. Start profiling if enabled
+if profiler_enabled:
+    omni_llm.start_profile(stages=[0])
+
+# Initialize generator
+omni_generator = omni_llm.generate(prompts, sampling_params_list, py_generator=args.py_generator)
+
+total_requests = len(prompts)
+processed_count = 0
+
+# Main Processing Loop
+for stage_outputs in omni_generator:
+
+    # ... [Output processing logic for text/audio would go here] ...
+
+    # Update count to track when to stop profiling
+    processed_count += len(stage_outputs.request_output)
+
+    # 2. Check if all requests are done to stop the profiler safely
+    if profiler_enabled and processed_count >= total_requests:
+        print(f"[Info] Processed {processed_count}/{total_requests}. Stopping profiler inside active loop...")
+
+        # Stop the profiler while workers are still active
+        omni_llm.stop_profile()
+
+        # Wait for traces to flush to disk
+        print("[Info] Waiting 30s for workers to write trace files to disk...")
+        time.sleep(30)
+        print("[Info] Trace export wait time finished.")
+
+omni_llm.close()
+```
+
+
+**Examples**:
+
+1. **Qwen2.5-Omni**:  [https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen2_5_omni/end2end.py](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen2_5_omni/end2end.py)
+
+2. **Qwen3-Omni**:   [https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen3_omni/end2end.py](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen3_omni/end2end.py)
+
+
+### 3. Profiling diffusion models
+
+Diffusion profiling is End-to-End, capturing encoding, denoising loops, and decoding.
+
+**CLI Usage:**
+```python
+
+python image_to_video.py \
+    --model Wan-AI/Wan2.2-I2V-A14B-Diffusers \
+    --image qwen-bear.png \
+    --prompt "A cat playing with yarn, smooth motion" \
+    \
+    # Minimize Spatial Dimensions (Optional but helpful):
+    #    Drastically reduces memory usage so the profiler doesn't
+    #    crash due to overhead, though for accurate performance
+    #    tuning you often want target resolutions.
+    --height 48 \
+    --width 64 \
+    \
+    # Minimize Temporal Dimension (Frames):
+    #    Video models process 3D tensors (Time, Height, Width).
+    #    Reducing frames to the absolute minimum (2) keeps the
+    #    tensor size small, ensuring the trace file doesn't become
+    #    multi-gigabytes in size.
+    --num-frames 2 \
+    \
+    # Minimize Iteration Loop (Steps):
+    #    This is the most critical setting for profiling.
+    #    Diffusion models run the same loop X times.
+    #    Profiling 2 steps gives you the exact same performance
+    #    data as 50 steps, but saves minutes of runtime and
+    #    prevents the trace viewer from freezing.
+    --num-inference-steps 2 \
+    \
+    --guidance-scale 5.0 \
+    --guidance-scale-high 6.0 \
+    --boundary-ratio 0.875 \
+    --flow-shift 12.0 \
+    --fps 16 \
+    --output i2v_output.mp4
+
+```
+
+**Examples**:
+
+1. **Qwen image edit**:  [https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/image_to_image/image_edit.py](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/image_to_image/image_edit.py)
+
+2. **Wan-AI/Wan2.2-I2V-A14B-Diffusers**:   [https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/image_to_video](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/image_to_video)
+
+> **Note:**
+As of now, asynchronous (online) profiling is not fully supported in vLLM-Omni. While start_profile() and stop_profile() methods exist, they are only reliable in offline inference scripts (e.g., the provided end2end.py examples). Do not use them in server-mode or streaming scenarios—traces may be incomplete or fail to flush.
+
+### 4. Analyzing Omni Traces
+
+Output files are saved to your configured ```VLLM_TORCH_PROFILER_DIR```.
+
+**Output**
+**Chrome Trace** (```.json.gz```): Visual timeline of kernels and stages. Open in Perfetto UI.
+
+**Viewing Tools:**
+
+- [Perfetto](https://ui.perfetto.dev/)(recommended)
+- ```chrome://tracing```(Chrome only)
+
+**Note**: vLLM-Omni reuses the PyTorch Profiler infrastructure from vLLM. See the official vLLM profiler documentation:  [vLLM Profiling Guide](https://docs.vllm.ai/en/stable/contributing/profiling/)
@@ -1,266 +0,0 @@
-# Test File Structure and Style Guide
-
-To ensure project maintainability and sustainable development, we encourage contributors to submit test code (unit tests, system tests, or end-to-end tests) alongside their code changes. This document outlines the guidelines for organizing and naming test files.
-
-## Test Types
-
-### Unit Tests and System Tests
-For unit tests and system tests, we strongly recommend placing test files in the same directory structure as the source code being tested, using the naming convention `test_*.py`.
-
-### End-to-End (E2E) Tests for Models
-End-to-end tests verify the complete functionality of a system or component. For our project, the E2E tests for different omni models are organized into two subdirectories:
-
- **`tests/e2e/offline_inference/`**: Tests for offline inference modes (e.g., Qwen3Omni offline inference)
-
- **`tests/e2e/online_serving/`**: Tests for online serving scenarios (e.g., API server tests)
-
-**Example:** The test file for `vllm_omni/entrypoints/omni_llm.py` should be located at `tests/entrypoints/test_omni_llm.py`.
-
-## Test Directory Structure
-
-The ideal directory structure mirrors the source code organization:
-
-```
-vllm_omni/                          tests/
-├── config/                    →    ├── config/
-│   └── model.py                    │   └── test_model.py
-│
-├── core/                      →    ├── core/
-│   └── sched/                      │   └── sched/                    # Maps to core/sched/
-│       ├── omni_ar_scheduler.py    │       ├── test_omni_ar_scheduler.py
-│       ├── omni_generation_scheduler.py │  ├── test_omni_generation_scheduler.py
-│       └── output.py               │       └── test_output.py
-│
-├── diffusion/                 →    ├── diffusion/
-│   ├── diffusion_engine.py         │   ├── test_diffusion_engine.py
-│   ├── omni_diffusion.py           │   ├── test_omni_diffusion.py
-│   ├── attention/                  │   ├── attention/                # Maps to diffusion/attention/
-│   │   └── backends/               │   │   └── test_*.py
-│   ├── models/                     │   ├── models/                   # Maps to diffusion/models/
-│   │   ├── qwen_image/             │   │   ├── qwen_image/
-│   │   │   └── ...                 │   │   │   └── test_*.py
-│   │   └── z_image/                │   │   └── z_image/
-│   │       └── ...                 │   │       └── test_*.py
-│   └── worker/                     │   └── worker/                   # Maps to diffusion/worker/
-│       └── ...                     │       └── test_*.py
-│
-├── distributed/               →    ├── distributed/
-│   └── ...                         │   └── test_*.py
-│
-├── engine/                    →    ├── engine/
-│   ├── processor.py                │   ├── test_processor.py
-│   └── output_processor.py         │   └── test_output_processor.py
-│
-├── entrypoints/               →    ├── entrypoints/
-│   ├── omni_llm.py                 │   ├── test_omni_llm.py          # UT: OmniLLM core logic (mocked)
-│   ├── omni_stage.py               │   ├── test_omni_stage.py         # UT: OmniStage logic
-│   ├── omni.py                     │   ├── test_omni.py               # E2E: Omni class (offline inference)
-│   ├── async_omni.py               │   ├── test_async_omni.py         # E2E: AsyncOmni class
-│   ├── cli/                        │   ├── cli/                       # Maps to entrypoints/cli/
-│   │   └── ...                     │   │   └── test_*.py
-│   └── openai/                     │   └── openai/                     # Maps to entrypoints/openai/
-│       ├── api_server.py           │       ├── test_api_server.py     # E2E: API server (online serving)
-│       └── serving_chat.py         │       └── test_serving_chat.py
-│
-├── inputs/                    →    ├── inputs/
-│   ├── data.py                     │   ├── test_data.py
-│   ├── parse.py                    │   ├── test_parse.py
-│   └── preprocess.py               │   └── test_preprocess.py
-│
-├── model_executor/            →    ├── model_executor/
-│   ├── layers/                     │   ├── layers/
-│   │   └── mrope.py                │   │   └── test_mrope.py
-│   ├── model_loader/               │   ├── model_loader/
-│   │   └── weight_utils.py         │   │   └── test_weight_utils.py
-│   ├── models/                     │   ├── models/
-│   │   ├── qwen2_5_omni/           │   │   ├── qwen2_5_omni/
-│   │   │   ├── qwen2_5_omni_thinker.py │ │   │   ├── test_qwen2_5_omni_thinker.py  # UT
-│   │   │   ├── qwen2_5_omni_talker.py │ │   │   ├── test_qwen2_5_omni_talker.py  # UT
-│   │   │   └── qwen2_5_omni_token2wav.py │ │   │   └── test_qwen2_5_omni_token2wav.py  # UT
-│   │   └── qwen3_omni/             │   │   └── qwen3_omni/
-│   │       └── ...                 │   │       └── test_*.py
-│   ├── stage_configs/              │   └── stage_configs/             # Configuration tests (if needed)
-│   │   └── ...                     │       └── test_*.py
-│   └── stage_input_processors/     │   └── stage_input_processors/
-│       └── ...                     │       └── test_*.py
-│
-├── sample/                    →    ├── sample/
-│   └── ...                         │   └── test_*.py
-│
-├── utils/                     →    ├── utils/
-│   └── platform_utils.py           │   └── test_platform_utils.py
-│
-├── worker/                    →    ├── worker/
-    ├── gpu_ar_worker.py            │   ├── test_gpu_ar_worker.py
-    ├── gpu_generation_worker.py    │   ├── test_gpu_generation_worker.py
-    ├── gpu_model_runner.py         │   ├── test_gpu_model_runner.py
-    └── npu/                        │   └── npu/                       # Maps to worker/npu/
-        └── ...                     │       └── test_*.py
-│
-└── e2e/                       →    ├── e2e/                # End-to-end scenarios (no 1:1 source mirror)
-                                    ├── online_serving/       # Full-stack online serving flows
-                                    │   └── (empty for now)
-                                    └── offline_inference/    # Full offline inference flows
-                                        ├── test_qwen2_5_omni.py     # Moved from multi_stages/
-                                        ├── test_qwen3_omni.py       # Moved from multi_stages_h100/
-                                        ├── test_t2i_model.py  # Moved from single_stage/
-                                        └── stage_configs/           # Shared stage configs
-                                            ├── qwen2_5_omni_ci.yaml
-                                            └── qwen3_omni_ci.yaml
-```
-
-
-
-### Naming Conventions
-
- **Unit/System Tests**: Use `test_<module_name>.py` format. Example: `omni_llm.py` → `test_omni_llm.py`
-
- **E2E Tests**: Place in `tests/e2e/offline_inference/` or `tests/e2e/online_serving/` with descriptive names. Example: `tests/e2e/offline_inference/test_qwen3_omni.py`, `tests/e2e/offline_inference/test_diffusion_model.py`
-
-### Best Practices
-
-1. **Mirror Source Structure**: Test directories should mirror the source code structure
-2. **Test Type Indicators**: Use comments to indicate test types (UT for unit tests, E2E for end-to-end tests)
-3. **Shared Resources**: Place shared test configurations (e.g., CI configs) in appropriate subdirectories
-4. **Consistent Naming**: Follow the `test_*.py` naming convention consistently across all test files
-
-
-## Test codes requirements
-
-### Coding style
-
-1. **File header**: Add SPDX license header to all test files
-2. **Imports**: Pls don't use manual `sys.path` modifications, use standard imports instead.
-3. **Test type differentiation**:
-
-      - Unit tests: Maintain mock style
-      - E2E tests for models: Consider using OmniRunner uniformly, avoid decorators
-
-4. **Documentation**: Add docstrings to all test functions
-5. **Environment variables**: Set uniformly in `conftest.py` or at the top of files
-6. **Type annotations**: Add type annotations to all test function parameters
-7. **Resources**, Using pytest tag to specify the computation resources the test required.
-
-### Template
-#### E2E - Online serving
-
-```python
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Online E2E smoke test for an omni model (video,text,audio → audio).
-"""
-from pathlib import Path
-
-import pytest
-import openai
-
-
-# Optional: set process start method for workers
-os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
-
-models = ["{your model name}"] #Edit here to load your model
-stage_configs = [str(Path(__file__).parent / "stage_configs" / {your model yaml})] #Edit here to load your model yaml
-test_params = [(model, stage_config) for model in models for stage_config in stage_configs]
-
-#OmniServer，Used to start the vllm-omni server
-class OmniServer:
-    xxx
-
-
-@pytest.fixture
-def omni_server(request):
-    model, stage_config_path = request.param
-    with OmniServer(model, ["--stage-configs-path", stage_config_path]) as server:
-        yield server
-
-
-#handle request message
-@pytest.fixture(scope="session")
-def base64_encoded_video() -> str:
-    xxx
-
-@pytest.fixture(scope="session")
-def dummy_messages_from_video_data(video_data_url: str, content_text: str) -> str:
-    xxx
-
-@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_video_to_audio(
-    client: openai.OpenAI,
-    omni_server,
-    base64_encoded_video: str,
-) -> None:
-    #set message
-    video_data_url = f"data:video/mp4;base64, {base64_encoded_video}"
-    messages = dummy_messages_from_video_data(video_data_url)
-
-    #send request
-    chat_completion = client.chat.completions.create(
-        model=omni_server.model,
-        messages=messages,
-    )
-
-    #verify text output
-    text_choice = chat_completion.choices[0]
-    assert text_choice.finish_reason == "length"
-
-    #verify audio output
-    audio_choice = chat_completion.choices[1]
-    audio_message = audio_choice.message
-    if hasattr(audio_message, "audio") and audio_message.audio:
-        assert audio_message.audio.data is not None
-        assert len(audio_message.audio.data) > 0
-```
-
-#### E2E - Offline inference
-```python
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Offline E2E smoke test for an omni model (video → audio).
-"""
-
-import os
-from pathlib import Path
-
-import pytest
-from vllm.assets.video import VideoAsset
-
-from ..multi_stages.conftest import OmniRunner
-
-# Optional: set process start method for workers
-os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
-
-models = ["{your model name}"] #Edit here to load your model
-stage_configs = [str(Path(__file__).parent / "stage_configs" / {your model yaml})] #Edit here to load your model yaml
-
-# Create parameter combinations for model and stage config
-test_params = [(model, stage_config) for model in models for stage_config in stage_configs]
-
-# function name: test_{input_modality}_to_{output_modality}
-# modality candidate: text, image, audio, video, mixed_modalities
-@pytest.mark.gpu_mem_high  # requires high-memory GPU node
-@pytest.mark.parametrize("test_config", test_params)
-def test_video_to_audio(omni_runner: type[OmniRunner], model: str) -> None:
-    """Offline inference: video input, audio output."""
-    model, stage_config_path = test_config
-    with omni_runner(model, seed=42, stage_configs_path=stage_config_path) as runner:
-        # Prepare inputs
-        video = VideoAsset(name="sample", num_frames=4).np_ndarrays
-
-        outputs = runner.generate_multimodal(
-            prompts="Describe this video briefly.",
-            videos=video,
-        )
-
-        # Minimal assertions: got outputs and at least one audio result
-        assert outputs
-        has_audio = any(o.final_output_type == "audio" for o in outputs)
-        assert has_audio
-```
-
-## Checklist before submitting your test files
-
-1. The file is saved in an appropriate place and the file name is clear.
-2. The coding style follows the requirements outlined above.
-3. For e2e model test, please ensure the test is configured under the `./buildkite/` folder.
@@ -64,13 +64,13 @@ According to analysis for current popular open-source models, most of them have

 ## Key Components

-| Component | Description |
-|-----------|-------------|
-| **OmniRouter** | provide an intelligent router for Omni-modality requests dispatch |
-| **EntryPoints** | define the APIs for offline/online serving (APIServer, Omni/AsyncOmni) and provide the OmniStage abstraction for different AR/DiT stages |
-| **AR** | adapted for omni-modality models while inheriting efficient features from vLLM, such as cache management |
-| **Diffusion** | natively implemented and optimized using acceleration components |
-| **OmniConnector** | supports fully disaggregation based on E/P/D/G (Encoding/Processing/Decoding/Generation) disaggregation across stages |
+| Component         | Description                                                                                                                              |
+| ----------------- | ---------------------------------------------------------------------------------------------------------------------------------------- |
+| **OmniRouter**    | provide an intelligent router for Omni-modality requests dispatch                                                                        |
+| **EntryPoints**   | define the APIs for offline/online serving (APIServer, Omni/AsyncOmni) and provide the OmniStage abstraction for different AR/DiT stages |
+| **AR**            | adapted for omni-modality models while inheriting efficient features from vLLM, such as cache management                                 |
+| **Diffusion**     | natively implemented and optimized using acceleration components                                                                         |
+| **OmniConnector** | supports fully disaggregation based on E/P/D/G (Encoding/Processing/Decoding/Generation) disaggregation across stages                    |

 Disaggregated stages are managed through configuration, such as in the Qwen3-Omni example, where stages like Thinker, Talker, and Code2wav are defined as separate OmniStage instances with specific resources and input/output type.

@@ -192,4 +192,4 @@ curl -sS -X POST http://localhost:8091/v1/chat/completions \
 }
 ```

-For more usages, please refer to [examples](../user_guide/examples/).
+For more usages, please refer to [examples](https://github.com/vllm-project/vllm-omni/tree/main/examples).
@@ -1,27 +1,35 @@
 # Disaggregated Inference for Omni-Modality Models

-This guide explains how to configure and use distributed connectors (vllm_omni/distributed/connectors) in vllm-omni for multi-stage pipelines.
+This guide explains how to configure and use distributed connectors
+(`vllm_omni/distributed/omni_connectors`) in vllm-omni for multi-stage pipelines.

-## 1. Overview
+Backend-specific setup lives in separate docs:
+
+- [SharedMemoryConnector](omni_connectors/shared_memory_connector.md)
+- [MooncakeConnector](omni_connectors/mooncake_connector.md)
+- [YuanrongConnector](omni_connectors/yuanrong_connector.md)
+
+## Overview

 Connectors enable data transfer between pipeline stages (e.g., Thinker -> Talker).
-Currently supported connectors operate in **D2H2D (Device-to-Host-to-Device)** mode:
-1. **SharedMemoryConnector**: Uses system shared memory.
-2. **MooncakeConnector**: Uses [Mooncake](https://github.com/kvcache-ai/Mooncake).
+Current connectors operate in D2H2D (device to host to device) mode.

-*   **SharedMemoryConnector (Default)**: Zero-copy (on host), lowest latency. Best for **single-node** deployments. Auto-configured if no connectors are specified.
-*   **MooncakeConnector**: TCP/RDMA based. Best for **multi-node** distributed deployments. Requires a Mooncake Master service.
+## Connector Choices

-## 2. API Design
+| Use Case | Recommended Connector | Notes |
+| :--- | :--- | :--- |
+| Single node | SharedMemoryConnector | Auto-configured if no connector is specified. |
+| Multi node (Mooncake) | MooncakeConnector | Requires Mooncake Master + metadata server. |
+| Multi node (Yuanrong) | YuanrongConnector | Requires Yuanrong Datasystem + etcd. |

-The connector system is built around the `OmniConnectorBase` abstraction, which decouples data transport from stage logic.
+## Core API

-### Core Interface
+The connector system is built around `OmniConnectorBase`.

 ```python
 class OmniConnectorBase(ABC):
    @abstractmethod
-    def put(self, from_stage: str, to_stage: str, request_id: str, data: Any) -> tuple[bool, int, Optional[dict]]:
+    def put(self, from_stage: str, to_stage: str, put_key: str, data: Any) -> tuple[bool, int, Optional[dict]]:
        """
        Store data.
        Returns: (success, serialized_size, metadata)
@@ -29,145 +37,23 @@ class OmniConnectorBase(ABC):
        pass

    @abstractmethod
-    def get(self, from_stage: str, to_stage: str, request_id: str, metadata: Optional[dict] = None) -> Optional[tuple[Any, int]]:
+    def get(self, from_stage: str, to_stage: str, get_key: str, metadata: Optional[dict] = None) -> Optional[tuple[Any, int]]:
        """
        Retrieve data.
-        Args: metadata - Transport-specific handles returned by put() (e.g., SHM name).
+        Args: metadata - transport-specific handles returned by put() (e.g., SHM name).
        Returns: (object, serialized_size)
        """
        pass
 ```

-### Key Concept: Metadata Passing
-Unlike a pure key-value store, some connectors (like `SharedMemoryConnector`) generate transient resources (e.g., a shared memory block name) during `put()`. This `metadata` **must be passed via the control plane** (e.g., HTTP headers, queue messages) from the producer stage to the consumer stage so `get()` can locate the data.
+### Metadata Passing

-## 3. Backends & Use Cases
+Some connectors (e.g., SharedMemoryConnector) generate transient resources during `put()`.
+This `metadata` must be passed through the control plane so `get()` can locate the data.

-### 3.1 SharedMemoryConnector
-**Best for:** Single-node, high-performance IPC.
+## Configuration Model

-*   **Mechanism:**
-    *   **Small Payloads (< Threshold)**: Data is serialized and passed directly "inline" within the `metadata` dictionary. This avoids the overhead of creating SHM blocks for tiny messages.
-    *   **Large Payloads (>= Threshold)**: Data is written to a named System Shared Memory block. The block name is returned in `metadata`.
-*   **Configuration:**
-    *   `shm_threshold_bytes`: Size in bytes to switch from inline to SHM (default: 64KB).
-
-### 3.2 MooncakeConnector
-**Best for:** Multi-node distributed inference.
-
-*   **Mechanism:** Uses Mooncake's distributed KVCache store.
-    *   **Data Plane**: TCP or RDMA for high-bandwidth transfer.
-    *   **Control Plane**: Uses a centralized Mooncake Master and Metadata Server.
-    *   **Keying**: Deterministic keys based on `request_id/from_stage_to_stage`.
-*   **Requirements**: Requires a running Mooncake Master service.
-
-## 4. Relationship with vLLM
-
-vLLM provides specialized distributed mechanisms for specific artifacts:
-*   **KV Transfer** (`vllm.distributed.kv_transfer`): Optimized for transferring KV caches between prefill and decode instances (using NCCL, Mooncake, etc.).
-*   **EC Transfer** (`vllm.distributed.ec_transfer`): Optimized for sharing encoder embeddings.
-*   **Device Communicators** (`vllm.distributed.device_communicators`): Low-level primitives (NCCL, SHM) for Tensor/Pipeline Parallelism.
-
-`vllm-omni` complements this by introducing a **Generalized Connector Abstraction** (`OmniConnector`) for multimodal pipelines. While vLLM's connectors are artifact-specific, `vllm-omni`:
-
-1.  **Unifies Transport**: Provides a single API (`put`/`get`) to transport *any* stage artifact (Input Embeddings, Hidden States, Audio/Image Tensors, KV Cache, Final Output) between arbitrary pipeline stages (e.g., AudioEncoder -> LLM -> AudioGenerator).
-2.  **Extends Connectivity**: Enables flexible construction of complex DAGs (Directed Acyclic Graphs) where stages can run in the same process, same node, or across nodes, using the most appropriate backend (SHM, Mooncake, etc.) for each edge.
-3.  **Wraps & Adapts**: Can internally utilize vLLM's specialized `kv_transfer` for KV paths while using generic transports (SHM/Mooncake) for other data types, presenting a consistent interface to the application layer.
-
-## 5. Installation (Mooncake)
-
-If using `MooncakeConnector`, install the library first:
-
-```bash
-# For CUDA-enabled systems (Recommended)
-pip install mooncake-transfer-engine
-
-# For non-CUDA systems
-pip install mooncake-transfer-engine-non-cuda
-```
-
-## 6. Using MooncakeConnector
-
-### 6.1 Start Mooncake Master
-
-Start the master service on your primary node:
-
-```bash
-# if you use mooncake SSD storage
-mkdir -p ./mc_storage
-
-mooncake_master \
-  --rpc_port=50051 \
-  --enable_http_metadata_server=true \
-  --http_metadata_server_host=0.0.0.0 \
-  --http_metadata_server_port=8080 \
-  --metrics_port=9003 \
-  --root_fs_dir=./mc_storage/ \
-  --cluster_id=mc-local-1 &
-```
-
-### 6.2 Configuration (YAML)
-
-Edit your stage config (e.g., `qwen2_5_omni.yaml`).
-
-**Step 1: Define Connector in Global Runtime**
-
-```yaml
-runtime:
-  connectors:
-    connector_of_mooncake:
-      name: MooncakeConnector
-      extra:
-        host: "127.0.0.1"           # Local Worker IP
-        metadata_server: "http://<MASTER_IP>:8080/metadata"
-        master: "<MASTER_IP>:50051"
-        segment: 512000000          # 512MB segment
-        localbuf: 64000000          # 64MB buffer
-        proto: "tcp"                # "tcp" or "rdma"
-    ```
-
-**Mooncake Configuration Parameters:**
-
-*   **host**: The hostname or IP address of the local machine (worker). Mooncake uses this to register itself in the metadata server so other nodes can find it.
-*   **metadata_server**: The URL of the metadata server. This is used for service discovery and connection establishment (e.g., exchanging QP information for RDMA).
-*   **master**: The address of the Mooncake Master Server (e.g., `<MASTER_IP>:50051`). This is used for global state management and control plane operations.
-*   **segment**: The size of the global memory segment in bytes (default: ~512MB). This defines the shared memory region accessible by Mooncake for data transfer.
-*   **localbuf**: The size of the local buffer in bytes (default: ~64MB). Used for local data buffering during transfer operations.
-*   **proto**: The transport protocol to use. Options:
-    *   `tcp`: Standard TCP/IP (easier setup, universal compatibility).
-    *   `rdma`: Remote Direct Memory Access (higher performance, requires RDMA-capable hardware).
-
-For more details, refer to the [Mooncake Repository](https://github.com/kvcache-ai/Mooncake).
-
-    **Step 2: Reference in Stages**
-
-Explicitly link stages using `input_connectors` and `output_connectors`:
-
-```yaml
-stage_args:
-  - stage_id: 0
-    # ...
-    output_connectors:
-      to_stage_1: connector_of_mooncake
-
-  - stage_id: 1
-    # ...
-    input_connectors:
-      from_stage_0: connector_of_mooncake
-```
-
-## 7. Using SharedMemoryConnector (Auto-Mode)
-
-**Best for single-node.**
-
-The system will automatically create `SharedMemoryConnector`s for any pipeline edge that does not have an explicit connector defined. This is inferred from:
-1.  `runtime.edges` list in the config.
-2.  `engine_input_source` dependencies defined in `stage_args`.
-
-### Threshold Configuration
-By default, payloads larger than **64KB** (default threshold) are transferred via shared memory, while smaller ones use the control queue (inline).
-
-To adjust this threshold (e.g., to 1GB), add the following to your `runtime.connectors`:
+Define connectors in runtime:

 ```yaml
 runtime:
@@ -175,26 +61,46 @@ runtime:
    connector_of_shared_memory:
      name: SharedMemoryConnector
      extra:
-        shm_threshold_bytes: 1024 # 1KB threshold
+        shm_threshold_bytes: 65536
 ```

-## 8. Summary
+Wire stages to connectors:

-| Use Case | Recommended Connector | Configuration |
-| :--- | :--- | :--- |
-| **Single Node** | `SharedMemoryConnector` | **None** (Automatic) or Custom Threshold |
-| **Multi Node** | `MooncakeConnector` | Explicit YAML + Mooncake Master |
+```yaml
+stage_args:
+  - stage_id: 0
+    output_connectors:
+      to_stage_1: connector_of_shared_memory

-## 9. Operational Notes (important)
+  - stage_id: 1
+    input_connectors:
+      from_stage_0: connector_of_shared_memory
+```

- **Fail-fast config validation**: the loader raises if any expected edge is missing a connector. Define `input_connectors`/`output_connectors` or rely on auto-SHM filling; otherwise startup aborts.
- **Missing payloads halt the stage**: workers expect connector payloads; if metadata or connector config is missing, the stage raises and stops. Verify connector wiring and metadata propagation before production.
+If a pipeline edge has no explicit connector, the system auto-creates a
+SharedMemoryConnector for that edge.

-## 10. Future Roadmap: Device-to-Device (D2D) Transport
+## Relationship with vLLM

-The current implementations (`SharedMemoryConnector`, `MooncakeConnector`) utilize a **D2H2D (Device-to-Host-to-Device)** data path. Tensors are moved to CPU memory (Host) for transport, which incurs PCIe overhead.
+vLLM provides specialized distributed mechanisms for specific artifacts:

-As outlined in the design RFC, future versions will introduce **D2D (Device-to-Device)** connectors:
+- KV Transfer (`vllm.distributed.kv_transfer`): optimized for KV caches.
+- EC Transfer (`vllm.distributed.ec_transfer`): optimized for encoder embeddings.
+- Device Communicators (`vllm.distributed.device_communicators`): low-level primitives (NCCL, SHM).

-*   **Goal**: Direct GPU-to-GPU transfer (via NCCL, UCX, or IPC) to minimize latency for large tensor payloads.
-*   **Mechanism**: The `OmniConnector` API allows `put()` to initiate a transfer and return a lightweight handle (metadata) via the control plane, while the heavy payload flows directly between devices.
+vllm-omni complements this with a generalized connector abstraction:
+
+1. Unifies transport via a single `put`/`get` API for any stage artifact.
+2. Enables DAG-style pipelines across processes or nodes with per-edge transports.
+3. Can wrap vLLM-specific transfers for KV paths while keeping a consistent interface.
+
+## Operational Notes
+
+- Fail-fast config validation: missing expected edges cause startup failures.
+- Missing payloads halt stages: verify connector wiring and metadata propagation.
+
+## Future Roadmap: D2D Transport
+
+Current connectors use D2H2D paths. Future versions will introduce direct
+device-to-device connectors (NCCL, UCX, IPC) to reduce latency for large
+tensor payloads.
@@ -1,47 +0,0 @@
-## Multi-Request Streaming (MRS) on a Single Machine
-
-### 1. Background & Scope
- All processing runs on a single physical machine with multi-process, per-stage workers. No proxy or network transport involved.
- Current alignment with vllm-omni: `OmniLLM` supports multiple stages (`OmniStage`). GPU runners already expose streamable steps (prefill/decoding/diffusion), but the entry layer still collects lists and lacks intra-stage streaming and window scheduling.
- Goal: implement multi-stage, multi-request streaming (MRS) locally. Each stage outputs segments; downstream stages stitch and trigger compute based on configured windows. Shared memory and zero-copy strategies reduce data movement overhead.
-
-### 2. Key Constraints
- Multi-process per stage: each stage is an independent process with a while loop; device visibility can be configured (`CUDA_VISIBLE_DEVICES`/`torch.cuda.set_device`).
- Simple IPC (copy-based): use `multiprocessing.Queue`/Pipe for inter-process communication with CPU copies/serialization; do not rely on CUDA IPC/SHM zero-copy in this version.
- Cross-stage pipeline: different stages can process different requests concurrently (e.g., stage A handles request 1 while stage B handles request 0).
-
-### 3. Architecture Overview
- Processes & IPC queues
-  - Each "sub-stage" is an OS process (worker). The loop: take from input_queue → compute → put to output_queue.
-  - Inter-stage connection via IPC: copy-based `multiprocessing.Queue` passing dict payloads; use shared memory for large objects.
-  - Each link is SPSC (single-producer/single-consumer): the upstream is the orchestrator and the downstream is a single stage process; queues are unbounded (maxsize=0) on the orchestrator side.
- Device visibility
-  - Each stage sets `CUDA_VISIBLE_DEVICES` or calls `torch.cuda.set_device` to bind to GPU sets.
-  - A stage may use multiple GPUs internally (TP/PP/DP) but presents as a single stage unit.
- Simplified IPC: copy-based queues/pipes for data transfer; zero-copy is future work.
- Pipeline progression: when a stage finishes a request, it enqueues outputs to the downstream stage; if downstream is idle, it starts immediately.
- Scheduling
-  - A downstream stage triggers only after the upstream completes the request.
-  - Windowed segmentation/stitched triggering is not implemented; intra-stage streaming is not provided.
-
-### 4. IPC Implementation (simplified: copy-based)
- Use `multiprocessing.Queue`/Pipe for inter-process communication (control + data).
- Data is serialized/copied via CPU; no CUDA IPC/SHM zero-copy in this version.
- Backpressure: queues are unbounded; pressure manifests as compute-rate differences. Optional SHM reduces large-object transfer cost; RX/decoding overhead is recorded for observability.
-
-### 5. Scheduling & Cancellation (simplified)
- Pipeline: when a stage finishes a request, it enqueues to the next stage; that stage immediately pulls the next request from its input queue, enabling cross-stage concurrency.
- Cancellation/timeout: explicit cancellation/timeouts are not provided; graceful shutdown uses a `None` sentinel sent to each stage input queue.
-
-#### Short sequence example (req0/req1, stage A→B)
-1) t0: stage A handles req0
-2) t1: req0 completes on A → enters B; A immediately starts req1
-3) t2: B handles req0 while A handles req1 (parallel across stages)
-
-### 6. Integration Points (by file)
- `vllm_omni/entrypoints/omni.py` (Orchestrator)
-  - Class `Omni` orchestrates multi-process stages; constructs `OmniStage` instances in parallel and spawns per-stage workers.
-  - Spawns stage processes per config (set `CUDA_VISIBLE_DEVICES`/`torch.cuda.set_device`), creates control/data channels, builds simple full-trigger flow.
-  - Stats/logging are disabled by default; per-stage and orchestrator stats are only written when explicitly enabled.
-  - Manages process lifecycle: start/wait for readiness, graceful shutdown; forwards results between stages using copy-based IPC and optional SHM.
-  - Stage readiness: each stage emits `{"type": "stage_ready"}` after initialization; the orchestrator waits for all stages or times out and logs diagnostic suggestions.
@@ -0,0 +1,74 @@
+# MooncakeConnector
+
+## When to Use
+
+Best for multi-node distributed inference using Mooncake.
+
+## Installation
+
+```bash
+# For CUDA-enabled systems (recommended)
+pip install mooncake-transfer-engine
+
+# For non-CUDA systems
+pip install mooncake-transfer-engine-non-cuda
+```
+
+## Start Mooncake Master
+
+```bash
+# If you use Mooncake SSD storage
+mkdir -p ./mc_storage
+
+mooncake_master \
+  --rpc_port=50051 \
+  --enable_http_metadata_server=true \
+  --http_metadata_server_host=0.0.0.0 \
+  --http_metadata_server_port=8080 \
+  --metrics_port=9003 \
+  --root_fs_dir=./mc_storage/ \
+  --cluster_id=mc-local-1 &
+```
+
+## Configuration
+
+Define the connector in runtime:
+
+```yaml
+runtime:
+  connectors:
+    connector_of_mooncake:
+      name: MooncakeConnector
+      extra:
+        host: "127.0.0.1"
+        metadata_server: "http://<MASTER_IP>:8080/metadata"
+        master: "<MASTER_IP>:50051"
+        segment: 512000000
+        localbuf: 64000000
+        proto: "tcp"
+```
+
+Wire stages to the connector:
+
+```yaml
+stage_args:
+  - stage_id: 0
+    output_connectors:
+      to_stage_1: connector_of_mooncake
+
+  - stage_id: 1
+    input_connectors:
+      from_stage_0: connector_of_mooncake
+```
+
+Parameters:
+
+- host: local worker IP registered in the metadata server.
+- metadata_server: metadata server URL for discovery and setup.
+- master: Mooncake Master address.
+- segment: global memory segment size in bytes.
+- localbuf: local buffer size in bytes.
+- proto: transport protocol ("tcp" or "rdma").
+
+For more details, refer to the
+[Mooncake repository](https://github.com/kvcache-ai/Mooncake).
@@ -0,0 +1,27 @@
+# SharedMemoryConnector
+
+## When to Use
+
+Best for single-node deployments where stages run on the same host. It is
+auto-configured when no explicit connector is specified for an edge.
+
+## How It Works
+
+- Small payloads (< threshold): serialized and passed inline in metadata.
+- Large payloads (>= threshold): stored in shared memory; the SHM name is
+  returned in metadata.
+
+## Configuration
+
+```yaml
+runtime:
+  connectors:
+    connector_of_shared_memory:
+      name: SharedMemoryConnector
+      extra:
+        shm_threshold_bytes: 65536
+```
+
+## Notes
+
+- Auto-mode uses SharedMemoryConnector if no connector is declared for an edge.
@@ -0,0 +1,101 @@
+# YuanrongConnector
+
+## When to Use
+
+Best for multi-node distributed inference using Yuanrong Datasystem.
+
+## Mechanism
+
+Uses Yuanrong Datasystem's distributed KV store (`datasystem.kv_client`).
+
+- Data Plane: TCP or RDMA for high-bandwidth transfer.
+- Control Plane: Yuanrong Datasystem workers and etcd.
+- Keying: deterministic keys based on `put_key` (often composed as `request_id:fromStage_toStage`).
+
+## Installation
+
+```bash
+pip install openyuanrong-datasystem
+```
+
+## Start etcd
+
+```bash
+# Download and install etcd (v3.5.12 or higher)
+ETCD_VERSION="v3.5.12"
+ETCD_ARCH="linux-arm64"
+wget https://github.com/etcd-io/etcd/releases/download/${ETCD_VERSION}/etcd-${ETCD_VERSION}-${ETCD_ARCH}.tar.gz
+tar -xvf etcd-${ETCD_VERSION}-${ETCD_ARCH}.tar.gz
+cd etcd-${ETCD_VERSION}-${ETCD_ARCH}
+sudo cp etcd etcdctl /usr/local/bin/
+
+# Start etcd
+etcd \
+  --name etcd-single \
+  --data-dir /tmp/etcd-data \
+  --listen-client-urls http://0.0.0.0:2379 \
+  --advertise-client-urls http://0.0.0.0:2379 \
+  --listen-peer-urls http://0.0.0.0:2380 \
+  --initial-advertise-peer-urls http://0.0.0.0:2380 \
+  --initial-cluster etcd-single=http://0.0.0.0:2380 &
+
+# Verify etcd is running
+etcdctl --endpoints "127.0.0.1:2379" put key "value"
+etcdctl --endpoints "127.0.0.1:2379" get key
+```
+
+For production environments, refer to the
+[official etcd clustering documentation](https://etcd.io/docs/current/op-guide/clustering/).
+
+## Start Datasystem Worker
+
+```bash
+# Replace ${ETCD_IP} with etcd node IP, ${WORKER_IP} with local node IP
+dscli start -w \
+  --worker_address "${WORKER_IP}:31501" \
+  --etcd_address "${ETCD_IP}:2379" \
+  --shared_memory_size_mb 20480
+```
+
+To stop the worker:
+
+```bash
+dscli stop --worker_address "${WORKER_IP}:31501"
+```
+
+## Configuration
+
+Define the connector in runtime:
+
+```yaml
+runtime:
+  connectors:
+    connector_of_yuanrong:
+      name: YuanrongConnector
+      extra:
+        host: "127.0.0.1"
+        port: 31501
+        get_sub_timeout_ms: 1000
+```
+
+Wire stages to the connector:
+
+```yaml
+stage_args:
+  - stage_id: 0
+    output_connectors:
+      to_stage_1: connector_of_yuanrong
+
+  - stage_id: 1
+    input_connectors:
+      from_stage_0: connector_of_yuanrong
+```
+
+Parameters:
+
+- host: datasystem worker host.
+- port: datasystem worker port.
+- get_sub_timeout_ms: get timeout in milliseconds (0 for no timeout).
+
+For more details, refer to the
+[Yuanrong Datasystem repository](https://atomgit.com/openeuler/yuanrong-datasystem).
@@ -1,14 +1,17 @@
 # Distributed utils

 This directory (vllm_omni/distributed/ray_utils) contains utilities for distributed execution in vllm-omni, supporting both **Ray** and **Multiprocessing** backends.
-
-## 1. Ray Utils
+## 1. Installation
+```bash
+pip install "ray[default]"
+```
+## 2. Ray Utils

 The `ray_utils` module provides helper functions for managing Ray clusters and actors, which is essential for:
 *   **Multi-node deployment**: Running pipeline stages across different physical machines.
 *   **Resource management**: Efficient GPU/CPU allocation.

-### 1.1 Basic Usage
+### 2.1 Basic Usage

 To use the Ray backend, specify `worker_backend="ray"` when initializing the engine.

@@ -21,7 +24,7 @@ vllm serve Qwen/Qwen2.5-Omni-7B \
  --ray-address auto
 ```

-### 1.2 Cluster Setup
+### 2.2 Cluster Setup

 **Step 1: Start Head Node**
 Run this on your primary machine:
@@ -38,7 +41,7 @@ ray start --address=<HEAD_NODE_IP>:6399
 > **Tip**: For a complete cluster setup script, refer to the vLLM example:
 > [run_cluster.sh](https://github.com/vllm-project/vllm/blob/main/examples/online_serving/run_cluster.sh)

-### 1.3 Distributed Connector Support
+### 2.3 Distributed Connector Support

 When running on Ray, the system automatically adapts its communication strategy:

@@ -46,16 +49,11 @@ When running on Ray, the system automatically adapts its communication strategy:
 *   **Same-Node**: Can still use `SharedMemoryConnector` for efficiency, or Ray's native object store (plasma).
 *   **SHM threshold default differs**: when `worker_backend="ray"`, the SharedMemoryConnector default threshold is set to `sys.maxsize`, which forces payloads to go inline (no SHM). Override `shm_threshold_bytes` in the connector config if you want SHM for Ray runs.

-### 1.4 Internal Helpers
+### 2.4 Internal Helpers

 *   **`initialize_ray_cluster`**: Connects to an existing Ray cluster or starts a local one.

-## 2. Troubleshooting
+## 3. Troubleshooting

 *   **Connection Issues**: Ensure the Ray head node is accessible and ports (default 6399 in this example) are open.
 *   **Version Mismatch**: Ensure all nodes run the same version of Ray and Python.
-
-### Installation
-```bash
-pip install "ray[default]"
-```
@@ -9,7 +9,6 @@ This section contains design documents and architecture specifications for vLLM-
 ## Feature Design Documents

 - [Disaggregated Inference](feature/disaggregated_inference.md)
- [Multi-Request Streaming](feature/multi_request_streaming.md)
 - [Ray-based Execution](feature/ray_based_execution.md)

 ## Module Design Documents
@@ -1,3 +1,7 @@
+---
+toc_depth: 4
+---
+
 # Diffusion Module Architecture Design

 The vLLM-Omni diffusion module (`vllm_omni/diffusion`) is a high-performance inference engine for diffusion models, designed with a modular architecture that separates concerns across multiple components. It provides efficient execution for non-autoregressive generation tasks such as image and video generation.
@@ -5,7 +9,7 @@ The vLLM-Omni diffusion module (`vllm_omni/diffusion`) is a high-performance inf
 This document describes the architecture design of the diffusion module, including the diffusion engine, scheduler, worker, diffusion pipeline, and acceleration components.

 <p align="center">
-   <img src="https://github.com/user-attachments/assets/cde1ebea-4006-4d30-8283-15a5721a09d2" alt="vLLM-Omni Diffusion Module Components" width="60%">
+   <img src="https://github.com/user-attachments/assets/cde1ebea-4006-4d30-8283-15a5721a09d2" alt="vLLM-Omni Diffusion Module Components" width="80%">
 </p>
 <p align="center">
  <em> Main Components of the Diffusion Module </em>
@@ -20,10 +24,10 @@ This document describes the architecture design of the diffusion module, includi
 - [Worker](#3-worker)
 - [Diffusion Pipeline](#4-diffusion-pipeline)
 - [Acceleration Components](#5-acceleration-components)
-   - [Attention Backends](#51-attention-backends)
-   - [Parallel Attention](#52-parallel-attention)
-   - [Cache Backends](#53-cache-backends)
-   - [Parallel Strategies](#54-parallel-strategies)
+    - [Attention Backends](#51-attention-backends)
+    - [Parallel Attention](#52-parallel-attention)
+    - [Cache Backends](#53-cache-backends)
+    - [Parallel Strategies](#54-parallel-strategies)
 - [Data Flow](#6-data-flow)

 ---
@@ -65,8 +69,11 @@ class DiffusionEngine:
 ```

 **Key Features**:
+
 - **Pre/Post Processing**: Registers model-specific pre-processing and post-processing functions via registry pattern
+
 - **Worker Management**: Launches and manages multiple worker processes (one per GPU)
+
 - **Process Isolation**: Uses multiprocessing for true parallelism

 #### 1.2 Worker Launch Process
@@ -86,8 +93,11 @@ def _launch_workers(self, broadcast_handle):
 ```

 **Design Decisions**:
+
 - **Spawn Method**: Ensures clean state for each worker (no shared memory issues)
+
 - **Pipe Communication**: Uses `mp.Pipe` for initialization handshake
+
 - **Device Selection**: Each worker is assigned a specific GPU (`cuda:{rank}`)

 #### 1.3 Request Processing Flow
@@ -106,8 +116,11 @@ def step(self, requests: list[OmniDiffusionRequest]):
 ```

 **Flow**:
+
 1. **Pre-processing**: Applies model-specific transformations
+
 2. **Scheduling**: Delegates to scheduler for distribution
+
 3. **Post-processing**: Converts raw outputs to final format (e.g., PIL images)

 ---
@@ -139,8 +152,11 @@ class Scheduler:
 ```

 **Communication Pattern**:
+
 - **Broadcast Queue**: One-to-many communication (scheduler → all workers)
+
 - **Result Queue**: One-to-one communication (rank 0 → scheduler)
+
 - **Shared Memory**: Uses `MessageQueue` (ZMQ-based) for efficient IPC

 #### 2.2 Request Distribution
@@ -156,8 +172,11 @@ def add_req(self, requests: list[OmniDiffusionRequest]) -> DiffusionOutput:
 ```

 **Design Features**:
+
 - **Broadcast Model**: All workers receive the same request (for tensor parallelism)
+
 - **Single Response**: Only rank 0 sends results back (avoids duplicate outputs)
+
 - **Synchronous**: Blocks until result is received (can be made async)

 #### 2.3 Singleton Pattern
@@ -176,8 +195,11 @@ scheduler = Scheduler()
 ```

 **Benefits**:
+
 - **Single Point of Control**: Ensures consistent state
+
 - **Easy Access**: Global `scheduler` instance accessible everywhere
+
 - **Resource Management**: Centralized queue management

 ---
@@ -212,12 +234,19 @@ class WorkerProc:
 ```

 **Initialization Steps**:
+
 1. **IPC Setup**: Creates ZMQ context and message queues
+
 2. **Distributed Environment Setup**: Initializes PyTorch distributed communication
-     - For CUDA GPUs: Uses NCCL (fast GPU communication)
-     - For NPU: Uses HCCL (Huawei Collective Communications Library)
-     - For other devices: Uses appropriate backend (GLOO, MCCL, etc.)
+
+    - For CUDA GPUs: Uses NCCL (fast GPU communication)
+
+    - For NPU: Uses HCCL (Huawei Collective Communications Library)
+
+    - For other devices: Uses appropriate backend (GLOO, MCCL, etc.)
+
 3. **Model Loading**: Loads diffusion pipeline on assigned GPU
+
 4. **Cache Setup**: Enables cache backend if configured.

 #### 3.2 GPU Worker
@@ -253,8 +282,11 @@ class GPUWorker:
 ```

 **Key Features**:
+
 - **Tensor Parallelism**: Supports multi-GPU tensor parallelism via PyTorch distributed
+
 - **Model Loading**: Uses `DiffusersPipelineLoader` for efficient weight loading
+
 - **Cache Integration**: Enables cache backends (TeaCache, cache-dit, etc.) transparently

 #### 3.3 Worker Busy Loop
@@ -283,9 +315,13 @@ def worker_busy_loop(self):
 ```

 **Execution Flow**:
+
 1. **Receive**: Dequeues unified messages from shared memory queue
+
 2. **Route**: Handles different message types (generation, RPC, shutdown)
+
 3. **Execute**: Runs forward pass through pipeline for generation requests
+
 4. **Respond**: Sends results back (rank 0 for generation, specified rank for RPC)

 #### 3.4 Model Execution
@@ -315,10 +351,14 @@ The model execution leverages multiple parallelism strategies that are transpare
 from vllm_omni.diffusion.distributed.parallel_state import (
    get_sp_group, get_dp_group, get_cfg_group, get_pp_group
 )
+```

 **Optimizations**:
+
 - **Cache Refresh**: Clears cache state before each generation for clean state
+
 - **Context Management**: Forward context ensures parallel groups are available during execution
+
 - **Single Request**: Currently processes one request at a time (batching TODO)

 ---
@@ -327,66 +367,9 @@ from vllm_omni.diffusion.distributed.parallel_state import (

 **Location**: `vllm_omni/diffusion/models/*/pipeline_*.py`

-### Architecture
-
 The pipeline is the **model-specific implementation** that orchestrates the diffusion process. Different models (QwenImage, Wan2.2, Z-Image) have their own pipeline implementations.

-### Key Components (QwenImage Example)
-
-#### 4.1 Pipeline Structure
-
-```python
-class QwenImagePipeline(nn.Module):
-    def __init__(self, od_config: OmniDiffusionConfig):
-        # Core components
-        self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(...)
-        self.text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained(...)
-        self.vae = AutoencoderKLQwenImage.from_pretrained(...)
-        self.transformer = QwenImageTransformer2DModel(od_config=od_config)
-        self.tokenizer = Qwen2Tokenizer.from_pretrained(...)
-
-        # Cache backend (set by worker)
-        self._cache_backend = None
-```
-
-**Component Responsibilities**:
- **Text Encoder**: Encodes prompts into embeddings
- **Transformer**: Main denoising network (DiT-style)
- **VAE**: Encodes/decodes between pixel and latent space
- **Scheduler**: Manages timestep scheduling
-
-#### 4.2 Forward Pass
-
-```python
-def forward(self, req: OmniDiffusionRequest, ...) -> DiffusionOutput:
-    # 1. Encode prompts
-    prompt_embeds, prompt_embeds_mask = self.encode_prompt(prompt)
-
-    # 2. Prepare latents
-    latents = self.prepare_latents(batch_size, num_channels, height, width, ...)
-
-    # 3. Prepare timesteps
-    timesteps, num_inference_steps = self.prepare_timesteps(...)
-
-    # 4. Diffusion loop
-    latents = self.diffuse(
-        prompt_embeds, prompt_embeds_mask,
-        negative_prompt_embeds, negative_prompt_embeds_mask,
-        latents, img_shapes, txt_seq_lens,
-        timesteps, do_true_cfg, guidance, true_cfg_scale,
-    )
-
-    # 5. Decode latents
-    image = self.vae.decode(latents)
-
-    return DiffusionOutput(output=image)
-```
-
-The main steps of the forward pass are referred from `diffusers`.
-
-#### 4.3 Diffusion Loop
-
-The multi-step diffusion loop is usually the most time-consuming part during the overall inference process.
+Most pipeline implementation are referred from `diffusers`. The multi-step diffusion loop is usually the most time-consuming part during the overall inference process, which is defined by the `diffuse` function in the pipeline class. An example is as follows:

 ```python
 def diffuse(self, ...):
@@ -416,11 +399,14 @@ def diffuse(self, ...):
 ```

 **Key Features**:
+
 - **CFG Support**: Handles classifier-free guidance with separate forward passes
+
 - **Cache Branching**: Uses `cache_branch` parameter for cache-aware execution
+
 - **True CFG**: Implements advanced CFG with norm preservation

-To learn how to add a new diffusion pipeline, please view [Adding Diffusion Model](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/model/adding_diffusion_model)
+To learn more about the diffusion pipeline and how to add a new diffusion pipeline, please view [Adding Diffusion Model](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/model/adding_diffusion_model)

 ---

@@ -448,9 +434,13 @@ class Attention(nn.Module):
 ```

 **Available Backends**:
+
 - **FlashAttention**: Optimized CUDA kernel (FA2/FA3) - memory efficient via tiling
+
 - **SDPA**: PyTorch's scaled dot-product attention - default, cross-platform
+
 - **SageAttention**: Sparse attention implementation from SageAttention library
+
 - **AscendAttention**: NPU-optimized attention for Ascend hardware

 These backends provide the **kernel implementations** for attention computation. For attention-level sequence parallelism strategies (Ring Attention, Ulysses), see [Parallel Attention](#52-parallel-attention).
@@ -470,16 +460,25 @@ def get_attn_backend(head_size: int) -> type[AttentionBackend]:
 ```

 **Selection Priority**:
+
 1. **Environment Variable**: `DIFFUSION_ATTENTION_BACKEND` for manual override
-   - Valid values: `FLASH_ATTN`, `TORCH_SDPA`, `SAGE_ATTN`, `ASCEND`
-   - Example: `export DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN`
+
+    - Valid values: `FLASH_ATTN`, `TORCH_SDPA`, `SAGE_ATTN`, `ASCEND`
+
+    - Example: `export DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN`
+
 2. **Automatic Fallback**: Falls back to SDPA if selected backend unavailable
+
 3. **Hardware Detection**: Can select based on device type (NPU, CUDA, etc.)

 **Backend Availability**:
+
 - **SDPA**: Always available (PyTorch built-in)
+
 - **FlashAttention**: Requires `flash-attn` package installed
+
 - **SageAttention**: Requires `sage-attention` package (from THU-ML GitHub)
+
 - **AscendAttention**: Only available on Ascend NPU hardware

 #### Attention Backend Registry
@@ -563,8 +562,11 @@ Parallel attention strategies implement **Sequence Parallelism (SP) at the atten
 **Key Distinction**: Unlike AttentionBackend (which provides kernel implementations), ParallelAttentionStrategy provides communication patterns for multi-GPU attention parallelism. These strategies implement the `ParallelAttentionStrategy` interface and use AttentionBackend implementations internally.

 Both Ring Attention and Ulysses are forms of Sequence Parallelism (SP) that:
+
 - Split the sequence dimension across GPUs
+
 - Contribute to `sequence_parallel_size` (via `ring_degree` and `ulysses_degree`)
+
 - Work at the attention layer level (not model/pipeline level)

 #### Ulysses Sequence Parallelism (USP)
@@ -574,7 +576,9 @@ Both Ring Attention and Ulysses are forms of Sequence Parallelism (SP) that:
 USP is a sequence-parallel attention strategy that splits attention computation across multiple GPUs by distributing both the sequence dimension and attention heads. It uses **all-to-all communication** to efficiently parallelize attention for very long sequences. Specifically, it uses **all-to-all** collective operations to redistribute Q/K/V tensors before attention computation and gather results afterward.

 Ulysses splits attention computation in two dimensions:
+
 1. **Sequence Dimension**: Splits the sequence length across GPUs
+
 2. **Head Dimension**: Splits attention heads across GPUs

 **Configuration**: `ulysses_degree` contributes to `sequence_parallel_size`
@@ -601,9 +605,13 @@ class RingParallelAttention:
 ```

 **Integration**:
+
 - Ring Attention is activated when `ring_degree > 1` in parallel config
+
 - It's selected by `build_parallel_attention_strategy()` in the attention layer
+
 - The `Attention` layer routes to `_run_ring_attention()` when Ring is enabled
+
 - Works alongside attention backends: Ring handles communication, backends handle computation

 **Configuration**: `ring_degree` contributes to `sequence_parallel_size`
@@ -611,8 +619,11 @@ class RingParallelAttention:
 #### Relationship with AttentionBackend

 Parallel attention strategies (Ring, Ulysses) work **on top of** AttentionBackend implementations:
+
 - They use AttentionBackend for the actual attention computation (FlashAttention, SDPA, etc.)
+
 - They handle the multi-GPU communication/parallelization layer
+
 - They implement `ParallelAttentionStrategy` interface (not `AttentionBackend`)

 For general parallelism strategies (Data Parallelism, Tensor Parallelism, Pipeline Parallelism), see [Parallel Strategies](#54-parallel-strategies).
@@ -651,8 +662,11 @@ class CacheBackend(ABC):
 ```

 **Design Pattern**:
+
 - **Abstract Base Class**: Defines contract for all cache backends
+
 - **Pipeline-based**: Works with pipeline instances (not just transformers)
+
 - **State Management**: Provides refresh mechanism for clean state between generations

 #### Available Backends
@@ -686,9 +700,13 @@ class TeaCacheBackend(CacheBackend):
 ```

 **TeaCache Features**:
+
 - **Timestep-aware**: Caches based on timestep embedding similarity
+
 - **Adaptive**: Dynamically decides when to reuse cached computations
+
 - **CFG-aware**: Handles positive/negative branches separately
+
 - **Custom Hook System**: Uses a custom forward interception mechanism (via `HookRegistry`) that wraps the module's `forward` method, allowing transparent integration without modifying model code

 **2. Cache-DiT Backend**
@@ -710,9 +728,13 @@ class CacheDiTBackend(CacheBackend):
 ```

 **Cache-DiT Features**:
+
 - **DBCache**: Dynamic block caching with configurable compute blocks
+
 - **SCM**: Step Computation Masking for additional speedup
+
 - **TaylorSeer**: Advanced calibration for cache accuracy
+
 - **Dual-transformer Support**: Handles models like Wan2.2 with two transformers

 #### Cache Backend Selector
@@ -748,9 +770,13 @@ def get_cache_backend(
 ```

 **Usage Flow**:
+
 1. **Selection**: `get_cache_backend()` returns appropriate backend instance
+
 2. **Enable**: `backend.enable(pipeline)` called during worker initialization
+
 3. **Refresh**: `backend.refresh(pipeline, num_inference_steps)` called before each generation
+
 4. **Check**: `backend.is_enabled()` verifies cache is active

 ### 5.4 Parallel Strategies
@@ -762,23 +788,35 @@ def get_cache_backend(
 The system supports multiple orthogonal parallelism strategies:

 **Sequence Parallelism (SP)**
+
 - **Purpose**: Split sequence dimension across GPUs
+
 - **Attention-level SP**: Ring Attention and Ulysses (USP) implement SP at the attention layer level
-  - See [Parallel Attention](#52-parallel-attention) for details
-  - Configuration: `ulysses_degree` × `ring_degree` = `sequence_parallel_size`
+
+    - See [Parallel Attention](#52-parallel-attention) for details
+
+    - Configuration: `ulysses_degree` × `ring_degree` = `sequence_parallel_size`
+
 - **Use Case**: Very long sequences (e.g., high-resolution images)

 **Data Parallelism (DP)**
+
 - **Purpose**: Replicate model across GPUs, split batch
+
 - **Use Case**: Batch processing, throughput optimization

 **Tensor Parallelism (TP)** (Experimental)
+
 - **Purpose**: Split model weights across GPUs
+
 - **Implementation**: Uses vLLM's tensor parallel groups
+
 - **Use Case**: Large models that don't fit on single GPU

 **CFG Parallelism**  (under development)
+
 - **Purpose**: Parallelize Classifier-Free Guidance (positive/negative prompts)
+
 - **Infrastructure**: CFG parallel groups are initialized and available via `get_cfg_group()`

 #### Parallel Group Management
@@ -0,0 +1,70 @@
+# vLLM-Omni ComfyUI Integration
+
+vLLM-Omni offers a ComfyUI integration on top of its online serving API.
+It can send model inference requests to either a locally running vLLM-Omni service or a remote one.
+
+## Requirement
+
+- Python 3.12 or above
+- [ComfyUI installed](https://docs.comfy.org/installation/system_requirements)
+- [vLLM-Omni installed](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/) on either the same device or another device discoverable via the internet.
+- No need to install additional packages apart from those already required by ComfyUI.
+
+!!! tip
+    If you run both ComfyUI and vLLM-Omni on the same device, you can create separate virtual environments and use different Python versions for them.
+
+
+## Installation
+
+Copy the `apps/ComfyUI-vLLM-Omni` folder to the `custom_nodes` subfolder of your ComfyUI installation. Your directory should look like `ComfyUI/custom_nodes/ComfyUI-vLLM-Omni`.
+
+If you are running ComfyUI during copying, you should restart ComfyUI to load this extension.
+
+!!! tip
+    You can use utility websites such as https://download-directory.github.io/ to download a subdirectory of a repo. Also checkout community discussions (e.g., https://stackoverflow.com/questions/7106012/download-a-single-folder-or-directory-from-a-github-repository) for more info.
+
+On the device and virtual environment you run ComfyUI, launch ComfyUI with
+```bash
+cd ComfyUI
+
+# The regular way
+python main.py
+
+# If you are mainly using this node, launch it faster with
+python main.py --cpu
+```
+
+On the device and virtual environment you run vLLM-Omni, start a model service with
+```bash
+vllm serve The_Model_ID_to_Serve --omni --port 8000
+```
+
+Check **ComfyUI's sidebar -> Node Library**. There should be a new folder named **vLLM-Omni**.
+If no, check your shell running the ComfyUI process. There may be some error messages before the line `Import times for custom nodes:` and the line `To see the GUI go to: http://127.0.0.1:8188`.
+
+## Quickstart
+
+This extension offers the following nodes based on the output modalities:
+
+- **Generate Image** for text-to-image and image-to-image tasks
+- **Multimodality Comprehension** for multimodality-to-text and multimodality-to-audio tasks
+- **TTS** and **TTS Voice Clone** for TTS tasks
+
+This extension also offers example workflows (at **ComfyUI sidebar -> Templates -> vLLM-Omni**)
+
+!!! info
+    The node UI and feature designs are intended to match vLLM-Omni online serving interfaces. It cannot offer more than what the interfaces support.
+
+To build a simple workflow yourself,
+
+- Drag a generation node onto the canvas.
+- Depending on your need, grab built-in multimedia file loader nodes, such as **image->Load Image**, **image->video->Load Video**, **audio->Load Audio**
+- Depending on your need, grab built-in multimedia file preview nodes, such as **image->Preview Image**, **image->video->Save Video**, **audio->Preview Audio**. For text output, you can install [ComfyUI-Custom-Scripts plugin](https://github.com/pythongosssss/ComfyUI-Custom-Scripts/) and grab its **utils->Show Text 🐍** node.
+- If you want to tune sampling parameters, grab corresponding nodes from **vLLM-Omni-> Sampling Params**.
+    - For multi-stage models, you can connect multiple **AR Sampling Params** and **Diffusion Sampling Params** nodes to a **Multi-Stage Sampling Params List** node, and connect this node to the generation node.
+    - For some multi-stage models like BAGEL, [only one stage's sampling parameters are exposed and tunable via vLLM-Omni's online serving API](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/examples/online_serving/bagel/). Thus, these models are treated as single-stage ones. Please check the vLLM-Omni documentation on how to correctly set each model's sampling parameters.
+    - For multi-stage models where all stages are either autoregression or diffusion, you can also connect only a single Sampling Params node, indicating that this set of sampling parameters will be used for all stages.
+
+## Examples & Screenshots
+
+Please read the [ComfyUI integration's README](https://github.com/vllm-project/vllm-omni/tree/main/apps/ComfyUI-vLLM-Omni) for more info.
@@ -0,0 +1,151 @@
+# Custom Pipeline Extension Guide
+
+Transformer already support Custom Pipeline via
+https://github.com/huggingface/diffusers/blob/main/docs/source/en/using-diffusers/custom_pipeline_overview.md
+
+This guide demonstrates how to use the newly added features for extending vLLM-Omni's diffusion pipeline with custom functionality.
+
+## Overview
+
+Three main features enable custom pipeline extension:
+
+1. **`WorkerWrapperBase`**: A wrapper class that enables dynamic worker extension with custom functionality
+2. **`load_format`**: A parameter that controls how diffusion models are loaded, including support for custom pipelines
+3. **`CustomPipelineWorkerExtension`**: An extension class that enables pipeline re-initialization with custom implementations
+
+## Features
+
+### WorkerWrapperBase
+
+`WorkerWrapperBase` is a wrapper class that creates `DiffusionWorker` instances with optional extension support. It enables dynamic inheritance, allowing you to add custom methods and functionality to workers without modifying the base worker class.
+
+**Key capabilities:**
+- Dynamic worker class extension via `worker_extension_cls`
+- Support for custom pipeline initialization via `custom_pipeline_args`
+- Method delegation to underlying worker
+- Attribute access forwarding
+
+**Location:** `vllm_omni/diffusion/worker/diffusion_worker.py`
+
+### load_format Parameter
+
+The `load_format` parameter controls how diffusion models are loaded. It supports the following values:
+
+- **`"default"`**: Standard model loading using the model registry (default behavior)
+- **`"custom_pipeline"`**: Load a custom pipeline class specified by `custom_pipeline_name`
+- **`"dummy"`**: Skip model loading (useful for testing or when pipeline will be initialized separately)
+
+**Location:** `vllm_omni/diffusion/model_loader/diffusers_loader.py`
+
+### CustomPipelineWorkerExtension
+
+`CustomPipelineWorkerExtension` is a mixin class that extends `DiffusionWorker` with the ability to re-initialize the pipeline with a custom implementation.
+
+**Key method:**
+- `re_init_pipeline(custom_pipeline_args)`: Re-initializes the pipeline with custom arguments, properly cleaning up the old pipeline first
+
+**Location:** `vllm_omni/diffusion/worker/diffusion_worker.py`
+
+## Usage Example
+
+### Step 1: Create a Custom Pipeline
+
+Create a custom pipeline class that extends an existing pipeline. In this example, we extend `QwenImageEditPipeline` to add trajectory tracking:
+
+```python
+# custom_pipeline.py
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image_edit import QwenImageEditPipeline
+import torch
+
+class CustomPipeline(QwenImageEditPipeline):
+    def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
+        super().__init__(od_config=od_config, prefix=prefix)
+
+    def forward(self, req, prompt=None, negative_prompt=None, **kwargs):
+        # Call parent's forward to get normal output
+        output = super().forward(req=req, prompt=prompt, negative_prompt=negative_prompt, **kwargs)
+
+        # Add custom trajectory data
+        actual_num_steps = req.sampling_params.num_inference_steps or kwargs.get('num_inference_steps', 50)
+        output.trajectory_timesteps = torch.linspace(1000, 0, actual_num_steps, dtype=torch.float32)
+        output.trajectory_latents = torch.randn(actual_num_steps, 1, 16, 64, 64, dtype=torch.float32)
+
+        return output
+```
+
+### Step 2: Use the Custom Pipeline with Omni
+
+Initialize the `Omni` engine with custom pipeline configuration:
+
+```python
+from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+# Initialize with custom pipeline
+omni = Omni(
+    model="Qwen/Qwen-Image-Edit",
+    diffusion_load_format="dummy",  # Skip initial loading
+    custom_pipeline_args={
+        "pipeline_class": "custom_pipeline.CustomPipeline"
+    },
+)
+
+# Generate with the custom pipeline
+outputs = omni.generate(
+    ...
+)
+
+# Access custom trajectory data
+output = outputs[0].request_output[0]
+print(f"Trajectory timesteps shape: {output.metrics['trajectory_timesteps'].shape}")
+print(f"Trajectory latents shape: {output.latents.shape}")
+```
+
+### Step 3: Run the Example
+
+The example provided in this directory demonstrates the complete workflow:
+
+```bash
+cd examples/offline_inference/custom_pipeline/image_to_image
+
+# Run with custom pipeline
+python image_edit.py \
+    --model Qwen/Qwen-Image-Edit-2511 \
+    --image cherry_blossom.jpg \
+    --prompt "Let this mascot dance under the moon, surrounded by floating stars" \
+    --output output_image_edit.png \
+    --num_inference_steps 10
+```
+
+## Advanced Usage
+
+### Custom Worker Extension
+
+You can create custom worker extensions to add new methods beyond pipeline re-initialization:
+
+```python
+from typing import Any
+from vllm_omni.diffusion.worker.diffusion_worker import CustomPipelineWorkerExtension
+
+class MyCustomExtension(CustomPipelineWorkerExtension):
+    def custom_method(self):
+        """Your custom worker method."""
+        return "custom_result"
+
+    def another_method(self, data: Any):
+        """Another custom method."""
+        # Access worker internals via self
+        return self.model_runner.some_operation(data)
+
+omni = Omni(
+    model="Qwen/Qwen-Image-Edit",
+    diffusion_load_format="dummy",
+    custom_pipeline_args={
+        "pipeline_class": "custom_pipeline.CustomPipeline"
+    },
+    worker_extension_cls=MyCustomExtension,
+    # Note: worker_extension_cls is an internal parameter
+    # CustomPipelineWorkerExtension will automatically init pipeline when custom_pipeline_args is provided
+)
+```
@@ -0,0 +1,39 @@
+# Sleep Mode
+
+vLLM-Omni’s **Sleep Mode** allows you to temporarily release most GPU memory used by a model—such as model weights and key-value (KV) caches (for autoregressive models)—**without stopping the server or unloading the Docker container**.
+
+This feature is inherited from [vLLM’s Sleep Mode](https://blog.vllm.ai/2025/10/26/sleep-mode.html), which provides zero-reload model switching for multi-model serving.  
+
+It is especially useful in **RLHF**, **training**, or **cost-saving scenarios**, where GPU resources must be freed between inference workloads.
+
+---
+
+## Omni Model
+
+Omni model inherit the feature from vLLM' Sleep Mode
+
+This means:
+
+- Support both Level 1 and Level 2 sleep, allow to release and reset both model weights and KV Cache
+
+## Diffusion Model Extension
+
+We added Sleep Mode support for **diffusion models**, which previously lacked this functionality.  
+In diffusion pipelines, this currently only offloads **model weight memory**, as these models typically do not use KV caches.
+
+This means:
+
+- Diffusion models can now enter Level 1 sleep.
+- Pipeline states (e.g., noise schedulers, buffers) remain intact after waking.
+- Useful for releasing VRAM between image generation or training cycles.
+
+---
+
+## Enable sleep mode
+To enable sleep mode, set the `enable_sleep_mode` in `engine_args` to `True`
+
+
+Example:
+```python
+omni = Omni(model=...,enable_sleep_mode=True)
+```
@@ -26,6 +26,8 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr

 ### Pre-built wheels

+Note: Pre-built wheels are currently only available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0. For the latest version, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
+
 === "NVIDIA CUDA"

    --8<-- "docs/getting_started/installation/gpu/cuda.inc.md:pre-built-wheels"
@@ -49,15 +51,18 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr

 ## Set up using Docker

+### Pre-built images
+
+=== "NVIDIA CUDA"
+
+    --8<-- "docs/getting_started/installation/gpu/cuda.inc.md:pre-built-images"
+
+=== "AMD ROCm"
+
+    --8<-- "docs/getting_started/installation/gpu/rocm.inc.md:pre-built-images"

 ### Build your own docker image

 === "AMD ROCm"

    --8<-- "docs/getting_started/installation/gpu/rocm.inc.md:build-docker"
-
-### Pre-built images
-
-=== "NVIDIA CUDA"
-
-    --8<-- "docs/getting_started/installation/gpu/cuda.inc.md:pre-built-images"
@@ -17,12 +17,10 @@ Therefore, it is recommended to install vLLM and vLLM-Omni with a **fresh new**
 # --8<-- [start:pre-built-wheels]

 #### Installation of vLLM
-Note: Pre-built wheels are currently only available for vLLM-Omni 0.11.0rc1. For the latest version, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
-

 vLLM-Omni is built based on vLLM. Please install it with command below.
 ```bash
-uv pip install vllm==0.11.0 --torch-backend=auto
+uv pip install vllm==0.16.0 --torch-backend=auto
 ```

 #### Installation of vLLM-Omni
@@ -36,13 +34,13 @@ uv pip install vllm-omni
 # --8<-- [start:build-wheel-from-source]

 #### Installation of vLLM
-If you do not need to modify source code of vLLM, you can directly install the stable 0.12.0 release version of the library
+If you do not need to modify source code of vLLM, you can directly install the stable 0.16.0 release version of the library

 ```bash
-uv pip install vllm==0.12.0 --torch-backend=auto
+uv pip install vllm==0.16.0 --torch-backend=auto
 ```

-The release 0.12.0 of vLLM is based on PyTorch 2.9.0 which requires CUDA 12.9 environment.
+The release 0.14.0 of vLLM is based on PyTorch 2.9.0 which requires CUDA 12.9 environment.

 #### Installation of vLLM-Omni
 Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -58,11 +56,11 @@ If you want to check, modify or debug with source code of vLLM, install the libr
 ```bash
 git clone https://github.com/vllm-project/vllm.git
 cd vllm
-git checkout v0.12.0
+git checkout v0.16.0
 ```
 Set up environment variables to get pre-built wheels. If there are internet problems, just download the whl file manually. And set `VLLM_PRECOMPILED_WHEEL_LOCATION` as your local absolute path of whl file.
 ```bash
-export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.12.0/vllm-0.12.0-cp38-abi3-manylinux_2_31_x86_64.whl
+export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.16.0/vllm-0.16.0-cp38-abi3-manylinux_2_31_x86_64.whl
 ```
 Install vllm with command below (If you have no existing PyTorch).
 ```bash
@@ -93,7 +91,7 @@ docker run --runtime nvidia --gpus 2 \
    --env "HF_TOKEN=$HF_TOKEN" \
    -p 8091:8091 \
    --ipc=host \
-    vllm/vllm-omni:v0.11.0rc1 \
+    vllm/vllm-omni:v0.14.0 \
    --model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 8091
 ```

@@ -9,10 +9,60 @@ vLLM-Omni current recommends the steps in under setup through Docker Images.

 # --8<-- [start:pre-built-wheels]

+#### Installation of vLLM
+
+vLLM-Omni is built based on vLLM. Please install it with command below.
+```bash
+uv pip install vllm==0.14.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.14.0/rocm700
+```
+
+#### Installation of vLLM-Omni
+
+```bash
+# we need to add --no-build-isolation as the torch
+# is not obtained from pypi, we have to install using the
+# torch installed in our environment
+uv pip install vllm-omni
+
+# Optional if want to run Qwen3 TTS
+uv pip uninstall onnxruntime # should be removed before we can install onnxruntime-rocm
+uv pip install onnxruntime-rocm sox
+```
+
 # --8<-- [end:pre-built-wheels]

 # --8<-- [start:build-wheel-from-source]

+#### Installation of vLLM
+If you do not need to modify source code of vLLM, you can directly install the stable 0.14.0 release version of the library
+
+```bash
+uv pip install vllm==0.14.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.14.0/rocm700
+```
+
+The release 0.14.0 of vLLM requires ROCm 7.0 environment.
+
+#### Installation of vLLM-Omni
+Since vllm-omni is rapidly evolving, it's recommended to install it from source
+```bash
+git clone https://github.com/vllm-project/vllm-omni.git
+cd vllm-omni
+VLLM_OMNI_TARGET_DEVICE=rocm uv pip install -e .
+# OR
+uv pip install -e . --no-build-isolation
+```
+
+<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:
+
+```bash
+git clone https://github.com/vllm-project/vllm.git
+cd vllm
+git checkout v0.14.0
+python3 -m pip install -r requirements/rocm.txt
+python3 setup.py develop
+```
+
 # --8<-- [end:build-wheel-from-source]

 # --8<-- [start:build-docker]
@@ -23,19 +73,30 @@ vLLM-Omni current recommends the steps in under setup through Docker Images.
 DOCKER_BUILDKIT=1 docker build -f docker/Dockerfile.rocm -t vllm-omni-rocm .
 ```

-If you want to specify which GPU Arch to build for to cutdown build time:
-
-```bash
-DOCKER_BUILDKIT=1 docker build \
-  -f docker/Dockerfile.rocm \
-  --build-arg PYTORCH_ROCM_ARCH="gfx942;gfx950" \
-  -t vllm-omni-rocm .
-```
-
 #### Launch the docker image

+##### Launch with OpenAI API Server
+
 ```
-docker run -it \
+docker run --rm \
+--group-add=video \
+--ipc=host \
+--cap-add=SYS_PTRACE \
+--security-opt seccomp=unconfined \
+--device /dev/kfd \
+--device /dev/dri \
+-v ~/.cache/huggingface:/root/.cache/huggingface \
+--env "HF_TOKEN=$HF_TOKEN" \
+-p 8091:8091 \
+--ipc=host \
+vllm-omni-rocm \
+--model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 8091
+```
+
+##### Launch with interactive session for development
+
+```
+docker run --rm -it \
 --network=host \
 --group-add=video \
 --ipc=host \
@@ -44,12 +105,51 @@ docker run -it \
 --device /dev/kfd \
 --device /dev/dri \
 -v <path/to/model>:/app/model \
-vllm-omni-rocm \
-bash
+-v ~/.cache/huggingface:/root/.cache/huggingface \
+--entrypoint bash \
+vllm-omni-rocm
 ```

 # --8<-- [end:build-docker]

 # --8<-- [start:pre-built-images]

+vLLM-Omni offers an official docker image for deployment. These images are built on top of vLLM docker images and available on Docker Hub as [vllm/vllm-omni-rocm](https://hub.docker.com/r/vllm/vllm-omni-rocm/tags). The version of vLLM-Omni indicates which release of vLLM it is based on.
+
+#### Launch vLLM-Omni Server
+Here's an example deployment command that has been verified on 2 x MI300's:
+```bash
+docker run --rm \
+  --group-add=video \
+  --ipc=host \
+  --cap-add=SYS_PTRACE \
+  --security-opt seccomp=unconfined \
+  --device /dev/kfd \
+  --device /dev/dri \
+  -v <path/to/model>:/app/model \
+  -v ~/.cache/huggingface:/root/.cache/huggingface \
+  --env "HF_TOKEN=$HF_TOKEN" \
+  -p 8091:8091 \
+  vllm/vllm-omni-rocm:v0.14.0 \
+  --model Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
+```
+
+#### Launch an interactive terminal with prebuilt docker image.
+If you want to run in dev environment you can launch the docker image as follows:
+```bash
+docker run --rm -it \
+  --network=host \
+  --group-add=video \
+  --ipc=host \
+  --cap-add=SYS_PTRACE \
+  --security-opt seccomp=unconfined \
+  --device /dev/kfd \
+  --device /dev/dri \
+  -v <path/to/model>:/app/model \
+  -v ~/.cache/huggingface:/root/.cache/huggingface \
+  --env "HF_TOKEN=$HF_TOKEN" \
+  --entrypoint bash \
+  vllm/vllm-omni-rocm:v0.14.0
+```
+
 # --8<-- [end:pre-built-images]
@@ -16,7 +16,13 @@ vLLM-Omni supports NPU through the vLLM Ascend Plugin (vllm-ascend). This is a c

 ## Installation

-### Recommended
+### Set up using Docker
+
+=== "NPU"
+
+    --8<-- "docs/getting_started/installation/npu/npu.inc.md:pre-built-images"
+
+### Build wheel from source

 === "NPU"

@@ -13,10 +13,10 @@ export DEVICE0=/dev/davinci0
 export DEVICE1=/dev/davinci1
 # Update the vllm-ascend image
 # Atlas A2:
-# export IMAGE=quay.io/ascend/vllm-ascend:v0.11.0rc2
+# export IMAGE=quay.io/ascend/vllm-ascend:v0.14.0rc1
 # Atlas A3:
-# export IMAGE=quay.io/ascend/vllm-ascend:v0.11.0rc2-a3
-export IMAGE=quay.io/ascend/vllm-ascend:v0.11.0rc2
+# export IMAGE=quay.io/ascend/vllm-ascend:v0.14.0rc1-a3
+export IMAGE=quay.io/ascend/vllm-ascend:v0.14.0rc1
 docker run --rm \
    --name vllm-omni-npu \
    --shm-size=1g \
@@ -36,9 +36,11 @@ docker run --rm \

 # Inside the container, install vLLM-Omni from source
 cd /vllm-workspace
-git clone -b v0.11.0rc1 https://github.com/vllm-project/vllm-omni.git
+git clone -b v0.14.0 https://github.com/vllm-project/vllm-omni.git
+
 cd vllm-omni
-pip install -v -e .
+VLLM_OMNI_TARGET_DEVICE=npu pip install -v -e .
+# OR pip install -v -e . --no-build-isolation
 export VLLM_WORKER_MULTIPROC_METHOD=spawn
 ```

@@ -46,4 +48,65 @@ The default workdir is `/workspace`, with vLLM, vLLM-Ascend and vLLM-Omni code p

 For other installation methods (pip installation, building from source, custom Docker builds), please refer to the [vllm-ascend installation guide](https://docs.vllm.ai/projects/ascend/en/latest/installation.html).

+We are keeping [issue #997](https://github.com/vllm-project/vllm-omni/issues/997) up to date with the aligned versions of vLLM, vLLM-Ascend, and vLLM-Omni, and also outlining the Q1 roadmap there.
+
 # --8<-- [end:installation]
+
+# --8<-- [start:pre-built-images]
+
+`vllm-ascend` offers Docker images for deployment. You can just pull the **prebuilt image** from the image repository [ascend/vllm-ascend](https://quay.io/repository/ascend/vllm-ascend?tab=tags) and run it with bash.
+
+Supported images as following.
+
+| image name | Hardware | OS |
+|-|-|-|
+| image-tag | Atlas A2 | Ubuntu |
+| image-tag-a3 | Atlas A3 | Ubuntu |
+
+vLLM-Omni offers Docker images for Ascend NPU deployment. You can just pull the **prebuilt image** from the image repository [ascend/vllm-ascend](https://quay.io/repository/ascend/vllm-ascend?tab=tags) and run it with bash.
+
+Here's an example deployment command that has been verified on 6 x NPUs:
+
+```bash
+# Atlas A2:
+# export IMAGE=quay.io/ascend/vllm-omni:v0.14.0
+# Atlas A3:
+# export IMAGE=quay.io/ascend/vllm-omni:v0.14.0-a3
+export IMAGE=quay.io/ascend/vllm-omni:v0.14.0
+docker run --rm \
+    --name vllm-omni-npu \
+    --shm-size=1g \
+    --device /dev/davinci0 \
+    --device /dev/davinci1 \
+    --device /dev/davinci2 \
+    --device /dev/davinci3 \
+    --device /dev/davinci4 \
+    --device /dev/davinci5 \
+    --device /dev/davinci_manager \
+    --device /dev/devmm_svm \
+    --device /dev/hisi_hdc \
+    -v /usr/local/dcmi:/usr/local/dcmi \
+    -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
+    -v /usr/local/Ascend/driver/lib64/:/usr/local/Ascend/driver/lib64/ \
+    -v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info \
+    -v /etc/ascend_install.info:/etc/ascend_install.info \
+    -v /root/.cache:/root/.cache \
+    -p 8000:8000 \
+    -it $IMAGE bash
+```
+
+!!! tip
+    You can use this docker image to serve models the same way you would with in vLLM! To do so, make sure you overwrite the default entrypoint (`vllm serve --omni`) which works only for models supported in the vLLM-Omni project.
+
+Or build IMAGE from **source code**:
+
+```bash
+git clone https://github.com/vllm-project/vllm-omni.git
+cd vllm-omni
+# A2
+docker build -t vllm-omni-dev-image:latest -f ./docker/Dockerfile.npu .
+# A3
+# docker build -t vllm-omni-dev-image:latest -f ./docker/Dockerfile.npu.a3 .
+```
+
+# --8<-- [end:pre-built-images]
@@ -17,7 +17,13 @@ For installation on GPU from source:
 ```bash
 uv venv --python 3.12 --seed
 source .venv/bin/activate
-uv pip install vllm==0.12.0 --torch-backend=auto
+
+# On CUDA
+uv pip install vllm==0.16.0 --torch-backend=auto
+
+# On ROCm
+uv pip install vllm==0.16.0 --extra-index-url https://wheels.vllm.ai/rocm/0.16.0/rocm700
+
 git clone https://github.com/vllm-project/vllm-omni.git
 cd vllm-omni
 uv pip install -e .
@@ -40,6 +46,46 @@ if __name__ == "__main__":
    images[0].save("coffee.png")
 ```

+You can pass a list of prompts and wait for them to process altogether, shown below.
+
+!!! info
+
+    However, it is not currently recommended to do so
+    because not all models support batch inference,
+    and batch requesting mostly does not provide significant performance improvement (despite the impression that it does).
+    This feature is primarily for the sake of interface compatibility with vLLM and to allow for future improvements.
+
+```python
+from vllm_omni.entrypoints.omni import Omni
+
+if __name__ == "__main__":
+    omni = Omni(
+        model="Tongyi-MAI/Z-Image-Turbo",
+        # stage_configs_path="./stage-config.yaml",  # See below
+    )
+    prompts = [
+        "a cup of coffee on a table",
+        "a toy dinosaur on a sandy beach",
+        "a fox waking up in bed and yawning",
+    ]
+    omni_outputs = omni.generate(prompts)
+    for i_prompt, prompt_output in enumerate(omni_outputs):
+        this_request_output = prompt_output.request_output[0]
+        this_images = this_request_output.images
+        for i_image, image in enumerate(this_images):
+            image.save(f"p{i_prompt}-img{i_image}.jpg")
+            print("saved to", f"p{i_prompt}-img{i_image}.jpg")
+            # saved to p0-img0.jpg
+            # saved to p1-img0.jpg
+            # saved to p2-img0.jpg
+```
+
+!!! info
+
+    For diffusion pipelines, the stage config field `stage_args.[].runtime.max_batch_size` is 1 by default, and the input
+    list is sliced into single-item requests before feeding into the diffusion pipeline. For models that do internally support
+    batched inputs, you can [modify this configuration](../configuration/stage_configs.md) to let the model accept a longer batch of prompts.
+
 For more usages, please refer to [offline inference](../user_guide/examples/offline_inference/qwen2_5_omni.md)

 ## Online Serving with OpenAI-Completions API
@@ -34,10 +34,10 @@ CATEGORIES = {
        "name": "Core",
        "description": "Core scheduling and caching components.",
    },
-    "model_executor": {
-        "name": "Model Executor",
-        "description": "Model execution components.",
-    },
+    # "model_executor": {
+    #     "name": "Model Executor",
+    #     "description": "Model execution components.",
+    # },
    "config": {
        "name": "Configuration",
        "description": "Configuration classes.",
@@ -20,10 +20,13 @@ th {
 | `Qwen3OmniMoeForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct` |
 | `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B` |
 | `BagelForConditionalGeneration` | BAGEL (DiT-only) | `ByteDance-Seed/BAGEL-7B-MoT` |
+| `HunyuanImage3ForCausalMM` | HunyuanImage3.0 (DiT-only) | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` |
 | `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` |
+| `QwenImagePipeline` | Qwen-Image-2512 | `Qwen/Qwen-Image-2512` |
 | `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` |
+| `GlmImagePipeline` | GLM-Image | `zai-org/GLM-Image` |
 |`ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` |
 | `WanPipeline` | Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` |
 | `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` |
@@ -31,6 +34,12 @@ th {
 |`LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` |
 |`LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` |
 |`StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` |
+|`Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` |
+|`FluxPipeline` | FLUX.1-dev | `black-forest-labs/FLUX.1-dev` |
+|`StableAudioPipeline` | Stable-Audio-Open | `stabilityai/stable-audio-open-1.0` |
+|`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` |
+|`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` |
+|`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-0.6B-Base` |


 ## List of Supported Models for NPU
@@ -44,6 +53,7 @@ th {

 | Architecture | Models | Example HF Models |
 |--------------|--------|-------------------|
+| `Qwen3OmniMoeForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct` |
 | `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B`|
 | `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` |
 | `QwenImagePipeline` | Qwen-Image-2512 | `Qwen/Qwen-Image-2512` |
@@ -52,3 +62,8 @@ th {
 | `QwenImageLayeredPipeline` | Qwen-Image-Layered | `Qwen/Qwen-Image-Layered` |
 | `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2511 | `Qwen/Qwen-Image-Edit-2511` |
 |`ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` |
+|`LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` |
+|`Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` |
+|`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` |
+|`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` |
+|`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-0.6B-Base` |
@@ -0,0 +1,205 @@
+# Image Edit API
+
+vLLM-Omni provides an OpenAI DALL-E compatible API for image edit using diffusion models.
+
+Each server instance runs a single model (specified at startup via `vllm serve <model> --omni`).
+
+## Quick Start
+
+### Start the Server
+
+For example...
+
+```bash
+# Qwen-Image
+vllm serve Qwen/Qwen-Image-Edit-2511 --omni --port 8000
+
+
+### Generate Images
+
+**Using curl:**
+
+```bash
+curl -s -D >(grep -i x-request-id >&2) \
+  -o >(jq -r '.data[0].b64_json' | base64 --decode > gift-basket.png) \
+  -X POST "http://localhost:8000/v1/images/edits" \
+  -F "model=xxx" \
+  -F "image=@./xx.png" \
+  -F "prompt='this bear is wearing sportwear. holding a basketball, and bending one leg.'" \
+  -F "size=1024x1024" \
+  -F "output_format=png"
+```
+
+
+**Using OpenAI SDK:**
+
+```python
+import base64
+from openai import OpenAI
+from pathlib import Path
+client = OpenAI(
+    api_key="None",
+    base_url="http://localhost:8000/v1"
+)
+
+input_image_url = "https://vllm-public-assets.s3.us-west-2.amazonaws.com/omni-assets/qwen-bear.png"
+
+result = client.images.edit(
+    image=[],
+    model="Qwen-Image-Edit-2511",
+    prompt="Change the bears in the two input images into walking together.",
+    size='512x512',
+    stream=False,
+    output_format='jpeg',
+    # url格式
+    extra_body={
+        "url": [input_image_url1,input_image_url],
+        "num_inference_steps": 50,
+        "guidance_scale": 1,
+        "seed": 777,
+    }
+)
+
+image_base64 = result.data[0].b64_json
+image_bytes = base64.b64decode(image_base64)
+
+# Save the image to a file
+with open("edit_out_http.jpeg", "wb") as f:
+    f.write(image_bytes)
+```
+
+## API Reference
+
+### Endpoint
+
+```
+POST /v1/images/edits
+Content-Type: multipart/form-data
+```
+
+### Request Parameters
+
+#### OpenAI Standard Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `prompt` | string | **required** | A text description of the desired image |
+| `model` | string | server's model | Model to use (optional, should match server if specified) |
+| `image` | string or array | **required** | The image(s) to edit. |
+| `n` | integer | 1 | Number of images to generate (1-10) |
+| `size` | string | "auto" | Image dimensions in WxH format (e.g., "1024x1024", "512x512"), when set to auto, it decide size from first input image. |
+| `response_format` | string | "b64_json" | Response format (only "b64_json" supported) |
+| `user` | string | null | User identifier for tracking |
+| `output_format` | string | "png" | The format in which the generated images are returned. Must be one of "png", "jpg", "jpeg", "webp". |
+| `output_compression` | integer | 100 | The compression level (0-100%) for the generated images. |
+| `background` | string or null | "auto" | Allows to set transparency for the background of the generated image(s).
+
+#### vllm-omni Extension Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `url` | string or array | None | The image(s) to edit. |
+| `negative_prompt` | string | null | Text describing what to avoid in the image |
+| `num_inference_steps` | integer | model defaults | Number of diffusion steps |
+| `guidance_scale` | float | model defaults | Classifier-free guidance scale (typically 0.0-20.0) |
+| `true_cfg_scale` | float | model defaults | True CFG scale (model-specific parameter, may be ignored if not supported) |
+| `seed` | integer | null | Random seed for reproducibility |
+
+### Response Format
+
+```json
+{
+  "created": 1701234567,
+  "data": [
+    {
+      "b64_json": "<base64-encoded PNG>",
+      "url": null,
+      "revised_prompt": null
+    }
+  ],
+  "output_format": null,
+  "size": null,
+}
+```
+
+## Examples
+
+### Multiple Images input
+
+```bash
+curl -s -D >(grep -i x-request-id >&2) \
+  -o >(jq -r '.data[0].b64_json' | base64 --decode > gift-basket.png) \
+  -X POST "http://localhost:8000/v1/images/edits" \
+  -F "model=xxx" \
+  -F "image=@xx.png" \
+  -F "image=@xx.png"
+  -F "prompt='this bear is wearing sportwear. holding a basketball, and bending one leg.'" \
+  -F "size=1024x1024" \
+  -F "output_format=png"
+```
+
+
+## Parameter Handling
+
+The API passes parameters directly to the diffusion pipeline without model-specific transformation:
+
+- **Default values**: When parameters are not specified, the underlying model uses its own defaults
+- **Pass-through design**: User-provided values are forwarded directly to the diffusion engine
+- **Minimal validation**: Only basic type checking and range validation at the API level
+
+### Parameter Compatibility
+
+The API passes parameters directly to the diffusion pipeline without model-specific validation.
+
+- Unsupported parameters may be silently ignored by the model
+- Incompatible values will result in errors from the underlying pipeline
+- Recommended values vary by model - consult model documentation
+
+**Best Practice:** Start with the model's recommended parameters, then adjust based on your needs.
+
+## Error Responses
+
+### 400 Bad Request
+
+Invalid parameters (e.g., model mismatch):
+
+```json
+{
+  "detail": "Invalid size format: '1024x'. Expected format: 'WIDTHxHEIGHT' (e.g., '1024x1024')."
+}
+```
+
+### 422 Unprocessable Entity
+
+Validation errors (missing required fields):
+
+```json
+{
+  "detail": "Field 'image' or 'url' is required"
+}
+```
+
+## Troubleshooting
+
+### Server Not Running
+
+```bash
+# Check if server is responding
+curl -X http://localhost:8000/v1/images/edit \
+  -F "prompt='test'"
+```
+
+### Out of Memory
+
+If you encounter OOM errors:
+1. Reduce image size: `"size": "512x512"`
+2. Reduce inference steps: `"num_inference_steps": 25`
+
+## Development
+
+Enable debug logging to see prompts and generation details:
+
+```bash
+vllm serve Qwen/Qwen-Image-Edit-2511 --omni \
+  --uvicorn-log-level debug
+```
@@ -0,0 +1,267 @@
+# Speech API
+
+vLLM-Omni provides an OpenAI-compatible API for text-to-speech (TTS) generation using Qwen3-TTS models.
+
+Each server instance runs a single model (specified at startup via `vllm serve <model> --omni`).
+
+## Quick Start
+
+### Start the Server
+
+```bash
+# CustomVoice model (predefined speakers)
+vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
+    --stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts.yaml \
+    --omni \
+    --port 8091 \
+    --trust-remote-code \
+    --enforce-eager
+```
+
+### Generate Speech
+
+**Using curl:**
+
+```bash
+curl -X POST http://localhost:8091/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Hello, how are you?",
+        "voice": "vivian",
+        "language": "English"
+    }' --output output.wav
+```
+
+**Using Python:**
+
+```python
+import httpx
+
+response = httpx.post(
+    "http://localhost:8091/v1/audio/speech",
+    json={
+        "input": "Hello, how are you?",
+        "voice": "vivian",
+        "language": "English",
+    },
+    timeout=300.0,
+)
+
+with open("output.wav", "wb") as f:
+    f.write(response.content)
+```
+
+**Using OpenAI SDK:**
+
+```python
+from openai import OpenAI
+
+client = OpenAI(base_url="http://localhost:8091/v1", api_key="none")
+
+response = client.audio.speech.create(
+    model="Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
+    voice="vivian",
+    input="Hello, how are you?",
+)
+
+response.stream_to_file("output.wav")
+```
+
+## API Reference
+
+### Endpoint
+
+```
+POST /v1/audio/speech
+Content-Type: application/json
+```
+
+### Request Parameters
+
+#### OpenAI Standard Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `input` | string | **required** | The text to synthesize into speech |
+| `model` | string | server's model | Model to use (optional, should match server if specified) |
+| `voice` | string | "vivian" | Speaker name (e.g., vivian, ryan, aiden) |
+| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus |
+| `speed` | float | 1.0 | Playback speed (0.25-4.0) |
+
+#### vLLM-Omni Extension Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `task_type` | string | "CustomVoice" | TTS task type: CustomVoice, VoiceDesign, or Base |
+| `language` | string | "Auto" | Language (see supported languages below) |
+| `instructions` | string | "" | Voice style/emotion instructions |
+| `max_new_tokens` | integer | 2048 | Maximum tokens to generate |
+
+**Supported languages:** Auto, Chinese, English, Japanese, Korean, German, French, Russian, Portuguese, Spanish, Italian
+
+#### Voice Clone Parameters (Base task)
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `ref_audio` | string | null | Reference audio (URL or base64 data URL) |
+| `ref_text` | string | null | Transcript of reference audio |
+| `x_vector_only_mode` | bool | null | Use speaker embedding only (no ICL) |
+
+### Response Format
+
+Returns binary audio data with appropriate `Content-Type` header (e.g., `audio/wav`).
+
+### Voices Endpoint
+
+```
+GET /v1/audio/voices
+```
+
+Lists available voices for the loaded model.
+
+```json
+{
+    "voices": ["aiden", "dylan", "eric", "ono_anna", "ryan", "serena", "sohee", "uncle_fu", "vivian"]
+}
+```
+
+## Examples
+
+### CustomVoice with Style Instruction
+
+```bash
+curl -X POST http://localhost:8091/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "I am so excited!",
+        "voice": "vivian",
+        "instructions": "Speak with great enthusiasm"
+    }' --output excited.wav
+```
+
+### VoiceDesign (Natural Language Voice Description)
+
+```bash
+# Start server with VoiceDesign model first
+vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign \
+    --stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts.yaml \
+    --omni \
+    --port 8091 \
+    --trust-remote-code \
+    --enforce-eager
+```
+
+```bash
+curl -X POST http://localhost:8091/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Hello world",
+        "task_type": "VoiceDesign",
+        "instructions": "A warm, friendly female voice with a gentle tone"
+    }' --output designed.wav
+```
+
+### Base (Voice Cloning)
+
+```bash
+# Start server with Base model first
+vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base \
+    --stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts.yaml \
+    --omni \
+    --port 8091 \
+    --trust-remote-code \
+    --enforce-eager
+```
+
+```bash
+curl -X POST http://localhost:8091/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Hello, this is a cloned voice",
+        "task_type": "Base",
+        "ref_audio": "https://example.com/reference.wav",
+        "ref_text": "Original transcript of the reference audio"
+    }' --output cloned.wav
+```
+
+## Supported Models
+
+| Model | Task Type | Description |
+|-------|-----------|-------------|
+| `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | CustomVoice | Predefined speaker voices with optional style control |
+| `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | VoiceDesign | Natural language voice style description |
+| `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | Base | Voice cloning from reference audio |
+| `Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice` | CustomVoice | Smaller/faster variant |
+| `Qwen/Qwen3-TTS-12Hz-0.6B-Base` | Base | Smaller/faster variant for voice cloning |
+
+## Error Responses
+
+### 400 Bad Request
+
+Invalid parameters:
+
+```json
+{
+    "error": {
+        "message": "Input text cannot be empty",
+        "type": "BadRequestError",
+        "param": null,
+        "code": 400
+    }
+}
+```
+
+### 404 Not Found
+
+Model not found:
+
+```json
+{
+    "error": {
+        "message": "The model `xxx` does not exist.",
+        "type": "NotFoundError",
+        "param": "model",
+        "code": 404
+    }
+}
+```
+
+## Troubleshooting
+
+### "TTS model did not produce audio output"
+
+Ensure you're using the correct model variant for your task type:
+- CustomVoice task → CustomVoice model
+- VoiceDesign task → VoiceDesign model
+- Base task → Base model
+
+### Server Not Running
+
+```bash
+# Check if server is responding
+curl http://localhost:8091/v1/audio/voices
+```
+
+### Out of Memory
+
+If you encounter OOM errors:
+1. Use smaller model variant: `Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice`
+2. Reduce `--gpu-memory-utilization`
+
+### Unsupported Speaker
+
+Use `/v1/audio/voices` to list available voices for the loaded model.
+
+## Development
+
+Enable debug logging:
+
+```bash
+vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
+    --stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts.yaml \
+    --omni \
+    --port 8091 \
+    --trust-remote-code \
+    --enforce-eager \
+    --uvicorn-log-level debug
+```
@@ -22,7 +22,7 @@ A: At first, you can check current [issues](https://github.com/vllm-project/vllm

 > Q: Does vLLM-Omni support AWQ or any other quantization?

-A: vLLM-Omni partitions model into several stages. For AR stages, it will reuse main logic of LLMEngine in vLLM. So current quantization supported in vLLM should be also supported in vLLM-Omni for them. But systematic verification is ongoing. For quantization for DiffusionEngine, we are working on it. Please stay tuned and welcome contribution!
+A: We plan to introduce GGUF FP8 prequantized models and online FP8 quantization in version 0.16.0. Support for other quantization types will follow in future releases. For details, please see our [Q1 quantization roadmap](https://github.com/vllm-project/vllm-omni/issues/1057).

 > Q: Does vLLM-Omni support multimodal streaming input and output?

@@ -1,274 +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
-
-
-The following table shows which models are currently supported by parallelism method:
-
-### ImageGen
-
-| Model | Model Identifier | Ulysses-SP | Ring-SP |
-|-------|------------------|-----------|---------|
-| **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-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` | ❌ | ❌ |
-
-### VideoGen
-
-| Model | Model Identifier | Ulysses-SP | Ring-SP |
-|-------|------------------|-----------|---------|
-| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ❌ | ❌ |
-
-### Sequence Parallelism
-
-#### Ulysses-SP
-
-##### 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.diffusion.data import DiffusionParallelConfig
-ulysses_degree = 2
-
-omni = Omni(
-    model="Qwen/Qwen-Image",
-    parallel_config=DiffusionParallelConfig(ulysses_degree=2)
-)
-
-outputs = omni.generate(prompt="A cat sitting on a windowsill", 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 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
-```
-
-##### 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.diffusion.data import DiffusionParallelConfig
-ring_degree = 2
-
-omni = Omni(
-    model="Qwen/Qwen-Image",
-    parallel_config=DiffusionParallelConfig(ring_degree=2)
-)
-
-outputs = omni.generate(prompt="A cat sitting on a windowsill", 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`.
-
-##### Offline Inference
-
-```python
-from vllm_omni import Omni
-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(prompt="A cat sitting on a windowsill", 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 |
-
-
-##### How to parallelize a new model
-
-If a diffusion model has been deployed in vLLM-Omni and supports single-card inference, you can refer to the following instructions to parallelize it with [Ulysses-SP](https://arxiv.org/pdf/2309.14509).
-
-This section uses **Qwen-Image** (`QwenImageTransformer2DModel`) as the reference implementation. Qwen-Image is a **dual-stream** transformer (text + image) that performs **joint attention** across the concatenated sequences. Because of that, when enabling sequence parallel you typically:
-
- Chunk **image tokens** (`hidden_states`) across SP ranks along the **sequence dimension**.
- Keep **text embeddings** (`encoder_hidden_states`) **replicated** on all SP ranks for correctness (unless you implement full joint-SP semantics and explicitly split text too).
- Chunk **image RoPE freqs** to match the chunked image tokens.
- `all_gather` the final output back along the sequence dimension.
-
-First, add the sequence-parallel helpers and (for Qwen-Image) the forward-context flag. Then, in the transformer's `forward()`:
-
- Chunk `hidden_states` by SP world size.
- Set `get_forward_context().split_text_embed_in_sp = False` because Qwen-Image uses joint attention and we would use a full text embedding in each rank.
- After `pos_embed`, chunk `img_freqs` on `dim=0` (token axis) to match the chunked image tokens; only chunk `txt_freqs` if you explicitly split text embeddings in SP.
-
-Taking `vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py` as an example:
-
-```diff
-from vllm_omni.diffusion.distributed.parallel_state import (
-    get_sequence_parallel_rank,
-    get_sequence_parallel_world_size,
-    get_sp_group,
-)
-from vllm_omni.diffusion.forward_context import get_forward_context
-
-class QwenImageTransformer2DModel(...):
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor = None,
-        encoder_hidden_states_mask: torch.Tensor = None,
-        timestep: torch.LongTensor = None,
-        img_shapes: list[tuple[int, int, int]] | None = None,
-        txt_seq_lens: list[int] | None = None,
-        ...
-    ):
-+       if self.parallel_config.sequence_parallel_size > 1:
-+           # Chunk image tokens along sequence dimension.
-+           hidden_states = torch.chunk(
-+               hidden_states,
-+               get_sequence_parallel_world_size(),
-+               dim=-2,
-+           )[get_sequence_parallel_rank()]
-+
-+           # Qwen-Image uses *dual-stream* (text + image) and runs *joint attention*.
-+           # Text embeddings should be replicated across SP ranks for correctness.
-+           get_forward_context().split_text_embed_in_sp = False
-
-        hidden_states = self.img_in(hidden_states)
-
-        ...
-        image_rotary_emb = self.pos_embed(img_shapes, txt_seq_lens, device=hidden_states.device)
-
-+       def get_rotary_emb_chunk(freqs):
-+           return torch.chunk(freqs, get_sequence_parallel_world_size(), dim=0)[
-+               get_sequence_parallel_rank()
-+           ]
-+
-+       if self.parallel_config.sequence_parallel_size > 1:
-+           img_freqs, txt_freqs = image_rotary_emb
-+           img_freqs = get_rotary_emb_chunk(img_freqs)
-+           if get_forward_context().split_text_embed_in_sp:
-+               txt_freqs = get_rotary_emb_chunk(txt_freqs)
-+           image_rotary_emb = (img_freqs, txt_freqs)
-```
-
-Next, at the end of the `forward()` function, call `get_sp_group().all_gather` to gather the chunked outputs across devices and concatenate them along the sequence dimension (matching the earlier `dim=-2` chunking):
-
-```diff
-class QwenImageTransformer2DModel(...):
-    def forward(...):
-        ...
-        hidden_states = self.norm_out(hidden_states, temb)
-        output = self.proj_out(hidden_states)
-
-+       if self.parallel_config.sequence_parallel_size > 1:
-+           output = get_sp_group().all_gather(output, dim=-2)
-        return Transformer2DModelOutput(sample=output)
-```
-
-Finally, you can set the parallel configuration and pass it to `Omni` and start parallel inference with:
-```diff
-from vllm_omni import Omni
-+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(prompt="A cat sitting on a windowsill", num_inference_steps=50)
-```
@@ -18,6 +18,7 @@ Enable cache-dit acceleration by simply setting `cache_backend="cache_dit"`. Cac

 ```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(
@@ -27,7 +28,7 @@ omni = Omni(

 images = omni.generate(
    "a beautiful landscape",
-    num_inference_steps=50,
+    OmniDiffusionSamplingParams(num_inference_steps=50),
 )
 ```

@@ -69,7 +70,7 @@ vllm serve Qwen/Qwen-Image --omni --port 8091 \

 ## Acceleration Methods

-For comprehensive illustration, please view cache-dit [User_Guide](https://github.com/vipshop/cache-dit/blob/main/docs/User_Guide.md)
+For comprehensive illustration, please view cache-dit [User_Guide](https://cache-dit.readthedocs.io/en/latest/user_guide/OVERVIEWS/)

 ### 1. DBCache (Dual Block Cache)

@@ -187,8 +188,8 @@ 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
+    --cache-backend cache_dit \
+    --num-inference-steps 50
 ```


@@ -222,6 +223,6 @@ To test another model, you can modify `--model` with the target model identifier

 ## Additional Resources

- [Cache-DiT User Guide](https://github.com/vipshop/cache-dit/blob/main/docs/User_Guide.md)
- [Cache-DiT Benchmark](https://github.com/vipshop/cache-dit/tree/main/bench)
- [DBCache Technical Details](https://github.com/vipshop/cache-dit/blob/main/docs/DBCache.md)
+- [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,143 @@
+# CPU Offloading for Diffusion Models
+
+## Overview
+
+vLLM-Omni provides two offloading strategies to reduce GPU memory usage for diffusion models:
+
+1. **Model-level (Sequential) Offloading**: Mutual exclusion between DiT model and encoder - only one is on GPU at a time.
+2. **Layerwise (Blockwise) Offloading**: Keeps only one transformer block on GPU at a time with compute-memory overlap.
+
+Both strategies use pinned memory for faster CPU-GPU transfers. The strategies are **mutually exclusive** for now - if both are enabled, layerwise takes priority.
+
+
+## Model-level (Sequential) Offloading
+
+### How It Works
+
+Model-level offloading implements mutual exclusion between DiT transformer and encoder modules using pre forward hooks:
+
+- **When encoders run**: DiT transformer is offloaded to CPU
+- **When DiT runs**: Encoders are offloaded to CPU
+- **VAE**: Stays resident on GPU
+
+Before each module's forward pass, the hook automatically moves it to GPU while offloading the other module group to CPU. Transfers use pinned memory for speed.
+
+### Usage
+
+**Python API:**
+```python
+from vllm_omni import Omni
+
+m = Omni(model="Wan-AI/Wan2.2-T2V-A14B-Diffusers", enable_cpu_offload=True)
+```
+
+**CLI:**
+```bash
+vllm-omni serve diffusion Wan-AI/Wan2.2-T2V-A14B-Diffusers --enable-cpu-offload
+```
+
+### Limitations
+- Cold start latency increases
+- Adds overhead from CPU-GPU transfers between encoder and denoising phases
+- Support single GPU only for now
+
+
+## Layerwise (Blockwise) Offloading
+
+### How It Works
+
+Layerwise offloading keeps only one transformer block on GPU at a time.
+
+As each block completes, the next block is prefetched to GPU while the current block is freed. The pre and forward hooks utilized by layerwise offloading apply a separate CUDA stream (`copy_stream`) to overlap weight transfer with computation, and retain flattened tensors in pinned CPU memory for block parameters re-materialization. Encoders, VAE, and non-block DiT modules (embeddings, norms) always stay on GPU.
+
+**Execution Flow:**
+
+| Block | Pre-forward Hook | Forward | Post-forward Hook |
+|-------|------------------|---------|-------------------|
+| block-0 | Prefetch block-1 (async) | Compute block-0 | Free block-0 |
+| block-1 | Prefetch block-2 (async) | Compute block-1 | Free block-1 |
+| ... | ... | ... | ... |
+| block-(n-1) | **Prefetch block-0** (async) | Compute block-(n-1) | Free block-(n-1) |
+
+Each transformer block has a `LayerwiseOffloadHook` that prefetches the next block before forward and frees the current block after forward.
+
+Layerwise offloading is primarily recommended for large **video generation models** where the compute cost per block is high enough to effectively overlap with memory prefetch operations. For example, Wan2.2 T2V and I2V pipelines.
+
+### Usage
+
+**Python API:**
+```python
+from vllm_omni import Omni
+
+# Text-to-video
+m = Omni(model="Wan-AI/Wan2.2-T2V-A14B-Diffusers", enable_layerwise_offload=True)
+
+# Or image-to-video
+m = Omni(model="Wan-AI/Wan2.2-I2V-A14B-Diffusers", enable_layerwise_offload=True)
+```
+
+**CLI:**
+```bash
+# Text-to-video
+vllm-omni serve diffusion Wan-AI/Wan2.2-T2V-A14B-Diffusers --enable-layerwise-offload
+
+# Or image-to-video
+vllm-omni serve diffusion Wan-AI/Wan2.2-I2V-A14B-Diffusers --enable-layerwise-offload
+```
+
+### To Support a Model
+
+Models must define the blocks attribute name for layerwise offloading:
+
+```python
+class WanTransformer3DModel(nn.Module):
+    _layerwise_offload_blocks_attr = "blocks"  # Attribute name containing transformer blocks
+
+    def __init__(self):
+        self.blocks = nn.ModuleList([...])  # Transformer blocks
+```
+
+### Limitations
+- Cold start latency increases because of
+    1) components are loaded to CPU first at the very first during initialization,
+    2) weight consolidation and pinning
+- Performance depends on compute cost and H2D bandwidth as well
+- Support single GPU only for now
+
+
+### Implementation Notes
+
+**Module Discovery**
+
+The offloader automatically discovers pipeline components:
+
+- **DiT modules**: `transformer`, `transformer_2`, `dit`
+- **Encoders**: `text_encoder`, `text_encoder_2`, `text_encoder_3`, `image_encoder`
+- **VAE**: `vae`
+
+**Hook System**
+
+Both strategies use vLLM-Omni's hook registry system (`HookRegistry` and `ModelHook`) to register pre/post forward callbacks on modules, enabling automatic swapping without modifying model code.
+
+**Backend Architecture**
+
+```
+OffloadBackend (base class)
+├── ModelLevelOffloadBackend → uses SequentialOffloadHook
+└── LayerWiseOffloadBackend → uses LayerwiseOffloadHook
+```
+
+Factory function `get_offload_backend()` selects the appropriate backend based on configuration.
+
+
+## Supported Models
+
+| Architecture | Example Models | DiT Class | Model-Level Offload | Layerwise Offload | Blocks Attr (Layerwise specific) |
+|--------------|----------------|-----------|---------------------|-------------------|-------------|
+| Wan22Pipeline | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | `"blocks"` |
+| Wan22I2VPipeline | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | `"blocks"` |
+| QwenImagePipeline | `Qwen/Qwen-Image` | `QwenImageTransformer2DModel` | ✓ | ✓ | `"transformer_blocks"` |
+
+**Notes:**
+- Model-Level Offloading is expected to be supported by all common diffusion models (DiT and encoders) naturally
+- Layerwise Offloading requires DiT class to define `_layerwise_offload_blocks_attr` pointing to transformer blocks
@@ -0,0 +1,332 @@
+# 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.
+
+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 |
+|--------------------------|--------------------------------------|:----------:|:-------:|:------------:|:---------------:|:------------------:|
+| **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-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`   |     ❌      |    ❌    |      ❌       |        ✅        |         ❌          |
+| **FLUX.2-klein**         | `black-forest-labs/FLUX.2-klein-4B`  |     ❌      |    ❌    |      ❌       |        ✅        |         ❌          |
+| **FLUX.1-dev**           | `black-forest-labs/FLUX.1-dev`       |     ❌      |    ❌    |      ✅       |        ✅        |         ❌          |
+
+!!! 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 |
+|-------|------------------|------------|---------|--------------------------|
+| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ✅ | ✅ | ✅ |
+
+### 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** (currently: `Tongyi-MAI/Z-Image-Turbo`).
+    - 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
+
+##### 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.
+
+##### 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
+```
+
+##### 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`.
+
+##### 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
+```
				`@@ -0,0 +1 @@`
				`# noqa: N999 # This is not a python library intended to be imported`
				`@@ -0,0 +1 @@`
				{"id":"1c99f525-0a37-45ba-a28a-7df7c3af66b4","revision":0,"last_node_id":12,"last_link_id":14,"nodes":[{"id":1,"type":"VLLMOmniComprehension","pos":[1191.2177053682556,144.66829928181377],"size":[400,268],"flags":{},"order":8,"mode":0,"inputs":[{"localized_name":"image","name":"image","shape":7,"type":"IMAGE","link":1},{"localized_name":"video","name":"video","shape":7,"type":"VIDEO","link":2},{"localized_name":"audio","name":"audio","shape":7,"type":"AUDIO","link":3},{"localized_name":"sampling_params","name":"sampling_params","shape":7,"type":"SAMPLING_PARAMS","link":14},{"localized_name":"url","name":"url","type":"STRING","widget":{"name":"url"},"link":null},{"localized_name":"model","name":"model","type":"STRING","widget":{"name":"model"},"link":null},{"localized_name":"prompt","name":"prompt","type":"STRING","widget":{"name":"prompt"},"link":null},{"localized_name":"output_text","name":"output_text","type":"BOOLEAN","widget":{"name":"output_text"},"link":null},{"localized_name":"output_audio","name":"output_audio","type":"BOOLEAN","widget":{"name":"output_audio"},"link":null},{"localized_name":"use_audio_in_video","name":"use_audio_in_video","type":"BOOLEAN","widget":{"name":"use_audio_in_video"},"link":null}],"outputs":[{"localized_name":"text_response","name":"text_response","type":"STRING","links":[8]},{"localized_name":"audio_response","name":"audio_response","type":"AUDIO","links":[9]}],"properties":{"Node name for S&R":"VLLMOmniComprehension"},"widgets_values":["http://localhost:8000/v1","Qwen/Qwen2.5-Omni-7B","",true,true,true]},{"id":3,"type":"LoadVideo","pos":[729.5984141255855,-198.631920454299],"size":[282.798828125,233.0743408203125],"flags":{},"order":0,"mode":0,"inputs":[{"localized_name":"file","name":"file","type":"COMBO","widget":{"name":"file"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"IMAGEUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"VIDEO","name":"VIDEO","type":"VIDEO","links":[2]}],"properties":{"Node name for S&R":"LoadVideo"},"widgets_values":["draw.mp4","image"]},{"id":4,"type":"LoadAudio","pos":[729.8037086965753,99.86963519703949],"size":[282.798828125,136],"flags":{},"order":1,"mode":0,"inputs":[{"localized_name":"audio","name":"audio","type":"COMBO","widget":{"name":"audio"},"link":null},{"localized_name":"audioUI","name":"audioUI","type":"AUDIO_UI","widget":{"name":"audioUI"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"AUDIOUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"AUDIO","name":"AUDIO","type":"AUDIO","links":[3]}],"properties":{"Node name for S&R":"LoadAudio"},"widgets_values":["Megan-Fox.mp3",null,null]},{"id":5,"type":"VLLMOmniARSampling","pos":[510.3517536642828,658.073751009259],"size":[270,178],"flags":{},"order":2,"mode":0,"inputs":[{"localized_name":"max_tokens","name":"max_tokens","type":"INT","widget":{"name":"max_tokens"},"link":null},{"localized_name":"temperature","name":"temperature","type":"FLOAT","widget":{"name":"temperature"},"link":null},{"localized_name":"top_p","name":"top_p","type":"FLOAT","widget":{"name":"top_p"},"link":null},{"localized_name":"repetition_penalty","name":"repetition_penalty","type":"FLOAT","widget":{"name":"repetition_penalty"},"link":null},{"localized_name":"seed","name":"seed","type":"INT","widget":{"name":"seed"},"link":null}],"outputs":[{"localized_name":"AR sampling params","name":"AR sampling params","type":"SAMPLING_PARAMS","links":[12]}],"properties":{"Node name for S&R":"VLLMOmniARSampling"},"widgets_values":[100,1,1,1,-1,"randomize"]},{"id":7,"type":"VLLMOmniARSampling","pos":[503.33235181647115,419.34158016181806],"size":[270,178],"flags":{},"order":3,"mode":0,"inputs":[{"localized_name":"max_tokens","name":"max_tokens","type":"INT","widget":{"name":"max_tokens"},"link":null},{"localized_name":"temperature","name":"temperature","type":"FLOAT","widget":{"name":"temperature"},"link":null},{"localized_name":"top_p","name":"top_p","type":"FLOAT","widget":{"name":"top_p"},"link":null},{"localized_name":"repetition_penalty","name":"repetition_penalty","type":"FLOAT","widget":{"name":"repetition_penalty"},"link":null},{"localized_name":"seed","name":"seed","type":"INT","widget":{"name":"seed"},"link":null}],"outputs":[{"localized_name":"AR sampling params","name":"AR sampling params","type":"SAMPLING_PARAMS","links":[5,13]}],"properties":{"Node name for S&R":"VLLMOmniARSampling"},"widgets_values":[100,1,1,1,-1,"randomize"]},{"id":8,"type":"VLLMOmniSamplingParamsList","pos":[820.6056617389042,426.38372037182273],"size":[263.066015625,66],"flags":{},"order":7,"mode":0,"inputs":[{"localized_name":"param1","name":"param1","type":"SAMPLING_PARAMS","link":5},{"localized_name":"param2","name":"param2","shape":7,"type":"SAMPLING_PARAMS","link":12},{"localized_name":"param3","name":"param3","shape":7,"type":"SAMPLING_PARAMS","link":13}],"outputs":[{"localized_name":"param list","name":"param list"Line truncated
				`@@ -0,0 +1 @@`
				{"id":"91f75acc-8040-40f6-865a-2e8a7cfd6672","revision":0,"last_node_id":11,"last_link_id":23,"nodes":[{"id":3,"type":"PreviewImage","pos":[1281.8455167767304,-69.02638461454333],"size":[305.6628685610174,307.0336884172169],"flags":{},"order":4,"mode":0,"inputs":[{"localized_name":"images","name":"images","type":"IMAGE","link":22}],"outputs":[],"properties":{"Node name for S&R":"PreviewImage"},"widgets_values":[]},{"id":11,"type":"VLLMOmniGenerateImage","pos":[816.0962515873642,-137.51112584854445],"size":[400,278],"flags":{},"order":3,"mode":0,"inputs":[{"localized_name":"image","name":"image","shape":7,"type":"IMAGE","link":21},{"localized_name":"mask","name":"mask","shape":7,"type":"MASK","link":null},{"localized_name":"sampling_params","name":"sampling_params","shape":7,"type":"SAMPLING_PARAMS","link":23},{"localized_name":"url","name":"url","type":"STRING","widget":{"name":"url"},"link":null},{"localized_name":"model","name":"model","type":"STRING","widget":{"name":"model"},"link":null},{"localized_name":"prompt","name":"prompt","type":"STRING","widget":{"name":"prompt"},"link":null},{"localized_name":"negative_prompt","name":"negative_prompt","type":"STRING","widget":{"name":"negative_prompt"},"link":null},{"localized_name":"width","name":"width","type":"INT","widget":{"name":"width"},"link":null},{"localized_name":"height","name":"height","type":"INT","widget":{"name":"height"},"link":null}],"outputs":[{"localized_name":"image","name":"image","type":"IMAGE","links":[22]}],"properties":{"Node name for S&R":"VLLMOmniGenerateImage"},"widgets_values":["http://localhost:8000/v1","Qwen/Qwen-Image-Edit","Put this figure in a realistic mountain view","",512,512]},{"id":8,"type":"VLLMOmniDiffusionSampling","pos":[478.59266934006774,183.67711984955648],"size":[284.205078125,202],"flags":{},"order":0,"mode":0,"inputs":[{"localized_name":"n","name":"n","type":"INT","widget":{"name":"n"},"link":null},{"localized_name":"num_inference_steps","name":"num_inference_steps","type":"INT","widget":{"name":"num_inference_steps"},"link":null},{"localized_name":"guidance_scale","name":"guidance_scale","type":"FLOAT","widget":{"name":"guidance_scale"},"link":null},{"localized_name":"true_cfg_scale","name":"true_cfg_scale","type":"FLOAT","widget":{"name":"true_cfg_scale"},"link":null},{"localized_name":"vae_use_slicing","name":"vae_use_slicing","type":"BOOLEAN","widget":{"name":"vae_use_slicing"},"link":null},{"localized_name":"seed","name":"seed","type":"INT","widget":{"name":"seed"},"link":null}],"outputs":[{"localized_name":"diffusion sampling params","name":"diffusion sampling params","type":"SAMPLING_PARAMS","links":[23]}],"properties":{"Node name for S&R":"VLLMOmniDiffusionSampling"},"widgets_values":[4,50,1,1,false,42,"fixed"]},{"id":10,"type":"MarkdownNote","pos":[227.99306819575278,-231.24143306069843],"size":[240.3326478922474,136.3505820791881],"flags":{},"order":2,"mode":0,"inputs":[],"outputs":[],"title":"Note: Task and Input","properties":{},"widgets_values":["vLLM-Omni nodes are categorized based on the output modality. The \"Generate Image\" node supports both text-to-image generation or image-to-image generation (a.k.a. image editing). The node will route to the correct endpoint depending on whether an input image is present or not."],"color":"#432","bgcolor":"#000"},{"id":4,"type":"LoadImage","pos":[496.31859627609606,-229.71277089860084],"size":[270,314],"flags":{},"order":1,"mode":0,"inputs":[{"localized_name":"image","name":"image","type":"COMBO","widget":{"name":"image"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"IMAGEUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"IMAGE","name":"IMAGE","type":"IMAGE","links":[21]},{"localized_name":"MASK","name":"MASK","type":"MASK","links":null}],"properties":{"Node name for S&R":"LoadImage"},"widgets_values":["example.png","image"]}],"links":[[21,4,0,11,0,"IMAGE"],[22,11,0,3,0,"IMAGE"],[23,8,0,11,2,"SAMPLING_PARAMS"]],"groups":[{"id":1,"title":"Input","bounding":[213.1706010087147,-313.20095750667554,560.8246471437027,407.4246532472182],"color":"#3f789e","font_size":24,"flags":{}}],"config":{},"extra":{"workflowRendererVersion":"LG","ds":{"scale":1.1469075819486894,"offset":[276.8427259980847,549.4885226647582]}},"version":0.4}
				`@@ -0,0 +1 @@`
				{"id":"c0a11241-f0e6-4f5c-8488-b3d344a9df3d","revision":0,"last_node_id":7,"last_link_id":3,"nodes":[{"id":4,"type":"PreviewAudio","pos":[690.1759063763212,226.76444957981545],"size":[270,88],"flags":{},"order":5,"mode":0,"inputs":[{"localized_name":"audio","name":"audio","type":"AUDIO","link":2},{"localized_name":"audioUI","name":"audioUI","type":"AUDIO_UI","widget":{"name":"audioUI"},"link":null}],"outputs":[],"properties":{"Node name for S&R":"PreviewAudio"},"widgets_values":[]},{"id":5,"type":"LoadAudio","pos":[-83.61093919682321,103.80510898465587],"size":[282.798828125,136],"flags":{},"order":0,"mode":0,"inputs":[{"localized_name":"audio","name":"audio","type":"COMBO","widget":{"name":"audio"},"link":null},{"localized_name":"audioUI","name":"audioUI","type":"AUDIO_UI","widget":{"name":"audioUI"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"AUDIOUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"AUDIO","name":"AUDIO","type":"AUDIO","links":[3]}],"properties":{"Node name for S&R":"LoadAudio"},"widgets_values":["Megan-Fox.mp3",null,null]},{"id":3,"type":"VLLMOmniVoiceClone","pos":[248.63466280064262,168.13112893706375],"size":[400,306],"flags":{},"order":4,"mode":0,"inputs":[{"localized_name":"ref_audio","name":"ref_audio","type":"AUDIO","link":3},{"localized_name":"model_specific_params","name":"model_specific_params","shape":7,"type":"TTS_PARAMS","link":1},{"localized_name":"url","name":"url","type":"STRING","widget":{"name":"url"},"link":null},{"localized_name":"model","name":"model","type":"STRING","widget":{"name":"model"},"link":null},{"localized_name":"input","name":"input","type":"STRING","widget":{"name":"input"},"link":null},{"localized_name":"voice","name":"voice","type":"STRING","widget":{"name":"voice"},"link":null},{"localized_name":"response_format","name":"response_format","type":"COMBO","widget":{"name":"response_format"},"link":null},{"localized_name":"speed","name":"speed","type":"FLOAT","widget":{"name":"speed"},"link":null},{"localized_name":"ref_text","name":"ref_text","type":"STRING","widget":{"name":"ref_text"},"link":null},{"localized_name":"x_vector_only_mode","name":"x_vector_only_mode","type":"BOOLEAN","widget":{"name":"x_vector_only_mode"},"link":null}],"outputs":[{"localized_name":"audio","name":"audio","type":"AUDIO","links":[2]}],"properties":{"Node name for S&R":"VLLMOmniVoiceClone"},"widgets_values":["http://localhost:8000/v1","Qwen/Qwen3-TTS-12Hz-1.7B-Base","Someone just spilled a cup of coffee on my jacket this morning in the subway. Now I have to wear this stained jacket for a whole day!","","mp3",1,"",true]},{"id":6,"type":"MarkdownNote","pos":[255.52667408769662,27.26842775454861],"size":[381.58204108637847,88],"flags":{},"order":2,"mode":0,"inputs":[],"outputs":[],"title":"Note: TTS Nodes","properties":{},"widgets_values":["Apart from the Voice Cloning node, there is also another TTS node for simple text-to-speech tasks (without extra inputs for reference audio)."],"color":"#432","bgcolor":"#000"},{"id":7,"type":"MarkdownNote","pos":[-196.99453003423977,598.3693693727904],"size":[382.4539509209102,88],"flags":{},"order":3,"mode":0,"inputs":[],"outputs":[],"title":"Note: Model-Specific Parameters","properties":{},"widgets_values":["TTS models often require some tailor-made parameters. If you need to customize these parameters, grab one that matches the model you are requesting from \"TTS Params\" subfolder."],"color":"#432","bgcolor":"#000"},{"id":2,"type":"VLLMOmniQwenTTSParams","pos":[-207.8925310560497,342.7942867650348],"size":[400,200],"flags":{},"order":1,"mode":0,"inputs":[{"localized_name":"task_type","name":"task_type","type":"COMBO","widget":{"name":"task_type"},"link":null},{"localized_name":"language","name":"language","type":"COMBO","widget":{"name":"language"},"link":null},{"localized_name":"instructions","name":"instructions","type":"STRING","widget":{"name":"instructions"},"link":null},{"localized_name":"max_new_tokens","name":"max_new_tokens","type":"INT","widget":{"name":"max_new_tokens"},"link":null}],"outputs":[{"localized_name":"Qwen TTS params","name":"Qwen TTS params","type":"TTS_PARAMS","links":[1]}],"properties":{"Node name for S&R":"VLLMOmniQwenTTSParams"},"widgets_values":["Base","Auto","Super angry",2048]}],"links":[[1,2,0,3,1,"TTS_PARAMS"],[2,3,0,4,0,"AUDIO"],[3,5,0,3,0,"AUDIO"]],"groups":[{"id":1,"title":"Model-Specific Parameters","bounding":[-228.38328407738007,264.42790274716276,443.9887293164586,439.6853171248897],"color":"#3f789e","font_size":24,"flags":{}}],"config":{},"extra":{"workflowRendererVersion":"LG","ds":{"scale":1.1469075819486911,"offset":[898.8819468322333,258.7179979718389]}},"version":0.4}