[Bugfix] Fix diffusion hf_config fallback (#4093 )

Signed-off-by: KexiongYu <yukexiong1@huawei.com>
2026-06-03 20:57:14 +08:00
617 changed files with 12200 additions and 47193 deletions
@@ -56,88 +56,6 @@ check_run_all_label() {
    fi
 }

-is_docs_only_change() {
-    local file_path
-    local has_any=0
-
-    while IFS= read -r file_path; do
-        [[ -z "${file_path}" ]] && continue
-        has_any=1
-
-        if [[ "${file_path}" == docs/* ]]; then
-            continue
-        fi
-        if [[ "${file_path}" == *.md ]]; then
-            continue
-        fi
-        if [[ "${file_path}" == "mkdocs.yml" ]]; then
-            continue
-        fi
-        return 1
-    done
-
-    [[ "${has_any}" -eq 1 ]]
-}
-
-resolve_skip_ci() {
-    local is_pr_build=0
-    local files
-    local base_branch base_ref
-
-    if [[ "${BUILDKITE_PULL_REQUEST:-false}" != "false" && -n "${BUILDKITE_PULL_REQUEST:-}" ]]; then
-        is_pr_build=1
-    fi
-
-    if [[ "${is_pr_build}" -eq 1 ]]; then
-        base_branch="${BUILDKITE_PULL_REQUEST_BASE_BRANCH:-main}"
-        if ! git rev-parse --verify "origin/${base_branch}" >/dev/null 2>&1; then
-            echo "resolve_skip_ci: origin/${base_branch} not found locally; trying fetch" >&2
-            git fetch --depth=200 origin "${base_branch}" >/dev/null 2>&1 || true
-        fi
-
-        base_ref=""
-        if git rev-parse --verify "origin/${base_branch}" >/dev/null 2>&1; then
-            base_ref="origin/${base_branch}"
-        elif git rev-parse --verify "${base_branch}" >/dev/null 2>&1; then
-            base_ref="${base_branch}"
-        else
-            echo "resolve_skip_ci: cannot resolve PR base ${base_branch}; skip-ci=0" >&2
-            echo -n 0
-            return 0
-        fi
-
-        if ! files="$(git diff --name-only "${base_ref}...${BUILDKITE_COMMIT}" 2>/dev/null)"; then
-            echo "resolve_skip_ci: failed to compute PR changed files; skip-ci=0" >&2
-            echo -n 0
-            return 0
-        fi
-    elif [[ "${BUILDKITE_BRANCH:-}" == "main" ]]; then
-        if ! git rev-parse --verify "${BUILDKITE_COMMIT}^" >/dev/null 2>&1; then
-            echo "resolve_skip_ci: commit has no parent on main; skip-ci=0" >&2
-            echo -n 0
-            return 0
-        fi
-        if ! files="$(git diff --name-only "${BUILDKITE_COMMIT}^..${BUILDKITE_COMMIT}" 2>/dev/null)"; then
-            echo "resolve_skip_ci: failed to compute main changed files; skip-ci=0" >&2
-            echo -n 0
-            return 0
-        fi
-    else
-        echo "resolve_skip_ci: not PR/main build; skip-ci=0" >&2
-        echo -n 0
-        return 0
-    fi
-
-    if is_docs_only_change <<< "${files}"; then
-        echo "resolve_skip_ci: docs-only change detected; skip-ci=1" >&2
-        echo -n 1
-        return 0
-    fi
-
-    echo "resolve_skip_ci: non-doc changes detected; skip-ci=0" >&2
-    echo -n 0
-}
-
 if [[ -z "${COV_ENABLED:-}" ]]; then
    COV_ENABLED=0
 fi
@@ -173,24 +91,8 @@ upload_pipeline() {

    cd .buildkite

-    # Select test definition file: merge suite for main, ready suite for PRs.
-    # For debugging, DEBUG_TEST_YAML can override the selection — accepts
-    # "merge" or "ready" (case-insensitive).
-    if [[ -n "${DEBUG_TEST_YAML:-}" ]]; then
-        case "${DEBUG_TEST_YAML,,}" in
-            merge)
-                TEST_YAML="test-amd-merge.yml"
-                ;;
-            ready)
-                TEST_YAML="test-amd-ready.yaml"
-                ;;
-            *)
-                echo "ERROR: DEBUG_TEST_YAML must be 'merge' or 'ready', got '$DEBUG_TEST_YAML'" >&2
-                exit 1
-                ;;
-        esac
-        echo "DEBUG_TEST_YAML override: using $TEST_YAML"
-    elif [[ $BUILDKITE_BRANCH == "main" ]]; then
+    # Select test definition file: merge suite for main, ready suite for PRs
+    if [[ $BUILDKITE_BRANCH == "main" ]]; then
        TEST_YAML="test-amd-merge.yml"
    else
        TEST_YAML="test-amd-ready.yaml"
@@ -238,13 +140,30 @@ fi
 # Early exit start: skip pipeline if conditions are met
 # ----------------------------------------------------------------------

-# Match CUDA Buildkite skip-ci behavior for docs/markdown-only changes.
+# skip pipeline if all changed files are under docs/
 if [[ "${DOCS_ONLY_DISABLE}" != "1" ]]; then
-    SKIP_CI="$(resolve_skip_ci)"
-    if [[ "${SKIP_CI}" == "1" ]]; then
-        echo "[docs-only] Docs/markdown-only changes detected. Exiting before pipeline upload."
-        exit 0
+  if [[ -n "${file_diff:-}" ]]; then
+    docs_only=1
+    # Robust iteration over newline-separated file_diff
+    while IFS= read -r f; do
+      [[ -z "$f" ]] && continue
+      # **Policy:** only skip if *every* path starts with docs/
+      if [[ "$f" != docs/* ]]; then
+        docs_only=0
+        break
+      fi
+    done < <(printf '%s\n' "$file_diff" | tr ' ' '\n' | tr -d '\r')
+
+    if [[ "$docs_only" -eq 1 ]]; then
+      buildkite-agent annotate ":memo: CI skipped — docs/** only changes detected
+
+\`\`\`
+${file_diff}
+\`\`\`" --style "info" || true
+      echo "[docs-only] All changes are under docs/. Exiting before pipeline upload."
+      exit 0
    fi
+  fi
 fi

 # ----------------------------------------------------------------------
@@ -0,0 +1,20 @@
+steps:
+  - label: "Build and upload wheel"
+    key: "build-wheel"
+    agents:
+      queue: cpu_queue_release
+    commands:
+      - "curl -LsSf https://astral.sh/uv/install.sh | sh"
+      - 'export PATH="$HOME/.local/bin:$PATH"'
+      - "uv venv --python=3.12 && source .venv/bin/activate"
+      - "uv pip install --upgrade build"
+      - "python3 -m build"
+      - "bash .buildkite/scripts/upload-nightly-wheels.sh"
+
+  - label: "Generate and upload wheel indices"
+    depends_on: "build-wheel"
+    allow_dependency_failure: true
+    agents:
+      queue: small_cpu_queue_release
+    commands:
+      - "bash .buildkite/scripts/generate-and-upload-nightly-index.sh"
@@ -10,7 +10,7 @@ steps:
          DOCKER_BUILDKIT: "1"
          # Buildkite will automatically replace this with the actual commit hash
          VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}"
-          VLLM_VERSION: "v0.21.0"
+          VLLM_VERSION: "v0.20.0"
        priority: 100
        timeout_in_minutes: 60
        soft_fail: false
@@ -1,110 +0,0 @@
-# steps:
-#   - label: ":buildkit: Build and Push NPU Test Image - A3"
-#     key: image-build-a3
-#     agents:
-#       queue: "ascend-a3"
-#       resource_class: "npu-2"
-#     plugins:
-#       - kubernetes:
-#           metadata:
-#             annotations:
-#               vault.hashicorp.com/agent-init-first: "true"
-#               vault.hashicorp.com/agent-inject: "true"
-#               vault.hashicorp.com/agent-inject-perms-ca.pem: "0400"
-#               vault.hashicorp.com/agent-inject-perms-cert.pem: "0400"
-#               vault.hashicorp.com/agent-inject-perms-config.json: "0400"
-#               vault.hashicorp.com/agent-inject-perms-key.pem: "0400"
-#               vault.hashicorp.com/agent-inject-secret-ca.pem: internal/data/ascend/buildkitd
-#               vault.hashicorp.com/agent-inject-secret-cert.pem: internal/data/ascend/buildkitd
-#               vault.hashicorp.com/agent-inject-secret-config.json: internal/data/ascend/buildkitd
-#               vault.hashicorp.com/agent-inject-secret-key.pem: internal/data/ascend/buildkitd
-#               vault.hashicorp.com/agent-inject-template-ca.pem: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.RootCA }}\n{{- end }}"
-#               vault.hashicorp.com/agent-inject-template-cert.pem: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.ClientCaCert }}\n{{- end }}"
-#               vault.hashicorp.com/agent-inject-template-config.json: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.dockerConfig }}\n{{- end }}"
-#               vault.hashicorp.com/agent-inject-template-key.pem: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.ClientCaKey }}\n{{- end }}"
-#               vault.hashicorp.com/agent-pre-populate-only: "true"
-#               vault.hashicorp.com/agent-run-as-group: "1000"
-#               vault.hashicorp.com/agent-run-as-user: "1000"
-#               vault.hashicorp.com/agent-service-account-token-volume-name: token-vol
-#               vault.hashicorp.com/role: ascend-gha-runners
-#               vault.hashicorp.com/secret-volume-path: /home/user/.docker/
-#               vault.hashicorp.com/tls-skip-verify: "true"
-#           podSpecPatch:
-#             volumes:
-#             - name: token-vol
-#               projected:
-#                 defaultMode: 420
-#                 sources:
-#                 - serviceAccountToken:
-#                     audience: api
-#                     expirationSeconds: 600
-#                     path: token
-#     env:
-#       VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}"
-#       IMAGE_NAME: "vllm-omni-ci-npu"
-#       IMAGE_REGISTRY: "swr.cn-southwest-2.myhuaweicloud.com/modelfoundry"
-#       BUILDKITD_ADDR: "tcp://buildkitd-service.buildkitd:1234"
-#     command: |
-#       set -ex
-
-#       echo "--- Building and pushing NPU Test Image"
-#       echo "Image: $${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG}"
-#       echo "buildkitd address: $${BUILDKITD_ADDR}"
-
-#       git config --global url."https://gh-proxy.test.osinfra.cn/https://github.com/".insteadOf "https://github.com/"
-
-#       if ! command -v buildctl &> /dev/null; then
-#         echo "Installing buildctl..."
-#         mkdir -p /tmp/buildkit
-#         BUILDKIT_VERSION="v0.29.0"
-#         wget -q "https://gh-proxy.test.osinfra.cn/https://github.com/moby/buildkit/releases/download/$${BUILDKIT_VERSION}/buildkit-$${BUILDKIT_VERSION}.linux-arm64.tar.gz" -O /tmp/buildkit.tar.gz
-#         tar -xzf /tmp/buildkit.tar.gz -C /tmp/buildkit
-#         cp /tmp/buildkit/bin/buildctl /usr/local/bin/
-#       fi
-
-#       export DOCKER_CONFIG=/home/user/.docker
-#       buildctl \
-#         --addr="$${BUILDKITD_ADDR}" \
-#         --tlscacert=/home/user/.docker/ca.pem \
-#         --tlscert=/home/user/.docker/cert.pem \
-#         --tlskey=/home/user/.docker/key.pem \
-#         build \
-#         --frontend dockerfile.v0 \
-#         --local context=. \
-#         --local dockerfile=./docker \
-#         --opt filename=Dockerfile.npu.ci.a3 \
-#         --secret id=dockerconfig,src=/home/user/.docker/config.json \
-#         --output type=image,name=$${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG},push=true \
-#         --progress=plain
-
-#       echo "--- Image pushed successfully"
-#       echo "$${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG}"
-
-#   # - label: "Upload Ready Pipeline - A3"
-#   #   depends_on: image-build-a3
-#   #   key: upload-ready-pipeline-a3
-#   #   if: build.branch != "main" && build.pull_request.labels includes "npu-test"
-#   #   agents:
-#   #     queue: "ascend-a3"
-#   #     resource_class: "npu-2"
-#   #   command: |
-#   #     buildkite-agent pipeline upload .buildkite/test-npu-ready.yaml
-
-  # # L4 Test — main+NIGHTLY=1 (scheduled), or PR with specific label (e.g. add label then Rebuild)
-  # - label: "Upload Nightly Pipeline"
-  #   depends_on: image-build
-  #   key: upload-nightly-pipeline
-  #   if: >-
-  #     (build.branch == "main" && build.env("NIGHTLY") == "1") ||
-  #     (build.branch != "main" && (
-  #       build.pull_request.labels includes "nightly-test" ||
-  #       build.pull_request.labels includes "omni-test" ||
-  #       build.pull_request.labels includes "tts-test" ||
-  #       build.pull_request.labels includes "diffusion-x2iat-test" ||
-  #       build.pull_request.labels includes "diffusion-x2v-test"
-  #     ))
-  #   commands:
-  #     - buildkite-agent pipeline upload .buildkite/test-npu-nightly.yaml
-  #   agents:
-  #     queue: "ascend-a3"
-  #     resource_class: "npu-2"
@@ -1,112 +0,0 @@
-steps:
-  - label: ":buildkit: Build and Push NPU Test Image - A2"
-    key: image-build
-    agents:
-      queue: "ascend-a2b3"
-      resource_class: "npu-1"
-    plugins:
-      - kubernetes:
-          metadata:
-            annotations:
-              vault.hashicorp.com/agent-init-first: "true"
-              vault.hashicorp.com/agent-inject: "true"
-              vault.hashicorp.com/agent-inject-perms-ca.pem: "0400"
-              vault.hashicorp.com/agent-inject-perms-cert.pem: "0400"
-              vault.hashicorp.com/agent-inject-perms-config.json: "0400"
-              vault.hashicorp.com/agent-inject-perms-key.pem: "0400"
-              vault.hashicorp.com/agent-inject-secret-ca.pem: internal/data/ascend/buildkitd
-              vault.hashicorp.com/agent-inject-secret-cert.pem: internal/data/ascend/buildkitd
-              vault.hashicorp.com/agent-inject-secret-config.json: internal/data/ascend/buildkitd
-              vault.hashicorp.com/agent-inject-secret-key.pem: internal/data/ascend/buildkitd
-              vault.hashicorp.com/agent-inject-template-ca.pem: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.RootCA }}\n{{- end }}"
-              vault.hashicorp.com/agent-inject-template-cert.pem: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.ClientCaCert }}\n{{- end }}"
-              vault.hashicorp.com/agent-inject-template-config.json: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.dockerConfig }}\n{{- end }}"
-              vault.hashicorp.com/agent-inject-template-key.pem: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.ClientCaKey }}\n{{- end }}"
-              vault.hashicorp.com/agent-pre-populate-only: "true"
-              vault.hashicorp.com/agent-run-as-group: "1000"
-              vault.hashicorp.com/agent-run-as-user: "1000"
-              vault.hashicorp.com/agent-service-account-token-volume-name: token-vol
-              vault.hashicorp.com/role: ascend-gha-runners
-              vault.hashicorp.com/secret-volume-path: /home/user/.docker/
-              vault.hashicorp.com/tls-skip-verify: "true"
-          podSpecPatch:
-            volumes:
-            - name: token-vol
-              projected:
-                defaultMode: 420
-                sources:
-                - serviceAccountToken:
-                    audience: api
-                    expirationSeconds: 600
-                    path: token
-    env:
-      VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}"
-      IMAGE_NAME: "vllm-omni-ci-npu"
-      IMAGE_REGISTRY: "swr.cn-southwest-2.myhuaweicloud.com/modelfoundry"
-      BUILDKITD_ADDR: "tcp://buildkitd-service.buildkitd:1234"
-    command: |
-      set -ex
-
-      echo "--- Building and pushing NPU Test Image"
-      echo "Image: $${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG}"
-      echo "buildkitd address: $${BUILDKITD_ADDR}"
-
-      if ! command -v buildctl &> /dev/null; then
-        echo "Installing buildctl..."
-        mkdir -p /tmp/buildkit
-        BUILDKIT_VERSION="v0.29.0"
-        wget -q "https://github.com/moby/buildkit/releases/download/$${BUILDKIT_VERSION}/buildkit-$${BUILDKIT_VERSION}.linux-arm64.tar.gz" -O /tmp/buildkit.tar.gz
-        tar -xzf /tmp/buildkit.tar.gz -C /tmp/buildkit
-        cp /tmp/buildkit/bin/buildctl /usr/local/bin/
-      fi
-
-      export DOCKER_CONFIG=/home/user/.docker
-      buildctl \
-        --addr="$${BUILDKITD_ADDR}" \
-        --tlscacert=/home/user/.docker/ca.pem \
-        --tlscert=/home/user/.docker/cert.pem \
-        --tlskey=/home/user/.docker/key.pem \
-        build \
-        --frontend dockerfile.v0 \
-        --local context=. \
-        --local dockerfile=./docker \
-        --opt filename=Dockerfile.npu.ci \
-        --secret id=dockerconfig,src=/home/user/.docker/config.json \
-        --output type=image,name=$${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG},push=true \
-        --progress=plain
-
-      echo "--- Image pushed successfully"
-      echo "$${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG}"
-
-  - label: "Upload Ready Pipeline"
-    depends_on: image-build
-    key: upload-ready-pipeline
-    if: >-
-      build.branch != "main" && (
-        build.pull_request.labels includes "npu-test" ||
-        build.pull_request.labels includes "ready"
-      )
-    agents:
-      queue: "ascend-a2b3"
-      resource_class: "npu-1"
-    command: |
-      buildkite-agent pipeline upload .buildkite/test-npu-ready.yaml
-
-  # L4 Test — main+NIGHTLY=1 (scheduled), or PR with specific label (e.g. add label then Rebuild)
-  - label: "Upload Nightly Pipeline"
-    depends_on: image-build
-    key: upload-nightly-pipeline
-    if: >-
-      (build.branch == "main" && build.env("NIGHTLY") == "1") ||
-      (build.branch != "main" && (
-        build.pull_request.labels includes "nightly-test" ||
-        build.pull_request.labels includes "omni-test" ||
-        build.pull_request.labels includes "tts-test" ||
-        build.pull_request.labels includes "diffusion-x2iat-test" ||
-        build.pull_request.labels includes "diffusion-x2v-test"
-      ))
-    commands:
-      - buildkite-agent pipeline upload .buildkite/test-npu-nightly.yaml
-    agents:
-      queue: "ascend-a2b3"
-      resource_class: "npu-1"
@@ -1,38 +0,0 @@
-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 --progress=plain --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"
-
-  # L2 Test
-  - label: "Upload Ready Pipeline"
-    depends_on: image-build
-    key: upload-ready-pipeline
-    commands:
-      - buildkite-agent pipeline upload .buildkite/test-ready.yml
-    agents:
-      queue: "cpu_queue_premerge"
-
-  # L3 Test
-  - label: "Upload Merge Pipeline"
-    depends_on: image-build
-    key: upload-merge-pipeline
-    commands:
-      - buildkite-agent pipeline upload .buildkite/test-merge.yml
-    agents:
-      queue: "cpu_queue_premerge"
-
-  # L4 Test
-  - label: "Upload Nightly Pipeline"
-    depends_on: image-build
-    key: upload-nightly-pipeline
-    if: build.env("NIGHTLY") == "1"
-    commands:
-      - buildkite-agent pipeline upload .buildkite/test-nightly.yml
-    agents:
-      queue: "cpu_queue_premerge"
@@ -1,110 +1,38 @@
 steps:
-  - input: "Provide Release version here"
-    id: input-release-version
-    fields:
-      - text: "What is the release version?"
-        key: release-version
-
-  - label: "Build and upload wheel"
-    depends_on: ~
-    key: "build-wheel"
-    agents:
-      queue: cpu_queue_release
+  - label: ":docker: Build image"
+    key: image-build
    commands:
-      - "curl -LsSf https://astral.sh/uv/install.sh | sh"
-      - 'export PATH="$HOME/.local/bin:$PATH"'
-      - "uv venv --python=3.12 && source .venv/bin/activate"
-      - "uv pip install --upgrade build"
-      - "python3 -m build"
-      - "bash .buildkite/scripts/upload-nightly-wheels.sh"
-
-  - label: "Generate and upload wheel indices"
-    depends_on: "build-wheel"
-    allow_dependency_failure: true
+      - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
+      - "docker build --progress=plain --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: small_cpu_queue_release
+      queue: "cpu_queue_premerge"
+
+  # L2 Test
+  - label: "Upload Ready Pipeline"
+    depends_on: image-build
+    key: upload-ready-pipeline
    commands:
-      - "bash .buildkite/scripts/generate-and-upload-nightly-index.sh"
-
-  - block: "Upload release wheels to PyPI"
-    depends_on: "build-wheel"
-    key: block-upload-release-wheels
-
-  - label: "Upload release wheels to PyPI"
-    depends_on:
-      - block-upload-release-wheels
-      - input-release-version
-    key: upload-release-wheels
+      - buildkite-agent pipeline upload .buildkite/test-ready.yml
    agents:
-      queue: small_cpu_queue_release
+      queue: "cpu_queue_premerge"
+
+  # L3 Test
+  - label: "Upload Merge Pipeline"
+    depends_on: image-build
+    key: upload-merge-pipeline
    commands:
-      - "bash .buildkite/scripts/publish-release-wheels-pypi.sh"
-
-  - block: "Unblock to build release Docker images"
-    depends_on: ~
-    key: block-build-release-images
-
-  - group: "Build release Docker images"
-    key: "build-release-images"
-    depends_on: block-build-release-images
-    allow_dependency_failure: true
-    steps:
-      - label: "Build release image - x86_64"
-        key: build-release-image-x86
-        depends_on: ~
-        agents:
-          queue: cpu_queue_release
-        commands:
-          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-          - |
-            DOCKER_BUILDKIT=1 docker build \
-              --build-arg max_jobs=16 \
-              --tag public.ecr.aws/q9t5s3a7/vllm-omni-release-repo:$BUILDKITE_COMMIT-$(uname -m) \
-              --progress plain \
-              -f docker/Dockerfile.cuda .
-          - "docker push public.ecr.aws/q9t5s3a7/vllm-omni-release-repo:$BUILDKITE_COMMIT-$(uname -m)"
-          # re-tag to default image tag and push, just in case arm64 build fails
-          - "docker tag public.ecr.aws/q9t5s3a7/vllm-omni-release-repo:$BUILDKITE_COMMIT-$(uname -m) public.ecr.aws/q9t5s3a7/vllm-omni-release-repo:$BUILDKITE_COMMIT"
-          - "docker push public.ecr.aws/q9t5s3a7/vllm-omni-release-repo:$BUILDKITE_COMMIT"
-        env:
-          DOCKER_BUILDKIT: "1"
-
-      - label: "Build release image - aarch64"
-        key: build-release-image-arm64
-        depends_on: ~
-        agents:
-          queue: arm64_cpu_queue_release
-        commands:
-          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-          - |
-            DOCKER_BUILDKIT=1 docker build \
-              --build-arg max_jobs=16 \
-              --tag public.ecr.aws/q9t5s3a7/vllm-omni-release-repo:$BUILDKITE_COMMIT-$(uname -m) \
-              --progress plain \
-              -f docker/Dockerfile.cuda .
-          - "docker push public.ecr.aws/q9t5s3a7/vllm-omni-release-repo:$BUILDKITE_COMMIT-$(uname -m)"
-        env:
-          DOCKER_BUILDKIT: "1"
-
-  - block: "Publish release images to DockerHub"
-    key: block-publish-release-images
-    depends_on:
-      - build-release-image-x86
-      - build-release-image-arm64
-      - input-release-version
-
-  - label: "Publish release images to DockerHub"
-    depends_on:
-      - block-publish-release-images
-    key: publish-release-images-dockerhub
+      - buildkite-agent pipeline upload .buildkite/test-merge.yml
    agents:
-      queue: small_cpu_queue_release
+      queue: "cpu_queue_premerge"
+
+  # L4 Test
+  - label: "Upload Nightly Pipeline"
+    depends_on: image-build
+    key: upload-nightly-pipeline
+    if: build.env("NIGHTLY") == "1"
    commands:
-      - "bash .buildkite/scripts/publish-release-images.sh"
-    plugins:
-      - docker-login#v3.0.0:
-          username: vllmbot
-          password-env: DOCKERHUB_TOKEN
-    env:
-      DOCKER_BUILDKIT: "1"
-      DOCKERHUB_USERNAME: "vllmbot"
+      - buildkite-agent pipeline upload .buildkite/test-nightly.yml
+    agents:
+      queue: "cpu_queue_premerge"
@@ -52,6 +52,5 @@ time timeout -k 30 30m docker run \
    pip install tblib==3.1.0
    cd /workspace/vllm-omni
    pytest -v -s -m "core_model and xpu and B60"
-    export VLLM_XPU_USE_SAMPLER_KERNEL=0    # NOTE: Remove this after vLLM v0.21.1 is merged. Fixes Qwen2-5 omni-expansion tests.
    pytest -v -s -m "advanced_model and xpu and B60"
 '
@@ -1,50 +0,0 @@
-#!/bin/bash
-# SPDX-License-Identifier: Apache-2.0
-#
-# Publish release Docker images from ECR to DockerHub.
-# Pulls per-arch images, tags with `latest` and versioned tags, pushes them,
-# then creates and pushes a multi-arch manifest.
-
-set -euo pipefail
-
-RELEASE_VERSION=$(buildkite-agent meta-data get release-version --default "" | sed 's/^v//')
-if [ -z "${RELEASE_VERSION}" ]; then
-  echo "ERROR: release-version metadata not set"
-  exit 1
-fi
-
-COMMIT="$BUILDKITE_COMMIT"
-DOCKERHUB_REPO="vllm/vllm-omni"
-ECR_REPO="public.ecr.aws/q9t5s3a7/vllm-omni-release-repo"
-
-echo "========================================"
-echo "Publishing release images v${RELEASE_VERSION}"
-echo "  Commit: ${COMMIT}"
-echo "========================================"
-
-# Login to ECR to pull staging images
-aws ecr-public get-login-password --region us-east-1 | \
-  docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
-
-docker pull "${ECR_REPO}:${COMMIT}-x86_64"
-docker pull "${ECR_REPO}:${COMMIT}-aarch64"
-
-docker tag "${ECR_REPO}:${COMMIT}-x86_64" "${DOCKERHUB_REPO}:latest-x86_64"
-docker tag "${ECR_REPO}:${COMMIT}-x86_64" "${DOCKERHUB_REPO}:v${RELEASE_VERSION}-x86_64"
-docker push "${DOCKERHUB_REPO}:latest-x86_64"
-docker push "${DOCKERHUB_REPO}:v${RELEASE_VERSION}-x86_64"
-
-docker tag "${ECR_REPO}:${COMMIT}-aarch64" "${DOCKERHUB_REPO}:latest-aarch64"
-docker tag "${ECR_REPO}:${COMMIT}-aarch64" "${DOCKERHUB_REPO}:v${RELEASE_VERSION}-aarch64"
-docker push "${DOCKERHUB_REPO}:latest-aarch64"
-docker push "${DOCKERHUB_REPO}:v${RELEASE_VERSION}-aarch64"
-
-docker manifest rm "${DOCKERHUB_REPO}:latest" || true
-docker manifest rm "${DOCKERHUB_REPO}:v${RELEASE_VERSION}" || true
-docker manifest create "${DOCKERHUB_REPO}:latest" "${DOCKERHUB_REPO}:latest-x86_64" "${DOCKERHUB_REPO}:latest-aarch64"
-docker manifest create "${DOCKERHUB_REPO}:v${RELEASE_VERSION}" "${DOCKERHUB_REPO}:v${RELEASE_VERSION}-x86_64" "${DOCKERHUB_REPO}:v${RELEASE_VERSION}-aarch64"
-docker manifest push "${DOCKERHUB_REPO}:latest"
-docker manifest push "${DOCKERHUB_REPO}:v${RELEASE_VERSION}"
-
-echo ""
-echo "Successfully published release images for v${RELEASE_VERSION}"
@@ -1,52 +0,0 @@
-#!/usr/bin/env bash
-# SPDX-License-Identifier: Apache-2.0
-
-set -e
-
-BUCKET="vllm-wheels"
-SUBPATH="omni/$BUILDKITE_COMMIT"
-S3_COMMIT_PREFIX="s3://$BUCKET/$SUBPATH/"
-
-RELEASE_VERSION=$(buildkite-agent meta-data get release-version | sed 's/^v//')
-if [[ -z "$RELEASE_VERSION" ]]; then
-  echo "[FATAL] release-version metadata not set."
-  exit 1
-fi
-echo "Release version: $RELEASE_VERSION"
-
-if [[ -z "$PYPI_TOKEN" ]]; then
-  echo "[FATAL] PYPI_TOKEN is not set."
-  exit 1
-else
-  export TWINE_USERNAME="__token__"
-  export TWINE_PASSWORD="$PYPI_TOKEN"
-fi
-
-set -x
-
-if ! command -v uv &> /dev/null; then
-  curl -LsSf https://astral.sh/uv/install.sh | sh
-  export PATH="$HOME/.local/bin:$PATH"
-fi
-
-uv venv --python 3.12 /tmp/vllm-omni-release-env
-source /tmp/vllm-omni-release-env/bin/activate
-uv pip install twine
-
-DIST_DIR=/tmp/vllm-omni-release-dist
-mkdir -p "$DIST_DIR"
-
-echo "Downloading wheels from S3:"
-aws s3 ls "$S3_COMMIT_PREFIX"
-aws s3 cp --recursive --exclude "*" --include "vllm_omni-${RELEASE_VERSION}*.whl" --exclude "*dev*" "$S3_COMMIT_PREFIX" "$DIST_DIR"
-ls -la "$DIST_DIR"
-
-PYPI_WHEEL_FILES=$(find "$DIST_DIR" -name "vllm_omni-${RELEASE_VERSION}*.whl" -not -name "*+*")
-if [[ -z "$PYPI_WHEEL_FILES" ]]; then
-  echo "[FATAL] No wheels found for version ${RELEASE_VERSION}"
-  exit 1
-fi
-
-python3 -m twine check $PYPI_WHEEL_FILES
-python3 -m twine upload --non-interactive --verbose $PYPI_WHEEL_FILES
-echo "Wheels uploaded to PyPI"
@@ -1,18 +1,10 @@
 #!/usr/bin/env bash
-# Evaluate docs-only skip-ci and upload continuation steps from the given pipeline YAML
+# Evaluate docs-only skip-ci and upload continuation steps from the same `.buildkite/pipeline.yml`
 # (YAML document after the first `---`). Buildkite `if` is evaluated at upload time.
-#
-# Usage: upload_pipeline_with_skip_ci.sh [pipeline_yaml]
-#   pipeline_yaml: path relative to repo root or absolute (default: .buildkite/pipeline.yml)
 set -euo pipefail

 ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
-PIPELINE_ARG="${1:-.buildkite/pipeline.yml}"
-if [[ "${PIPELINE_ARG}" = /* ]]; then
-  PIPELINE_YML="${PIPELINE_ARG}"
-else
-  PIPELINE_YML="${ROOT}/${PIPELINE_ARG}"
-fi
+PIPELINE_YML="${ROOT}/.buildkite/pipeline.yml"

 # Prints a single digit to stdout: 1 = skip image CI, 0 = run. Logs go to stderr.
 is_docs_only_change() {
@@ -29,7 +21,7 @@ is_docs_only_change() {
    if [[ "${file_path}" == *.md ]]; then
      continue
    fi
-    if [[ "${file_path}" == "mkdocs.yml" ]]; then
+    if [[ "${file_path}" == "mkdocs.yaml" ]]; then
      continue
    fi
    return 1
@@ -112,13 +104,12 @@ import pathlib
 path = pathlib.Path(os.environ["PIPELINE_YML"])
 text = path.read_text(encoding="utf-8")
 sep = "\n---\n"
-# Two supported layouts:
-#   - multi-doc: doc 1 = bootstrap (loaded by Buildkite), doc 2 (after `---`) = uploaded steps.
-#   - single-doc: whole file is the uploaded steps; caller (UI init step) is the bootstrap.
-if sep in text:
-    _, continuation = text.split(sep, 1)
-else:
-    continuation = text
+if sep not in text:
+    raise SystemExit(
+        "upload_pipeline_with_skip_ci: .buildkite/pipeline.yml must contain a '\\n---\\n' separator "
+        "(document 1 = bootstrap, document 2 = uploaded steps)"
+    )
+_, continuation = text.split(sep, 1)

 skip = os.environ.get("SKIP_CI") == "1"
 # When docs-only skip-ci: skip default CI image, but still build for L4 nightly (PR label nightly-test or
@@ -135,7 +126,7 @@ if skip:
 else:
    rep = "'true'"
    ready_rep = "'build.branch != \"main\" && build.pull_request.labels includes \"ready\"'"
-    merge_rep = "'(build.branch == \"main\" && build.env(\"NIGHTLY\") != \"1\" && build.env(\"WEEKLY\") != \"1\") || (build.branch != \"main\" && build.pull_request.labels includes \"merge-test\")'"
+    merge_rep = "'(build.branch == \"main\" && build.env(\"NIGHTLY\") != \"1\") || (build.branch != \"main\" && build.pull_request.labels includes \"merge-test\")'"
 rendered = (
    continuation
    .replace("__IMAGE_BUILD_IF__", rep)
@@ -9,7 +9,16 @@ steps:
    - export VLLM_ROCM_USE_AITER=0
    # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
    # that is not supported on CPU
-    - "timeout 60m pytest -v -s -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
+    - "timeout 20m pytest -v -s -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
+
+- label: "Diffusion Model Test"
+  agent_pool: mi325_2
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export GPU_ARCHS=gfx942
+    - timeout 20m pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "advanced_model and diffusion" --run-level "advanced_model"

 - label: "Diffusion Images API LoRA E2E"
  agent_pool: mi325_1
@@ -17,7 +26,7 @@ steps:
  mirror_hardwares: [amdproduction]
  grade: Blocking
  commands:
-    - timeout 20m pytest -s -v tests/diffusion/lora/test_images_generations_lora.py
+    - timeout 20m pytest -s -v tests/e2e/online_serving/test_images_generations_lora.py

 - label: "Diffusion Model CPU offloading Test"
  agent_pool: mi325_1
@@ -30,9 +39,9 @@ steps:
    - |
      timeout 20m bash -c '
        set +e
-        pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py
+        pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
        EXIT1=\$?
-        pytest -s -v tests/diffusion/offloader/test_diffusion_layerwise_offload.py
+        pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
        EXIT2=\$?
        exit \$((EXIT1 | EXIT2))
      '
@@ -67,7 +76,7 @@ steps:
  commands:
    - export VLLM_LOGGING_LEVEL=DEBUG
    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - timeout 20m pytest -s -v tests/diffusion/distributed/test_sequence_parallel.py
+    - timeout 20m pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py
    - timeout 20m pytest -s -v tests/diffusion/distributed/test_ulysses_uaa_perf.py

 # merge-only tests
@@ -77,15 +86,7 @@ steps:
  mirror_hardwares: [amdproduction]
  grade: Blocking
  commands:
-    - pytest -s -v tests/diffusion/distributed/test_zimage_parallelism.py
-
- label: "Diffusion Model Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - timeout 30m pytest -s -v tests/e2e/offline_inference/test_z_image.py
+    - pytest -s -v tests/e2e/offline_inference/test_zimage_parallelism.py

 - label: "Diffusion GPU Worker Test"
  agent_pool: mi325_2
@@ -123,6 +124,7 @@ steps:
  grade: Blocking
  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 tests/e2e/online_serving/test_qwen3_omni.py tests/e2e/online_serving/test_mimo_audio.py -m "advanced_model" --run-level "advanced_model"

 - label: "Qwen3-TTS CustomVoice E2E Test"
@@ -153,6 +155,17 @@ steps:
        pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py tests/e2e/offline_inference/test_qwen3_tts_base.py -m "advanced_model" --run-level "advanced_model"
      '

+- label: "Diffusion Image Edit Test"
+  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
+    - timeout 20m pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py
+
 # TODO: Bagel test on ROCm is very unstable. @tjtanaa
 # Need to debug before reneable numerical changes across large PRs
 # # split Bagel Model Test with H100 (Real Weights) into three tests
@@ -163,10 +176,11 @@ steps:
 #   grade: Blocking
 #   commands:
 #     - export GPU_ARCHS=gfx942
+#     - export VLLM_TEST_CLEAN_GPU_MEMORY=1
 #     - export VLLM_LOGGING_LEVEL=DEBUG
 #     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
 #     - export VLLM_ROCM_USE_AITER_RMSNORM=0
-#     - timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel.py -m "advanced_model" --run-level "advanced_model" -k "shared_memory" -k "rocm"
+#     - timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py -m "advanced_model" --run-level "advanced_model" -k "shared_memory" -k "rocm"

 # - label: "Bagel Img2Img Model Test (2/3)"
 #   agent_pool: mi325_1
@@ -175,10 +189,11 @@ steps:
 #   grade: Blocking
 #   commands:
 #     - export GPU_ARCHS=gfx942
+#     - export VLLM_TEST_CLEAN_GPU_MEMORY=1
 #     - export VLLM_LOGGING_LEVEL=DEBUG
 #     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
 #     - export VLLM_ROCM_USE_AITER_RMSNORM=0
-#     - timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel.py -m "advanced_model" --run-level "advanced_model" -k "rocm"
+#     - timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py -m "advanced_model" --run-level "advanced_model" -k "rocm"

 # - label: "Bagel Online Serving Test (3/3)"
 #   agent_pool: mi325_1
@@ -187,23 +202,12 @@ steps:
 #   grade: Blocking
 #   commands:
 #     - export GPU_ARCHS=gfx942
+#     - export VLLM_TEST_CLEAN_GPU_MEMORY=1
 #     - export VLLM_IMAGE_FETCH_TIMEOUT=60
 #     - export VLLM_LOGGING_LEVEL=DEBUG
 #     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
 #     - export VLLM_ROCM_USE_AITER_RMSNORM=0
-#     - timeout 40m pytest -s -v tests/e2e/online_serving/test_bagel.py -m "advanced_model" --run-level "advanced_model" -k "rocm"
-
- label: "Diffusion · Wan22 Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - |
-      timeout 40m bash -c "
-        export VLLM_IMAGE_FETCH_TIMEOUT=60
-        pytest -s -v tests/e2e/offline_inference/test_wan22_t2v.py tests/e2e/online_serving/test_wan22_t2v.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
-      "
+#     - timeout 40m pytest -s -v tests/e2e/online_serving/test_bagel_online.py -m "advanced_model" --run-level "advanced_model" -k "rocm"

 - label: "Voxtral-TTS E2E Test"
  agent_pool: mi325_1
@@ -9,7 +9,7 @@ steps:
    - export VLLM_ROCM_USE_AITER=0
    # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
    # that is not supported on CPU
-    - "timeout 60m pytest -v -s -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
+    - "timeout 20m pytest -vvvv -s -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"

 - label: "Voxtral TTS CUDA Unit Test"
  agent_pool: mi325_1
@@ -19,13 +19,21 @@ steps:
  commands:
    - timeout 10m pytest -s -v tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py

+- label: "Diffusion Model Test"
+  agent_pool: mi325_1
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - timeout 30m pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "core_model and diffusion" --run-level "core_model"
+
 - label: "Diffusion Batching Test"
  agent_pool: mi325_1
  depends_on: amd-build
  mirror_hardwares: [amdproduction]
  grade: Blocking
  commands:
-    - timeout 20m pytest -s -v tests/diffusion/batching/test_diffusion_batching.py -m "core_model and diffusion" --run-level "core_model"
+    - timeout 20m pytest -s -v tests/e2e/offline_inference/test_qwen_image_diffusion_batching.py -m "core_model and diffusion" --run-level "core_model"

 - label: "Custom Pipeline Test"
  agent_pool: mi325_1
@@ -46,9 +54,9 @@ steps:
    - |
      timeout 20m bash -c '
        set +e
-        pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py
+        pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
        EXIT1=\$?
-        pytest -s -v tests/diffusion/offloader/test_diffusion_layerwise_offload.py
+        pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
        EXIT2=\$?
        exit \$((EXIT1 | EXIT2))
      '
@@ -65,6 +73,17 @@ steps:
 #     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
 #     - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_expansion.py -m "advanced_model and diffusion and L4" --run-level advanced_model

+- label: "AudioX Online Test"
+  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
+    - timeout 20m pytest -s -v tests/e2e/online_serving/test_audiox_online.py
+
 - label: "Diffusion Cache Backend Test"
  agent_pool: mi325_1
  depends_on: amd-build
@@ -84,7 +103,7 @@ steps:
  commands:
    - export VLLM_LOGGING_LEVEL=DEBUG
    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - timeout 20m pytest -s -v tests/diffusion/distributed/test_sequence_parallel.py -m core_model
+    - timeout 20m pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py -m core_model

 - label: "Diffusion GPU Worker Test"
  agent_pool: mi325_2
@@ -94,27 +113,6 @@ steps:
  commands:
    - timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py

- label: "Diffusion · Model Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - timeout 15m pytest -sv tests/diffusion/models/ -m 'core_model and cuda' --run-level "core_model"
-
- label: "Diffusion · Wan22 Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - |
-      timeout 40m bash -c "
-        export VLLM_TEST_CLEAN_GPU_MEMORY=1
-        export VLLM_IMAGE_FETCH_TIMEOUT=60
-        pytest -s -v tests/e2e/online_serving/test_wan22_t2v.py -m 'core_model' --run-level 'core_model'
-      "
-
 - label: "Engine Test"
  agent_pool: mi325_1
  depends_on: amd-build
@@ -135,6 +133,7 @@ steps:
 #   commands:
 #     - export VLLM_LOGGING_LEVEL=DEBUG
 #     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+#     - export VLLM_TEST_CLEAN_GPU_MEMORY=1
 #     - timeout 10m pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
 #     - timeout 20m pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "core_model" --run-level "core_model"

@@ -176,6 +175,17 @@ steps:
        pytest -s -v tests/e2e/offline_inference/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
      '

+- label: "Diffusion Image Edit Test"
+  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
+    - timeout 20m pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py
+
 # TODO: Bagel test on ROCm is very unstable. @tjtanaa
 # Need to debug before reneable numerical changes across large PRs
 # - label: "Bagel Text2Img Model Test"
@@ -185,10 +195,11 @@ steps:
 #   grade: Blocking
 #   commands:
 #     - export GPU_ARCHS=gfx942
+#     - export VLLM_TEST_CLEAN_GPU_MEMORY=1
 #     - export VLLM_LOGGING_LEVEL=DEBUG
 #     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
 #     - export VLLM_ROCM_USE_AITER_RMSNORM=0
-#     - timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel.py -m "core_model" --run-level "core_model" -k "rocm"
+#     - timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py -m "core_model" --run-level "core_model" -k "rocm"

 # - label: "Bagel Img2Img Model Test"
 #   agent_pool: mi325_1
@@ -197,10 +208,11 @@ steps:
 #   grade: Blocking
 #   commands:
 #     - export GPU_ARCHS=gfx942
+#     - export VLLM_TEST_CLEAN_GPU_MEMORY=1
 #     - export VLLM_LOGGING_LEVEL=DEBUG
 #     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
 #     - export VLLM_ROCM_USE_AITER_RMSNORM=0
-#     - timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel.py -m "core_model" --run-level "core_model" -k "rocm"
+#     - timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py -m "core_model" --run-level "core_model" -k "rocm"

 # - label: "Bagel Online Serving Test"
 #   agent_pool: mi325_1
@@ -209,11 +221,12 @@ steps:
 #   grade: Blocking
 #   commands:
 #     - export GPU_ARCHS=gfx942
+#     - export VLLM_TEST_CLEAN_GPU_MEMORY=1
 #     - export VLLM_IMAGE_FETCH_TIMEOUT=60
 #     - export VLLM_LOGGING_LEVEL=DEBUG
 #     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
 #     - export VLLM_ROCM_USE_AITER_RMSNORM=0
-#     - timeout 40m pytest -s -v tests/e2e/online_serving/test_bagel.py -m "core_model" --run-level "core_model" -k "rocm"
+#     - timeout 40m pytest -s -v tests/e2e/online_serving/test_bagel_online.py -m "core_model" --run-level "core_model" -k "rocm"

 - label: "CosyVoice3-TTS E2E Test"
  agent_pool: mi325_1
@@ -223,5 +236,5 @@ steps:
  commands:
    - |
      timeout 20m bash -c '
-        pytest -s -v tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py -m "full_model" --run-level "core_model"
+        pytest -s -v tests/e2e/online_serving/test_cosyvoice3_tts.py -m "core_model" --run-level "core_model"
      '
@@ -8,7 +8,7 @@ steps:
  grade: Blocking
  commands:
    - export GPU_ARCHS=gfx942
-    - pytest -s -v tests/e2e/offline_inference/test_z_image.py
+    - pytest -s -v tests/e2e/offline_inference/test_t2i_model.py

 - label: "Diffusion Images API LoRA E2E"
  timeout_in_minutes: 30
@@ -20,7 +20,7 @@ steps:
    - export GPU_ARCHS=gfx942
    - export VLLM_LOGGING_LEVEL=DEBUG
    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - pytest -s -v tests/diffusion/lora/test_images_generations_lora.py
+    - pytest -s -v tests/e2e/online_serving/test_images_generations_lora.py

 - label: "Diffusion Model CPU offloading Test"
  timeout_in_minutes: 20
@@ -32,8 +32,8 @@ steps:
    - export GPU_ARCHS=gfx942
    - export VLLM_LOGGING_LEVEL=DEBUG
    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py
-    - pytest -s -v tests/diffusion/offloader/test_diffusion_layerwise_offload.py
+    - pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
+    - pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py

 - label: "Diffusion Cache Backend Test"
  timeout_in_minutes: 15
@@ -57,7 +57,7 @@ steps:
    - export GPU_ARCHS=gfx942
    - export VLLM_LOGGING_LEVEL=DEBUG
    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - pytest -s -v tests/diffusion/distributed/test_sequence_parallel.py -m core_model
+    - pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py -m core_model

 - label: "Diffusion Tensor Parallelism Test"
  timeout_in_minutes: 20
@@ -67,7 +67,7 @@ steps:
    - export GPU_ARCHS=gfx942
    - export VLLM_LOGGING_LEVEL=DEBUG
    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - pytest -s -v tests/diffusion/distributed/test_zimage_parallelism.py
+    - pytest -s -v tests/e2e/offline_inference/test_zimage_parallelism.py

 - label: "Diffusion GPU Worker Test"
  timeout_in_minutes: 20
@@ -100,11 +100,25 @@ steps:
  commands:
    - export VLLM_LOGGING_LEVEL=DEBUG
    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - export VLLM_TEST_CLEAN_GPU_MEMORY="1"
    - export VLLM_ROCM_USE_AITER=0
    - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
    - pytest -s -v tests/e2e/online_serving/test_qwen3_omni.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
+
+
 - label: "Omni Sleep Mode Test"
  timeout_in_minutes: 40
  agent_pool: mi325_2
@@ -115,4 +129,5 @@ steps:
    - export GPU_ARCHS=gfx942
    - export VLLM_LOGGING_LEVEL=DEBUG
    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - pytest -s -v tests/entrypoints/test_omni_sleep_mode.py -m "advanced_model and omni and MI325" --run-level "advanced_model"
+    - export VLLM_TEST_CLEAN_GPU_MEMORY="1"
+    - pytest -s -v tests/e2e/offline_inference/test_omni_sleep_mode.py -m "advanced_model and omni and MI325" --run-level "advanced_model"
@@ -1,6 +1,5 @@
 env:
  VLLM_WORKER_MULTIPROC_METHOD: spawn
-  VLLM_USE_FLASHINFER_MOE_FP16: 0
  HF_HUB_DOWNLOAD_TIMEOUT: 300
  HF_HUB_ETAG_TIMEOUT: 60

@@ -14,7 +13,7 @@ steps:
      - label: ":full_moon: Omni · Function Test with H100"
        timeout_in_minutes: 90
        commands:
-          - pytest -sv tests/e2e --run-level "full_model" -m "full_model and H100 and omni" --ignore=tests/e2e/accuracy --ignore=tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py
+          - pytest -s -v tests/e2e/ -m "full_model and H100 and omni" --run-level "full_model" --ignore=tests/e2e/accuracy
        agents:
          queue: "mithril-h100-pool"
        plugins:
@@ -170,104 +169,14 @@ steps:
                      path: /mnt/hf-cache
                      type: DirectoryOrCreate

-      - label: ":full_moon: Omni · Perf Test · No Async Chunk"
-        key: nightly-omni-performance-no-async-chunk
+      - label: ":full_moon: Omni · Perf Test"
+        key: nightly-omni-performance
        timeout_in_minutes: 180
        commands:
          - export BENCHMARK_DIR=tests/dfx/perf/results
          - |
            set +e
-            pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen3_omni_no_async_chunk.json
-            EXIT=$$?
-            buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
-            exit $$EXIT
-        agents:
-          queue: "mithril-h100-pool"
-        plugins:
-          - kubernetes:
-              podSpec:
-                containers:
-                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-                    resources:
-                      limits:
-                        nvidia.com/gpu: 2
-                    volumeMounts:
-                      - name: devshm
-                        mountPath: /dev/shm
-                      - name: hf-cache
-                        mountPath: /root/.cache/huggingface
-                    env:
-                      - name: HF_HOME
-                        value: /root/.cache/huggingface
-                      - name: HF_TOKEN
-                        valueFrom:
-                          secretKeyRef:
-                            name: hf-token-secret
-                            key: token
-                nodeSelector:
-                  node.kubernetes.io/instance-type: gpu-h100-sxm
-                volumes:
-                  - name: devshm
-                    emptyDir:
-                      medium: Memory
-                  - name: hf-cache
-                    hostPath:
-                      path: /mnt/hf-cache
-                      type: DirectoryOrCreate
-
-      - label: ":full_moon: Omni · Perf Test · Async Chunk"
-        key: nightly-omni-performance-async-chunk
-        timeout_in_minutes: 180
-        commands:
-          - export BENCHMARK_DIR=tests/dfx/perf/results
-          - |
-            set +e
-            pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json
-            EXIT=$$?
-            buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
-            exit $$EXIT
-        agents:
-          queue: "mithril-h100-pool"
-        plugins:
-          - kubernetes:
-              podSpec:
-                containers:
-                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-                    resources:
-                      limits:
-                        nvidia.com/gpu: 2
-                    volumeMounts:
-                      - name: devshm
-                        mountPath: /dev/shm
-                      - name: hf-cache
-                        mountPath: /root/.cache/huggingface
-                    env:
-                      - name: HF_HOME
-                        value: /root/.cache/huggingface
-                      - name: HF_TOKEN
-                        valueFrom:
-                          secretKeyRef:
-                            name: hf-token-secret
-                            key: token
-                nodeSelector:
-                  node.kubernetes.io/instance-type: gpu-h100-sxm
-                volumes:
-                  - name: devshm
-                    emptyDir:
-                      medium: Memory
-                  - name: hf-cache
-                    hostPath:
-                      path: /mnt/hf-cache
-                      type: DirectoryOrCreate
-
-      - label: ":full_moon: Omni · Perf Test · vLLM Text"
-        key: nightly-omni-performance-vllm-text
-        timeout_in_minutes: 180
-        commands:
-          - export BENCHMARK_DIR=tests/dfx/perf/results
-          - |
-            set +e
-            pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen3_omni_vllm_text.json
+            pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_omni.json
            EXIT=$$?
            buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
            exit $$EXIT
@@ -309,7 +218,7 @@ steps:
      - label: ":full_moon: Omni · Multi-Replica Startup Test with 4x H100"
        timeout_in_minutes: 45
        commands:
-          - pytest -s -v tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py -m "full_model" --run-level "core_model"
+          - pytest -s -v tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py -m "core_model" --run-level "core_model"
        agents:
          queue: "mithril-h100-pool"
        plugins:
@@ -353,7 +262,7 @@ steps:
        timeout_in_minutes: 90
        commands:
          - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
-          - pytest -s -v tests/e2e/ -m "full_model and L4 and tts" --run-level "full_model" --ignore=tests/e2e/accuracy
+          - pytest -s -v tests/e2e/ -m "full_model and L4 and omni" -k "not test_qwen2_5" --run-level "full_model" --ignore=tests/e2e/accuracy
        agents:
          queue: "gpu_1_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
        plugins:
@@ -389,7 +298,7 @@ steps:
                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
                    resources:
                      limits:
-                        nvidia.com/gpu: 2
+                        nvidia.com/gpu: 1
                    volumeMounts:
                      - name: devshm
                        mountPath: /dev/shm
@@ -423,65 +332,13 @@ steps:
      build.pull_request.labels includes "nightly-test" ||
      build.pull_request.labels includes "diffusion-x2iat-test"
    steps:
-      # Use explicit file shards so single-card cases do not reserve 2+ GPUs,
-      # and Bagel feature coverage is not duplicated via a broad tests/e2e sweep.
-      - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100 · Single-GPU"
+      - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100"
        timeout_in_minutes: 120
        commands:
-          - >-
-            pytest -sv
-            tests/e2e/online_serving/test_qwen_image_expansion.py
-            tests/e2e/online_serving/test_qwen_image_edit_expansion.py
-            tests/e2e/online_serving/test_qwen_image_layered_expansion.py
-            tests/e2e/online_serving/test_longcat_image_expansion.py
-            tests/e2e/online_serving/test_longcat_image_edit_expansion.py
-            tests/e2e/online_serving/test_flux_2_dev_expansion.py
-            tests/e2e/online_serving/test_bagel_expansion.py
-            -m "full_model and diffusion and H100 and not distributed_cuda"
-            --run-level "full_model"
-        agents:
-          queue: "mithril-h100-pool"
-        plugins:
-          - kubernetes:
-              podSpec:
-                containers:
-                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-                    resources:
-                      limits:
-                        nvidia.com/gpu: 1
-                    volumeMounts:
-                      - name: devshm
-                        mountPath: /dev/shm
-                      - name: hf-cache
-                        mountPath: /root/.cache/huggingface
-                    env:
-                      - name: HF_HOME
-                        value: /root/.cache/huggingface
-                      - name: HF_TOKEN
-                        valueFrom:
-                          secretKeyRef:
-                            name: hf-token-secret
-                            key: token
-                nodeSelector:
-                  node.kubernetes.io/instance-type: gpu-h100-sxm
-                volumes:
-                  - name: devshm
-                    emptyDir:
-                      medium: Memory
-                  - name: hf-cache
-                    hostPath:
-                      path: /mnt/hf-cache
-                      type: DirectoryOrCreate
-
-      - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100 · Multi-GPU Qwen-Image"
-        timeout_in_minutes: 120
-        commands:
-          - >-
-            pytest -sv
-            tests/e2e/online_serving/test_qwen_image_expansion.py
-            tests/e2e/online_serving/test_qwen_image_edit_expansion.py
-            -m "full_model and diffusion and H100 and distributed_cuda"
-            --run-level "full_model"
+          - pytest -sv tests/e2e/ -k "not test_wan and not test_bagel_expansion and not hunyuan" -m "full_model and diffusion and H100" --run-level "full_model" --ignore=tests/e2e/accuracy
+          # revoke after bagel optimization finished
+          # Keep Bagel expansion and multi-replica tests in their dedicated H100 jobs below.
+          # - pytest -sv tests/e2e/ -k "not test_wan and not test_bagel_expansion and not hunyuan" -m "full_model and diffusion and H100" --run-level "full_model" --ignore=tests/e2e/accuracy --ignore=tests/e2e/online_serving/test_bagel_multi_replicas.py
        agents:
          queue: "mithril-h100-pool"
        plugins:
@@ -516,59 +373,10 @@ steps:
                      path: /mnt/hf-cache
                      type: DirectoryOrCreate

-      - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100 · Multi-GPU Layered/LongCat/FLUX"
-        timeout_in_minutes: 90
-        commands:
-          - >-
-            pytest -sv
-            tests/e2e/online_serving/test_qwen_image_layered_expansion.py
-            tests/e2e/online_serving/test_longcat_image_expansion.py
-            tests/e2e/online_serving/test_longcat_image_edit_expansion.py
-            tests/e2e/online_serving/test_flux_2_dev_expansion.py
-            -m "full_model and diffusion and H100 and distributed_cuda"
-            --run-level "full_model"
-        agents:
-          queue: "mithril-h100-pool"
-        plugins:
-          - kubernetes:
-              podSpec:
-                containers:
-                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-                    resources:
-                      limits:
-                        nvidia.com/gpu: 2
-                    volumeMounts:
-                      - name: devshm
-                        mountPath: /dev/shm
-                      - name: hf-cache
-                        mountPath: /root/.cache/huggingface
-                    env:
-                      - name: HF_HOME
-                        value: /root/.cache/huggingface
-                      - name: HF_TOKEN
-                        valueFrom:
-                          secretKeyRef:
-                            name: hf-token-secret
-                            key: token
-                nodeSelector:
-                  node.kubernetes.io/instance-type: gpu-h100-sxm
-                volumes:
-                  - name: devshm
-                    emptyDir:
-                      medium: Memory
-                  - name: hf-cache
-                    hostPath:
-                      path: /mnt/hf-cache
-                      type: DirectoryOrCreate
-
-      - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100 · Multi-GPU Bagel"
+      - label: ":full_moon: Diffusion X2I(&A&T) · Bagel Function Test with H100"
        timeout_in_minutes: 120
        commands:
-          - >-
-            pytest -sv
-            tests/e2e/online_serving/test_bagel_expansion.py
-            -m "full_model and diffusion and H100 and distributed_cuda"
-            --run-level "full_model"
+          - pytest -sv tests/e2e/online_serving/test_bagel_expansion.py -m "full_model and diffusion and H100" --run-level "full_model"
        agents:
          queue: "mithril-h100-pool"
        plugins:
@@ -603,6 +411,44 @@ steps:
                      path: /mnt/hf-cache
                      type: DirectoryOrCreate

+      # revoke after bagel optimization finished
+      # - label: ":full_moon: Diffusion X2I(&A&T) · BAGEL Multi-Replica with H100"
+      #   timeout_in_minutes: 120
+      #   commands:
+      #     - pytest -s -v tests/e2e/online_serving/test_bagel_multi_replicas.py -m "full_model and diffusion and H100" --run-level "full_model"
+      #   agents:
+      #     queue: "mithril-h100-pool"
+      #   plugins:
+      #     - kubernetes:
+      #         podSpec:
+      #           containers:
+      #             - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+      #               resources:
+      #                 limits:
+      #                   nvidia.com/gpu: 4
+      #               volumeMounts:
+      #                 - name: devshm
+      #                   mountPath: /dev/shm
+      #                 - name: hf-cache
+      #                   mountPath: /root/.cache/huggingface
+      #               env:
+      #                 - name: HF_HOME
+      #                   value: /root/.cache/huggingface
+      #                 - name: HF_TOKEN
+      #                   valueFrom:
+      #                     secretKeyRef:
+      #                       name: hf-token-secret
+      #                       key: token
+      #           nodeSelector:
+      #             node.kubernetes.io/instance-type: gpu-h100-sxm
+      #           volumes:
+      #             - name: devshm
+      #               emptyDir:
+      #                 medium: Memory
+      #             - name: hf-cache
+      #               hostPath:
+      #                 path: /mnt/hf-cache
+      #                 type: DirectoryOrCreate

      - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with L4"
        timeout_in_minutes: 60
@@ -625,6 +471,7 @@ steps:
      - label: ":full_moon: Diffusion X2I(&A&T) · Doc Test"
        timeout_in_minutes: 60
        commands:
+          - export VLLM_TEST_CLEAN_GPU_MEMORY="1"
          - pytest -s -v tests/examples/*/test_text_to_image.py -m "full_model and example and H100" --run-level "full_model"
        agents:
          queue: "mithril-h100-pool"
@@ -660,6 +507,87 @@ steps:
                      path: /mnt/hf-cache
                      type: DirectoryOrCreate

+      - label: ":full_moon: Diffusion X2I(&A&T) · GEBench Accuracy Test"
+        timeout_in_minutes: 60
+        commands:
+          - pytest -s -v tests/e2e/accuracy/test_gebench_h100_smoke.py --run-level full_model --gebench-model Qwen/Qwen-Image-2512 --accuracy-judge-model QuantTrio/Qwen3-VL-30B-A3B-Instruct-AWQ --accuracy-gpu 0 --gebench-port 8093 --accuracy-workers 1
+          - buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gebench_qwen-image-2512/summary*.json"
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 1
+                    volumeMounts:
+                      - name: devshm
+                        mountPath: /dev/shm
+                      - name: hf-cache
+                        mountPath: /root/.cache/huggingface
+                    env:
+                      - name: HF_HOME
+                        value: /root/.cache/huggingface
+                      - name: HF_TOKEN
+                        valueFrom:
+                          secretKeyRef:
+                            name: hf-token-secret
+                            key: token
+                nodeSelector:
+                  node.kubernetes.io/instance-type: gpu-h100-sxm
+                volumes:
+                  - name: devshm
+                    emptyDir:
+                      medium: Memory
+                  - name: hf-cache
+                    hostPath:
+                      path: /mnt/hf-cache
+                      type: DirectoryOrCreate
+
+      - label: ":full_moon: Diffusion X2I(&A&T) · GEdit-Bench Accuracy Test"
+        timeout_in_minutes: 60
+        commands:
+          - pytest -s -v tests/e2e/accuracy/test_gedit_bench_h100_smoke.py --run-level full_model --gedit-model Qwen/Qwen-Image-Edit --accuracy-judge-model QuantTrio/Qwen3-VL-30B-A3B-Instruct-AWQ --accuracy-gpu 0 --gedit-port 8093 --gedit-samples-per-group 20 --accuracy-workers 1
+          - buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gedit_scores_qwen-image-edit/qwen-image-edit_all_all_vie_score_*.csv"
+          - buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gedit_scores_qwen-image-edit/qwen-image-edit_all_all_summary_*.json"
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 1
+                    volumeMounts:
+                      - name: devshm
+                        mountPath: /dev/shm
+                      - name: hf-cache
+                        mountPath: /root/.cache/huggingface
+                    env:
+                      - name: HF_HOME
+                        value: /root/.cache/huggingface
+                      - name: VLLM_HTTP_TIMEOUT_KEEP_ALIVE
+                        value: "120"
+                      - name: HF_TOKEN
+                        valueFrom:
+                          secretKeyRef:
+                            name: hf-token-secret
+                            key: token
+                nodeSelector:
+                  node.kubernetes.io/instance-type: gpu-h100-sxm
+                volumes:
+                  - name: devshm
+                    emptyDir:
+                      medium: Memory
+                  - name: hf-cache
+                    hostPath:
+                      path: /mnt/hf-cache
+                      type: DirectoryOrCreate
+
      - label: ":full_moon: Diffusion X2I(&A&T) · Accuracy Test"
        timeout_in_minutes: 180
        commands:
@@ -701,116 +629,28 @@ steps:
                      path: /mnt/hf-cache
                      type: DirectoryOrCreate

-      - label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image"
-        key: nightly-diffusion-x2iat-performance-qwen-image
-        timeout_in_minutes: 180
+      - label: ":full_moon: Diffusion X2I(&A&T) · Perf Test"
+        key: nightly-diffusion-x2iat-performance
+        timeout_in_minutes: 360
        commands:
          - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
          - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
          - export CACHE_DIT_VERSION=1.3.0
+          # [HACK]: run upload in the same command block as pytest.
+          # Because `exit` aborts the entire commands list.
          - |
            set +e
            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json
-            EXIT=$$?
-            buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
-            buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
-            exit $$EXIT
-        agents:
-          queue: "mithril-h100-pool"
-        plugins:
-          - kubernetes:
-              podSpec:
-                containers:
-                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-                    resources:
-                      limits:
-                        nvidia.com/gpu: 4
-                    volumeMounts:
-                      - name: devshm
-                        mountPath: /dev/shm
-                      - name: hf-cache
-                        mountPath: /root/.cache/huggingface
-                    env:
-                      - name: HF_HOME
-                        value: /root/.cache/huggingface
-                      - name: HF_TOKEN
-                        valueFrom:
-                          secretKeyRef:
-                            name: hf-token-secret
-                            key: token
-                nodeSelector:
-                  node.kubernetes.io/instance-type: gpu-h100-sxm
-                volumes:
-                  - name: devshm
-                    emptyDir:
-                      medium: Memory
-                  - name: hf-cache
-                    hostPath:
-                      path: /mnt/hf-cache
-                      type: DirectoryOrCreate
-
-      - label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image-Edit"
-        key: nightly-diffusion-x2iat-performance-qwen-image-edit
-        timeout_in_minutes: 180
-        commands:
-          - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
-          - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
-          - export CACHE_DIT_VERSION=1.3.0
-          - |
-            set +e
+            EXIT1=$$?
            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json
-            EXIT=$$?
-            buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
-            buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
-            exit $$EXIT
-        agents:
-          queue: "mithril-h100-pool"
-        plugins:
-          - kubernetes:
-              podSpec:
-                containers:
-                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-                    resources:
-                      limits:
-                        nvidia.com/gpu: 4
-                    volumeMounts:
-                      - name: devshm
-                        mountPath: /dev/shm
-                      - name: hf-cache
-                        mountPath: /root/.cache/huggingface
-                    env:
-                      - name: HF_HOME
-                        value: /root/.cache/huggingface
-                      - name: HF_TOKEN
-                        valueFrom:
-                          secretKeyRef:
-                            name: hf-token-secret
-                            key: token
-                nodeSelector:
-                  node.kubernetes.io/instance-type: gpu-h100-sxm
-                volumes:
-                  - name: devshm
-                    emptyDir:
-                      medium: Memory
-                  - name: hf-cache
-                    hostPath:
-                      path: /mnt/hf-cache
-                      type: DirectoryOrCreate
-
-      - label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image-Edit-2509"
-        key: nightly-diffusion-x2iat-performance-qwen-image-edit-2509
-        timeout_in_minutes: 180
-        commands:
-          - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
-          - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
-          - export CACHE_DIT_VERSION=1.3.0
-          - |
-            set +e
+            EXIT2=$$?
            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_2509_vllm_omni.json
-            EXIT=$$?
+            EXIT3=$$?
+            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json
+            EXIT4=$$?
            buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
            buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
-            exit $$EXIT
+            exit $$((EXIT1 | EXIT2 | EXIT3 | EXIT4))
        agents:
          queue: "mithril-h100-pool"
        plugins:
@@ -845,54 +685,6 @@ steps:
                      path: /mnt/hf-cache
                      type: DirectoryOrCreate

-      - label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image-Layered"
-        key: nightly-diffusion-x2iat-performance-qwen-image-layered
-        timeout_in_minutes: 120
-        commands:
-          - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
-          - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
-          - export CACHE_DIT_VERSION=1.3.0
-          - |
-            set +e
-            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json
-            EXIT=$$?
-            buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
-            buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
-            exit $$EXIT
-        agents:
-          queue: "mithril-h100-pool"
-        plugins:
-          - kubernetes:
-              podSpec:
-                containers:
-                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-                    resources:
-                      limits:
-                        nvidia.com/gpu: 1
-                    volumeMounts:
-                      - name: devshm
-                        mountPath: /dev/shm
-                      - name: hf-cache
-                        mountPath: /root/.cache/huggingface
-                    env:
-                      - name: HF_HOME
-                        value: /root/.cache/huggingface
-                      - name: HF_TOKEN
-                        valueFrom:
-                          secretKeyRef:
-                            name: hf-token-secret
-                            key: token
-                nodeSelector:
-                  node.kubernetes.io/instance-type: gpu-h100-sxm
-                volumes:
-                  - name: devshm
-                    emptyDir:
-                      medium: Memory
-                  - name: hf-cache
-                    hostPath:
-                      path: /mnt/hf-cache
-                      type: DirectoryOrCreate
-
      - label: ":full_moon: Diffusion X2I(&A&T) · HunyuanImage3 · DiT Perf Test"
        key: nightly-hunyuan-image3-performance
        soft_fail: true
@@ -959,7 +751,7 @@ steps:
      - label: ":full_moon: Diffusion X2V · Function Test"
        timeout_in_minutes: 90
        commands:
-          - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py tests/e2e/online_serving/test_wan_2_1_vace_expansion.py tests/e2e/online_serving/test_hunyuan_video_15_expansion.py  -m "full_model and cuda" --run-level "full_model"
+          - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py tests/e2e/online_serving/test_wan_2_1_vace_expansion.py tests/e2e/online_serving/test_hunyuan_video_15_expansion.py  --run-level "full_model"
        agents:
          queue: "mithril-h100-pool"
        plugins:
@@ -998,7 +790,7 @@ steps:
        timeout_in_minutes: 180
        commands:
          - pytest -s -v tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py -m full_model --run-level full_model
-          - pytest -s -v tests/e2e/accuracy/test_ltx2_3_video_similarity.py -m full_model --run-level full_model
+          - pytest -s -v tests/e2e/accuracy/test_ltx2_3_video_similarity.py -m advanced_model --run-level advanced_model
          - pytest -s -v tests/e2e/accuracy/test_diffusers_backend_similarity.py -k '2v_matches_diffusers' -m full_model --run-level full_model
        agents:
          queue: "mithril-h100-pool"
@@ -1081,65 +873,6 @@ steps:
                      path: /mnt/hf-cache
                      type: DirectoryOrCreate

-  - group: ":card_index_dividers: Diffusion · Quantization Test"
-    depends_on: upload-nightly-pipeline
-    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" || build.pull_request.labels includes "diffusion-x2iat-test"
-    steps:
-      - label: ":full_moon: Quantization Test with H100"
-        timeout_in_minutes: 90
-        commands:
-          - timeout 30m pytest -sv tests/diffusion/quantization -m 'full_model and cuda and H100' --run-level "full_model"
-        agents:
-          queue: "mithril-h100-pool"
-        plugins:
-          - kubernetes:
-              podSpec:
-                containers:
-                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-                    resources:
-                      limits:
-                        nvidia.com/gpu: 1
-                    volumeMounts:
-                      - name: devshm
-                        mountPath: /dev/shm
-                      - name: hf-cache
-                        mountPath: /root/.cache/huggingface
-                    env:
-                      - name: HF_HOME
-                        value: /root/.cache/huggingface
-                      - name: HF_TOKEN
-                        valueFrom:
-                          secretKeyRef:
-                            name: hf-token-secret
-                            key: token
-                nodeSelector:
-                  node.kubernetes.io/instance-type: gpu-h100-sxm
-                volumes:
-                  - name: devshm
-                    emptyDir:
-                      medium: Memory
-                  - name: hf-cache
-                    hostPath:
-                      path: /mnt/hf-cache
-                      type: DirectoryOrCreate
-
-      - label: ":full_moon: Quantization Test with L4"
-        timeout_in_minutes: 60
-        commands:
-          - pytest -sv tests/diffusion/quantization -m 'full_model and cuda and L4' --run-level "full_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"
-                - "HF_TOKEN"
-              volumes:
-                - "/fsx/hf_cache:/fsx/hf_cache"
-
  - label: ":bar_chart: Testcase Statistics"
    key: nightly-testcase-statistics
    timeout_in_minutes: 120
@@ -1157,7 +890,7 @@ steps:
              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
                resources:
                  limits:
-                    nvidia.com/gpu: 1
+                    nvidia.com/gpu: 2
                volumeMounts:
                  - name: devshm
                    mountPath: /dev/shm
@@ -1186,14 +919,9 @@ steps:
  - label: ":email: Nightly Collection & Email"
    key: nightly-perf-distribution
    depends_on:
-      - nightly-omni-performance-no-async-chunk
-      - nightly-omni-performance-async-chunk
-      - nightly-omni-performance-vllm-text
+      - nightly-omni-performance
      - nightly-tts-performance
-      - nightly-diffusion-x2iat-performance-qwen-image
-      - nightly-diffusion-x2iat-performance-qwen-image-edit
-      - nightly-diffusion-x2iat-performance-qwen-image-edit-2509
-      - nightly-diffusion-x2iat-performance-qwen-image-layered
+      - nightly-diffusion-x2iat-performance
      - nightly-diffusion-x2v-performance
      - nightly-testcase-statistics
    if: build.env("NIGHTLY") == "1"
@@ -1202,13 +930,8 @@ steps:
      - export DEFAULT_INPUT_DIR=tests/dfx/perf/results
      - export DEFAULT_OUTPUT_DIR=tests/dfx/perf/results
      - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-tts-performance
-      - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-no-async-chunk
-      - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-async-chunk
-      - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-vllm-text
-      - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image
-      - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-edit
-      - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-edit-2509
-      - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-layered
+      - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance
+      - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance
      - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2v-performance
      - buildkite-agent artifact download "tests/dfx/perf/results/*.html" . --step nightly-testcase-statistics
      - python tools/nightly/generate_nightly_perf_excel.py
@@ -1,24 +0,0 @@
-steps:
-  - group: ":card_index_dividers: Diffusion X2V Model Test"
-    key: nightly-diffusion-x2v-group
-    depends_on: upload-nightly-pipeline
-    if: >-
-      build.env("NIGHTLY") == "1" ||
-      build.pull_request.labels includes "nightly-test" ||
-      build.pull_request.labels includes "diffusion-x2v-test"
-    steps:
-      - label: ":full_moon: Diffusion X2V · Function Test"
-        agents:
-          queue: "ascend-a2b3"
-          resource_class: "npu-4"
-        image: "swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/vllm-omni-ci-npu:${BUILDKITE_COMMIT}"
-        plugins:
-          - kubernetes:
-              podSpecPatch:
-                imagePullSecrets:
-                  - name: swr-secret
-        env:
-          VLLM_WORKER_MULTIPROC_METHOD: spawn
-          HF_TOKEN: "${HF_TOKEN}"
-        commands:
-          - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k npu_i2v --run-level "full_model"
@@ -1,21 +0,0 @@
-steps:
-  - group: ":card_index_dividers: TTS Model Test"
-    key: ready-tts-group
-    depends_on: upload-ready-pipeline
-    steps:
-      - label: ":sound: TTS · Qwen3-TTS CustomVoice Test"
-        agents:
-          queue: "ascend-a2b3"
-          resource_class: "npu-1"
-        image: "swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/vllm-omni-ci-npu:${BUILDKITE_COMMIT}"
-        plugins:
-          - kubernetes:
-              podSpecPatch:
-                imagePullSecrets:
-                  - name: swr-secret
-        env:
-          VLLM_WORKER_MULTIPROC_METHOD: spawn
-          VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
-          HF_TOKEN: "${HF_TOKEN}"
-        commands:
-          - pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m 'core_model' --run-level 'core_model'
@@ -1,6 +1,5 @@
 env:
  VLLM_WORKER_MULTIPROC_METHOD: spawn
-  VLLM_USE_FLASHINFER_MOE_FP16: 0
  HF_HUB_DOWNLOAD_TIMEOUT: 300
  HF_HUB_ETAG_TIMEOUT: 60

@@ -84,65 +83,3 @@ steps:
                hostPath:
                  path: /mnt/hf-cache
                  type: DirectoryOrCreate
-
-  # Invalid HTTP / WS payloads vs live Omni servers (Issue #3649 tracks known skips).
-  - group: "Reliability Test - Invalid parameters Test"
-    depends_on: upload-weekly-pipeline
-    if: build.env("WEEKLY") == "1" || build.pull_request.labels includes "weekly-test"
-    steps:
-      - label: "Invalid parameters Test · H100"
-        timeout_in_minutes: 120
-        commands:
-          - pytest -s -v tests/dfx/reliability/invalid_param_test/ -m "slow and H100"
-        agents:
-          queue: "mithril-h100-pool"
-        plugins:
-          - kubernetes:
-              podSpec:
-                containers:
-                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-                    resources:
-                      limits:
-                        nvidia.com/gpu: 2
-                    volumeMounts:
-                      - name: devshm
-                        mountPath: /dev/shm
-                      - name: hf-cache
-                        mountPath: /root/.cache/huggingface
-                    env:
-                      - name: HF_HOME
-                        value: /root/.cache/huggingface
-                      - name: HF_TOKEN
-                        valueFrom:
-                          secretKeyRef:
-                            name: hf-token-secret
-                            key: token
-                nodeSelector:
-                  node.kubernetes.io/instance-type: gpu-h100-sxm
-                volumes:
-                  - name: devshm
-                    emptyDir:
-                      medium: Memory
-                  - name: hf-cache
-                    hostPath:
-                      path: /mnt/hf-cache
-                      type: DirectoryOrCreate
-
-      # Stable Audio Open, Qwen3-TTS speech/batch/voices paths marked ``L4``.
-      - label: "Invalid parameters Test · L4"
-        timeout_in_minutes: 120
-        commands:
-          - pytest -s -v tests/dfx/reliability/invalid_param_test/ -m "slow and L4"
-        agents:
-          queue: "gpu_1_queue"
-        plugins:
-          - docker#v5.2.0:
-              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-              always-pull: true
-              propagate-environment: true
-              shm-size: "8gb"
-              environment:
-                - "HF_HOME=/fsx/hf_cache"
-                - "HF_TOKEN"
-              volumes:
-                - "/fsx/hf_cache:/fsx/hf_cache"
@@ -515,15 +515,15 @@ After adding parallelism support, update:

 ### Step 11: Add CPU Offload Support

-Implement `SupportsComponentDiscovery` on your pipeline class to enable
+Implement `SupportsModuleOffload` on your pipeline class to enable
 `--enable-cpu-offload` and `--enable-layerwise-offload`. The protocol
 declares which submodules the offloader should manage:

 ```python
 from typing import ClassVar
-from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
+from vllm_omni.diffusion.models.interface import SupportsModuleOffload

-class YourPipeline(nn.Module, SupportsComponentDiscovery):
+class YourPipeline(nn.Module, SupportsModuleOffload):
    _dit_modules: ClassVar[list[str]] = ["transformer"]
    _encoder_modules: ClassVar[list[str]] = ["text_encoder"]
    _vae_modules: ClassVar[list[str]] = ["vae"]
@@ -540,7 +540,7 @@ class YourPipeline(nn.Module, SupportsComponentDiscovery):
 All attribute names support dotted paths for nested submodules
 (e.g. `"pipe.transformer"`, `"bagel.time_embedder"`).

-Pipelines without `SupportsComponentDiscovery` fall back to scanning
+Pipelines without `SupportsModuleOffload` fall back to scanning
 well-known attribute names (`transformer`, `text_encoder`, `vae`,
 etc.), which fails for non-standard names.

@@ -92,7 +92,6 @@ class TPSelfAttention(nn.Module):
            softmax_scale=1.0 / (self.head_dim ** 0.5),
            causal=False,
            num_kv_heads=self.to_qkv.num_kv_heads,  # Local KV heads per GPU
-            role="self",
        )

    def forward(self, x):
@@ -61,7 +61,6 @@ class SelfAttentionBlock(nn.Module):
            softmax_scale=1.0 / (self.head_dim ** 0.5),
            causal=False,
            num_kv_heads=num_heads,
-            role="self",
        )

    def forward(self, x, attn_mask=None):
@@ -76,24 +75,6 @@ class SelfAttentionBlock(nn.Module):
        return self.to_out(out)
 ```

-**Declaring the attention role.** Always pass a `role` string to `Attention(...)`. It is the lookup key users target with `--diffusion-attention-config.per_role.<role>.backend` to swap kernels per site without touching model code.
-
-| Convention | When to use |
-|---|---|
-| `role="self"` | Q/K/V come from the same hidden state |
-| `role="cross"` | K/V come from a separate `encoder_hidden_states`; pair with `skip_sequence_parallel=True` if K/V is replicated across SP ranks |
-
-For multi-modal sites that don't fit `self` / `cross`, use a dot-namespaced role and provide `role_category` so user config can fall back to the generic category:
-
-```python
-self.audio_to_video_attn = Attention(
-    num_heads=num_heads, head_size=self.head_dim,
-    softmax_scale=1.0 / (self.head_dim ** 0.5), causal=False,
-    role="mymodel.audio_to_video",
-    role_category="cross",
-)
-```
-
 #### Fused QKV with TP (Advanced)

 For tensor parallelism, use vLLM's parallel linear layers:
@@ -123,7 +104,6 @@ class TPSelfAttention(nn.Module):
            softmax_scale=1.0 / (self.head_dim ** 0.5),
            causal=False,
            num_kv_heads=num_heads,
-            role="self",
        )
 ```

@@ -68,7 +68,7 @@ out = out.reshape(batch, seq_len, -1)
 ### Different Output vs Diffusers Reference

 **Common causes**:
-1. **Attention backend difference**: FlashAttention vs SDPA may produce slightly different results. Force SDPA via `--diffusion-attention-backend TORCH_SDPA` (or `DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA`) to match diffusers; for finer-grained per-role overrides see `--diffusion-attention-config.per_role.<role>.backend`
+1. **Attention backend difference**: FlashAttention vs SDPA may produce slightly different results. Set `DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA` to match diffusers
 2. **Float precision**: vLLM-Omni may use bfloat16 where diffusers uses float32 for some operations
 3. **Missing normalization**: Check all LayerNorm/RMSNorm are preserved
 4. **Scheduler rounding**: Some schedulers have numerical sensitivity
@@ -1,488 +0,0 @@
---
-name: diffusion-perf-opt
-description: Diagnose and optimize vLLM Omni diffusion workloads, especially Wan/Qwen/Flux-style image and video generation. Use when Codex is asked to analyze profiling traces, choose parallel strategies, inspect torch profiler trace.json or trace.json.gz timelines, estimate optimization ROI, investigate GPU idle/free bubbles, compare USP/CFG/HSDP/VAE parallelism, or design operator/host/quantization optimizations for vLLM Omni.
---
-
-# vLLM Omni Performance Optimization
-
-Use this skill to run a disciplined optimization loop for vLLM Omni diffusion workloads. Keep two ideas separate: real performance baselines are collected with low overhead, while torch profiler traces are diagnostic artifacts and may distort latency.
-
-## First Questions
-
-Before proposing changes, ask for the optimization scene if it is not already known:
-
- GPU model, card count, topology, and whether NVLink is present.
- Model and pipeline, for example Wan2.2 I2V A14B.
- User workload: resolution, frames, steps, batch/concurrency, CFG scales, prompt/image inputs.
- Three runnable commands for each model/workload, or scripts that generate them:
-  - **Server startup command**: exact `vllm serve` command, environment variables, model path, port, parallelism flags, profiler flags, and precision/compile settings.
-  - **User request command**: exact single-request client command (for example curl or Python client) with fixed prompt/media/seed/size/steps. Use this to validate correctness and collect per-request stage timings.
-  - **User benchmark command**: exact repeatable benchmark command or script with warmup count, measured iteration count, concurrency/batch policy, output directory, and summary format.
- Current enabled strategies: USP/SP, CFG parallel, HSDP/FSDP, VAE patch parallel, torch.compile, profiler options.
- Optimization target: latency, throughput, memory, cost, or quality-preserving speed.
- Precision/quality tolerance: bf16/fp8/quantization/sparsity/approximate attention allowed or not.
-
-## Workflow
-
-1. **Freeze the measurement protocol and commands.**
-   - Before analyzing parallel strategies or traces, establish the three commands for the target model/workload: server startup, single user request, and user benchmark.
-   - Prefer checked-in scripts or one-off shell scripts over free-form commands in chat. The scripts should make workload variables explicit, including model path, port/base URL, prompt/media inputs, resolution, frames, steps, seed, warmup, measured iterations, concurrency, and output path.
-   - Measurement reports must preserve the concrete process, not just the final numbers. For each tested configuration, write down the server startup command, the user request command, the user benchmark command or polling loop, the final server response metrics, and the client-side timing/HTTP result.
-   - Present all duration and latency values in **milliseconds (ms)** in tables and summaries. If an API returns seconds, convert to ms before reporting. Keep raw units only inside quoted raw logs or source snippets.
-   - Keep profiler modes separate in the command set:
-     - Baseline/benchmark commands should avoid torch profiler and stack collection.
-     - Diagnostic commands may enable torch profiler after a baseline identifies a bottleneck.
-   - If any of the three commands is missing, ask for it or create a proposed script before proceeding.
-   - Make clear which command is authoritative for correctness validation and which command is authoritative for performance numbers.
-
-2. **Collect a real baseline.**
-   - Disable PyTorch profiler and stack collection.
-   - Disable or fix torch.compile state so A/B is fair.
-   - Avoid `--enforce-eager` for production-speed baselines unless eager is the target.
-   - Prefer `--log-stats` and `--enable-diffusion-pipeline-profiler` for low-overhead stage timing. PR 3069 has the relevant metrics/log-stats changes; if local code does not include them, fetch or cherry-pick the minimal metrics changes rather than merging unrelated PR drift.
-   - Run warmup requests and exclude first-request lazy init.
-
-3. **Model the parallel strategy before testing.**
-   - Estimate compute and communication for self-attention, cross-attention, FFN, CFG, VAE encode/decode, and HSDP.
-   - Select a small candidate matrix rather than exhaustively testing everything.
-   - Typical 2-card candidates: `USP=2`, `CFG=2`, `USP=1/HSDP on-off`, VAE parallel on-off if memory allows.
-   - Typical 4-card candidates: `USP=4`, `CFG=2 x USP=2`, `USP=2 x HSDP`, VAE parallel on-off.
-   - Typical 8-card candidates for official CFG-enabled video diffusion models:
-     - Primary candidate: `CFG=2 x USP=4` with VAE patch parallel across all 8 ranks.
-     - Compare against: `CFG=1 x USP=8` to test whether larger sequence parallel groups beat CFG branch parallelism.
-     - Isolate VAE parallelism: keep the DiT strategy fixed and compare VAE patch parallel on/off or different VAE patch sizes.
-     - If long-video `diffuse` remains dominant and Ulysses all-to-all is suspected, test a hybrid sequence strategy such as `CFG=2 x USP=2 x Ring=2`.
-     - Test HSDP off only if the model fits without it; treat HSDP primarily as a memory strategy until A/B proves otherwise.
-   - Prefer USP/SP for long video token sequences; prefer CFG parallel when CFG doubles transformer forwards and sequence length is modest.
-   - Convert workload shape to latent/patch token counts before choosing candidates. For video models, record frames, latent frames, latent HxW, patch size, approximate token count, and which stages should scale with the token count.
-
-4. **Search for the best parallel configuration.**
-   - Start with a small matrix that answers one question per comparison:
-     - `CFG parallel vs larger USP`: compare `CFG=2 x USP=world/2` against `CFG=1 x USP=world` for CFG-enabled workloads.
-     - `VAE patch parallel value`: compare the best DiT strategy with VAE patch parallel enabled and disabled.
-     - `HSDP cost`: compare HSDP on/off only when both configurations fit memory.
-     - `Ulysses vs Ring`: test a Ring/Ulysses hybrid only after long-sequence `diffuse` is confirmed dominant.
-   - For each configuration, create or record a stable config id, for example `A_cfg2_usp4_vaepp8_hsdp_tiling`.
-   - For each config id, capture:
-     - exact server command and environment variables,
-     - observed distributed setup from logs, such as SP groups, CFG groups, HSDP shard/replicate sizes, VAE patch size,
-     - exact request command for every scenario,
-     - final server response metrics and client-side elapsed time,
-     - output artifact paths and any failed/empty responses.
-   - Use one warmup request per scenario, then at least three measured repeats for the shortlisted config. For exploratory matrix pruning, one measured request is acceptable only if the margin is large; label it as one-shot.
-   - Compare configurations by stage, not only end-to-end latency. A configuration can improve `diffuse` while hurting `vae.decode`; record both effects.
-   - Select the best config only after checking the target metric, dominant stages, memory headroom, and output correctness.
-
-5. **Run targeted A/B tests.**
-   - Change one variable per test.
-   - Keep model, input, seed, request parameters, GPU placement, and warmup policy fixed.
-   - Record latency, stage timings, memory, output quality, and logs.
-   - Report comparison tables in ms. Include at least end-to-end client time, server `inference_time`, server stage generation time if available, `vae.encode`, `diffuse`, `vae.decode`, and peak memory.
-
-6. **Collect diagnostic trace only after narrowing hypotheses.**
-   - Use torch profiler for a small number of requests.
-   - Run two separate diagnostic traces instead of mixing concerns:
-     - **Operator/shape trace**: enable `torch_profiler_record_shapes=True` and keep stack collection disabled. Use this to rank CUDA kernels, NCCL collectives, attention/MLP/norm/RoPE work, and shape-specific hot operators.
-     - **Host-stack trace**: enable `torch_profiler_with_stack=True` and normally keep shape collection disabled. Use this to map CPU/Python host gaps, synchronization points, scheduler paths, and request handling overhead.
-   - Keep both trace commands and reports separate from baseline/benchmark commands. Torch profiler latency is diagnostic only and must not be used as the final latency claim.
-   - Prefer profiling only the narrowed dominant scenario, for example the highest-resolution/video-length case where `diffuse` dominates.
-   - If profiler endpoints are available, run one warmup request first, then call `/start_profile`, run one profiled request, and call `/stop_profile`. This keeps model initialization and warmup out of the diagnostic trace.
-   - Start by analyzing rank 0 only. Expand to more ranks only if rank 0 suggests imbalance, unclear GPU idle/free bubbles, high NCCL wait, server timing mismatch, CFG branch imbalance, or USP group stragglers.
-   - Diagnostic reports must be written to disk and preserve: server command, profiler config, warmup/request/polling commands, trace artifact paths, rank analyzed, analyzer output or summary, and the decision about whether additional ranks are necessary.
-   - Analyze both rank-level balance and device-level free bubbles when additional ranks are opened.
-
-7. **Analyze host, communication, and operators.**
-   - Find GPU idle/free intervals and map each large gap to the enclosing CPU/Python code.
-   - Separate real GPU idle from profiler overhead such as CUPTI `Command Buffer Full`.
-   - Compare NCCL kernel time to user annotations; annotations can overcount nested intervals.
-   - Rank operator work by total CUDA time and by repeated small-kernel launch count.
-
-8. **Produce an optimization plan.**
-   - Classify candidates as P0/P1/P2.
-   - For each candidate, state necessity, expected benefit, implementation path, validation plan, and quality risk.
-   - Do not implement high-risk operator rewrites before proving the operator is a bottleneck for the target shapes.
-   - End the plan with a user-facing candidate selection table. The assistant
-     should not automatically choose a risky optimization just because it is
-     technically possible. Present the options clearly and let the user decide
-     which item is worth implementing based on latency target, engineering
-     budget, memory headroom, and quality tolerance.
-   - Organize the plan by optimization layer, not by a flat list of ideas:
-     - host/runtime optimization,
-     - measurement/benchmark reliability,
-     - parallelism and communication,
-     - VAE encode/decode and media pre/post processing,
-     - operator fusion and layout cleanup,
-     - attention main-path optimization,
-     - algorithmic, precision, or approximation changes.
-   - For each layer, tie every candidate back to evidence from baseline
-     metrics, diagnostic trace, source code, or output quality requirements.
-     Do not list generic optimizations without a trace or workload reason.
-
-## Priority Rules
-
- **P0:** low risk, likely useful, or required for trustworthy measurement. Examples: real baseline, warmup, targeted parallel A/B, disabling avoidable `empty_cache`, scheduler coefficient caching.
- **P1:** meaningful code changes with contained risk. Examples: cross-attention KV caching, VAE gather/broadcast reduction, AdaLayerNorm/RMSNorm/RoPE fusion after trace evidence.
- **P2:** high implementation or quality risk. Examples: FA to LA replacement, custom Triton/CUDA fused kernels, FP8/quantization, sparsity/Rainfusion-style acceleration.
-
-Every implemented optimization needs A/B validation and quality regression. A/B means same workload and hardware before/after. Quality regression means checking generated image/video stability, artifacts, temporal flicker, and seed behavior when precision or approximate kernels change.
-
-## Optimization Layers
-
-After baseline, parallel-search, and diagnostic traces, summarize optimization
-opportunities by layer. This is the core of the performance analysis: the goal
-is to connect evidence to a scoped implementation and a validation plan.
-
-### Host and Runtime Optimization
-
-Purpose: remove CPU/Python stalls, synchronization points, allocator overhead,
-and request-path overhead that leave GPU lanes empty.
-
-Evidence to look for:
-
- High `idle_pct` or large `GAP` blocks in `trace_analyzer.py`.
- Host-stack trace lines such as `torch.cuda.empty_cache`,
-  `cudaStreamSynchronize`, `cudaDeviceSynchronize`, Python locks, scheduler
-  waits, image/video preprocessing, or repeated small allocation paths.
- Difference between client wall-clock time and server `inference_time_s`.
-
-Typical candidates:
-
- Make avoidable `torch.cuda.empty_cache()` optional or guard it by memory
-  headroom.
- Cache scheduler coefficients, timesteps, masks, or other tiny repeated CPU
-  computations when the request shape/steps are fixed.
- Remove avoidable host-device synchronizations and blocking logging/stat calls.
- Move expensive preprocessing out of the critical path or cache fixed prompt,
-  image, and transform work for benchmark scenarios.
- Ensure benchmark scripts record client-side elapsed time, HTTP status, output
-  path, and server response metrics.
-
-Priority guidance:
-
- Usually P0 when the change is measurement reliability or an obvious removable
-  synchronization.
- Usually P1 when it changes scheduling, memory lifetime, or request execution
-  order.
-
-Validation:
-
- Re-run non-profiler baseline with same workload and seed.
- Confirm peak memory headroom if disabling cache cleanup.
- Confirm generated output exists and quality/seed behavior is unchanged.
-
-### Parallelism and Communication Optimization
-
-Purpose: choose the right decomposition for CFG branches, sequence tokens, model
-weights, VAE tiles, and rank topology.
-
-Evidence to look for:
-
- Baseline A/B across `CFG`, `USP/SP`, `Ring`, `HSDP/FSDP`, and VAE patch
-  parallelism.
- Stage timing shifts: `diffuse`, `vae.encode`, `vae.decode`, and server
-  end-to-end.
- NCCL kernel time from trace, not only `user_annotation` time.
- Rank imbalance across SP group ranks or CFG branch ranks.
- Memory headroom and OOM risk.
-
-Typical candidates:
-
- `CFG=2 x USP=world/2` versus `CFG=1 x USP=world` for CFG-enabled models.
- VAE patch parallel on/off or patch size tuning.
- HSDP/FSDP on/off only if both configurations fit memory.
- Ulysses versus Ulysses+Ring only after long-sequence `diffuse` is confirmed
-  dominant and all-to-all is suspected.
- Rank mapping/topology changes if all-rank traces show stragglers or NCCL wait.
- Buffer reuse or preallocation for FSDP/HSDP all-gather paths.
-
-Priority guidance:
-
- P0/P1 for configuration-only changes with strong measured wins.
- P1 for buffer reuse or rank mapping changes.
- P2 for invasive distributed algorithm changes.
-
-Validation:
-
- Measure by stage and memory, not only end-to-end.
- Use one-variable A/B with identical prompt/media/seed/shape/steps.
- When communication is suspected, compare rank0-3 in one USP group and rank0
-  versus rank4 across CFG branches for `CFG=2 x USP=4`.
-
-### VAE and Media Pipeline Optimization
-
-Purpose: reduce encode/decode, tiling, split/gather, and media conversion time.
-
-Evidence to look for:
-
- Large `vae.encode` or `vae.decode` in low-overhead stage timings.
- Host-stack gaps in VAE tile split, gather, merge, broadcast, or image/video
-  transforms.
- VAE kernels or cuDNN convolution in operator trace.
- Whether every rank needs the final decoded tensor.
-
-Typical candidates:
-
- Keep VAE patch parallel enabled when it has clear measured benefit.
- Reduce VAE gather/broadcast to only ranks that need the final media output.
- Reuse tile metadata, split buffers, or gather buffers.
- Evaluate bf16/autocast behavior for VAE only with visual quality checks.
- Avoid redundant image conversion, resize, or tensor construction in repeated
-  benchmark runs.
-
-Priority guidance:
-
- P0/P1 if VAE is a large share of the target workload or if a host gap is
-  obvious and low risk.
- Lower priority when `diffuse` dominates and VAE is already patch-parallelized.
-
-Validation:
-
- Compare `vae.encode`, `vae.decode`, server end-to-end, and peak memory.
- Check output video integrity, artifacts, flicker, and seed stability.
-
-### Operator Fusion and Layout Cleanup
-
-Purpose: reduce high-frequency small kernels, memory bandwidth pressure, layout
-conversions, and launch overhead in transformer and VAE blocks.
-
-Evidence to look for:
-
- Top operator tables showing many `aten::copy_`, `aten::cat`,
-  `split_with_sizes_copy`, `aten::add`, `aten::mul`, `aten::div`, norm,
-  activation, RoPE, or reshape/layout kernels.
- `ops_rankN.xlsx` `by_shape` sheet showing repeated small shapes inside the
-  same block path.
- Trace lanes showing many short kernels between larger GEMM/attention kernels.
- Source code patterns with repeated elementwise chains or layout conversions.
-
-Typical fusion targets:
-
- AdaLayerNorm / RMSNorm / LayerNorm plus scale/shift fusion.
- RoPE fusion with Q/K layout preparation when shapes are stable.
- Residual add, scale, gate, and elementwise chains.
- MLP gate/up/down cleanup, such as fusing activation and multiply around
-  GEMM outputs when feasible.
- QKV projection and reshape/split/cat path cleanup.
- Attention pre/post layout cleanup to avoid unnecessary copies, cats, and
-  splits around sequence parallel all-to-all.
-
-Priority guidance:
-
- P1 when implemented with existing PyTorch/Triton/local helper patterns and
-  validated against exact outputs or tolerances.
- P2 when it requires custom CUDA/Triton kernels, changes numerics, or touches
-  attention math directly.
-
-Validation:
-
- First prove the operator family is material for the target shape.
- Use non-profiler A/B for latency and stage timing.
- Use quality regression checks for generated video stability.
- Check compile behavior and graph breaks if using `torch.compile`.
-
-### Attention Main-Path Optimization
-
-Purpose: address the dominant self-attention cost when FlashAttention or other
-attention kernels dominate CUDA time.
-
-Evidence to look for:
-
- Operator trace where FlashAttention/SDPA kernels dominate total CUDA time.
- Attention shape from `ops_rankN.xlsx` `by_shape`, model code, or trace
-  metadata.
- Whether attention cost scales with latent frames, latent H/W, patch size, or
-  CFG duplication.
- Layout/copy/all-to-all work around attention.
-
-Typical candidates:
-
- Verify the attention backend and shape are on the intended fast kernel path.
- Compare supported attention backends only with identical workload and quality
-  settings.
- Reduce attention input size by safe model/config choices when allowed:
-  latent resolution, frame count, patching, boundary ratio, or windowing.
- Remove avoidable layout conversions before/after attention.
- Reuse condition-side KV or other static inputs if the model structure allows.
- Consider custom kernels, sparse/window/linear attention, or approximation only
-  after quality risk is accepted.
-
-Priority guidance:
-
- P1 for backend/config/layout changes with preserved math.
- P2 for approximate attention, sparsity, custom kernels, or any change that can
-  alter quality/temporal consistency.
-
-Validation:
-
- Always include output quality and seed behavior checks.
- Compare `diffuse`, server end-to-end, peak memory, and attention kernel time
-  in diagnostic traces if needed.
-
-### Algorithmic, Precision, and Approximation Optimization
-
-Purpose: reduce mathematical work or precision cost beyond local code cleanup.
-
-Evidence to look for:
-
- A single operator family dominates even after low-risk runtime, parallel, and
-  fusion work.
- Memory bandwidth or compute utilization suggests precision or quantization
-  could matter.
- The user explicitly allows quality-preserving or approximate methods.
-
-Typical candidates:
-
- FP8/quantization for transformer or selected projections.
- Sparsity or Rainfusion-style acceleration.
- Reduced steps, scheduler changes, distillation, or caching across frames.
- Approximate attention or linear attention.
-
-Priority guidance:
-
- Usually P2 because quality, numerics, and implementation risk are high.
-
-Validation:
-
- Requires strict A/B, visual quality review, temporal flicker checks, seed
-  stability, and possibly human evaluation.
-
-### Interpolation, Super-Resolution, and E2E Pipeline Optimization
-
-Purpose: optimize the whole user-visible video product, not only the base
-diffusion invocation. Some deployments trade base-model latency against
-post-processing, interpolation, or super-resolution stages.
-
-Evidence to look for:
-
- E2E latency breakdown across base generation, interpolation, super-resolution,
-  encoding, storage, and response streaming.
- Fast/slow GPU or fast/slow stage analysis across multiple cards and pipeline
-  stages.
- User quality target: resolution, FPS, temporal smoothness, and acceptable
-  post-processing artifacts.
-
-Typical candidates:
-
- Add or optimize a frame interpolation stage when it reduces required base
-  model frames for the same perceived FPS.
- Add or optimize a super-resolution model when generating lower base
-  resolution plus SR is faster for the target quality.
- Analyze E2E2 pipeline behavior: client request, service scheduling, diffusion,
-  VAE/media, post-process, file write, and response.
- Identify fast/slow cards or stages and rebalance pipeline placement.
-
-Priority guidance:
-
- P1 when using proven interpolation/SR components without changing diffusion
-  math.
- P2 when quality risk is high or the pipeline adds significant operational
-  complexity.
-
-Validation:
-
- Measure E2E wall-clock, per-stage server timings, output FPS/resolution,
-  artifacts, flicker, and user-visible quality.
-
-### Optimization Candidate Library
-
-Use this table as a compact menu, not as automatic recommendations. Pick items
-only when baseline metrics, trace evidence, source inspection, or quality
-tolerance supports them.
-
-| Layer | Candidate | Evidence | Priority | Validation focus |
-|---|---|---|---|---|
-| Measurement | Freeze server/request/benchmark commands | Missing or drifting commands | P0 | Repeatable non-profiler A/B |
-| Measurement | Separate baseline and diagnostic profiler runs | Profiler used for latency claims | P0 | Low-overhead stage timings |
-| Host/runtime | Guard or remove avoidable `empty_cache` | Host-stack gaps or sync stalls | P0/P1 | Latency, peak memory, OOM safety |
-| Host/runtime | Cache scheduler coefficients/timesteps | Repeated tiny CPU/GPU work | P0/P1 | Same seed/output, stage timing |
-| Host/runtime | Reduce framework scheduling overhead | Client time exceeds server time | P1 | E2E latency, throughput |
-| Parallel | `CFG=2 x USP=world/2` vs `USP=world` | CFG doubles forward work | P0/P1 | `diffuse`, NCCL, memory |
-| Parallel | Tune VAE patch parallelism | VAE encode/decode is material | P0/P1 | VAE time, output correctness |
-| Parallel | HSDP on/off or buffer reuse | HSDP affects memory/all-gather | P1 | Memory, latency, OOM risk |
-| Parallel | Ulysses vs Ulysses+Ring | Long sequence all-to-all suspected | P1/P2 | Rank balance, NCCL kernels |
-| Cross-attn | Disable SP for short condition tokens | Cross-attn comm exceeds compute | P1 | `diffuse`, correctness |
-| VAE/media | Reduce VAE gather/broadcast | Rank traces show VAE wait | P1 | Rank balance, output file |
-| VAE/media | Reuse tile metadata/buffers | Tile split/merge host gaps | P1 | `vae.encode/decode`, memory |
-| VAE/media | VAE bf16/autocast | VAE float kernels are slow | P1/P2 | Artifacts, flicker, seed stability |
-| Operator fusion | AdaLayerNorm/LayerNorm fusion | Norm plus scale/shift kernels | P1 | Numeric tolerance, latency |
-| Operator fusion | RMSNorm fusion | Many small RMSNorm kernels | P1 | Numeric tolerance, latency |
-| Operator fusion | RoPE cache/fuse/layout cleanup | RoPE copy/reshape kernels | P1/P2 | Kernel count, correctness |
-| Operator layout | QKV or attention layout cleanup | Copy/cat/split around attention | P1 | Copy kernels, compile behavior |
-| Attention | Verify backend fast path | FA/SDPA dominates trace | P1 | `diffuse`, attention kernels |
-| Attention | FA to LA or selected-head LA | Attention remains dominant | P2 | Quality, temporal stability |
-| Precision | Transformer FP8/quantization | Compute/bandwidth bound and allowed | P2 | Quality, speed, stability |
-| Sparsity | Rainfusion-style acceleration | DiT compute remains dominant | P2 | Prompt diversity, quality |
-| Pipeline | Frame interpolation | Fewer base frames can meet FPS | P1/P2 | E2E latency, motion artifacts |
-| Pipeline | Super-resolution | Lower base res plus SR may win | P1/P2 | Detail quality, artifacts |
-| E2E | Fast/slow-card analysis | Multi-card stragglers | P0/P1 | Per-rank/stage wall-clock |
-
-### Optimization Plan Template
-
-Use this table shape when reporting the next work items:
-
-| Priority | Layer | Candidate | Evidence | Expected benefit | Implementation path | Validation | Quality risk |
-|---|---|---|---|---|---|---|---|
-| P0 | Host/runtime | Guard `empty_cache` | Host-stack gap points to `torch.cuda.empty_cache` | Small latency reduction, less idle | Add config/env guard | Non-profiler A/B, memory check | Low |
-| P1 | Operator fusion | RMSNorm/AdaLayerNorm fusion | High-frequency norm/elementwise kernels | Lower launch/bandwidth overhead | Use existing fusion helper or targeted Triton | A/B + output check | Medium |
-| P1/P2 | Attention | Attention layout/backend investigation | FA kernel dominates CUDA time | Potentially large | Inspect shapes/backend and remove layout copies | A/B + trace + quality | Medium/high |
-
-Then present a short selection prompt using the same rows:
-
-```text
-Which candidate should we implement next?
-
-1. P0 Host/runtime: guard empty_cache
-   - Expected benefit: small but low-risk latency reduction.
-   - Risk: possible memory increase/OOM if memory headroom is insufficient.
-
-2. P1 Operator fusion: inspect by_shape and implement first norm/RoPE/layout fusion
-   - Expected benefit: medium if high-frequency small kernels are confirmed.
-   - Risk: numerical/compile/quality validation needed.
-
-3. P1/P2 Attention: FA/LA/backend/layout investigation
-   - Expected benefit: potentially large.
-   - Risk: high quality and implementation risk.
-```
-
-If the user has not chosen an item, default to explaining tradeoffs and asking
-which candidate to execute. Only proceed autonomously on low-risk P0 measurement
-or instrumentation fixes.
-
-## Analysis Helpers
-
-PyTorch profiler traces are Chrome/Perfetto-compatible JSON files, usually
-`trace_rankN.json` or `trace_rankN.json.gz`. They normally contain a top-level
-`traceEvents` list, though some exporters emit the raw event list directly.
-
-Use the checked-in analyzer from the repository root:
-
-```bash
-python3 .claude/skills/diffusion-perf-opt/scripts/trace_analyzer.py \
-  vllm_profile/.../trace_rank0.json.gz \
-  --min-gap-ms 5 \
-  --topn 20
-```
-
-For rank imbalance or communication questions, pass all relevant rank traces in
-one command. For host gaps, lower `--min-gap-ms` to `1` and use a host-stack
-trace. Read `gpu_span_s`, `busy_union_s`, `idle_union_s`, `idle_pct`, `GAP`
-blocks, `Top GPU/operator events by total duration`, and `Top NCCL-like events
-by category`. Treat `cat=user_annotation` NCCL ranges as enclosing annotations;
-prefer `cat=kernel` or `cat=gpu_user_annotation` for real device work.
-
-The analyzer summarizes timing only. It does not parse tensor shapes, attribute
-overlap to individual streams, prove quality, or provide final latency claims.
-Use `ops_rankN.xlsx` or PyTorch key averages for shape analysis, and re-test any
-optimization with non-profiler baseline commands.
-
-Read `references/optimization-playbook.md` when drafting the optimization table or comparing candidate techniques.
-
-## vLLM Omni Heuristics
-
- Cross-attention usually should not use USP/SP when text/image condition token count is much smaller than latent video tokens. Confirm via trace; in Wan2.2 I2V, self-attention dominates cross-attention.
- VAE bf16/autocast is often worthwhile but requires visual quality checks.
- VAE patch parallel can help decode/encode but may add gather/merge/broadcast overhead. Check whether all ranks need the final decoded tensor.
- HSDP/FSDP is primarily a memory strategy. If the model fits without it, run an on/off latency comparison.
- Scheduler work can create small host/device gaps; cache tiny solve coefficients when timesteps/order are known.
- `torch.cuda.empty_cache()` can prevent OOM but creates synchronization/idle. Make it optional if memory headroom is sufficient.
- `Command Buffer Full` in profiler output is profiler overhead, not a model optimization target.
@@ -1,4 +0,0 @@
-interface:
-  display_name: "vLLM Omni Perf Opt"
-  short_description: "Profile and optimize vLLM Omni diffusion performance."
-  default_prompt: "Analyze and optimize a vLLM Omni diffusion workload using baselines, traces, parallel strategy modeling, and prioritized optimization plans."
@@ -1,82 +0,0 @@
-# Optimization Playbook
-
-Use this reference when turning trace evidence into prioritized vLLM Omni optimizations.
-
-## Parallel Strategy
-
-Start from workload shape:
-
- Long video tokens, high resolution, many frames: self-attention and FFN dominate. USP/SP is usually valuable.
- CFG enabled with positive and negative prompts: transformer forward count doubles. CFG parallel may beat USP on small/moderate sequence lengths.
- Cross-attention: skip SP/USP when condition tokens are much smaller than latent tokens; extra all-to-all can exceed compute saved.
- HSDP/FSDP: use to fit memory first. If memory allows, compare without HSDP because all-gathers and lazy init add overhead.
- VAE parallel: useful when encode/decode is material. Check gather/merge/broadcast and whether every rank needs final decoded output.
-
-Candidate matrices:
-
- 2 cards: `USP=2`, `CFG=2`, HSDP on/off, VAE PP on/off.
- 4 cards: `USP=4`, `CFG=2 x USP=2`, `USP=2 x HSDP`, VAE PP on/off.
- 8 cards: add pipeline/stage parallel only after measuring orchestration and inter-stage transfer overhead.
-
-## Host And Free Bubble Analysis
-
-Look for device intervals where no kernel, memcpy, or memset is active.
-
-Common causes:
-
- First request FSDP/HSDP lazy init.
- Python preprocessing: image load, resize, `to_tensor`, tokenizer, request marshaling.
- `torch.cuda.empty_cache()` or explicit synchronizations.
- Scheduler small linear algebra such as tiny `torch.linalg.solve`.
- VAE tile split/merge Python loops.
- Profiler overhead such as CUPTI command buffer flush/full.
-
-Also inspect trace-external time:
-
- API request parsing and multipart upload.
- Engine/orchestrator dispatch.
- Worker RPC and queue wait.
- Output serialization, video write, base64, HTTP response.
-
-## Operator Optimization Candidates
-
-Evaluate necessity before implementation:
-
- **FA to LA / LA for head:** high potential only if attention is the main CUDA hotspot and target shapes fit LA well.
- **Cross-attention BSND/no transpose:** do only if transpose/copy around cross-attn is measurable.
- **AdaLayerNorm fusion:** fuse norm, scale/shift, gate, residual, and dtype conversion if repeated elementwise kernels are significant.
- **RMSNorm fusion:** useful for Q/K norms and VAE norms when small-kernel launch count is high.
- **FP32 LayerNorm to bf16:** quality-sensitive; run visual regression.
- **LA preprocess:** integrate only if preprocess overhead is cached or amortized.
- **RoPE move/cache/fuse:** useful when RoPE small kernels are visible around every self-attn block.
- **VAE RMSNorm replacement:** only after VAE internal trace shows norm as a real component of decode/encode.
- **Scheduler fusion/cache:** cache coefficient solves and fuse CFG combine with scheduler arithmetic when possible.
-
-Before writing a custom kernel, check:
-
- Local vLLM Omni implementations.
- Upstream vLLM kernels and attention backends.
- SGLang, FlashInfer, FlashAttention, and Triton examples.
- Shape constraints and layout expectations.
-
-## Quantization And Sparsity
-
-Treat these as high-risk optimizations:
-
- Rainfusion/sparse diffusion can be high reward but needs model-specific quality validation.
- VAE decode quantization can reduce decode time but usually has smaller end-to-end upside than transformer attention/FFN work.
- FP8 or approximate kernels require seed-stable visual regression, temporal flicker checks, and prompt diversity.
-
-## Output Format For Recommendations
-
-Return a table with:
-
- Priority: P0/P1/P2.
- Optimization.
- Evidence from trace or code.
- Expected benefit.
- Risk/quality impact.
- Implementation path.
- Validation command or A/B test.
-
-State clearly when a recommendation is a hypothesis requiring new trace with shapes.
@@ -1,144 +0,0 @@
-#!/usr/bin/env python3
-"""Summarize torch profiler Chrome traces for vLLM Omni diffusion optimization."""
-
-from __future__ import annotations
-
-import argparse
-import collections
-import gzip
-import json
-from pathlib import Path
-from typing import Any
-
-GPU_CATS = {"kernel", "gpu_memcpy", "gpu_memset"}
-CPU_CATS = {"python_function", "user_annotation", "cpu_op", "cuda_runtime", "cuda_driver"}
-
-
-def open_trace(path: Path):
-    if path.suffix == ".gz":
-        return gzip.open(path, "rt")
-    return path.open("rt")
-
-
-def event_name(event: dict[str, Any]) -> str:
-    return str(event.get("name", ""))
-
-
-def summarize_trace(path: Path, min_gap_us: float, topn: int) -> None:
-    with open_trace(path) as f:
-        data = json.load(f)
-
-    events = data if isinstance(data, list) else data.get("traceEvents", [])
-    gpu: list[tuple[float, float, float, str, str, Any, Any]] = []
-    cpu: list[tuple[float, float, float, str, str, Any, Any]] = []
-    by_gpu_name: dict[str, list[float]] = collections.defaultdict(lambda: [0, 0.0, 0.0])
-    by_name: dict[str, list[float]] = collections.defaultdict(lambda: [0, 0.0, 0.0])
-    nccl: dict[tuple[str, str], list[float]] = collections.defaultdict(lambda: [0, 0.0, 0.0])
-
-    for event in events:
-        dur = event.get("dur")
-        ts = event.get("ts")
-        if dur is None or ts is None or dur <= 0:
-            continue
-        cat = str(event.get("cat", ""))
-        name = event_name(event)
-        row = (float(ts), float(ts + dur), float(dur), name, cat, event.get("pid"), event.get("tid"))
-        stat = by_name[name]
-        stat[0] += 1
-        stat[1] += dur
-        stat[2] = max(stat[2], dur)
-        if cat in GPU_CATS:
-            gpu.append(row)
-            gstat = by_gpu_name[name]
-            gstat[0] += 1
-            gstat[1] += dur
-            gstat[2] = max(gstat[2], dur)
-        elif cat in CPU_CATS:
-            cpu.append(row)
-        if "nccl" in name.lower():
-            nstat = nccl[(cat, name)]
-            nstat[0] += 1
-            nstat[1] += dur
-            nstat[2] = max(nstat[2], dur)
-
-    print(f"\n== {path}")
-    print(f"events={len(events)} gpu_events={len(gpu)} cpu_events={len(cpu)}")
-    if not gpu:
-        print("No GPU events found.")
-        return
-
-    gpu.sort()
-    merged: list[list[Any]] = []
-    for start, end, dur, name, cat, pid, tid in gpu:
-        if not merged or start > merged[-1][1]:
-            merged.append([start, end, [(start, end, dur, name, cat, pid, tid)]])
-        else:
-            merged[-1][1] = max(merged[-1][1], end)
-            merged[-1][2].append((start, end, dur, name, cat, pid, tid))
-
-    span = max(end for _, end, *_ in gpu) - min(start for start, *_ in gpu)
-    busy = sum(end - start for start, end, _ in merged)
-    idle = span - busy
-    print(
-        f"gpu_span_s={span / 1e6:.3f} "
-        f"busy_union_s={busy / 1e6:.3f} "
-        f"idle_union_s={idle / 1e6:.3f} "
-        f"idle_pct={idle / span * 100:.2f}"
-    )
-
-    interesting_cpu = [
-        row
-        for row in cpu
-        if row[2] >= 1000
-        and (
-            row[4] in {"python_function", "user_annotation"}
-            or "cudaStreamSynchronize" in row[3]
-            or "cudaDeviceSynchronize" in row[3]
-            or "cudaLaunch" in row[3]
-            or "cudaMemcpy" in row[3]
-        )
-    ]
-
-    gaps = []
-    for idx in range(1, len(merged)):
-        gap_start = merged[idx - 1][1]
-        gap_end = merged[idx][0]
-        gap_dur = gap_end - gap_start
-        if gap_dur >= min_gap_us:
-            prev_event = max(merged[idx - 1][2], key=lambda x: x[1])
-            next_event = min(merged[idx][2], key=lambda x: x[0])
-            mid = (gap_start + gap_end) / 2
-            containers = [row for row in interesting_cpu if row[0] <= mid <= row[1]]
-            containers = sorted(containers, key=lambda x: x[2])[:8]
-            gaps.append((gap_dur, gap_start, gap_end, prev_event, next_event, containers))
-
-    print(f"gaps_ge_{min_gap_us / 1000:.3f}ms count={len(gaps)} sum_s={sum(g[0] for g in gaps) / 1e6:.3f}")
-    for gap_dur, gap_start, gap_end, prev_event, next_event, containers in sorted(gaps, reverse=True)[:topn]:
-        print(f"\nGAP {gap_dur / 1000:.3f} ms ts={gap_start:.0f}->{gap_end:.0f}")
-        print(f"  prev {prev_event[4]} {prev_event[2] / 1000:.3f} ms {prev_event[3][:160]}")
-        print(f"  next {next_event[4]} {next_event[2] / 1000:.3f} ms {next_event[3][:160]}")
-        for row in containers:
-            print(f"  in   {row[4]} {row[2] / 1000:.3f} ms {row[3][:180]}")
-
-    print("\nTop GPU/operator events by total duration:")
-    for name, (count, total, max_dur) in sorted(by_gpu_name.items(), key=lambda kv: kv[1][1], reverse=True)[:topn]:
-        print(f"  {int(count):8d} total={total / 1e6:9.3f}s max={max_dur / 1000:9.3f}ms {name[:180]}")
-
-    print("\nTop NCCL-like events by category:")
-    for (cat, name), (count, total, max_dur) in sorted(nccl.items(), key=lambda kv: kv[1][1], reverse=True)[:topn]:
-        print(f"  {int(count):8d} total={total / 1e6:9.3f}s max={max_dur / 1000:9.3f}ms cat={cat} {name[:160]}")
-
-
-def main() -> None:
-    parser = argparse.ArgumentParser(description=__doc__)
-    parser.add_argument("traces", nargs="+", type=Path, help="trace.json or trace.json.gz files")
-    parser.add_argument("--min-gap-ms", type=float, default=5.0, help="minimum GPU idle gap to print")
-    parser.add_argument("--topn", type=int, default=20, help="number of gaps/hotspots to print")
-    args = parser.parse_args()
-
-    for trace in args.traces:
-        summarize_trace(trace, args.min_gap_ms * 1000.0, args.topn)
-
-
-if __name__ == "__main__":
-    main()
@@ -18,9 +18,6 @@ include:

 - `add-diffusion-model`: guides integration of a new diffusion model into
  `vllm-omni`
- `diffusion-perf-opt`: guides diffusion model performance optimization,
-  including profiling traces, parallel strategies, stage timing analysis, and
-  benchmark-driven tuning
 - `add-omni-model`: covers addition of new omni-modality model support
 - `add-tts-model`: covers integration of new TTS models and related serving
  workflows
@@ -3,64 +3,15 @@
 # More info: https://docs.github.com/en/repositories/managing-your-repositorys-settings-and-features/customizing-your-repository/about-code-owners

 # Default owners (fallback reviewers for all files)
-* @Gaohan123 @hsliuustc0106 @ywang96
+# later, we can add more reviewers here
+* @hsliuustc0106

-# Repository operations
-/.github/ @hsliuustc0106
+# Assign reviewers by directory
+/.github/ @owner1
+/tests/ @owner1
+/docs/ @owner1

-# Documentation and Recipes
-/docs/ @hsliuustc0106 @Gaohan123 @david6666666
-/recipes @hsliuustc0106 @Gaohan123 @david6666666
-
-# Skills
-/.claude/ @hsliuustc0106
-
-# Tests
-/tests/ @yenuo26
-
-# CI
-/.buildkite/ @yenuo26 @congw729
-
-# Engine and model_executor
-/vllm_omni/core/ @tzhouam @yinpeiqi @fake0fan
-/vllm_omni/engine/ @tzhouam @yinpeiqi @fake0fan
-/vllm_omni/entrypoints/ @tzhouam @yinpeiqi @fake0fan
-/vllm_omni/model_executor/ @tzhouam @gcanlin
-/vllm_omni/worker/ @tzhouam @gcanlin
-
-# Diffusion
-/vllm_omni/diffusion/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @david6666666
-/vllm_omni/lora/ @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @yuanheng-zhao
-
-# Diffusion CPU offloading
-/vllm_omni/diffusion/offloader/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @yuanheng-zhao
-
-# Omni and TTS support
-/vllm_omni/transformers_utils/ @tzhouam @yuanheng-zhao
-/vllm_omni/model_executor/models/ @linyueqian @ZeldaHuang @princepride @yuanheng-zhao @amy-why-3459 @Bounty-hunter
-
-# Benchmarks
-/benchmarks/accuracy/ @david6666666
-/benchmarks/build_dataset/ @linyueqian @yuanheng-zhao
-/benchmarks/diffusion/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa
-/benchmarks/distributed/ @princepride
-/benchmarks/glm_image/ @princepride @JaredforReal
-/benchmarks/tts/ @linyueqian @ZeldaHuang @princepride
-
-# Usability and Observability
-/vllm_omni/deploy/ @tzhouam @lishunyang12
-/vllm_omni/metrics/ @lishunyang12
-/vllm_omni/profiler/ @gcanlin
-
-# ComfyUI
-/apps/ @fhfuih
-
-# Large Scale Serving
-/vllm_omni/distributed/ @princepride @yuanheng-zhao @xuechendi @natureofnature @fake0fan
-
-# Quantization and configuration
-/vllm_omni/config/ @lishunyang12
-/vllm_omni/quantization/ @david6666666 @Isotr0py @lishunyang12
-
-# Hardware plugins and accelerators
-/vllm_omni/platforms/ @gcanlin @tjtanaa @xuechendi
+# Example: reviewers for specific directories
+# /examples/offline_inference/ @offline-inference-team
+# /vllm_omni/core/ @core-team
+# /vllm_omni/model_executor/ @model-team
@@ -239,7 +239,6 @@ data/
 datasets/
 *.csv
 *.json
-!tests/dfx/perf/tests/*.json
 !apps/ComfyUI-vLLM-Omni/example_workflows/*.json
 *.jsonl
 *.parquet
@@ -9,19 +9,19 @@ 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://deepwiki.com/vllm-project/vllm-omni"><b>DeepWiki</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/111-L8zF7A1j_YI_cR8JsblofdScdRr2f/edit?usp=sharing&ouid=110473603432222024453&rtpof=true&sd=true"><b>Slides</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/1XJWgv79lORl8rbaVvp2d5Sqs6ZEBgAgj/edit?slide=id.p1#slide=id.p1"><b>Slides</b></a> |
 </p>


 ---

 *Latest News* 🔥
- [2026/05] We released [0.20.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.20.0) - refreshes the serving/runtime stack for large-scale omni workloads, and improves diffusion model performance, quantization, and hardware readiness across CUDA, ROCm, MUSA, NPU, and XPU backends.
 - [2026/03] We released [0.18.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.18.0) - strengthens the core runtime through a large entrypoint refactor and scheduler/runtime cleanups, expands unified quantization and diffusion execution, broadens multimodal model coverage, and improves production readiness across audio, omni, image, video, RL, and multi-platform deployments.
 - [2026/03] Check out our first public [project deepdive](https://youtu.be/sgwNfsNnR9I) at the vLLM Hong Kong Meetup!
 - [2026/03] **[vllm-omni-skills](https://github.com/hsliuustc0106/vllm-omni-skills)** is a community-driven collection of AI assistant skills that help developers work with vLLM-Omni more effectively. These skills can be used with popular agentic AI coding assistants like **Cursor IDE**, **Claude**, **Codex**, and more.
 - [2026/02] We released [0.16.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.16.0) - A major alignment + capability release that rebases onto **upstream vLLM v0.16.0** and significantly expands performance, distributed execution, and production readiness across **Qwen3-Omni / Qwen3-TTS**, **Bagel**, **MiMo-Audio**, **GLM-Image** and the **Diffusion (DiT) image/video stack**—while also improving platform coverage (CUDA / ROCm / NPU / XPU), CI quality, and documentation.
 - [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.

 ---
@@ -18,13 +18,3 @@ vd017|||The quarterly results exceed expectations across all major metrics.|A sh
 vd018|||Chapter one. The morning sun filtered gently through the forest canopy.|A smooth, rich male audiobook narrator voice, expressive and engaging.
 vd019|||To be or not to be, that is the question.|A theatrical female voice, dramatic and expressive, stage projection quality.
 vd020|||And that is all for tonight. Stay well out there, everyone.|A warm, velvety male late-night radio DJ voice, smooth and intimate.
-vd021|||Please fasten your seatbelts; we will be experiencing some turbulence shortly.|A calm, professional female flight attendant voice, reassuring and clear.
-vd022|||The fossil record suggests these creatures lived more than sixty million years ago.|A measured male documentary narrator voice, scientific and curious.
-vd023|||Hey kiddo, did you finish your homework before dinner today?|A gentle, caring father voice, warm and slightly tired after a long day.
-vd024|||Live from courtside, the home team has just taken a commanding lead.|A high-energy male sports broadcaster, urgent and engaged.
-vd025|||Please remain on the line while I transfer you to the next available agent.|A polite, neutral female call-center voice, professional and patient.
-vd026|||I have a dream that one day this nation will rise up and live out the true meaning of its creed.|A powerful, resonant male orator voice, slow and impassioned.
-vd027|||Mama always said life was like a box of chocolates; you never know what you're gonna get.|A folksy male voice with a Southern American accent, warm and reflective.
-vd028|||Ahoy there, matey! Hoist the colors and prepare to set sail.|A theatrical pirate male voice, gruff and boisterous.
-vd029|||Ladies and gentlemen, please welcome to the stage our keynote speaker.|A polished female event MC voice, gracious and articulate.
-vd030|||Once you've completed the form, click submit and we'll review it within two business days.|A friendly female product-onboarding voice, encouraging and clear.
@@ -18,13 +18,3 @@ smoke017|||In the distance, the mountains were shrouded in thick morning fog.
 smoke018|||Our company reported record revenue for the fourth quarter of the fiscal year.
 smoke019|||She explained the new policy in detail during the staff meeting this morning.
 smoke020|||The children laughed and played in the garden until the sun began to set.
-smoke021|||A long line of cars stretched along the coastal highway after the lighthouse opened.
-smoke022|||The chef recommended the seared scallops with a lemon butter sauce as the special.
-smoke023|||Our quarterly engineering review begins promptly at two in the afternoon on Thursday.
-smoke024|||Despite the heavy rain, the marathon runners kept a steady pace through downtown.
-smoke025|||She finished the crossword puzzle in under fifteen minutes while sipping her tea.
-smoke026|||The new release adds streaming output, faster startup, and better error reporting.
-smoke027|||His grandfather taught him how to whittle small wooden birds from cedar offcuts.
-smoke028|||Travelers were advised to check the timetable since trains were running on a holiday schedule.
-smoke029|||The startup launched a public beta after eighteen months of closed development.
-smoke030|||At night the desert sky filled with stars so bright the Milky Way was clearly visible.
@@ -30,8 +30,10 @@ python3 benchmarks/diffusion/diffusion_benchmark_serving.py \

 **Notes**

- By default, image tasks talk to `http://<host>:<port>/v1/chat/completions`; video tasks talk to `/v1/videos`.
+- The benchmark talks to `http://<host>:<port>/v1/chat/completions`.
 - If you run the server on another host or port, pass `--base-url` accordingly.
+- To compare diffusion multi-replica serving, restart the server with different
+  `runtime.num_replicas` values and keep the benchmark parameters unchanged.

 ## 2. Supported Datasets

@@ -93,11 +95,39 @@ You can point to your own trace using `--dataset-path`.

 ## 3. Benchmark Parameters

+### Multi-replica sweep
+
+For replica scaling, use the same model, prompt count, image/video size, and
+step count for every run. Change only the serving stage config:
+
+```yaml
+runtime:
+  devices: "0,1,2,3"
+  num_replicas: 4
+```
+
+Then run a fixed benchmark and write one JSON file per replica count:
+
+```bash
+python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
+  --base-url http://localhost:8099 \
+  --model Qwen/Qwen-Image \
+  --task t2i \
+  --dataset random \
+  --num-prompts 32 \
+  --max-concurrency 8 \
+  --request-rate inf \
+  --num-inference-steps 10 \
+  --width 512 --height 512 \
+  --output-file diffusion_replica_4.json
+```
+
+Repeat for `num_replicas=1`, `2`, `3`, and `4`, using matching `runtime.devices`.
+
 ### Basic flags

- `--base-url`: Server address; `--endpoint` selects the path appended to this base URL.
+- `--base-url`: Server address (the script calls `.../v1/chat/completions`).
 - `--model`: The OpenAI-compatible `model` field.
- `--endpoint`: API endpoint path. Leading `/` is optional, e.g. `/v1/videos` or `v1/videos`.
 - `--task`: Task type (e.g., `t2i`, `t2v`, `i2i`, `i2v`).
 - `--dataset`: Dataset mode (`vbench` / `trace` / `random`).
 - `--num-prompts`: Number of requests to send.
@@ -151,7 +151,7 @@ async def async_request_chat_completions(
    return output


-async def async_request_openai_image_generations(
+async def async_request_openai_images(
    input: RequestFuncInput,
    session: aiohttp.ClientSession,
    pbar: tqdm | None = None,
@@ -278,7 +278,7 @@ async def async_request_v1_videos(
                return output

        # invoke a poll request (GET /v1/videos/{video_id})
-        poll_interval = 0.2  # Unit(s)
+        poll_interval = 2.0  # Unit(s)
        timeout_seconds = 600.0
        deadline = time.perf_counter() + timeout_seconds
        job_url = f"{input.api_url}/{job_id}"
@@ -348,38 +348,12 @@ async def async_request_v1_videos(
    return output


-LEGACY_BACKEND_ENDPOINT_ALIASES = {
-    "vllm-omni": "/v1/chat/completions",
-    "openai": "/v1/images/generations",
-}
-
-
-def normalize_endpoint(value: str) -> str:
-    endpoint = str(value).strip()
-    if not endpoint:
-        raise ValueError("endpoint must not be empty.")
-    endpoint = LEGACY_BACKEND_ENDPOINT_ALIASES.get(
-        endpoint,
-        LEGACY_BACKEND_ENDPOINT_ALIASES.get(endpoint.lstrip("/"), endpoint),
-    )
-    if not endpoint.startswith("/"):
-        endpoint = f"/{endpoint}"
-    return endpoint
-
-
-def endpoint_filename_token(value: str) -> str:
-    token = normalize_endpoint(value).lstrip("/")
-    for bad in ("/", "\\", ":", "*", "?", '"', "<", ">", "|"):
-        token = token.replace(bad, "_")
-    return token or "endpoint"
-
-
 backends_function_mapping = {
    "2i": {
-        "/v1/chat/completions": (async_request_chat_completions, "/v1/chat/completions"),
-        "/v1/images/generations": (async_request_openai_image_generations, "/v1/images/generations"),
+        "vllm-omni": (async_request_chat_completions, "/v1/chat/completions"),
+        "openai": (async_request_openai_images, "/v1/images/generations"),
    },
    "2v": {
-        "/v1/videos": (async_request_v1_videos, "/v1/videos"),
+        "v1/videos": (async_request_v1_videos, "/v1/videos"),
    },
 }
@@ -6,31 +6,31 @@
 Benchmark online serving for diffusion models (Image/Video Generation).
 If you want to use i2v, i2i dataset, you should `uv pip install gdown` first

-Supports multiple endpoints:
-    - /v1/chat/completions: OpenAI chat-compatible image requests
-    - /v1/images/generations: OpenAI image generation requests
-    - /v1/videos: Async video jobs
+Supports multiple backends:
+    - vllm-omni: Uses /v1/chat/completions endpoint (default)
+    - openai: Uses /v1/images/generations endpoint
+    - v1/videos: Use /v1/videos endpoint

 Usage:
-    # Video (/v1/videos endpoint)
+    # Video (v1/videos backend)
    t2v:
    python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
-        --endpoint /v1/videos --dataset vbench --task t2v --num-prompts 10 \
+        --backend v1/videos --dataset vbench --task t2v --num-prompts 10 \
        --height 480 --width 640 --fps 16 --num-frames 80

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


-    # Image (/v1/chat/completions endpoint)
+    # Image (vllm-omni backend)
    t2i:
    python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
-        --endpoint /v1/chat/completions --dataset vbench --task t2i --num-prompts 10 \
+        --backend vllm-omni --dataset vbench --task t2i --num-prompts 10 \
        --height 1024 --width 1024

    python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
-        --endpoint /v1/chat/completions --dataset random --task t2i --num-prompts 1 \
+        --backend vllm-omni --dataset random --task t2i --num-prompts 1 \
        --max-concurrency 1 --enable-negative-prompt \
        --random-request-config '[
            {"width":512,"height":512,"num_inference_steps":20,"weight":0.15},
@@ -41,18 +41,18 @@ Usage:

    i2i:
    python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
-        --endpoint /v1/chat/completions --dataset vbench --task i2i --num-prompts 10
+        --backend vllm-omni --dataset vbench --task i2i --num-prompts 10

-    # Image (/v1/images/generations endpoint)
+    # Image (openai backend)
    t2i:
    python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
-        --endpoint /v1/images/generations --dataset vbench --task t2i --num-prompts 10 \
+        --backend openai --dataset vbench --task t2i --num-prompts 10 \
        --height 1024 --width 1024 --port 3000

    # Video (v1/videos)
    t2v:
    python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
-        --endpoint /v1/videos --dataset random --task t2v --num-prompts 1 \
+        --backend v1/videos --dataset random --task t2v --num-prompts 1 \
        --max-concurrency 1 --enable-negative-prompt \
        --random-request-config '[
            {"width":854,"height":480,"num_inference_steps":18,"num_frames":120,"fps":24,"weight":1}
@@ -68,7 +68,6 @@ import glob
 import json
 import logging
 import os
-import pathlib
 import random
 import tempfile
 import time
@@ -81,12 +80,7 @@ from typing import Any
 import aiohttp
 import numpy as np
 import requests
-from backends import (
-    RequestFuncInput,
-    RequestFuncOutput,
-    backends_function_mapping,
-    normalize_endpoint,
-)
+from backends import RequestFuncInput, RequestFuncOutput, backends_function_mapping
 from PIL import Image
 from tqdm.asyncio import tqdm

@@ -599,13 +593,6 @@ class RandomDataset(BaseDataset):
        if self.enable_negative_prompt:
            extra_body["negative_prompt"] = f"Negative prompt {idx} for benchmarking diffusion models"

-        request_fields = {
-            "width",
-            "height",
-            "num_frames",
-            "num_inference_steps",
-            "fps",
-        }
        params = {
            "width": self.args.width,
            "height": self.args.height,
@@ -615,11 +602,9 @@ class RandomDataset(BaseDataset):
        }
        if self._sampled_requests:
            profile = self._sampled_requests[idx]
-            params.update({key: value for key, value in profile.items() if key in request_fields})
-            extra_body.update({key: value for key, value in profile.items() if key not in request_fields})
-        prompt = self.args.random_prompt or f"Random prompt {idx} for benchmarking diffusion models"
+            params.update(profile)
        return RequestFuncInput(
-            prompt=prompt,
+            prompt=f"Random prompt {idx} for benchmarking diffusion models",
            api_url=self.api_url,
            model=self.model,
            seed=self.args.seed,
@@ -879,35 +864,6 @@ def calculate_metrics(
    return metrics


-def _response_suffix(endpoint: str) -> str:
-    if endpoint == "/v1/videos":
-        return ".mp4"
-    return ".json"
-
-
-def _save_success_responses(
-    outputs: list[RequestFuncOutput],
-    requests_list: list[RequestFuncInput],
-    output_dir: str,
-    endpoint: str,
-) -> list[str]:
-    path = pathlib.Path(output_dir).expanduser().resolve()
-    path.mkdir(parents=True, exist_ok=True)
-    saved_paths: list[str] = []
-    suffix = _response_suffix(endpoint)
-
-    for index, (request, output) in enumerate(zip(requests_list, outputs, strict=False)):
-        if not output.success:
-            continue
-        item_path = path / f"{index:05d}_{request.request_id}{suffix}"
-        if isinstance(output.response_body, bytes):
-            item_path.write_bytes(output.response_body)
-        else:
-            item_path.write_text(json.dumps(output.response_body, indent=2, sort_keys=True), encoding="utf-8")
-        saved_paths.append(str(item_path))
-    return saved_paths
-
-
 def wait_for_service(base_url: str, timeout: int = 120) -> None:
    print(f"Waiting for service at {base_url}...")
    start_time = time.time()
@@ -927,14 +883,6 @@ def wait_for_service(base_url: str, timeout: int = 120) -> None:
        time.sleep(1)


-def _default_endpoint_for_task(task: str) -> str:
-    if task in {"t2v", "i2v", "ti2v"}:
-        return "/v1/videos"
-    if task in {"t2i", "i2i", "ti2i"}:
-        return "/v1/chat/completions"
-    raise ValueError(f"Unsupported task for endpoint resolution: {task}")
-
-
 async def benchmark(args):
    # Construct base_url if not provided
    if args.base_url is None:
@@ -954,23 +902,18 @@ async def benchmark(args):
            f"Valid image tasks: {sorted(IMAGE_TASKS)}"
        )

-    raw_endpoint = args.endpoint if args.endpoint is not None else args.backend
-    if raw_endpoint is None:
-        raw_endpoint = _default_endpoint_for_task(args.task)
-    args.endpoint = normalize_endpoint(raw_endpoint)
+    valid_backends = sorted(backends_function_mapping[task_type].keys())

-    valid_endpoints = sorted(backends_function_mapping[task_type].keys())
-
-    if args.endpoint not in valid_endpoints:
+    if args.backend not in valid_backends:
        logger.error(
-            f"Invalid endpoint '{args.endpoint}' for task '{args.task}' (task type: '{task_type}').\n"
-            f"Valid endpoints for this task type: {valid_endpoints}\n"
-            f"Example usage: --task {args.task} --endpoint {valid_endpoints[0]}"
+            f"Invalid backend '{args.backend}' for task '{args.task}' (task type: '{task_type}').\n"
+            f"Valid backends for this task type: {valid_backends}\n"
+            f"Example usage: --task {args.task} --backend {valid_backends[0]}"
        )
-        raise ValueError("Endpoint validation failed. See log above for valid options.")
+        raise ValueError("Backend validation failed. See log above for valid options.")

-    # Setup API URL and request function based on endpoint.
-    request_func, api_url = backends_function_mapping[task_type][args.endpoint]
+    # Setup API URL and request function based on backend
+    request_func, api_url = backends_function_mapping[task_type][args.backend]
    api_url = f"{args.base_url}{api_url}"

    if args.dataset == "vbench":
@@ -1033,22 +976,15 @@ async def benchmark(args):
    metrics = calculate_metrics(outputs, total_duration, requests_list, args, args.slo)

    # Add configuration info to metrics for JSON output
-    metrics["endpoint"] = args.endpoint
+    metrics["backend"] = args.backend
    metrics["model"] = args.model
    metrics["dataset"] = args.dataset
    metrics["task"] = args.task
-    if args.save_response_dir:
-        metrics["saved_responses"] = _save_success_responses(
-            outputs,
-            requests_list,
-            args.save_response_dir,
-            args.endpoint,
-        )

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

    # Section 1: Configuration
-    print("{:<40} {:<15}".format("Endpoint:", args.endpoint))
+    print("{:<40} {:<15}".format("Backend:", args.backend))
    print("{:<40} {:<15}".format("Model:", args.model))
    print("{:<40} {:<15}".format("Dataset:", args.dataset))
    print("{:<40} {:<15}".format("Task:", args.task))
@@ -1111,17 +1047,12 @@ if __name__ == "__main__":
    parser.add_argument("--host", type=str, default="localhost", help="Server host.")
    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(
-        "--endpoint",
-        type=str,
-        default=None,
-        help=("Endpoint path to target. Leading '/' is optional, e.g. /v1/videos or v1/videos."),
-    )
    parser.add_argument(
        "--backend",
        type=str,
-        default=None,
-        help=argparse.SUPPRESS,
+        default="vllm-omni",
+        choices=["vllm-omni", "openai", "v1/videos"],
+        help="Backend to target the benchmark to.",
    )
    parser.add_argument(
        "--dataset",
@@ -1200,12 +1131,6 @@ 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(
-        "--save-response-dir",
-        type=str,
-        default=None,
-        help="Directory to save successful raw responses. Video responses are saved as MP4 files.",
-    )
    parser.add_argument(
        "--slo",
        action="store_true",
@@ -1241,12 +1166,6 @@ if __name__ == "__main__":
            '{"width":768,"height":768,"num_inference_steps":20,"weight":0.85}]'
        ),
    )
-    parser.add_argument(
-        "--random-prompt",
-        type=str,
-        default=None,
-        help="Prompt to use for every synthetic random request. Defaults to generated placeholder prompts.",
-    )
    parser.add_argument(
        "--num-input-images",
        type=int,
@@ -1,222 +0,0 @@
-# Helios-Distilled Serving Performance Dashboard
-
-This document describes how to deploy and benchmark **BestWishYsh/Helios-Distilled** using vLLM-Omni. It follows the same serving benchmark structure as the Qwen-Image and Wan2.2 dashboards.
-
---
-
-# 1. Overview
-
-Helios-Distilled is a text-to-video diffusion model. Helios uses an autoregressive chunking pattern that generates 33 frames per chunk. For stable comparisons, choose `num_frames` as a multiple of 33, such as 99.
-
-This document covers:
-
-* Service launch configuration
-* Benchmark scripts and usage
-* Dataset and workload settings
-* Performance and accuracy comparison methodology
-* Reproducibility guidelines
-
---
-
-# 2. Test Environment
-
-| Component | Specification |
-|------------|----------------|
-| GPU | NVIDIA H100 for CI accuracy/perf gate; local bring-up also used L20X |
-| Diffusion Attention Backend | FlashAttention |
-| Model | `BestWishYsh/Helios-Distilled` |
-
---
-
-# 3. Service Launch Configuration
-
-## 3.1 Basic Serving Command
-
-```bash
-vllm serve BestWishYsh/Helios-Distilled --omni \
-    --port 8091
-```
-
-## 3.2 Key Parameters
-
-| Parameter | Description |
-| --- | --- |
-| `--port` | Serving port |
-| `--max-num-seqs` | Maximum number of active sequences |
-| `--enable-diffusion-pipeline-profiler` | Optional stage-level profiling |
-| `--profiler-config` | Optional torch profiler configuration |
-
-Record these parameters when reporting performance results.
-
---
-
-# 4. Benchmark Scripts
-
-## 4.1 Online Serving Benchmark
-
-Use the shared diffusion serving benchmark, as with Qwen-Image and Wan2.2:
-
-```bash
-python benchmarks/diffusion/diffusion_benchmark_serving.py \
-    --base-url http://localhost:8091 \
-    --endpoint /v1/videos \
-    --model BestWishYsh/Helios-Distilled \
-    --dataset random \
-    --task t2v \
-    --num-prompts 1 \
-    --max-concurrency 1 \
-    --disable-tqdm \
-    --random-request-config '[
-        {
-            "width": 640,
-            "height": 384,
-            "num_inference_steps": 50,
-            "num_frames": 99,
-            "fps": 16,
-            "guidance_scale": 1.0,
-            "is_enable_stage2": true,
-            "pyramid_num_stages": 3,
-            "pyramid_num_inference_steps_list": [1, 1, 1],
-            "is_amplify_first_chunk": false,
-            "weight": 1
-        }
-    ]'
-```
-
-`--random-request-config` forwards standard fields such as `width`, `height`, `num_frames`, `fps`, and `num_inference_steps` as first-class request fields. Model-specific fields are forwarded in the request body, which is required for Helios stage-2 settings.
-
-## 4.2 Warmed Serving Benchmark
-
-For local bring-up, use the same shared serving benchmark with one warmup request and one measured request:
-
-```bash
-python benchmarks/diffusion/diffusion_benchmark_serving.py \
-    --base-url http://localhost:8091 \
-    --endpoint /v1/videos \
-    --model BestWishYsh/Helios-Distilled \
-    --dataset random \
-    --task t2v \
-    --num-prompts 1 \
-    --max-concurrency 1 \
-    --warmup-requests 1 \
-    --warmup-num-inference-steps 1 \
-    --random-prompt "A cat wearing sunglasses dances on a beach at sunset, cinematic lighting." \
-    --seed 42 \
-    --disable-tqdm \
-    --output-file metrics.json \
-    --save-response-dir responses \
-    --random-request-config '[
-        {
-            "width": 640,
-            "height": 384,
-            "num_inference_steps": 50,
-            "num_frames": 99,
-            "fps": 16,
-            "guidance_scale": 1.0,
-            "is_enable_stage2": true,
-            "pyramid_num_stages": 3,
-            "pyramid_num_inference_steps_list": [1, 1, 1],
-            "is_amplify_first_chunk": false,
-            "weight": 1
-        }
-    ]'
-```
-
-## 4.3 PR-vs-Main Accuracy and Performance Comparison
-
-The PR-vs-main accuracy/performance gate lives in the Helios accuracy pytest. It checks out `origin/main`, starts a server for each checkout, and drives both with `diffusion_benchmark_serving.py`:
-
-```bash
-pytest -s -v tests/e2e/accuracy/test_helios_distilled.py::test_helios_distilled_pr_matches_main_acc_perf \
-    -m 'full_model and benchmark' \
-    --run-level 'full_model'
-```
-
---
-
-# 5. Dataset & Workload Settings
-
-## 5.1 Recommended Evaluation Configurations
-
-### Dataset A (99 frames, 384x640)
-
-* Dataset: `random`
-* Task: t2v
-* Concurrency: 1
-* Frame count: 99 (`33 x 3`)
-* FPS: 16
-* Stage-2 steps: `[1, 1, 1]`
-
-```json
-[
-  {
-    "width": 640,
-    "height": 384,
-    "num_inference_steps": 50,
-    "num_frames": 99,
-    "fps": 16,
-    "guidance_scale": 1.0,
-    "is_enable_stage2": true,
-    "pyramid_num_stages": 3,
-    "pyramid_num_inference_steps_list": [1, 1, 1],
-    "is_amplify_first_chunk": false,
-    "weight": 1
-  }
-]
-```
-
---
-
-# 6. Performance Metrics
-
-| Metric | Description | Unit |
-| --- | --- | --- |
-| Mean Latency | Mean request latency | seconds |
-| P99 Latency | P99 request latency | seconds |
-| Generated FPS | `num_frames / measured_time_s` | frames/second |
-| Candidate/reference latency ratio | PR latency divided by reference latency | ratio |
-
---
-
-# 7. Accuracy Metrics
-
-The PR-vs-main gate compares encoded video frames extracted from MP4 outputs:
-
-| Metric | Threshold | Direction |
-| --- | --- | --- |
-| PSNR | `>= 35 dB` | Higher is better |
-| MAE | `<= 3.0 / 255` | Lower is better |
-| Candidate/reference latency ratio | `<= 1.03` | Lower is better |
-
-Local L20X bring-up for this PR produced:
-
-| Variant | Measured Latency (ms) | Generated FPS |
-| --- | --- | --- |
-| `origin/main` | 5317.95 | 18.62 |
-| Candidate PR | 5283.27 | 18.74 |
-
-Encoded-video similarity:
-
-| Metric | Value |
-| --- | --- |
-| PSNR | 38.66 dB |
-| MAE | 1.89 / 255 |
-| Cosine similarity | 0.99978 |
-
---
-
-# 8. Reproducibility Checklist
-
-To ensure consistent and comparable benchmark results:
-
-* Record GPU type
-* Record vLLM and vLLM-Omni versions
-* Record benchmark parameters: resolution, frame count, FPS, seed, concurrency, and stage-2 steps
-* Use frame counts that are multiples of 33
-* Keep prompt, seed, resolution, and sampling parameters identical when comparing PR vs main
-* Ensure no background workload on GPUs during testing
-* Use the same MP4 export and frame extraction path for both variants
-
---
-
-This document serves as the Helios-Distilled serving performance reference under vLLM-Omni.
@@ -58,7 +58,7 @@ Record these parameters when reporting performance results.

 ```bash
 python benchmarks/diffusion/diffusion_benchmark_serving.py \
-    --endpoint /v1/chat/completions \
+    --backend vllm-omni \
    --dataset <DATASET_NAME> \
    --task t2i \
    --num-prompts <N> \
@@ -71,7 +71,7 @@ python benchmarks/diffusion/diffusion_benchmark_serving.py \

 | Parameter              | Description                       |
 | ---------------------- | --------------------------------- |
-| `--endpoint`           | API endpoint (use `/v1/chat/completions`; leading `/` is optional) |
+| `--backend`            | Serving backend (use `vllm-omni`) |
 | `--dataset`            | Dataset name (`random` or custom) |
 | `--task`               | Task type (e.g., `t2i`)           |
 | `--num-prompts`        | Total number of requests          |
@@ -129,7 +129,7 @@ python benchmarks/diffusion/diffusion_benchmark_serving.py \

 ```bash
 python benchmarks/diffusion/diffusion_benchmark_serving.py \
-    --endpoint /v1/chat/completions \
+    --backend vllm-omni \
    --dataset random \
    --task t2i \
    --num-prompts 1 \
@@ -53,7 +53,7 @@ Record these parameters when reporting performance results.

 ```bash
 python benchmarks/diffusion/diffusion_benchmark_serving.py \
-    --endpoint /v1/videos \
+    --backend v1/videos \
    --dataset <DATASET_NAME> \
    --task t2v \
    --num-prompts <N> \
@@ -66,7 +66,7 @@ python benchmarks/diffusion/diffusion_benchmark_serving.py \

 | Parameter              | Description                       |
 | ---------------------- | --------------------------------- |
-| `--endpoint`           | API endpoint (use `/v1/videos`; leading `/` is optional) |
+| `--backend`            | Serving backend (use `v1/videos`) |
 | `--dataset`            | Dataset name (`random` or custom) |
 | `--task`               | Task type (e.g., `t2v`)           |
 | `--num-prompts`        | Total number of requests          |
@@ -120,7 +120,7 @@ python benchmarks/diffusion/diffusion_benchmark_serving.py \

 ```bash
 python benchmarks/diffusion/diffusion_benchmark_serving.py \
-    --endpoint /v1/videos \
+    --backend v1/videos \
    --dataset random \
    --task t2v \
    --num-prompts 1 \
@@ -148,11 +148,10 @@ def _build_omni_kwargs(args, quantization=None):
 def _generate_image(omni, args, prompt, seed):
    """Generate a single image and return (PIL.Image, time_seconds, memory_gib)."""
    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-    from vllm_omni.outputs import OmniRequestOutput
    from vllm_omni.platforms import current_omni_platform

    generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(seed)
-    torch.accelerator.reset_peak_memory_stats()
+    torch.cuda.reset_peak_memory_stats()
    start = time.perf_counter()
    outputs = omni.generate(
        {"prompt": prompt},
@@ -164,11 +163,10 @@ def _generate_image(omni, args, prompt, seed):
        ),
    )
    elapsed = time.perf_counter() - start
-    peak_mem = torch.accelerator.max_memory_allocated() / (1024**3)
+    peak_mem = torch.cuda.max_memory_allocated() / (1024**3)

-    req_out = OmniRequestOutput.unwrap_result(outputs)
-    if not req_out.images:
-        raise ValueError("Could not extract image output from result.")
+    first = outputs[0]
+    req_out = first.request_output[0] if hasattr(first, "request_output") else first
    img = req_out.images[0]
    return img, elapsed, peak_mem

@@ -180,7 +178,7 @@ def _generate_video(omni, args, prompt, seed):
    from vllm_omni.platforms import current_omni_platform

    generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(seed)
-    torch.accelerator.reset_peak_memory_stats()
+    torch.cuda.reset_peak_memory_stats()
    start = time.perf_counter()
    outputs = omni.generate(
        {"prompt": prompt, "negative_prompt": ""},
@@ -194,7 +192,7 @@ def _generate_video(omni, args, prompt, seed):
        ),
    )
    elapsed = time.perf_counter() - start
-    peak_mem = torch.accelerator.max_memory_allocated() / (1024**3)
+    peak_mem = torch.cuda.max_memory_allocated() / (1024**3)

    first = outputs[0]
    if hasattr(first, "request_output") and isinstance(first.request_output, list):
@@ -230,8 +228,8 @@ def _unload_omni(omni):
    del omni
    gc.collect()
    if torch.cuda.is_available():
-        torch.accelerator.empty_cache()
-        torch.accelerator.synchronize()
+        torch.cuda.empty_cache()
+        torch.cuda.synchronize()


 def run_benchmark(args):
@@ -614,8 +614,8 @@ Examples:
        if not torch.cuda.is_available():
            print("[ERROR] CUDA is not available but GPU mode was requested.")
            sys.exit(1)
-        if args.gpu_id >= torch.accelerator.device_count():
-            print(f"[ERROR] GPU {args.gpu_id} not available. Found {torch.accelerator.device_count()} GPUs.")
+        if args.gpu_id >= torch.cuda.device_count():
+            print(f"[ERROR] GPU {args.gpu_id} not available. Found {torch.cuda.device_count()} GPUs.")
            sys.exit(1)
        print(f"[INFO] Using GPU {args.gpu_id}: {torch.cuda.get_device_name(args.gpu_id)}")

@@ -63,7 +63,7 @@ def _find_repo_root(start: Path) -> Path:

 REPO_ROOT = _find_repo_root(Path(__file__).resolve())
 DEFAULT_STAGE_ASYNC = REPO_ROOT / "vllm_omni" / "model_executor" / "stage_configs" / "voxcpm_async_chunk.yaml"
-DEFAULT_STAGE_SYNC = REPO_ROOT / "vllm_omni" / "deploy" / "voxcpm.yaml"
+DEFAULT_STAGE_SYNC = REPO_ROOT / "vllm_omni" / "model_executor" / "stage_configs" / "voxcpm.yaml"

 logger = logging.getLogger(__name__)

@@ -1,5 +1,5 @@
 ARG VLLM_BASE_IMAGE=vllm/vllm-openai
-ARG VLLM_BASE_TAG=v0.21.0
+ARG VLLM_BASE_TAG=v0.20.0
 FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG}
 ARG APP_DIR=/workspace/vllm-omni
 WORKDIR ${APP_DIR}
@@ -1,48 +0,0 @@
-ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
-ARG VLLM_ASCEND_TAG=nightly-main
-FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
-
-ARG APP_DIR=/vllm-workspace/vllm-omni
-WORKDIR ${APP_DIR}
-COPY . .
-
-RUN apt-get update && \
-    apt-get install -y --no-install-recommends \
-        ffmpeg \
-        libgl1 \
-        libglib2.0-0 \
-        libgomp1 \
-        libsm6 \
-        libxext6 \
-        libxrender1 \
-        libxcb1 \
-        espeak-ng && \
-    rm -rf /var/cache/apt/* && \
-    rm -rf /var/lib/apt/lists/*
-
-RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
-    source /usr/local/Ascend/nnal/atb/set_env.sh && \
-    export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
-    python3 -m pip install -U pip && \
-    python3 -m pip install -v \
-      "pytest>=7.0.0" \
-      "pytest-asyncio>=0.21.0" \
-      "pytest-cov>=4.0.0" \
-      "pytest-mock>=3.10.0" \
-      "openpyxl>=3.0.0" \
-      "psutil>=7.2.0" \
-      "soundfile>=0.13.1" \
-      "pyttsx3>=2.99" \
-      "opencv-python>=4.12.0.88" \
-      "imageio[ffmpeg]>=0.6.0" \
-      --index-url http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple \
-      --trusted-host cache-service.nginx-pypi-cache.svc.cluster.local \
-      --default-timeout=60 && \
-    python3 -m pip install -v -e /vllm-workspace/vllm-omni/ --no-build-isolation \
-      --index-url http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple \
-      --trusted-host cache-service.nginx-pypi-cache.svc.cluster.local \
-      --default-timeout=60
-
-ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
-
-ENTRYPOINT []
@@ -1,48 +0,0 @@
-ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
-ARG VLLM_ASCEND_TAG=nightly-main-a3
-FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
-
-ARG APP_DIR=/vllm-workspace/vllm-omni
-WORKDIR ${APP_DIR}
-COPY . .
-
-RUN apt-get update && \
-    apt-get install -y --no-install-recommends \
-        ffmpeg \
-        libgl1 \
-        libglib2.0-0 \
-        libgomp1 \
-        libsm6 \
-        libxext6 \
-        libxrender1 \
-        libxcb1 \
-        espeak-ng && \
-    rm -rf /var/cache/apt/* && \
-    rm -rf /var/lib/apt/lists/*
-
-RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
-    source /usr/local/Ascend/nnal/atb/set_env.sh && \
-    export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
-    python3 -m pip install -U pip && \
-    python3 -m pip install -v \
-      "pytest>=7.0.0" \
-      "pytest-asyncio>=0.21.0" \
-      "pytest-cov>=4.0.0" \
-      "pytest-mock>=3.10.0" \
-      "openpyxl>=3.0.0" \
-      "psutil>=7.2.0" \
-      "soundfile>=0.13.1" \
-      "pyttsx3>=2.99" \
-      "opencv-python>=4.12.0.88" \
-      "imageio[ffmpeg]>=0.6.0" \
-      --index-url http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple \
-      --trusted-host cache-service.nginx-pypi-cache.svc.cluster.local \
-      --default-timeout=60 && \
-    python3 -m pip install -v -e /vllm-workspace/vllm-omni/ --no-build-isolation \
-      --index-url http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple \
-      --trusted-host cache-service.nginx-pypi-cache.svc.cluster.local \
-      --default-timeout=60
-
-ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
-
-ENTRYPOINT []
@@ -1,4 +1,4 @@
-ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.21.0
+ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.20.0
 FROM ${BASE_IMAGE} AS base

 # Declare a variable to know if we want to use the nightly build or the stable build.
@@ -88,7 +88,7 @@ ENV UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
 ENV UV_INDEX_STRATEGY="unsafe-best-match"
 ENV UV_LINK_MODE="copy"

-ARG VLLM_VERSION=v0.21.0
+ARG VLLM_VERSION=v0.20.0
 RUN git clone -b ${VLLM_VERSION} https://github.com/vllm-project/vllm
 WORKDIR /workspace/vllm

@@ -6,11 +6,10 @@ nav:
    - getting_started/installation/*
  - Serving:
    - OpenAI-Compatible API:
-      - Chat Completions API: serving/chat_completions_api.md
+      - Diffusion Chat API: serving/diffusion_chat_api.md
      - Image Generation: serving/image_generation_api.md
      - Image Edit: serving/image_edit_api.md
      - Text to Speech: serving/speech_api.md
-      - Videos: serving/videos_api.md
      - Audio Generation: serving/audio_generate_api.md
      - Streaming Video Input: serving/video_stream_api.md
  - Examples:
@@ -77,10 +76,8 @@ nav:
    - Quantization:
      - Overview: user_guide/quantization/overview.md
      - Online Quantization: user_guide/quantization/online.md
-      - Quantized KV Cache: user_guide/quantization/quantized_kvcache.md
      - FP8 W8A8: user_guide/quantization/fp8.md
      - Int8 W8A8: user_guide/quantization/int8.md
-      - ModelOpt: user_guide/quantization/modelopt.md
      - GGUF: user_guide/quantization/gguf.md
      - AutoRound: user_guide/quantization/autoround.md
      - msModelSlim: user_guide/quantization/msmodelslim.md
@@ -99,6 +99,7 @@ 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.cosyvoice3.config.CosyVoice3Config][]
 - [vllm_omni.model_executor.models.fish_speech.configuration_fish_speech.FishSpeechConfig][]
 - [vllm_omni.model_executor.models.fish_speech.configuration_fish_speech.FishSpeechFastARConfig][]
 - [vllm_omni.model_executor.models.fish_speech.configuration_fish_speech.FishSpeechSlowARConfig][]
@@ -115,14 +116,10 @@ Configuration classes.
 - [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][]
- [vllm_omni.transformers_utils.configs.cosyvoice3.CosyVoice3Config][]
 - [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Config][]
 - [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Qwen2_5_VLConfig][]
 - [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Qwen2_5_VLTextConfig][]
 - [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Qwen2_5_VLVisionConfig][]
- [vllm_omni.transformers_utils.configs.omnivoice.OmniVoiceConfig][]
- [vllm_omni.transformers_utils.configs.voxcpm.VoxCPMConfig][]
- [vllm_omni.transformers_utils.configs.voxcpm2.VoxCPM2Config][]

 ## Workers

@@ -152,8 +152,6 @@ vllm_omni/                                    tests/
 │   │                                            │   │   └── test_lora_manager.py        ✅
 │   ├── models/                                 │   ├── models/
 │   │   ├── qwen_image/                         │   │   ├── qwen_image/                 (e2e coverage)
-│   │   ├── ovis_image/                         │   │   ├── ovis_image/
-│   │   │   └── ...                             │   │   │   └── test_ovis_image.py     ✅
 │   │   ├── z_image/                            │   │   └── z_image/
 │   │   └── ...                                 │   │       └── test_zimage_tp_constraints.py  ✅
 │   └── worker/                                 │   └── worker/
@@ -234,13 +232,15 @@ vllm_omni/                                    tests/
                                               │   ├── test_qwen3_omni.py
                                               │   ├── test_qwen3_omni_expansion.py
                                               │   ├── test_mimo_audio.py
+                                               │   ├── test_image_gen_edit.py
                                               │   └── test_images_generations_lora.py
                                               └── offline_inference/                  ✅
                                                   ├── test_qwen2_5_omni.py
                                                   ├── test_qwen3_omni.py
                                                   ├── test_bagel_text2img.py
-                                                   ├── test_z_image.py
-                                                   ├── test_wan22.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
@@ -80,7 +80,7 @@ Our test scripts use the pytest framework. First, please use `git clone https://

    ```bash
    cd tests
-    pytest -s -v -m "full_model and (omni or tts) and H100" --run-level=full_model
+    pytest -s -v -m "full_model and omni and H100" --run-level=full_model
    ```

 === "L5 level"
@@ -91,16 +91,10 @@ Our test scripts use the pytest framework. First, please use `git clone https://
    cd tests

    # Stability: Qwen3-Omni
-    pytest -s -v dfx/stability/scripts/test_stability_qwen3_omni.py -m slow
+    pytest -s -v dfx/stability/scripts/test_stability_qwen3_omni.py

    # Stability: Wan2.2 (v1/videos diffusion benchmark loop)
-    pytest -s -v dfx/stability/scripts/test_stability_wan22.py -m slow
-
-    # Reliability: Qwen3-Omni
-    pytest -s -v dfx/reliability/test_reliability_qwen3_omni.py -m slow
-
-    # Reliability: Wan2.2
-    pytest -s -v dfx/reliability/test_reliability_wan22.py -m slow
+    pytest -s -v dfx/stability/scripts/test_stability_wan22.py

    ```

@@ -11,8 +11,7 @@ Defined in `pyproject.toml`:
 | `advanced_model`   | L3 tests (run on each merge to main)                 |
 | `full_model`       | L4 tests (run nightly) |
 | `diffusion`        | Diffusion model tests                                     |
-| `omni`             | Omni multimodal model tests |
-| `tts`              | TTS model tests     |
+| `omni`             | Omni model tests                                          |
 | `cache`            | Cache backend tests                                       |
 | `parallel`         | Parallelism/distributed tests                             |
 | `cpu`              | Tests that run on CPU                                     |
@@ -54,8 +54,6 @@ vllm_omni/                                    tests/
 │   │                                            │   │   └── test_lora_manager.py        ✅
 │   ├── models/                                 │   ├── models/
 │   │   ├── qwen_image/                         │   │   ├── qwen_image/                 (e2e coverage)
-│   │   ├── ovis_image/                         │   │   ├── ovis_image/
-│   │   │   └── ...                             │   │   │   └── test_ovis_image.py     ✅
 │   │   ├── z_image/                            │   │   └── z_image/
 │   │   └── ...                                 │   │       └── test_zimage_tp_constraints.py  ✅
 │   └── worker/                                 │   └── worker/
@@ -136,13 +134,15 @@ vllm_omni/                                    tests/
                                               │   ├── test_qwen3_omni.py
                                               │   ├── test_qwen3_omni_expansion.py
                                               │   ├── test_mimo_audio.py
+                                               │   ├── test_image_gen_edit.py
                                               │   └── test_images_generations_lora.py
                                               └── offline_inference/                  ✅
                                                   ├── test_qwen2_5_omni.py
                                                   ├── test_qwen3_omni.py
                                                   ├── test_bagel_text2img.py
-                                                   ├── test_z_image.py
-                                                   ├── test_wan22.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
@@ -213,6 +213,7 @@ 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
@@ -115,20 +115,17 @@ class YourAttentionBlock(nn.Module):
 from vllm_omni.diffusion.attention.layer import Attention
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata

-class YourSelfAttentionBlock(nn.Module):
+class YourAttentionBlock(nn.Module):
    def __init__(self, ...):
        super().__init__()

-        # Initialize vLLM-Omni's Attention layer.
-        # `role` lets users target this site with --diffusion-attention-config
-        # (e.g. --diffusion-attention-config.per_role.self.backend SAGE_ATTN).
+        # Initialize vLLM-Omni's Attention layer
        self.attn = Attention(
            num_heads=self.num_heads,
            head_size=self.head_dim,
            softmax_scale=1.0 / (self.head_dim ** 0.5),
            causal=False,  # Diffusion models typically use bidirectional attention
            num_kv_heads=self.num_kv_heads,
-            role="self",
        )

    def forward(self, hidden_states, encoder_hidden_states=None, attention_mask=None, ...):
@@ -136,6 +133,7 @@ class YourSelfAttentionBlock(nn.Module):
        # Create attention metadata
        attn_metadata = AttentionMetadata(attn_mask=attention_mask)
        hidden_states = self.attn(query, key, value, attn_metadata=attn_metadata)
+
 ```

 **Key Points:**
@@ -143,35 +141,8 @@ class YourSelfAttentionBlock(nn.Module):
 - **Attention layer initialization:** Done in `__init__`, not per-forward
 - **Tensor shapes:** vLLM-Omni `Attention` expects QKV to have `[B, seq, num_heads, head_dim]` shape
 - **AttentionMetadata:** Wraps attention mask and other metadata
- **Role:** Tag every `Attention` site with a `role` string so users can configure backends per role (see below)

-**Declaring attention roles**
-
-The `role` argument is a free-form string that identifies this attention site. Users can match it from `--diffusion-attention-config.per_role.<role>.*` to swap backends without touching model code. Two conventions cover the common cases:
-
-| Convention | When to use | Example |
-|---|---|---|
-| `"self"` | Q/K/V come from the same hidden state | DiT self-attention block |
-| `"cross"` | K/V come from a separate `encoder_hidden_states` | Text-conditioned cross-attention |
-
-For multi-modal or unusual sites, use a dot-namespaced role and pair it with `role_category` so it can fall back to the generic config when nothing model-specific is set:
-
-```python
-# A model-specific cross-attention site that user config can target
-# either as 'mymodel.audio_to_video' (exact) or as 'cross' (category fallback).
-self.audio_to_video_attn = Attention(
-    num_heads=self.num_heads,
-    head_size=self.head_dim,
-    softmax_scale=1.0 / (self.head_dim ** 0.5),
-    causal=False,
-    role="mymodel.audio_to_video",
-    role_category="cross",
-)
-```
-
-For cross-attention sites whose K/V are replicated across ranks (e.g. text encoder output), pass `skip_sequence_parallel=True` to opt this layer out of sequence-parallel sharding.
-
-**Attention backends:** When the user does not configure a backend, vLLM-Omni asks the current platform for its default (typically `FLASH_ATTN` on CUDA when available). Users override the default via `--diffusion-attention-backend`, the `DIFFUSION_ATTENTION_BACKEND` env var, or finer-grained `--diffusion-attention-config.per_role.*` flags. See [Diffusion Attention Backends](../../user_guide/diffusion/attention_backends.md) for the full configuration surface.
+**Attention backends:** vLLM-Omni automatically selects the attention backend given the environmental variable `DIFFUSION_ATTENTION_BACKEND`. The default attention backend is `FLASH_ATTN` for diffusion models.

 #### 1.3: Replace Imports and Utilities

@@ -1,356 +0,0 @@
-# Pipeline Parallel
-
-This section describes how to add Pipeline Parallelism (PP) to a diffusion pipeline. We use the Wan2.2 text-to-video and
-image-to-video pipelines as the reference implementations.
-
-## Implementation Checklist
-
-Adding Pipeline Parallel support requires:
-
-1. ✅ **Inherit the mixin** - add `PipelineParallelMixin` before `CFGParallelMixin` in the pipeline class
-2. ✅ **Make stage forward resumable** - support `intermediate_tensors` in the transformer forward path
-3. ✅ **Return the right object type** - `IntermediateTensors` on non-last PP ranks, final model output on the last rank
-4. ✅ **Use the standard denoising loop** - call `predict_noise_maybe_with_cfg()` and `scheduler_step_maybe_with_cfg()`
-5. ✅ **Test parity** - compare PP results against the single-GPU baseline
-
---
-
-## Table of Contents
-
- [Overview](#overview)
- [Step-by-Step Implementation](#step-by-step-implementation)
- [Testing](#testing)
- [Troubleshooting](#troubleshooting)
- [Reference Implementations](#reference-implementations)
-
---
-
-## Overview
-
-### What is Pipeline Parallelism?
-
-Pipeline Parallelism splits the denoising transformer into multiple sequential stages and places each stage on a
-different rank. Instead of every rank holding the full DiT, each PP rank owns only a slice of the layers.
-
-For each denoising step:
-
-1. Rank 0 starts the forward pass using the current latents.
-2. Each intermediate rank receives hidden states from the previous rank, runs its local layer slice, and forwards the
-   intermediate tensors downstream.
-3. The last PP rank produces the final noise prediction.
-4. The last PP rank applies the scheduler step and sends the updated latents back to rank 0 for the next timestep.
-
-This reduces per-rank model memory and enables larger diffusion transformers to run across multiple GPUs. It can also be
-combined with CFG-Parallel, where each PP pipeline carries one CFG branch.
-
-### Architecture
-
-vLLM-Omni provides `PipelineParallelMixin` to encapsulate the PP communication pattern for diffusion pipelines.
-
-| Method                            | Purpose                       | Automatic Behavior                                                              |
-|-----------------------------------|-------------------------------|---------------------------------------------------------------------------------|
-| `diffuse()`                       | Denoising loop boundary       | Wrapped by `PipelineParallelMixin` to flush pending async PP sends on exit      |
-| `predict_noise_maybe_with_cfg()`  | Predict noise with PP support | Runs partial forwards on non-last PP ranks, combines with CFG logic when needed |
-| `scheduler_step_maybe_with_cfg()` | Step scheduler with PP sync   | Runs scheduler on the last PP rank, returns updated latents to rank 0           |
-| `_sync_pp_send()`                 | Flush pending async sends     | Waits for outstanding `isend` handles before later collectives or decode        |
-
-`PipelineParallelMixin` is intentionally a pipeline-level abstraction. Your model-specific `predict_noise()` still
-defines how a local stage executes.
-
-### How It Works
-
-`predict_noise_maybe_with_cfg()` automatically switches between these modes:
-
- **PP disabled** (`pipeline_parallel_size == 1`):
-    - Falls back to `CFGParallelMixin.predict_noise_maybe_with_cfg()`
- **PP only** (`pipeline_parallel_size > 1`, `cfg_parallel_size == 1`):
-    - Rank 0 starts with the input latents
-    - Middle ranks receive `intermediate_tensors`, run their local layer range, and asynchronously send downstream
-    - The last rank returns the final noise prediction
-    - When CFG is enabled in this mode (sequential CFG), both the positive and negative branches run through the same
-      PP pipeline, doubling the communication volume per denoising step. Prefer PP + CFG-Parallel when
-      `cfg_parallel_size > 1` is available to avoid this overhead.
- **PP + CFG-Parallel** (`pipeline_parallel_size > 1`, `cfg_parallel_size > 1`):
-    - Each PP pipeline carries one CFG branch
-    - The last PP rank all-gathers across the CFG group
-    - CFG combination happens on every last PP rank in the CFG group, matching the non-PP CFG-parallel behavior
-
-`scheduler_step_maybe_with_cfg()` keeps the denoising loop consistent:
-
- **PP disabled**:
-    - Falls back to `scheduler_step_maybe_with_cfg()`
- **PP enabled**:
-    - Only the last PP rank has `noise_pred` and runs the scheduler step
-    - The resulting latents are sent back to rank 0
-    - Rank 0 receives an `AsyncLatents` wrapper that resolves only when the tensor is next consumed
-
-This asynchronous design avoids unnecessary blocking between denoising steps.
-
-When a pipeline class is defined, `PipelineParallelMixin` wraps its `diffuse()` method and calls `_sync_pp_send()` in a
-`finally` block after `diffuse()` returns or raises. This keeps model code free of explicit PP cleanup while still
-ensuring that the final non-blocking PP send completes before decode or later collectives.
-
---
-
-## Step-by-Step Implementation
-
-### Step 1: Inherit `PipelineParallelMixin` and `CFGParallelMixin`
-
-`PipelineParallelMixin` **requires** `CFGParallelMixin` and must appear before it in the class MRO. This is enforced at
-class definition time via `__init_subclass__`: defining a pipeline that inherits `PipelineParallelMixin` without
-`CFGParallelMixin`, or after `CFGParallelMixin`, raises a `TypeError` immediately on import.
-
-`PipelineParallelMixin` delegates noise prediction, CFG combination, and scheduler stepping to `CFGParallelMixin`, which
-supplies `predict_noise()`, `predict_noise_maybe_with_cfg()`, `scheduler_step_maybe_with_cfg()`, and
-`combine_cfg_noise()`.
-
-**Example:**
-
-```python
-from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
-from vllm_omni.diffusion.distributed.pipeline_parallel import PipelineParallelMixin
-import torch.nn as nn
-
-
-class YourPipeline(nn.Module, PipelineParallelMixin, CFGParallelMixin):
-    ...
-```
-
-The order matters: `PipelineParallelMixin` must be listed before `CFGParallelMixin` so calls to
-`predict_noise_maybe_with_cfg()` and `scheduler_step_maybe_with_cfg()` resolve to the PP-aware wrappers, while their
-`super()` calls delegate to the CFG implementation when PP is disabled or after the last PP stage.
-
-### Step 2: Make model forward and `predict_noise()` PP-aware
-
-The PP mixin injects `intermediate_tensors` into the normal `predict_noise()` call. Your model forward path must support
-two inputs:
-
- Normal input from rank 0, typically passed as `hidden_states` or `x`
- `intermediate_tensors` from upstream PP ranks
-
-The standard model forward pattern is:
-
-1. If `intermediate_tensors` is present, read the local hidden state from it.
-2. Run only this rank's layer slice.
-3. Return `IntermediateTensors(...)` on non-last PP ranks.
-4. Return the final model output on the last PP rank. `CFGParallelMixin.predict_noise()` already follows this contract
-   for common pipelines.
-
-**Minimal example:**
-
-```python
-from vllm.sequence import IntermediateTensors
-from vllm_omni.diffusion.distributed.parallel_state import get_pp_group
-
-
-def forward(self, hidden_states=None, intermediate_tensors=None, **kwargs):
-    if intermediate_tensors is not None:
-        hidden_states = intermediate_tensors["hidden_states"]
-
-    for i in range(self.start_layer, self.end_layer):
-        hidden_states = self.layers[i](hidden_states)
-
-    pp_group = get_pp_group()
-    if not pp_group.is_last_rank:
-        return IntermediateTensors({"hidden_states": hidden_states})
-    return (hidden_states,)
-```
-
-### Step 3: Partition the transformer layers
-
-The local module on each PP rank must expose only that rank's layer slice. In the reference tests and vLLM model
-implementations, this is typically done with vLLM utilities such as `make_layers(...)`, which fill missing ranges with
-`PPMissingLayer`. Besides partitioning the layers themselves, model authors should also wire
-`make_empty_intermediate_tensors_factory(...)` for intermediate tensor allocation and `is_pp_missing_parameter(...)` for
-PP-aware weight loading.
-
-To prepare a transformer for PP, implement the following pieces in order.
-
-#### 3.1 Split the transformer layers across PP ranks
-
-Each PP rank should own only its local layer range, typically exposed as `[start_layer, end_layer)`. In practice this is
-usually done with `make_layers(...)`, which constructs local layers and fills missing ranges with `PPMissingLayer`.
-
-The goal is:
-
- every PP rank knows its `[start_layer, end_layer)` range
- non-local layers do not execute on this rank
- the forward can resume from incoming `intermediate_tensors`
-
-By default, layers are distributed evenly across PP ranks using `get_pp_indices()`. When the number of layers is not
-divisible by the PP size, remaining layers are assigned to middle partitions to balance compute and memory. You can
-override this with the `VLLM_PP_LAYER_PARTITION` environment variable to specify an exact per-rank layer count:
-
-```bash
-# Example: 40 layers across 4 PP ranks, assigning 8 / 12 / 12 / 8 layers
-export VLLM_PP_LAYER_PARTITION=8,12,12,8
-```
-
-The value must be a comma-separated list of integers whose length equals `pipeline_parallel_size` and whose sum equals
-the total number of transformer layers.
-
-#### 3.2 Expose `make_empty_intermediate_tensors`
-
-Your transformer module should expose `self.make_empty_intermediate_tensors`, typically created with
-`make_empty_intermediate_tensors_factory(...)`.
-
-This is important because PP ranks need a consistent way to allocate placeholder `IntermediateTensors` with the expected
-keys and hidden dimension.
-
-**Example:**
-
-```python
-from vllm.model_executor.models.utils import make_empty_intermediate_tensors_factory
-
-self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
-    ["hidden_states"],
-    inner_dim,
-)
-```
-
-For Wan2.2, the intermediate payload between PP stages is the token sequence stored under `"hidden_states"`, so the
-factory is created with that key and the transformer hidden size.
-
-#### 3.3 Return `IntermediateTensors` on non-last PP ranks
-
-Your model `forward()` or custom `predict_noise()` implementation should consume `intermediate_tensors` on non-first
-ranks and return `IntermediateTensors(...)` on non-last ranks.
-
-This lets each PP stage resume from the upstream hidden states and pass the local result to the next stage.
-
-#### 3.4 Skip non-local weights during `load_weights()`
-
-When a model is PP-partitioned, many parameters in the checkpoint belong to layers that do not exist on the current
-rank. `load_weights()` must therefore skip parameters for missing PP stages using `is_pp_missing_parameter(...)`.
-
-If this is not done, weight loading will either fail or incorrectly try to load tensors into `PPMissingLayer`
-placeholders.
-
-The Wan2.2 transformer is the best reference here: it uses `is_pp_missing_parameter(...)` before loading both remapped
-and fused parameters.
-
-If your model has multiple transformer variants, PP still works as long as each selected transformer obeys the same
-contract.
-
-### Step 4: Use the standard denoising contract
-
-vLLM-Omni diffusion pipelines already route denoising through `diffuse()`, `predict_noise_maybe_with_cfg()`, and
-`scheduler_step_maybe_with_cfg()`. `PipelineParallelMixin` overrides those standard helpers when
-`pipeline_parallel_size > 1`, so model integrations should not add separate PP-specific helper names or manual post-loop
-synchronization.
-
-In PP mode:
-
- Non-last PP ranks return `None` from `predict_noise_maybe_with_cfg()` after sending `IntermediateTensors` downstream.
- The last PP rank returns the final prediction, and with CFG-Parallel every last PP rank in the CFG group receives the
-  combined result.
- Rank 0 receives `AsyncLatents` from `scheduler_step_maybe_with_cfg()`, which resolves only when the latents are
-  consumed.
- Pending non-blocking sends are flushed automatically when `diffuse()` exits.
-
---
-
-## Testing
-
-Use an offline inference script with `pipeline_parallel_size > 1`.
-
-```bash
-python examples/offline_inference/text_to_video/text_to_video.py \
--model=Wan-AI/Wan2.2-TI2V-5B-Diffusers \
--width=1280 \
--height=704 \
--guidance-scale=5.0 \
--prompt="Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage" \
--output=t2v_5B_pp2.mp4 \
--pipeline-parallel-size=2
-```
-
-For PP + CFG-Parallel together:
-
-```bash
-python examples/offline_inference/text_to_video/text_to_video.py \
--model=Wan-AI/Wan2.2-TI2V-5B-Diffusers \
--width=1280 \
--height=704 \
--guidance-scale=5.0 \
--prompt="Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage" \
--output=t2v_5B_pp2_cfg2.mp4 \
--pipeline-parallel-size=2 \
--cfg-parallel-size=2
-```
-
-**Verify:**
-
-1. The run completes without hangs at the PP stage boundary
-2. Output quality matches the non-PP baseline within normal numerical variance
-3. Peak memory per GPU drops compared with the single-rank model
-4. There are no outstanding communication errors before decode
-
---
-
-## Troubleshooting
-
-### Issue: `TypeError` on import — `CFGParallelMixin` missing
-
-**Symptoms:** Importing a pipeline that inherits `PipelineParallelMixin` raises:
-
-```
-TypeError: YourPipeline inherits PipelineParallelMixin but not CFGParallelMixin.
-```
-
-or:
-
-```
-TypeError: YourPipeline must inherit PipelineParallelMixin before CFGParallelMixin ...
-```
-
-**Cause:** `PipelineParallelMixin` enforces via `__init_subclass__` that the subclass also inherits `CFGParallelMixin`
-and lists `PipelineParallelMixin` first in the MRO.
-
-**Solution:** Add `CFGParallelMixin` to the pipeline's base classes after `PipelineParallelMixin`:
-
-```python
-from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
-from vllm_omni.diffusion.distributed.pipeline_parallel import PipelineParallelMixin
-
-
-class YourPipeline(nn.Module, PipelineParallelMixin, CFGParallelMixin):
-    ...
-```
-
-### Issue: Non-last PP ranks crash when calling `predict_noise`
-
-**Symptoms:** Shape or missing-input errors on ranks other than the first or last PP rank.
-
-**Cause:** The model forward path assumes direct input tensors and ignores `intermediate_tensors`.
-
-**Solution:** Update the transformer `forward()` or custom `predict_noise()` path to load hidden states from
-`intermediate_tensors` when present.
-
-### Issue: PP outputs differ from single-GPU baseline
-
-**Symptoms:** The PP run finishes but produces numerically inconsistent outputs.
-
-**Causes & Solutions:**
-
- **Local layer partitioning is wrong**
-    - Verify each rank runs only its `[start_layer, end_layer)` slice
- **Non-last ranks return a plain tensor instead of `IntermediateTensors`**
-    - Return `IntermediateTensors({...})` until the last PP stage
- **CFG branch wiring is incorrect**
-    - When CFG is enabled, confirm positive and negative kwargs are passed exactly as in the non-PP path
-
---
-
-## Reference Implementations
-
-Complete examples in the codebase:
-
-| Component               | Path                                                       | Notes                                                                                                                   |
-|-------------------------|------------------------------------------------------------|-------------------------------------------------------------------------------------------------------------------------|
-| `PipelineParallelMixin` | `vllm_omni/diffusion/distributed/pipeline_parallel.py`     | Core PP communication and scheduler helpers                                                                             |
-| `CFGParallelMixin`      | `vllm_omni/diffusion/distributed/cfg_parallel.py`          | Default `predict_noise()` tuple normalization and CFG helper fallback                                                   |
-| Wan2.2 transformer      | `vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py`  | Reference for layer partitioning, `IntermediateTensors`, `make_empty_intermediate_tensors`, and PP-aware weight loading |
-| Wan2.2 T2V pipeline     | `vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py`     | Reference PP + CFG integration for text-to-video                                                                        |
-| Wan2.2 I2V pipeline     | `vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py` | Reference PP + CFG integration for image-to-video                                                                       |
-| PP tests                | `tests/diffusion/distributed/test_pipeline_parallel.py`    | Baseline parity and async communication tests                                                                           |
@@ -455,7 +455,7 @@ To learn more about the diffusion pipeline and how to add a new diffusion pipeli

 #### Architecture

-The attention system uses a **role-aware backend selector pattern**. Each `Attention` site declares a semantic `role` (`"self"`, `"cross"`, or a model-specific string). The selector consults the user's `AttentionConfig` to pick a backend per role and falls back to the platform default when nothing is configured.
+The attention system uses a **backend selector pattern** that automatically chooses the optimal attention implementation based on hardware and model configuration.

 #### Backend Selection

@@ -463,90 +463,92 @@ The attention system uses a **role-aware backend selector pattern**. Each `Atten

 ```python
 class Attention(nn.Module):
-    def __init__(self, num_heads, head_size, causal, softmax_scale, *,
-                 role="self", role_category=None, ...):
-        # Resolve backend for this role from the active OmniDiffusionConfig
-        config = get_current_diffusion_config_or_none()
-        attention_config = config.attention_config if config is not None else None
-        attn_backend_cls, spec = get_attn_backend_for_role(
-            role=role,
-            head_size=head_size,
-            attention_config=attention_config,
-            role_category=role_category,
-        )
-        self.attn_impl_cls = attn_backend_cls.get_impl_cls()
-        self.attention = self.attn_impl_cls(..., backend_kwargs=spec.extra if spec else None)
+    def __init__(self, num_heads, head_size, causal, softmax_scale, ...):
+        # Auto-select backend
+        self.attn_backend = get_attn_backend(-1)
+        self.attn_impl_cls = self.attn_backend.get_impl_cls()
+        self.attention = self.attn_impl_cls(...)
 ```

 **Available Backends**:

- **FlashAttention** (`FLASH_ATTN`): Memory-efficient kernel that dispatches to FA2/FA3 on CUDA, AITER on ROCm, and `mindiesd` on Ascend NPU.
+- **FlashAttention**: Optimized CUDA kernel (FA2/FA3) - memory efficient via tiling

- **SDPA** (`TORCH_SDPA`): PyTorch's scaled dot-product attention — cross-platform fallback.
+- **SDPA**: PyTorch's scaled dot-product attention - default, cross-platform

- **SageAttention** (`SAGE_ATTN`): Quantized attention implementation from the SageAttention library.
+- **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).

 #### Backend Selection Mechanism

-Selection is driven by `AttentionConfig` on `OmniDiffusionConfig`, which carries a global `default` plus a `per_role` map of `AttentionSpec` entries. A role string is resolved in this order:
-
 ```python
-def get_attn_backend_for_role(role, head_size, attention_config=None, role_category=None):
-    # 1. attention_config.per_role[role]            — exact match
-    # 2. attention_config.per_role[role_category]   — category fallback
-    # 3. attention_config.default                   — global default
-    # 4. Platform default                           — hardware-specific
-    if attention_config is not None:
-        spec, source = attention_config.resolve_with_source(
-            role=role, role_category=role_category,
-        )
-        if spec is not None:
-            return load_backend(spec.backend), spec
-    return current_omni_platform.get_diffusion_attn_backend_cls(...), None
+def get_attn_backend(head_size: int) -> type[AttentionBackend]:
+    # Check environment variable
+    backend_name = os.environ.get("DIFFUSION_ATTENTION_BACKEND")
+
+    if backend_name:
+        return load_backend(backend_name.upper())
+
+    # Default to SDPA
+    return SDPABackend
 ```

 **Selection Priority**:

-1. **Per-role override** (`--diffusion-attention-config.per_role.<role>.backend`) — finest control; matched on the layer's exact `role` string.
+1. **Environment Variable**: `DIFFUSION_ATTENTION_BACKEND` for manual override

-2. **Role category fallback** (`--diffusion-attention-config.per_role.<category>.backend`) — used when an exact match is missing and the layer declared `role_category` (e.g. a `"mymodel.audio_to_video"` site can fall back to `"cross"`).
+    - Valid values: `FLASH_ATTN`, `TORCH_SDPA`, `SAGE_ATTN`, `ASCEND`

-3. **Global default** (`--diffusion-attention-backend FLASH_ATTN`, or `--diffusion-attention-config.default.backend`, or `DIFFUSION_ATTENTION_BACKEND` env var).
+    - Example: `export DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN`

-4. **Platform default** — `current_omni_platform.get_diffusion_attn_backend_cls(...)` picks the best-available kernel for the hardware.
+2. **Automatic Fallback**: Falls back to SDPA if selected backend unavailable

-`AttentionSpec.extra` is forwarded to the backend constructor as `backend_kwargs`, so backend-specific parameters (e.g. SparseBlock `block_size`) can be set without changing model code.
-
-For the user-facing CLI surface, see [Diffusion Attention Backends](../../user_guide/diffusion/attention_backends.md). For the role declaration contract on the model side, see [Adding a Diffusion Model](../../contributing/model/adding_diffusion_model.md).
+3. **Hardware Detection**: Can select based on device type (NPU, CUDA, etc.)

 **Backend Availability**:

 - **SDPA**: Always available (PyTorch built-in)

- **FlashAttention**: Requires `flash-attn` on CUDA, `aiter` on ROCm, or `mindiesd` on Ascend NPU
+- **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

-**Location**: `vllm_omni/diffusion/attention/selector.py` + `vllm_omni/platforms/<device>/platform.py`
+**Location**: `vllm_omni/diffusion/attention/selector.py`

-Backend resolution is delegated to the active platform: `current_omni_platform.get_diffusion_attn_backend_cls(selected_backend, head_size)` returns a fully-qualified class path that `_load_backend_cls` imports lazily, with the result cached by `(backend_name, head_size)`.
+The attention system uses a **registry pattern** to manage and dynamically load attention backends. This allows for easy extension and runtime selection of backends.
+
+
+**Registry Structure**:

 ```python
-@cache
-def _cached_get_backend_cls(backend_name: str | None, head_size: int) -> type[AttentionBackend]:
-    backend_cls_path = current_omni_platform.get_diffusion_attn_backend_cls(
-        selected_backend=backend_name,
-        head_size=head_size,
-    )
-    return _load_backend_cls(backend_cls_path)
+# Registry mapping backend names to their module paths and class names
+_BACKEND_CONFIG = {
+    "FLASH_ATTN": {
+        "module": "vllm_omni.diffusion.attention.backends.flash_attn",
+        "class": "FlashAttentionBackend",
+    },
+    "TORCH_SDPA": {
+        "module": "vllm_omni.diffusion.attention.backends.sdpa",
+        "class": "SDPABackend",
+    },
+    "SAGE_ATTN": {
+        "module": "vllm_omni.diffusion.attention.backends.sage_attn",
+        "class": "SageAttentionBackend",
+    },
+    "ASCEND": {
+        "module": "vllm_omni.diffusion.attention.backends.ascend_attn",
+        "class": "AscendAttentionBackend",
+    },
+}
 ```

-Each platform (`cuda`, `rocm`, `xpu`, `musa`, `npu`) maps backend names like `"FLASH_ATTN"`, `"TORCH_SDPA"`, `"SAGE_ATTN"` to the right backend class path for that hardware, including head-size compatibility checks. Passing `selected_backend=None` lets the platform pick its own default.
-
 #### Attention Backend Integration

 The `Attention` layer integrates backends through a unified interface. Here's how **FlashAttentionBackend** is integrated as an example:
@@ -34,7 +34,7 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr

 ### Pre-built wheels

-Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, and 0.21.0. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
+Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, and 0.18.0. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).

 === "NVIDIA CUDA"

@@ -20,7 +20,7 @@ Therefore, it is recommended to install vLLM and vLLM-Omni with a **fresh new**

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

 #### Installation of vLLM-Omni
@@ -39,13 +39,13 @@ uv pip install 'vllm-omni[demo]'
 # --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.21.0 release version of the library
+If you do not need to modify source code of vLLM, you can directly install the stable 0.20.0 release version of the library

 ```bash
-uv pip install vllm==0.21.0 --torch-backend=auto
+uv pip install vllm==0.20.0 --torch-backend=auto
 ```

-The 0.21.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
+The 0.20.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.

 #### Installation of vLLM-Omni
 Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -66,12 +66,12 @@ 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.21.0
+git checkout v0.20.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
-#For CUDA 13.0 (the default for v0.21.0; the wheel filename has no `+cu130` suffix)
-export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.21.0/vllm-0.21.0-cp38-abi3-manylinux_2_35_x86_64.whl
+#For CUDA 13.0 (the default for v0.20.0; the wheel filename has no `+cu130` suffix)
+export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.20.0/vllm-0.20.0-cp38-abi3-manylinux_2_35_x86_64.whl
 ```
 Install vllm with command below (If you have no existing PyTorch).
 ```bash
@@ -102,7 +102,7 @@ docker run --runtime nvidia --gpus 2 \
    --env "HF_TOKEN=$HF_TOKEN" \
    -p 8091:8091 \
    --ipc=host \
-    vllm/vllm-omni:v0.21.0 \
+    vllm/vllm-omni:v0.18.0 \
    --model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 8091
 ```

@@ -13,7 +13,7 @@ vLLM-Omni current recommends the steps in under setup through Docker Images.

 vLLM-Omni is built based on vLLM. Please install it with command below.
 ```bash
-uv pip install vllm==0.21.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.21.0/rocm722
+uv pip install vllm==0.20.0+rocm721 --extra-index-url https://wheels.vllm.ai/rocm/0.20.0/rocm721
 ```

 #### Installation of vLLM-Omni
@@ -34,13 +34,13 @@ uv pip install onnxruntime-rocm
 # --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.21.0 release version of the library
+If you do not need to modify source code of vLLM, you can directly install the stable 0.20.0 release version of the library

 ```bash
-uv pip install vllm==0.21.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.21.0/rocm722
+uv pip install vllm==0.20.0+rocm721 --extra-index-url https://wheels.vllm.ai/rocm/0.20.0/rocm721
 ```

-The pre-built 0.21.0 vLLM wheel targets ROCm 7.2.2. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
+The pre-built 0.20.0 vLLM wheel targets ROCm 7.2.1. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.

 #### Installation of vLLM-Omni
 Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -58,7 +58,7 @@ 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.21.0
+git checkout v0.20.0
 python3 -m pip install -r requirements/rocm.txt
 python3 setup.py develop
 ```
@@ -130,7 +130,7 @@ docker run --rm \
  -v ~/.cache/huggingface:/root/.cache/huggingface \
  --env "HF_TOKEN=$HF_TOKEN" \
  -p 8091:8091 \
-  vllm/vllm-omni-rocm:v0.21.0 \
+  vllm/vllm-omni-rocm:v0.20.0 \
  --model Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 ```

@@ -149,7 +149,7 @@ docker run --rm -it \
  -v ~/.cache/huggingface:/root/.cache/huggingface \
  --env "HF_TOKEN=$HF_TOKEN" \
  --entrypoint bash \
-  vllm/vllm-omni-rocm:v0.21.0
+  vllm/vllm-omni-rocm:v0.20.0
 ```

 # --8<-- [end:pre-built-images]
@@ -19,10 +19,10 @@ uv venv --python 3.12 --seed
 source .venv/bin/activate

 # On CUDA
-uv pip install vllm==0.21.0 --torch-backend=auto
+uv pip install vllm==0.20.0 --torch-backend=auto

 # On ROCm
-uv pip install vllm==0.21.0+rocm721 --extra-index-url https://wheels.vllm.ai/rocm/0.21.0/rocm721
+uv pip install vllm==0.20.0+rocm721 --extra-index-url https://wheels.vllm.ai/rocm/0.20.0/rocm721

 git clone https://github.com/vllm-project/vllm-omni.git
 cd vllm-omni
@@ -35,7 +35,7 @@ For additional installation methods — please see the [installation guide](inst
 !!! note
    It is important to install the same major & minor version of vLLM and vLLM Omni, otherwise things may not work as expected. If the versions are misaligned, you will see a warning when you import vLLM Omni.

-    If you are seeing strange behavior with the `vllm` command not handling the `--omni` flag correctly, you most likely have a version mismatch with vLLM < `0.21.0` and vLLM Omni `0.21.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue.
+    If you are seeing strange behavior with the `vllm` command not handling the `--omni` flag correctly, you most likely have a version mismatch with vLLM < `0.20.0` and vLLM Omni `0.20.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue.

 ## Offline Inference

@@ -64,11 +64,9 @@ th {
 | `Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | ✅︎ | ✅︎ | | |
 | `FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✅︎ | ✅︎ | | |
 | `DreamIDOmniPipeline` | DreamID-Omni | `XuGuo699/DreamID-Omni` | ✅︎ | ✅︎ | | |
-| `SenseNovaU1Pipeline` | SenseNova-U1 (DiT-only) | `SenseNova/SenseNova-U1-8B-MoT` | ✅︎ | | | |
 | `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | ✅︎ | ✅︎ | | |
 | `HunyuanVideo15ImageToVideoPipeline` | HunyuanVideo-1.5-I2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v` | ✅︎ | ✅︎ | | |
 | `VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✅︎ | ✅︎ | | |
-| `CovoAudioForConditionalGeneration` | Covo-Audio-Chat | `tencent/Covo-Audio-Chat` | ✅︎ | | | |
 |`DyninOmniForConditionalGeneration` | Dynin-Omni | `snu-aidas/Dynin-Omni` | ✅︎ | | | |
 | `ErnieImagePipeline` | ERNIE-Image | `baidu/ERNIE-Image`, `baidu/ERNIE-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 ✅︎ indicates the model is supported on that backend. Empty cells mean not listed as supported on that backend.
@@ -1,4 +1,4 @@
-# Chat Completions API
+# Diffusion Chat Completions API

 vLLM-Omni supports generating and editing images via the `/v1/chat/completions`
 endpoint using diffusion models. This page explains how to pass generation
@@ -531,18 +531,16 @@ for result in response.json()["results"]:

 All items are fanned out to `generate()` concurrently. The engine's stage worker automatically batches them up to the configured `max_batch_size` and queues the rest — no client-side throttling needed.

-For best throughput, set both stages' `max_num_seqs` above 1 via `--stage-overrides`. On the current Qwen3-TTS CustomVoice benchmark, stage 1 performed best at `max_num_seqs: 10`:
+For best throughput, set both stages' `max_num_seqs` to ≥4 via `--stage-overrides`:

 ```bash
 vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
    --omni --port 8091 --trust-remote-code --enforce-eager \
-    --stage-overrides '{"0":{"max_num_seqs":10,"gpu_memory_utilization":0.2},
-                        "1":{"max_num_seqs":10,"gpu_memory_utilization":0.2}}'
+    --stage-overrides '{"0":{"max_num_seqs":4,"gpu_memory_utilization":0.2},
+                        "1":{"max_num_seqs":4,"gpu_memory_utilization":0.2}}'
 ```

-The bundled `qwen3_tts.yaml` uses a multi-request default and lets stage 1 batch chunks across in-flight requests. For latency-sensitive deployments, avoid forcing stage 1 back to `max_num_seqs: 1`; benchmark before reducing it below `10`.
-
-The bundled config also sets `initial_codec_chunk_frames: 1`. This emits only the first audio chunk early for lower TTFA, then returns to the normal `codec_chunk_frames` window so Code2Wav does not repeatedly decode tiny overlapping chunks.
+The bundled `qwen3_tts.yaml` uses `max_num_seqs: 1` (single request) on both stages. Bumping to 4 yields roughly 4× throughput on the talker and lets stage 1 batch chunks across in-flight requests.

 ## Supported Models

@@ -1,152 +0,0 @@
-# Videos API
-
-vLLM-Omni provides an OpenAI-compatible video generation API for diffusion
-video models. The API supports asynchronous video jobs through `/v1/videos` and
-a synchronous benchmark-oriented endpoint through `/v1/videos/sync`.
-
-Each server instance runs a single model specified at startup with
-`vllm serve <model> --omni`.
-
-## Quick Start
-
-### Start the Server
-
-```bash
-vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --port 8091
-```
-
-### Create a Video Job
-
-```bash
-create_response=$(curl -s http://localhost:8091/v1/videos \
-  -F "prompt=A cinematic tracking shot of a mountain lake at sunrise" \
-  -F "width=1280" \
-  -F "height=720" \
-  -F "num_frames=80" \
-  -F "fps=16" \
-  -F "num_inference_steps=40")
-
-video_id=$(echo "${create_response}" | jq -r '.id')
-```
-
-### Poll and Download
-
-```bash
-curl -s "http://localhost:8091/v1/videos/${video_id}" | jq .
-curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o output.mp4
-```
-
-## API Reference
-
-### Endpoints
-
-| Endpoint | Method | Description |
-|----------|--------|-------------|
-| `/v1/videos` | `POST` | Create an asynchronous video generation job |
-| `/v1/videos/sync` | `POST` | Generate a video synchronously and return raw video bytes |
-| `/v1/videos/{video_id}` | `GET` | Retrieve job status and metadata |
-| `/v1/videos` | `GET` | List stored video jobs |
-| `/v1/videos/{video_id}/content` | `GET` | Download generated video content |
-| `/v1/videos/{video_id}` | `DELETE` | Delete a video job and stored output |
-
-### Request Parameters
-
-`POST /v1/videos` and `POST /v1/videos/sync` accept `multipart/form-data`.
-
-#### OpenAI-style fields
-
-| Parameter | Type | Default | Description |
-|-----------|------|---------|-------------|
-| `prompt` | string | **required** | Text prompt for video generation |
-| `model` | string | server's model | Optional model name |
-| `seconds` | string | null | Requested clip duration in seconds |
-| `size` | string | null | Requested output size in `WIDTHxHEIGHT` format |
-| `user` | string | null | Optional user identifier |
-
-#### vLLM-Omni extension fields
-
-| Parameter | Type | Default | Description |
-|-----------|------|---------|-------------|
-| `input_reference` | file | null | Uploaded reference image for image-to-video requests |
-| `image_reference` | string | null | JSON-encoded reference image payload; do not combine with `input_reference` |
-| `width` | integer | model default | Output video width |
-| `height` | integer | model default | Output video height |
-| `num_frames` | integer | model default | Number of generated frames |
-| `fps` | integer | model default | Output frames per second |
-| `num_inference_steps` | integer | model default | Number of diffusion steps |
-| `guidance_scale` | number | null | CFG guidance scale for the low-noise stage |
-| `guidance_scale_2` | number | null | CFG guidance scale for the high-noise stage |
-| `boundary_ratio` | number | null | Boundary split ratio for multi-stage denoising |
-| `flow_shift` | number | null | Scheduler flow-shift value |
-| `true_cfg_scale` | number | null | True CFG scale when supported by the model |
-| `seed` | integer | null | Random seed for reproducibility |
-| `negative_prompt` | string | null | Text describing what to avoid in the generated video |
-| `enable_frame_interpolation` | boolean | null | Enable post-generation frame interpolation |
-| `frame_interpolation_exp` | integer | null | Interpolation exponent; `1=2x`, `2=4x`, and so on |
-| `frame_interpolation_scale` | number | null | RIFE inference scale |
-| `frame_interpolation_model_path` | string | null | Local path or Hugging Face repo for the interpolation model |
-| `lora` | string | null | JSON-encoded LoRA configuration object |
-| `extra_params` | string | null | JSON-encoded object for additional model-specific parameters |
-
-### Create Response
-
-`POST /v1/videos` returns a job record:
-
-```json
-{
-  "id": "video-123",
-  "status": "queued",
-  "created_at": 1701234567
-}
-```
-
-The final content is available from `/v1/videos/{video_id}/content` after the
-job status becomes `completed`.
-
-### Synchronous Response
-
-`POST /v1/videos/sync` blocks until generation finishes and returns raw video
-bytes. It is useful for benchmarks and simple scripts that do not need job
-storage or polling.
-
-## Examples
-
-### Image-to-Video
-
-```bash
-curl -s http://localhost:8091/v1/videos \
-  -F "prompt=animate this image with subtle camera movement" \
-  -F "input_reference=@input.png" \
-  -F "width=1280" \
-  -F "height=720" \
-  -F "num_frames=80" \
-  -F "fps=16"
-```
-
-### Synchronous Generation
-
-```bash
-curl -X POST http://localhost:8091/v1/videos/sync \
-  -F "prompt=A small robot walking through a neon city" \
-  -F "width=854" \
-  -F "height=480" \
-  -F "num_frames=80" \
-  -F "fps=16" \
-  -o output.mp4
-```
-
-## Storage
-
-Set `VLLM_OMNI_STORAGE_PATH` to control where asynchronous video outputs are
-stored:
-
-```bash
-export VLLM_OMNI_STORAGE_PATH=/var/tmp/vllm-omni-videos
-```
-
-## Model-Specific Examples
-
-For complete text-to-video and image-to-video walkthroughs, see:
-
- [Text-to-Video](../user_guide/examples/online_serving/text_to_video.md)
- [Image-to-Video](../user_guide/examples/online_serving/image_to_video.md)
@@ -1,10 +1,10 @@
 # Diffusion Attention Backends

-This document describes the diffusion attention backends available in vLLM-Omni, how to select them globally and per-role, and how to use SageAttention.
+This document describes the diffusion attention backends available in vLLM-Omni, how to select them, and how to use SageAttention.

 ## Overview

-Diffusion attention backend selection is resolved in `vllm_omni.diffusion.attention.selector`. It looks up the backend from a structured `AttentionConfig` carried on `OmniDiffusionConfig` and falls back to the platform default when nothing is configured.
+Diffusion attention backend selection is controlled by the `DIFFUSION_ATTENTION_BACKEND` environment variable and resolved in `vllm_omni.diffusion.attention.selector`.

 This backend is used by diffusion attention layers such as the DiT attention in video and image generation models.

@@ -14,7 +14,7 @@ On CUDA, the practical choices today are:
 - `TORCH_SDPA`: PyTorch `scaled_dot_product_attention`.
 - `SAGE_ATTN`: SageAttention backend, if `sageattention` is installed.

-If no attention backend is configured, vLLM-Omni asks the current platform to choose the default. On CUDA, that normally means `FLASH_ATTN` when available, otherwise `TORCH_SDPA`.
+If `DIFFUSION_ATTENTION_BACKEND` is unset, vLLM-Omni asks the current platform to choose the default backend. On CUDA, that normally means `FLASH_ATTN` when available, otherwise `TORCH_SDPA`.

 ## Backend Options

@@ -24,77 +24,19 @@ If no attention backend is configured, vLLM-Omni asks the current platform to ch
 | `TORCH_SDPA` | Most conservative fallback. Useful for debugging or compatibility. |
 | `SAGE_ATTN` | Requires `sageattention`. Can improve performance on some workloads, but output quality must be validated model-by-model. |

-## Configuration
+## Selection Priority

-Diffusion attention backends can be configured three ways, in priority order:
+Diffusion attention backend selection follows this order:

-1. **`--diffusion-attention-config`** — structured per-role config (highest priority).
-2. **`--diffusion-attention-backend` / `DIFFUSION_ATTENTION_BACKEND` env var** — global shorthand that sets the default backend.
-3. **Platform default** — used when nothing is configured.
+1. `DIFFUSION_ATTENTION_BACKEND`
+2. Platform default

-`--diffusion-attention-backend` is shorthand for `--diffusion-attention-config.default.backend`. It may be combined with `--diffusion-attention-config.per_role.*` overrides, but is mutually exclusive with `--diffusion-attention-config.default.backend`.
-
-### Global default
-
-Set the default backend for every diffusion attention layer:
+Example:

 ```bash
-# CLI flag
-vllm-omni serve <model> --diffusion-attention-backend SAGE_ATTN
-
-# Environment variable (also recognized for backwards compatibility)
 export DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN
 ```

-### Per-role configuration
-
-Roles are free-form strings declared by each diffusion model. The two common categories are `"self"` and `"cross"`; model-specific roles (e.g. `"ltx2.audio_to_video"`) may also be declared. A role string is matched in this order:
-
-1. Exact `per_role[role]` match
-2. `per_role[role_category]` fallback (e.g. `"ltx2.audio_to_video"` → `"cross"`)
-3. `default`
-4. Platform default
-
-Use vLLM-style dotted flags or one JSON blob:
-
-```bash
-# Dotted flags
-vllm-omni serve <model> \
-    --diffusion-attention-config.default.backend FLASH_ATTN \
-    --diffusion-attention-config.per_role.cross.backend TORCH_SDPA
-
-# JSON
-vllm-omni serve <model> \
-    --diffusion-attention-config '{"default":{"backend":"FLASH_ATTN"},"per_role":{"cross":{"backend":"TORCH_SDPA"}}}'
-```
-
-Backends may also accept backend-specific parameters via `extra`:
-
-```bash
--diffusion-attention-config.per_role.self.backend SPARSE_BLOCK \
--diffusion-attention-config.per_role.self.extra.block_size 128
-```
-
-### Programmatic API
-
-When constructing `OmniDiffusionConfig` directly:
-
-```python
-from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, OmniDiffusionConfig
-
-config = OmniDiffusionConfig(
-    attention_config=AttentionConfig(
-        default=AttentionSpec(backend="FLASH_ATTN"),
-        per_role={
-            "cross": AttentionSpec(backend="TORCH_SDPA"),
-        },
-    ),
-    ...,
-)
-```
-
-A plain dict is also accepted and normalized to `AttentionConfig`.
-
 ## SageAttention Installation

 vLLM-Omni expects SageAttention to be installed into the same Python environment as vLLM-Omni.
@@ -143,21 +85,7 @@ DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN python examples/offline_inference/text_to_
    --output outputs/wan22_sage.mp4
 ```

-### Mixed backends across roles
-
-Use `FLASH_ATTN` for self-attention and `TORCH_SDPA` for cross-attention:
-
-```bash
-python examples/offline_inference/text_to_video/text_to_video.py \
-    --model Wan-AI/Wan2.2-TI2V-5B-Diffusers \
-    --prompt "A dog running across a field of golden wheat." \
-    --diffusion-attention-config.per_role.self.backend FLASH_ATTN \
-    --diffusion-attention-config.per_role.cross.backend TORCH_SDPA \
-    --tensor-parallel-size 2 \
-    --output outputs/wan22_mixed.mp4
-```
-
-### Compare against FlashAttention
+### Compare Against FlashAttention

 Unset the backend override, or explicitly use `FLASH_ATTN`:

@@ -38,15 +38,14 @@ vllm-omni serve diffusion Wan-AI/Wan2.2-T2V-A14B-Diffusers --enable-cpu-offload

 ### To Support a Model

-Implement the `SupportsComponentDiscovery` protocol to declare which
-submodules serve as pipeline components (used by offloading, HSDP
-sharding, and other framework features):
+Implement the `SupportsModuleOffload` protocol to declare which
+submodules participate in offloading:

 ```python
 from typing import ClassVar
-from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
+from vllm_omni.diffusion.models.interface import SupportsModuleOffload

-class MyPipeline(nn.Module, SupportsComponentDiscovery):
+class MyPipeline(nn.Module, SupportsModuleOffload):
    _dit_modules: ClassVar[list[str]] = ["transformer"]
    _encoder_modules: ClassVar[list[str]] = ["text_encoder", "vision_model"]
    _vae_modules: ClassVar[list[str]] = ["vae"]
@@ -159,7 +158,7 @@ class Flux2Transformer2DModel(nn.Module):
 The offloader discovers pipeline components in two ways:

 1. **Protocol-based** (preferred): If the pipeline implements
-    `SupportsComponentDiscovery`, its `_dit_modules`, `_encoder_modules`,
+    `SupportsModuleOffload`, its `_dit_modules`, `_encoder_modules`,
    `_vae_modules`, and `_resident_modules` class variables are used
    directly.  All attribute names support dotted paths (e.g.
    `"pipe.transformer"`, `"bagel.time_embedder"`) for nested submodules.
@@ -28,6 +28,7 @@ Frame interpolation is currently supported for:

 - `WanPipeline` (Wan2.2 text-to-video)
 - `WanImageToVideoPipeline`
+- `Wan22TI2VPipeline`

 ## Request Parameters

@@ -0,0 +1,80 @@
+# Diffusion Multi-Replica Serving
+
+Diffusion stages can be replicated in online serving by setting `runtime.num_replicas` on a `stage_type: diffusion` stage. Each replica owns an independent diffusion engine process and receives requests through round-robin routing with request affinity.
+
+## When to Use
+
+Use multi-replica serving when one diffusion process does not saturate the available GPUs, or when you want higher request throughput for many independent image/video generation requests. It improves concurrency throughput; it does not make a single request faster.
+
+## Configuration
+
+A replica consumes the devices required by that diffusion stage. For single-device diffusion pipelines, set `runtime.devices` to one device per replica:
+
+```yaml
+stages:
+  - stage_id: 0
+    stage_type: diffusion
+    runtime:
+      devices: "0,1,2,3"
+      num_replicas: 4
+    engine_args:
+      model: Qwen/Qwen-Image
+      parallel_config:
+        world_size: 1
+```
+
+For multi-device diffusion pipelines, `parallel_config.world_size` is treated as the number of devices per replica. For example, four 2-GPU replicas need eight devices:
+
+```yaml
+stages:
+  - stage_id: 0
+    stage_type: diffusion
+    runtime:
+      devices: "0,1,2,3,4,5,6,7"
+      num_replicas: 4
+    engine_args:
+      model: Wan-AI/Wan2.2-T2V-A14B-Diffusers
+      parallel_config:
+        world_size: 2
+```
+
+The runtime splits the device list into contiguous chunks per replica. With the second example, replica 0 uses `0,1`, replica 1 uses `2,3`, replica 2 uses `4,5`, and replica 3 uses `6,7`.
+
+## Serving
+
+Start the server with the stage config:
+
+```bash
+vllm serve Qwen/Qwen-Image \
+  --omni \
+  --port 8099 \
+  --stage-configs-path /path/to/diffusion_multi_replica.yaml
+```
+
+Then send normal OpenAI-compatible diffusion chat requests. The orchestrator routes new requests across replicas and keeps updates/abort calls on the same replica for each request.
+
+## Benchmarking
+
+Benchmark each replica count by changing `runtime.num_replicas` and `runtime.devices`, restarting the server, then running:
+
+```bash
+python benchmarks/diffusion/diffusion_benchmark_serving.py \
+  --base-url http://localhost:8099 \
+  --model Qwen/Qwen-Image \
+  --task t2i \
+  --dataset random \
+  --num-prompts 8 \
+  --request-rate inf \
+  --num-inference-steps 2 \
+  --width 256 --height 256 \
+  --output-file diffusion_replica_4.json
+```
+
+For a quick scaling sweep, test `num_replicas` values `1`, `2`, `3`, and `4` with the same prompt count and generation parameters. Compare throughput and end-to-end latency from the JSON outputs.
+
+## Notes and Limitations
+
+- Multi-replica serving is currently supported for diffusion stages. Single-stage mode still rejects non-diffusion replicas.
+- `runtime.devices` must provide exactly `num_replicas * devices_per_replica` devices.
+- Request-level latency may stay similar or slightly increase because each request still runs on one replica.
+- Throughput improves only when enough concurrent requests are issued and each replica has enough GPU memory.
@@ -4,14 +4,13 @@ This guide covers the parallelism methods in vLLM-Omni for speeding up diffusion

 ## Supported Methods

-| Method                                             | Description                                                                                                         |
-|----------------------------------------------------|---------------------------------------------------------------------------------------------------------------------|
-| **[Tensor Parallelism](tensor_parallel.md)**       | Shards DiT weights across GPUs to reduce per-GPU memory                                                             |
-| **[Sequence Parallelism](sequence_parallel.md)**   | Splits sequence dimension across GPUs (Ulysses-SP, Ring-Attention, or hybrid) for high-resolution images and videos |
-| **[CFG-Parallel](cfg_parallel.md)**                | Runs CFG positive/negative branches on separate GPUs for ~1.8x speedup on guided generation                         |
-| **[Pipeline Parallelism](pipeline_parallel.md)**   | Splits the denoising transformer block-wise across sequential GPU stages to reduce per-GPU model memory             |
-| **[VAE Patch Parallelism](vae_patch_parallel.md)** | Distributes VAE decode spatially across GPUs to reduce peak VAE memory                                              |
-| **[HSDP](hsdp.md)**                                | Shards full model weights via PyTorch FSDP2 to enable large-model inference on memory-constrained GPUs              |
-| **[Expert Parallelism](expert_parallel.md)**       | Shards MoE expert blocks across GPUs for MoE models (e.g. HunyuanImage3.0)                                          |
+| Method | Description |
+|--------|-------------|
+| **[Tensor Parallelism](tensor_parallel.md)** | Shards DiT weights across GPUs to reduce per-GPU memory |
+| **[Sequence Parallelism](sequence_parallel.md)** | Splits sequence dimension across GPUs (Ulysses-SP, Ring-Attention, or hybrid) for high-resolution images and videos |
+| **[CFG-Parallel](cfg_parallel.md)** | Runs CFG positive/negative branches on separate GPUs for ~1.8x speedup on guided generation |
+| **[VAE Patch Parallelism](vae_patch_parallel.md)** | Distributes VAE decode spatially across GPUs to reduce peak VAE memory |
+| **[HSDP](hsdp.md)** | Shards full model weights via PyTorch FSDP2 to enable large-model inference on memory-constrained GPUs |
+| **[Expert Parallelism](expert_parallel.md)** | Shards MoE expert blocks across GPUs for MoE models (e.g. HunyuanImage3.0) |

 See [Supported Models](../../diffusion_features.md#supported-models) for per-model compatibility.
@@ -1,205 +0,0 @@
-# Pipeline Parallelism Guide
-
-## Table of Content
-
- [Overview](#overview)
- [Quick Start](#quick-start)
- [Example Script](#example-script)
- [Configuration Parameters](#configuration-parameters)
- [Best Practices](#best-practices)
- [Troubleshooting](#troubleshooting)
- [Summary](#summary)
-
---
-
-## Overview
-
-Pipeline Parallelism splits the denoising transformer block-wise into sequential stages across GPUs. Each rank owns only
-part of the transformer, which reduces per-GPU model memory and enables larger diffusion models to run across multiple
-devices.
-
-It can also be combined with other distributed methods such as CFG-Parallel, Tensor Parallelism, and Sequence
-Parallelism.
-
-See supported models list in [Supported Models](../../diffusion_features.md#supported-models).
-
---
-
-## Quick Start
-
-### Basic Usage
-
-Simplest working example:
-
-```python
-from vllm_omni import Omni
-from vllm_omni.diffusion.data import DiffusionParallelConfig
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-omni = Omni(
-    model="Wan-AI/Wan2.2-TI2V-5B-Diffusers",
-    parallel_config=DiffusionParallelConfig(
-        pipeline_parallel_size=2,
-    ),
-)
-
-outputs = omni.generate(
-    {"prompt": "A cinematic drone shot over snowy mountains"},
-    OmniDiffusionSamplingParams(
-        num_inference_steps=40,
-        num_frames=81,
-        height=704,
-        width=1280,
-    ),
-)
-```
-
---
-
-## Example Script
-
-### Offline Inference
-
-Use python scripts under:
-
- `examples/offline_inference/text_to_video/text_to_video.py`
- `examples/offline_inference/image_to_video/image_to_video.py`
-
-Text-to-video example:
-
-```bash
-python examples/offline_inference/text_to_video/text_to_video.py \
--model=Wan-AI/Wan2.2-TI2V-5B-Diffusers \
--width=1280 \
--height=704 \
--guidance-scale=5.0 \
--prompt="Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage" \
--output=t2v_5B_pp2.mp4 \
--pipeline-parallel-size=2
-```
-
-Pipeline Parallelism can also be combined with CFG-Parallel:
-
-```bash
-python examples/offline_inference/text_to_video/text_to_video.py \
--model=Wan-AI/Wan2.2-TI2V-5B-Diffusers \
--width=1280 \
--height=704 \
--guidance-scale=5.0 \
--prompt="Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage" \
--output=t2v_5B_pp2_cfg2.mp4 \
--pipeline-parallel-size=2 \
--cfg-parallel-size=2
-```
-
-### Online Serving
-
-Enable Pipeline Parallelism in online serving:
-
-```bash
-# Default PP configuration
-vllm serve Wan-AI/Wan2.2-TI2V-5B-Diffusers --omni --port 8091 --pipeline-parallel-size 2
-
-# PP + CFG-Parallel
-vllm serve Wan-AI/Wan2.2-TI2V-5B-Diffusers --omni --port 8091 \
-  --pipeline-parallel-size 2 \
-  --cfg-parallel-size 2
-```
-
---
-
-## Configuration Parameters
-
-In `DiffusionParallelConfig`
-
-| Parameter                | Type | Default | Description                                                                                                      |
-|--------------------------|------|---------|------------------------------------------------------------------------------------------------------------------|
-| `pipeline_parallel_size` | int  | 1       | Number of pipeline-parallel stages. Set to a value greater than 1 to split the denoising transformer across GPUs |
-
-> [!NOTE]
-> Total GPU count is the product of all enabled distributed dimensions, for example
-`pipeline_parallel_size * cfg_parallel_size * tensor_parallel_size * ulysses_degree * ring_degree`.
-
-### Manual Layer Partitioning
-
-By default, transformer layers are distributed evenly across PP ranks. You can override this with the
-`VLLM_PP_LAYER_PARTITION` environment variable to assign a specific number of layers to each rank:
-
-```bash
-# Example: 40 layers across 4 PP ranks, assigning 8 / 12 / 12 / 8 layers
-export VLLM_PP_LAYER_PARTITION=8,12,12,8
-```
-
-The value must be a comma-separated list of integers whose length equals `pipeline_parallel_size` and whose sum equals
-the total number of transformer layers. This is useful when you want to balance memory or compute asymmetrically across
-ranks.
-
-
---
-
-## Best Practices
-
-### When to Use
-
-**Good for:**
-
- Large diffusion transformers that do not fit comfortably on one GPU
- Multi-GPU setups where reducing per-GPU model memory is more important than minimizing communication
- Combining with CFG-Parallel or other distributed methods on supported models
-
-**Not for:**
-
- Single GPU setups
- Models that do not support Pipeline Parallelism
-  (check [supported models](../../diffusion_features.md#supported-models))
- Very small models where inter-stage communication overhead may outweigh the benefit
-
-### Expected Behavior
-
-Pipeline Parallelism primarily reduces per-GPU memory usage by splitting the transformer into stage-local blocks.
-Depending on the model, topology, and resolution, it may also help execution fit into available hardware, but it is not
-primarily a latency optimization.
-
---
-
-## Troubleshooting
-
-### Common Issue 1: No benefit from Pipeline Parallelism
-
-**Symptoms**: PP is enabled, but latency does not improve or becomes slightly worse.
-
-**Solutions**:
-
-1. **Check your goal:**
-
-```python
-# PP is mainly for memory scaling, not guaranteed latency speedup
-parallel_config = DiffusionParallelConfig(pipeline_parallel_size=2)
-```
-
-2. **Check model support:**
-    - Verify your model in [supported models](../../diffusion_features.md#supported-models)
-    - PP is currently validated only on selected pipelines
-
-3. **Combine with other methods when appropriate:**
-    - PP can be combined with CFG-Parallel, Tensor Parallelism, or Sequence Parallelism on supported models
-
-### Common Issue 2: PP pipeline fails at import
-
-**Symptoms**: Importing a custom pipeline raises a `TypeError` about `CFGParallelMixin` or mixin order.
-
-**Solutions**:
-
-1. Inherit both mixins.
-2. Put `PipelineParallelMixin` before `CFGParallelMixin` in the class base list.
-3. Use `class YourPipeline(nn.Module, PipelineParallelMixin, CFGParallelMixin): ...` as the reference pattern.
-
---
-
-## Summary
-
-1. ✅ **Enable Pipeline Parallelism** - Set `pipeline_parallel_size > 1` in `DiffusionParallelConfig`
-2. ✅ **Use Supported Models** - Verify your model supports PP in
-   [supported models](../../diffusion_features.md#supported-models)
-3. ✅ **Combine When Needed** - PP can be combined with CFG-Parallel and other distributed methods on supported pipelines
-4. ✅ **Scale for Memory** - Use PP primarily to reduce per-GPU model memory and fit larger transformers
@@ -15,7 +15,7 @@ vLLM-Omni supports various advanced features for diffusion models:
 - Acceleration: **cache methods**, **parallelism methods**, **startup optimizations**
 - Memory optimization: **cpu offloading**, **quantization**
 - Extensions: **LoRA inference**, **frame interpolation**
- Execution modes: **step execution**
+- Execution modes: **step execution**, **multi-replica serving**

 ## Supported Features

@@ -35,15 +35,15 @@ Cache methods trade minimal quality for significant speedup. Quality loss is typ

 Parallelism methods distribute computation across GPUs without quality loss (mathematically equivalent to single-GPU).

-| Method                                                                 | Description                                                              | Best For                                                          |
-|------------------------------------------------------------------------|--------------------------------------------------------------------------|-------------------------------------------------------------------|
-| **[Ulysses-SP](diffusion/parallelism/sequence_parallel.md)**           | Sequence parallelism via all-to-all communication                        | High-resolution images (>1536px) or long videos with 2-8 GPUs     |
-| **[Ring-Attention](diffusion/parallelism/sequence_parallel.md)**       | Sequence parallelism via ring-based communication                        | Videos, very long sequences, memory-constrained, with 2-8 GPUs    |
-| **[CFG-Parallel](diffusion/parallelism/cfg_parallel.md)**              | Splits CFG positive/negative branches across devices                     | Image editing with CFG guidance (true_cfg_scale > 1) on 2 GPUs    |
-| **[Tensor Parallelism](diffusion/parallelism/tensor_parallel.md)**     | Shards model weights across devices                                      | Large models that don't fit in single GPU, with 2+ GPUs           |
-| **[Pipeline Parallelism](diffusion/parallelism/pipeline_parallel.md)** | Splits the denoising transformer block-wise across sequential GPU stages | Large diffusion transformers that need lower per-GPU model memory |
-| **[HSDP](diffusion/parallelism/hsdp.md)**                              | Weight sharding via FSDP2, redistributed on-demand at runtime            | Very large models (14B+) on limited VRAM, combinable with SP      |
-| **[Expert Parallelism](diffusion/parallelism/expert_parallel.md)**     | Shards MoE expert MLP blocks across devices                              | MoE diffusion models (e.g., HunyuanImage3.0)                      |
+| Method | Description | Best For |
+|--------|-------------|----------|
+| **[Ulysses-SP](diffusion/parallelism/sequence_parallel.md)** | Sequence parallelism via all-to-all communication | High-resolution images (>1536px) or long videos with 2-8 GPUs |
+| **[Ring-Attention](diffusion/parallelism/sequence_parallel.md)** | Sequence parallelism via ring-based communication | Videos, very long sequences, memory-constrained, with 2-8 GPUs |
+| **[CFG-Parallel](diffusion/parallelism/cfg_parallel.md)** | Splits CFG positive/negative branches across devices | Image editing with CFG guidance (true_cfg_scale > 1) on 2 GPUs |
+| **[Tensor Parallelism](diffusion/parallelism/tensor_parallel.md)** | Shards model weights across devices | Large models that don't fit in single GPU, with 2+ GPUs |
+| **[HSDP](diffusion/parallelism/hsdp.md)** | Weight sharding via FSDP2, redistributed on-demand at runtime | Very large models (14B+) on limited VRAM, combinable with SP |
+| **[Expert Parallelism](diffusion/parallelism/expert_parallel.md)** | Shards MoE expert MLP blocks across devices | MoE diffusion models (e.g., HunyuanImage3.0) |
+| **[Multi-Replica Serving](diffusion/multi_replica_serving.md)** | Runs multiple independent diffusion stage replicas and routes requests across them | Online serving throughput with many concurrent requests |

 #### Startup Optimization

@@ -106,32 +106,31 @@ The following tables show which models support each feature:

 ### ImageGen

-| Model                    | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀Pipeline-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
-|--------------------------|:---------:|:----------:|:---------------------:|:--------------:|:-----------------:|:-------------------:|:------:|:-------------------------:|:--------------------:|:--------------:|:----------------:|
-| **Bagel**                |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |
-| **FLUX.1-dev**           |     ✅     |     ✅      |           ❌           |       ✅        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ✅        |        ❌         |
-| **FLUX.1-schnell**       |     ❌     |     ✅      |           ❌           |       ✅        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ✅        |        ❌         |
-| **FLUX.2-klein**         |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ✅        |        ❌         |
-| **FLUX.1-Kontext-dev**   |     ❌     |     ❌      |           ❌           |       ❌        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ❌        |        ❌         |
-| **FLUX.2-dev**           |     ✅     |     ✅      |           ❌           |       ✅        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ❌        |        ❌         |
-| **GLM-Image**            |     ❌     |     ❌      |           ❌           |       ✅        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ❌        |        ❌         |
-| **HunyuanImage3**        |     ❌     |     ✅      |           ❌           |       ❌        |         ✅         |          ❌          |   ❌    |             ❌             |          ❌           |       ✅        |        ❌         |
-| **LongCat-Image**        |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
-| **LongCat-Image-Edit**   |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
-| **MagiHuman**            |     ❌     |     ❌      |           ❌           |       ❓        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
-| **MammothModa2(T2I)**    |     ❌     |     ❌      |           ❌           |       ❌        |         ❌         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
-| **Nextstep_1(T2I)**      |     ❓     |     ❓      |           ❌           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
-| **OmniGen2**             |     ❌     |     ✅      |           ❌           |       ❌        |         ✅         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
-| **Ovis-Image**           |     ❌     |     ✅      |           ❌           |       ✅        |         ❌         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
-| **Qwen-Image**           |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ✅        |        ✅         |
-| **Qwen-Image-2512**      |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ✅        |        ✅         |
-| **Qwen-Image-Edit**      |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
-| **Qwen-Image-Edit-2509** |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |        ✅ (decode)         |          ✅           |       ❌        |        ❌         |
-| **Qwen-Image-Layered**   |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
-| **SenseNova-U1**         |     ❌     |     ❌      |           ❌           |       ❌        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
-| **Stable-Diffusion3.5**  |     ❌     |     ✅      |           ❌           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
-| **Z-Image**              |     ✅     |     ✅      |           ✅           |       ❓        |   ✅ (TP=2 only)   |          ❌          |   ✅    |             ❌             |      ✅ (decode)      |       ✅        |        ❌         |
-| **ERNIE-Image**          |     ❌     |     ✅      |           ✅           |       ❓        |         ✅         |          ❌          |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |
+| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
+|-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|:----------------:|
+| **Bagel** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅  | ❌ | ❌ | ❌ |
+| **FLUX.1-dev** | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
+| **FLUX.1-schnell** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
+| **FLUX.2-klein** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
+| **FLUX.1-Kontext-dev** | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
+| **FLUX.2-dev** | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
+| **GLM-Image** | ❌ | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
+| **HunyuanImage3** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ |
+| **LongCat-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
+| **LongCat-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
+| **MagiHuman** | ❌ | ❌ | ❌ | ❓ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
+| **MammothModa2(T2I)** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
+| **Nextstep_1(T2I)** | ❓ | ❓ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
+| **OmniGen2** | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
+| **Ovis-Image** | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
+| **Qwen-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ✅ |
+| **Qwen-Image-2512** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ✅ |
+| **Qwen-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
+| **Qwen-Image-Edit-2509** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ❌ | ❌ |
+| **Qwen-Image-Layered** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
+| **Stable-Diffusion3.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ (decode) | ❌ | ❌ |
+| **Z-Image** | ✅ | ✅ | ✅ | ❓ | ✅ (TP=2 only) | ✅ | ❌ | ✅ (decode) | ✅ | ❌ |
+| **ERNIE-Image** | ❌ | ✅ | ✅ | ❓ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |

 > Notes:
 > 1. Nextstep_1(T2I) does not support cache acceleration methods such as TeaCache or Cache-DiT.
@@ -139,15 +138,15 @@ The following tables show which models support each feature:

 ### VideoGen

-| Model                        | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | Pipeline-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
-|------------------------------|:---------:|:----------:|:---------------------:|:--------------:|:-----------------:|:-----------------:|:------:|:-------------------------:|:--------------------:|:--------------:|:----------------:|
-| **Wan2.2**                   |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ✅         |   ✅    |             ✅             |  ✅ (encode/decode)   |       ❌        |        ❌         |
-| **Wan2.1-VACE**              |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
-| **LTX-2**                    |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |
-| **LTX-2.3**                  |     ❌     |     ✅      |           ✅           |       ❌        |         ✅         |         ❌         |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
-| **Helios**                   |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |
-| **HunyuanVideo-1.5 T2V I2V** |     ❌     |     ✅      |           ❌           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |      ✅ (decode)      |       ✅        |        ❌         |
-| **DreamID-Omni**             |     ❌     |     ❌      |           ❌           |       ✅        |         ❌         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |
+| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
+|-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|:----------------:|
+| **Wan2.2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (encode/decode) | ❌ | ❌ |
+| **Wan2.1-VACE** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
+| **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
+| **LTX-2.3** | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
+| **Helios** | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
+| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ❌ |
+| **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |

 **Frame Interpolation Support**

@@ -156,9 +155,9 @@ The following tables show which models support each feature:

 ### AudioGen

-| Model                 | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀Pipeline-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
-|-----------------------|:---------:|:----------:|:---------------------:|:--------------:|:-----------------:|:-------------------:|:------:|:-------------------------:|:--------------------:|:--------------:|:----------------:|
-| **Stable-Audio-Open** |     ✅     |     ❌      |           ❓           |       ❓        |         ❌         |          ❌          |   ✅    |             ✅             |          ❌           |       ✅        |        ❌         |
+| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
+|-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|:----------------:|
+| **Stable-Audio-Open** | ✅ | ❌ | ❓ | ❓ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ |


 ## Feature Compatibility
@@ -258,7 +257,7 @@ Measured on NVIDIA H800:

 **Parallelism Methods:**

- **[Parallelism Overview](diffusion/parallelism/overview.md)** - Tensor Parallelism, Sequence Parallelism, CFG Parallelism, Pipeline Parallelism, HSDP, and Expert Parallelism
+- **[Parallelism Overview](diffusion/parallelism/overview.md)** - Tensor Parallelism, Sequence Parallelism, CFG Parallelism, HSDP, and Expert Parallelism

 **Memory Optimization:**

@@ -101,7 +101,7 @@ guides.

 When using `/v1/chat/completions`, pass these inside `extra_body` in the curl
 JSON, or via the `extra_body` keyword argument in the OpenAI Python SDK (see the
-[Chat Completions API guide](../../../serving/chat_completions_api.md)).
+[Diffusion Chat API guide](../../../serving/diffusion_chat_api.md)).
 When using the dedicated [`/v1/images/generations`](../../../serving/image_generation_api.md)
 or [`/v1/images/edits`](../../../serving/image_edit_api.md) endpoints, pass
 the supported generation controls as top-level fields directly. For image
@@ -208,7 +208,7 @@ Wrap generation parameters inside `extra_body` in the request JSON:
    ```

    For details on how generation parameters are handled across different clients, see the
-    [Chat Completions API guide](../../../../serving/chat_completions_api.md).
+    [Diffusion Chat API guide](../../../../serving/diffusion_chat_api.md).

 ### Layered Image Generation (Qwen-Image-Layered)

@@ -361,7 +361,7 @@ Provide multiple images in `content` (order matters):

 When using `/v1/chat/completions`, pass these inside `extra_body` in the curl
 JSON, or via the `extra_body` keyword argument in the OpenAI Python SDK (see the
-[Chat Completions API guide](../../../../serving/chat_completions_api.md)).
+[Diffusion Chat API guide](../../../../serving/diffusion_chat_api.md)).
 When using the dedicated [`/v1/images/edits`](../../../../serving/image_edit_api.md)
 endpoint, pass the supported generation controls as top-level form fields
 directly. For image dimensions and count, use `size` and `n` rather than
@@ -216,7 +216,7 @@ Wrap generation parameters inside `extra_body` in the request JSON:
    ```

    For details on how generation parameters are handled across different clients, see the
-    [Chat Completions API guide](../../../../serving/chat_completions_api.md).
+    [Diffusion Chat API guide](../../../../serving/diffusion_chat_api.md).

 ### Multimodal Input (Text + Structured Content)

@@ -237,7 +237,7 @@ Wrap generation parameters inside `extra_body` in the request JSON:

 When using `/v1/chat/completions`, pass these inside `extra_body` in the curl
 JSON, or via the `extra_body` keyword argument in the OpenAI Python SDK (see the
-[Chat Completions API guide](../../../../serving/chat_completions_api.md)).
+[Diffusion Chat API guide](../../../../serving/diffusion_chat_api.md)).
 When using the dedicated [`/v1/images/generations`](../../../../serving/image_generation_api.md)
 endpoint, pass the supported generation controls as top-level JSON fields
 directly. For image dimensions and count, use `size` and `n` rather than
@@ -2,11 +2,11 @@

 ## Overview

-FP8 quantization converts BF16/FP16 weights to FP8 at model load time. Online
-activation scaling is the default and does not require calibration. Static
-activation scaling is supported when calibrated scale information is available.
-For ModelOpt-produced pre-quantized checkpoints, see
-[ModelOpt Quantization](modelopt.md).
+FP8 quantization converts BF16/FP16 weights to FP8 at model load time, or loads
+a checkpoint whose target stage already declares an FP8 quantization config.
+Online activation scaling is the default and does not require calibration.
+Static activation scaling is supported when calibrated scale information is
+available.

 Some architectures can quantize all linear layers. Others have
 quality-sensitive layers that should stay in BF16 through `ignored_layers`.
@@ -46,7 +46,7 @@ guide. FP8 on Ampere may use a weight-only path where available.

 | Model | Scope | Format | Status |
 |-------|-------|--------|--------|
-| Qwen3-Omni | Thinker language-model stage | [ModelOpt](modelopt.md) `quant_algo=FP8` | Tested for thinker memory reduction |
+| Qwen3-Omni | Thinker language-model stage | ModelOpt `quant_algo=FP8` | Tested for thinker memory reduction |
 | Qwen3-TTS | TTS language-model stage | Checkpoint config | Not validated |

 Audio encoder, vision encoder, talker, and code2wav stay in BF16 unless a
@@ -1,138 +0,0 @@
-# ModelOpt Quantization
-
-## Overview
-
-ModelOpt quantization loads checkpoints produced by NVIDIA ModelOpt. The
-quantized weights and scale tensors are generated before serving, so inference
-does not run online calibration or convert a BF16 checkpoint at startup.
-
-vLLM-Omni currently validates the ModelOpt FP8 checkpoint path for diffusion
-transformers. The loader auto-detects supported ModelOpt FP8 checkpoint configs
-and keeps non-transformer components, such as the tokenizer, scheduler, text
-encoder, and VAE, on the base checkpoint unless a model-specific guide says
-otherwise.
-
-!!! note
-    `--force-cutlass-fp8` is an explicit runtime override for diffusion
-    checkpoints that already carry a supported ModelOpt FP8 config. It does not
-    quantize BF16 checkpoints and it does not apply to online `--quantization
-    fp8`. The flag only takes effect for ModelOpt FP8 diffusion stages on CUDA
-    SM89+ devices; other platforms and non-ModelOpt FP8 paths fall back to the
-    normal vLLM kernel selection.
-
-## Supported ModelOpt Checkpoint Formats
-
-vLLM-Omni treats ModelOpt checkpoints as pre-quantized checkpoints. The
-checkpoint config must identify ModelOpt as the quantization method or producer,
-and the quantization algorithm must be one of the validated FP8 algorithms.
-
-| Checkpoint field | Supported value |
-|------------------|-----------------|
-| `method` / `quant_method` | `modelopt` |
-| `producer.name` | `modelopt` |
-| `quant_algo` | `FP8`, `FP8_PER_CHANNEL_PER_TOKEN` |
-
-Other ModelOpt algorithms, such as NVFP4, are not enabled by this diffusion
-FP8 path until they have separate model and quality validation.
-
-## Hardware Support
-
-| Device | Support |
-|--------|---------|
-| NVIDIA Blackwell GPU (SM 100+) | ✅ |
-| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ |
-| NVIDIA Ampere GPU (SM 80+) | ⭕ |
-| AMD ROCm | ⭕ |
-| Intel XPU | ⭕ |
-| Ascend NPU | ❌ |
-
-Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this guide.
-The optional CUTLASS FP8 runtime override requires CUDA SM89+.
-
-## Model Type Support
-
-### Diffusion Model
-
-| Model | HF checkpoint | Scope | Status |
-|-------|---------------|-------|--------|
-| Qwen-Image 2512 | `feizhai123/qwen-image-2512-modelopt-fp8-dynamic-all` | Diffusion transformer | Validated for ModelOpt FP8 checkpoints |
-| Z-Image | `feizhai123/z-image-modelopt-fp8-conservative` | Diffusion transformer | Validated for ModelOpt FP8 checkpoints |
-| FLUX.2-dev | `feizhai123/flux2-dev-modelopt-fp8` | Diffusion transformer | Validated for ModelOpt FP8 checkpoints |
-| FLUX.2-klein 4B | `feizhai123/flux2-klein-4b-modelopt-fp8` | Diffusion transformer | Validated for ModelOpt FP8 checkpoints |
-| HunyuanImage-3.0 | `feizhai123/hunyuan-image3-modelopt-fp8` | MoE diffusion transformer | Validated for ModelOpt FP8 checkpoints |
-| Wan2.2 | Not available | Diffusion transformer | Not validated |
-
-### Multi-Stage Omni/TTS Model
-
-| Model | Scope | Status |
-|-------|-------|--------|
-| Qwen3-Omni | Thinker language-model stage | ModelOpt FP8 checkpoint path |
-| Qwen3-TTS | TTS language-model stage | Not validated |
-
-Audio encoder, vision encoder, talker, and code2wav stages stay in BF16 unless
-a model-specific guide documents otherwise.
-
-### Multi-Stage Diffusion Model
-
-ModelOpt checkpoints must be routed to the stage whose checkpoint contains the
-ModelOpt `quantization_config`. BAGEL and GLM-Image are not listed as validated
-ModelOpt targets yet.
-
-## Configuration
-
-For pre-quantized ModelOpt FP8 checkpoints, no `--quantization fp8` flag is
-needed. The checkpoint config selects the ModelOpt path.
-
-Online serving:
-
-```bash
-vllm serve <modelopt-fp8-checkpoint> \
-  --omni \
-  --tensor-parallel-size <N> \
-  --force-cutlass-fp8
-```
-
-Offline inference:
-
-```bash
-python examples/offline_inference/text_to_image/text_to_image.py \
-  --model <modelopt-fp8-checkpoint> \
-  --tensor-parallel-size <N> \
-  --prompt "a red ceramic teapot on a wooden table" \
-  --height 1024 \
-  --width 1024 \
-  --num-inference-steps 20 \
-  --seed 42 \
-  --output outputs/modelopt_fp8.png
-```
-
-Python API:
-
-```python
-from vllm_omni import Omni
-
-omni = Omni(
-    model="<modelopt-fp8-checkpoint>",
-    tensor_parallel_size=2,
-    force_cutlass_fp8=True,
-)
-```
-
-## Parameters
-
-| Parameter | Type | Default | Description |
-|-----------|------|---------|-------------|
-| `force_cutlass_fp8` / `--force-cutlass-fp8` | bool | `False` | Force CUTLASS FP8 linear kernels for supported ModelOpt FP8 diffusion stages on CUDA SM89+ |
-
-## Validation and Notes
-
-1. Compare the ModelOpt FP8 checkpoint against the BF16 baseline with the same
-   prompt, resolution, seed, and inference steps.
-2. Use `tests/diffusion/quantization/test_quantization_quality.py` with
-   `VLLM_OMNI_QUALITY_CONFIGS` to validate local baseline and quantized model
-   paths.
-3. Report LPIPS, PSNR, MAE, throughput, latency, and peak memory when adding a
-   new validated ModelOpt diffusion checkpoint.
-4. Keep `--quantization fp8` for online FP8 from BF16 checkpoints; use this
-   ModelOpt path only when the checkpoint already contains ModelOpt FP8 weights
-   and scales.
@@ -1,187 +0,0 @@
-# W8A8 MXFP8 Quantization
-
-## Overview
-
-W8A8 MXFP8 (Microscaling FP8) quantizes both weights and activations to FP8
-using the OCP MX format: groups of 32 K-dimension elements share a single
-`float8_e8m0fnu` exponent scale. This gives better accuracy than channel-wise
-FP8 while keeping the same 8-bit weight footprint.
-
-This method supports two modes:
-
-| Mode | Description |
-|------|-------------|
-| **Online** | BF16 weights are quantized to MXFP8 at load time — no pre-processing needed |
-| **Offline** | msModelSlim-exported MXFP8 weights converted to diffusers format via `merge_mxfp8_checkpoint.py` — weights and scales are loaded directly from the preprocessed checkpoint |
-
-## Hardware Support
-
-| Device | Support |
-|--------|---------|
-| NVIDIA Blackwell GPU (SM 100+) | ⭕ |
-| NVIDIA Ada/Hopper GPU (SM 89+) | ⭕ |
-| NVIDIA Ampere GPU (SM 80+) | ⭕ |
-| AMD ROCm | ⭕ |
-| Intel XPU | ⭕ |
-| Ascend NPU (Atlas 950 A5) | ✅ |
-
-Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
-guide.
-
-## Model Type Support
-
-### Diffusion Model (Wan2.2)
-
-| Model | Mode | Notes |
-|-------|------|-------|
-| Wan2.2-T2V-A14B | Online + Offline | MoE cascade; quantizes two transformers (`transformer` + `transformer_2`) |
-| Wan2.2-I2V-A14B | Online + Offline | MoE cascade; quantizes two transformers (`transformer` + `transformer_2`) |
-| Wan2.2-TI2V-5B | Online + Offline | Single transformer |
-
-### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
-
-| Model | Status | Notes |
-|-------|--------|-------|
-| Qwen3-Omni | Not validated | — |
-| Qwen3-TTS | Not validated | — |
-
-### Multi-Stage Diffusion Model (BAGEL, GLM-Image)
-
-| Model | Status | Notes |
-|-------|--------|-------|
-| BAGEL | Not validated | — |
-| GLM-Image | Not validated | — |
-
-## Configuration
-
-### Online Mode
-
-Online mode requires no pre-processing. vLLM-Omni quantizes BF16 weights to
-MXFP8 at load time.
-
-Python API:
-
-```python
-from vllm_omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-omni = Omni(model="<your-model>", quantization="mxfp8")
-
-outputs = omni.generate(
-    "A cat sitting on a windowsill",
-    OmniDiffusionSamplingParams(num_inference_steps=50),
-)
-```
-
-CLI:
-
-```bash
-python text_to_video.py --model <your-model> --quantization mxfp8
-
-# Online serving
-vllm serve <your-model> --omni --quantization mxfp8
-```
-
-### Offline Mode
-
-Offline mode loads a pre-quantized checkpoint from msModelSlim. A preprocessing
-step converts the raw quantized output to the diffusers format expected by
-vLLM-Omni and injects the quantization config into `transformer/config.json` so
-that vLLM-Omni auto-detects the offline path without a `--quantization` flag.
-
-#### Step 1 — Quantize with msModelSlim
-
-```bash
-msmodelslim quant \
-  --model_path /path/to/Wan2.2-TI2V-5B-Diffusers \
-  --save_path  /path/to/wan2_2_ti2v_quantized_raw \
-  --device npu \
-  --model_type Wan2_2 \
-  --config_path /path/to/wan2_2_w8a8f8_mxfp.yaml \
-  --trust_remote_code True
-```
-
-After this step, `--save_path` contains the raw quantized safetensors files and
-a metadata JSON (`quant_model_description*.json`).
-
-For cascade MoE models (T2V-A14B, I2V-A14B), msModelSlim outputs two
-subdirectories: `high_noise_model/` and `low_noise_model/`.
-
-#### Step 2 — Preprocess with merge_mxfp8_checkpoint.py
-
-The script (`vllm_omni/quantization/tools/merge_mxfp8_checkpoint.py`):
-
-1. Copies the original diffusers model to `--output-path` (VAE, text encoder,
-   scheduler, etc. are preserved).
-2. Remaps tensor names from msModelSlim convention to diffusers convention.
-3. Saves the converted weights as `diffusion_pytorch_model.safetensors`.
-4. Copies the original `transformer/config.json` and injects
-   `quantization_config` so that vLLM-Omni auto-detects offline MXFP8.
-
-For cascade MoE models, steps 2–4 run separately for `high_noise_model/` →
-`transformer/` and `low_noise_model/` → `transformer_2/`.
-
-```bash
-python vllm_omni/quantization/tools/merge_mxfp8_checkpoint.py \
-  --model-type     Wan2.2-TI2V-5B \
-  --original-model /path/to/Wan2.2-TI2V-5B-Diffusers \
-  --quant-path     /path/to/wan2_2_ti2v_quantized_raw \
-  --output-path    /path/to/Wan2.2-TI2V-5B-MXFP8
-```
-
-| Argument | Description |
-|----------|-------------|
-| `--model-type` | Model variant: `Wan2.2-T2V-A14B`, `Wan2.2-I2V-A14B`, or `Wan2.2-TI2V-5B` |
-| `--original-model` | Root directory of the original BF16 diffusers model |
-| `--quant-path` | Root directory of the msModelSlim quantized output |
-| `--output-path` | Output directory for the merged model (created by the script) |
-
-The script outputs a complete diffusers model directory at `--output-path`,
-with each transformer subfolder containing:
-
- `diffusion_pytorch_model.safetensors` — converted FP8 weights
- `config.json` — original transformer config with `quantization_config` injected
- `quant_model_description.json` — renamed quantization metadata (reference only)
-
-#### Step 3 — Serve
-
-```bash
-python text_to_video.py --model /path/to/Wan2.2-TI2V-5B-MXFP8
-
-# Online serving
-vllm serve /path/to/Wan2.2-TI2V-5B-MXFP8 --omni
-```
-
-Python API:
-
-```python
-omni = Omni(model="/path/to/Wan2.2-TI2V-5B-MXFP8")
-```
-
-!!! note
-    No `--quantization` flag is needed for offline mode. The preprocessing
-    script injects `quantization_config` into each `transformer/config.json`,
-    which vLLM-Omni reads automatically to activate the offline MXFP8 method.
-
-## Parameters
-
-| Parameter | Type | Default | Description |
-|-----------|------|---------|-------------|
-| `method` | str | — | Must be `"mxfp8"` |
-| `is_checkpoint_mxfp8_serialized` | bool | `False` | `True` for offline pre-quantized checkpoints; auto-set from `config.json` when using the preprocessing script |
-| `ignored_layers` | list[str] | `[]` | Layer name substrings to keep in BF16 (e.g. `"to_out"` matches `blocks.0.attn1.to_out.0`) |
-
-## Validation and Notes
-
-1. Online mode quantizes BF16 weights at load time using
-   `npu_dynamic_mx_quant`. This adds a one-time overhead on the first load
-   but requires no checkpoint preparation.
-2. Offline mode loads FP8 weights directly from the checkpoint. Scales are
-   stored as `uint8` bytes in safetensors (same bit layout as
-   `float8_e8m0fnu`) and are reinterpreted at load time without a dtype
-   conversion.
-3. If the offline checkpoint was produced with the old `merge_mxfp8_checkpoint.py`
-   interface (arguments `--quant-dir`, `--orig-dir`, `--meta-json`,
-   `--output-dir`), regenerate it with the current script. The old script
-   wrote a separate `quantization_config.json` that is not read by vLLM-Omni;
-   the current script injects the config directly into `transformer/config.json`.
@@ -1,4 +1,4 @@
-# Quantization
+# Quantization

 vLLM-Omni exposes quantization through the unified `quantization_config`
 path. The same configuration entrypoint is used across diffusion-only models,
@@ -9,20 +9,19 @@ type has a different quantization scope.

 | Mode | Guide | Description | Methods |
 |------|-------|-------------|---------|
-| Online quantization | [Online Quantization](online.md) | vLLM-Omni computes quantized weights and scales while loading the model. | FP8 W8A8, Int8 W8A8, MXFP8 W8A8 |
-| Runtime attention quantization | [Quantized KV Cache](quantized_kvcache.md) | vLLM-Omni dynamically quantizes eligible diffusion Flash Attention tensors during inference. | FP8 FA |
-| Pre-quantized checkpoints | Method-specific guides | The checkpoint or an offline quantizer provides quantized weights and scales before serving. | ModelOpt, GGUF, AutoRound, msModelSlim, serialized Int8, offline MXFP8 |
+| Online quantization | [Online Quantization](online.md) | vLLM-Omni computes quantized weights and scales while loading the model. | FP8 W8A8, Int8 W8A8 |
+| Pre-quantized checkpoints | Method-specific guides | The checkpoint or an offline quantizer provides quantized weights and scales before serving. | GGUF, AutoRound, msModelSlim, serialized Int8 |

 ## Hardware Support

-| Device | FP8 W8A8 | Int8 W8A8 | ModelOpt | MXFP8 W8A8 | GGUF | AutoRound | msModelSlim |
-|--------|----------|-----------|----------|------------|------|-----------|-------------|
-| NVIDIA Blackwell GPU (SM 100+) | ✅ | ✅ | ✅ | ⭕ | ✅ | ✅ | ❌ |
-| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ | ✅ | ⭕ | ✅ | ✅ | ❌ |
-| NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ | ⭕ | ⭕ | ✅ | ✅ | ❌ |
-| AMD ROCm | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ❌ |
-| Intel XPU | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ✅ | ❌ |
-| Ascend NPU | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ |
+| Device | FP8 W8A8 | Int8 W8A8 | GGUF | AutoRound | msModelSlim |
+|--------|----------|-----------|------|-----------|-------------|
+| NVIDIA Blackwell GPU (SM 100+) | ✅ | ✅ | ✅ | ✅ | ❌ |
+| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ | ✅ | ✅ | ❌ |
+| NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ | ✅ | ✅ | ❌ |
+| AMD ROCm | ⭕ | ⭕ | ⭕ | ⭕ | ❌ |
+| Intel XPU | ⭕ | ⭕ | ⭕ | ✅ | ❌ |
+| Ascend NPU | ❌ | ✅ | ❌ | ❌ | ✅ |

 Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
 guide. FP8 on Ampere may use a weight-only path where available.
@@ -40,8 +39,6 @@ otherwise.
 |--------|-------|------|----------------|--------|
 | FP8 W8A8 | [FP8](fp8.md) | Online W8A8 or checkpoint FP8 | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image family and other DiT models |
 | Int8 W8A8 | [Int8](int8.md) | Online or serialized W8A8 | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image and Z-Image |
-| ModelOpt | [ModelOpt](modelopt.md) | Pre-quantized FP8 checkpoints | Qwen-Image, Z-Image, FLUX.2, HunyuanImage-3.0 | Validated for ModelOpt FP8 diffusion checkpoints |
-| MXFP8 W8A8 | [MXFP8](mxfp8.md) | Online W8A8 or offline pre-quantized | Wan2.2-T2V-A14B, I2V-A14B, TI2V-5B | Ascend NPU only; validated for Wan2.2 |
 | GGUF | [GGUF](gguf.md) | Pre-quantized transformer weights | Qwen-Image | Validated where a model-specific GGUF adapter exists |
 | AutoRound | [AutoRound](autoround.md) | Pre-quantized W4A16 checkpoints | FLUX.1-dev; Qwen-Image/Wan2.2 not validated | Checkpoint-driven |
 | msModelSlim | [msModelSlim](msmodelslim.md) | Pre-quantized Ascend checkpoints | Wan2.2 recipe; HunyuanImage-3.0 inference target | Ascend/NPU path |
@@ -55,9 +52,8 @@ in BF16 unless the model guide explicitly adds support.

 | Method | Guide | Scope | Example models | Status |
 |--------|-------|-------|----------------|--------|
-| ModelOpt | [ModelOpt](modelopt.md) | Thinker or language-model checkpoint config | Qwen3-Omni thinker | ModelOpt checkpoint path |
+| FP8 | [FP8](fp8.md) | Thinker or language-model checkpoint config | Qwen3-Omni thinker | ModelOpt checkpoint path |
 | Int8 | [Int8](int8.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
-| MXFP8 | [MXFP8](mxfp8.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
 | GGUF | [GGUF](gguf.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
 | AutoRound | [AutoRound](autoround.md) | Thinker or language-model checkpoint config | Qwen2.5-Omni, Qwen3-Omni | Supported through AutoRound checkpoints |
 | msModelSlim | [msModelSlim](msmodelslim.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
@@ -71,8 +67,6 @@ attached to the intended stage rather than applied globally.
 |--------|-------|-------|----------------|--------|
 | FP8 | [FP8](fp8.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Requires model-specific validation |
 | Int8 | [Int8](int8.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Requires model-specific validation |
-| ModelOpt | [ModelOpt](modelopt.md) | Checkpoint-defined diffusion stage | BAGEL, GLM-Image | Requires model-specific validation |
-| MXFP8 | [MXFP8](mxfp8.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Not validated |
 | GGUF | [GGUF](gguf.md) | Stage-specific transformer weights | BAGEL, GLM-Image | No validated adapter listed |
 | AutoRound | [AutoRound](autoround.md) | Checkpoint-defined stage | BAGEL, GLM-Image | No validated checkpoint listed |
 | msModelSlim | [msModelSlim](msmodelslim.md) | Ascend-generated stage weights | GLM-Image | Requires model-specific adaptation |
@@ -100,7 +94,7 @@ config = build_quant_config({

 | Component | Default quantized? | Notes |
 |-----------|--------------------|-------|
-| Diffusion transformer | Yes | Primary target for FP8, Int8, ModelOpt, GGUF, AutoRound, and msModelSlim |
+| Diffusion transformer | Yes | Primary target for FP8, Int8, GGUF, AutoRound, and msModelSlim |
 | Text encoder | No | Keep BF16 unless a method-specific guide documents support |
 | VAE | No | Keep BF16; storage-only paths are method-specific |
 | Scheduler/tokenizer | No | Loaded from the base model repository |
@@ -138,98 +132,3 @@ build_quant_config({"method": "gguf", "gguf_model": "/path/to/model.gguf"})
 build_quant_config({"transformer": {"method": "fp8"}, "vae": None})
 build_quant_config(None)
 ```
-
-## Output Similarity Comparison Tool
-
-Use `vllm_omni.quantization.tools.compare_diffusion_trajectory_similarity`
-to compare a reference diffusion run with a quantized candidate run using the
-same prompt, seed, resolution, scheduler settings, and inference steps. The
-tool compares final decoded images or video frames, and also reports generation
-latency and worker-reported peak memory when available.
-
-This is useful when validating whether online quantization, an offline
-pre-quantized checkpoint, or a new `ignored_layers` choice keeps generation
-quality close to the BF16 reference.
-
-### Online Quantization Example
-
-```bash
-python -m vllm_omni.quantization.tools.compare_diffusion_trajectory_similarity \
-  --task t2i \
-  --model Qwen/Qwen-Image \
-  --candidate-quantization fp8 \
-  --ignored-layers img_mlp \
-  --prompt "a cup of coffee on the table" \
-  --height 512 --width 512 \
-  --num-inference-steps 20 \
-  --seed 142 \
-  --output-json /tmp/qwen_image_fp8_similarity/result.json \
-  --save-output-dir /tmp/qwen_image_fp8_similarity/images \
-  --enforce-eager
-```
-
-### Offline Checkpoint Example
-
-Use `--candidate-model` when the candidate is already quantized or lives at a
-different model path:
-
-```bash
-python -m vllm_omni.quantization.tools.compare_diffusion_trajectory_similarity \
-  --task t2i \
-  --reference-model Qwen/Qwen-Image \
-  --candidate-model /path/to/qwen-image-fp8-checkpoint \
-  --prompt "a cup of coffee on the table" \
-  --height 512 --width 512 \
-  --num-inference-steps 20 \
-  --seed 142 \
-  --output-json /tmp/qwen_image_fp8_checkpoint_similarity/result.json
-```
-
-If the checkpoint does not include a loadable quantization config, pass one
-explicitly:
-
-```bash
--candidate-quantization-config-json '{"method":"fp8"}'
-```
-
-### Output Metrics
-
-The output JSON includes `output_metrics`, `reference_generation`, and
-`candidate_generation`.
-
-| Metric | Direction | Meaning |
-|--------|-----------|---------|
-| `cosine_similarity` | Higher is better | Vector direction similarity between output pixels or frames. Useful as a broad sanity check. |
-| `mae` | Lower is better | Mean absolute pixel or frame error. For decoded outputs, values are in uint8 pixel units. |
-| `mse` / `rmse` | Lower is better | Squared error and its square root. These penalize localized large differences more than `mae`. |
-| `max_abs` | Lower is better | Worst single-element absolute error. Treat it as an outlier/debug signal, not as a release gate. |
-| `l2` / `relative_l2` | Lower is better | Absolute and reference-normalized L2 distance. `relative_l2` is easier to compare across resolutions. |
-| `psnr_db` | Higher is better | Pixel-space signal-to-noise ratio in dB for uint8 images or frames. |
-| `avg_generation_time_s` | Lower is better | Average wall-clock generation time across measured runs. |
-| `max_peak_memory_mb` | Lower is better | Maximum worker-reported peak device memory across measured runs, when the worker reports it. |
-
-Recommended starting thresholds for same-seed diffusion comparisons:
-
-| Metric | Smoke threshold | Stricter target | Notes |
-|--------|-----------------|-----------------|-------|
-| `psnr_db` | `>= 20.0` | `>= 25.0` | Good for quick image or frame regression checks. |
-| `mae` | `<= 12.0` | `<= 6.0` | Interpreted in decoded uint8 pixel units. |
-| `cosine_similarity` | `>= 0.98` | `>= 0.995` | Less sensitive to global scale than L2-style metrics. |
-| `relative_l2` | `<= 0.20` | `<= 0.08` | Useful when comparing across prompts or resolutions. |
-
-These thresholds are heuristics. Tune them by model family, task, resolution,
-quantization method, and deployment tolerance. For release gating, pair the
-numeric report with visual inspection of saved reference and candidate outputs.
-
-The tool intentionally reports separate quality, latency, and memory metrics
-instead of a single consolidated similarity score. A single score can hide
-important tradeoffs, for example a candidate with good PSNR but a meaningful
-memory regression, or a candidate with low average error but localized visual
-artifacts. If you need a project-specific pass/fail gate, define it as an
-explicit policy over the individual metrics.
-
-Pixel-level metrics do not measure semantic consistency. For higher-cost
-evaluation, you can complement this report with a vision-language judge that
-describes the reference and candidate outputs and compares those descriptions.
-Keep that semantic check separate from this lightweight tool so users can
-choose whether the additional model cost and latency are appropriate.
@@ -1,125 +0,0 @@
-# Quantized KV Cache
-
-## Overview
-
-In DiT-based image and video generation, Flash Attention can take a large share
-of denoising time, especially for high-resolution or long-frame workloads.
-vLLM-Omni supports online FP8 quantization for eligible diffusion Flash
-Attention (FA) to reduce FA latency while keeping model weights in their
-original dtype.
-
-This feature is configured through `diffusion_kv_cache_dtype` on
-`OmniDiffusionConfig` (CLI: `--diffusion-kv-cache-dtype`). It is intentionally
-**not** the same as vLLM's `--kv-cache-dtype`, which controls autoregressive
-language-model KV cache storage and defaults to `"auto"`. Diffusion FA
-quantization uses the dedicated diffusion flags so omni serve does not inherit
-that default.
-
-In vLLM-Omni diffusion pipelines, this is a runtime FA path: Q/K/V tensors are
-dynamically quantized before the attention operator. It does not quantize model
-weights and is separate from [FP8 W8A8](fp8.md), [Int8 W8A8](int8.md), or
-pre-quantized checkpoint formats.
-
-If `diffusion_kv_cache_dtype` is not set, behavior is unchanged and attention
-runs in the native dtype.
-
-## Hardware Support
-
-| Device | FP8 FA |
-|--------|--------|
-| Ascend NPU | ✅ |
-| NVIDIA GPU | ❌ |
-| AMD ROCm | ❌ |
-| Intel XPU | ❌ |
-
-Legend: `✅` supported, `❌` unsupported.
-
-FP8 FA is currently implemented only for the NPU Flash Attention backend. Other
-backends do not support `diffusion_kv_cache_dtype="fp8"` for diffusion attention
-and fall back to native dtype execution.
-
-## Model Type Support
-
-### Diffusion Model
-
-| Model | Scope | Status | Notes |
-|-------|-------|--------|-------|
-| Wan2.2 | Eligible DiT full-attention FA on Ascend NPU | Tested | Compare quality and latency against a BF16 baseline before production use |
-| Other diffusion models | Eligible DiT full-attention FA on Ascend NPU | Not tested | You can try `diffusion_kv_cache_dtype="fp8"`; tune `diffusion_kv_cache_skip_steps` and `diffusion_kv_cache_skip_layers` when higher precision is needed |
-
-### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
-
-Not tested for FP8 FA. Treat any use as experimental unless a model-specific
-guide documents support.
-
-### Multi-Stage Diffusion Model (BAGEL, GLM-Image)
-
-Not tested. If the diffusion stage uses the same NPU Flash Attention backend,
-`diffusion_kv_cache_dtype` may apply in theory; validate quality and latency for
-each stage and model.
-
-## Configuration
-
-Offline diffusion example:
-
-```bash
-python examples/offline_inference/image_to_video/image_to_video.py \
-    --model <your-wan2.2-model> \
-    --prompt "A cat sitting on a surfboard at the beach" \
-    --height 1280 \
-    --width 720 \
-    --num-frames 61 \
-    --num-inference-steps 4 \
-    --ulysses-degree 4 \
-    --vae-patch-parallel-size 4 \
-    --diffusion-kv-cache-dtype fp8 \
-    --diffusion-kv-cache-skip-steps "0,1" \
-    --diffusion-kv-cache-skip-layers "0-2"
-```
-
-Online serving:
-
-```bash
-vllm serve <your-model> --omni --diffusion-kv-cache-dtype fp8
-```
-
-Stage config:
-
-```yaml
-stage_args:
-  - stage_id: 0
-    stage_type: diffusion
-    engine_args:
-      model_stage: dit
-      diffusion_kv_cache_dtype: "fp8"
-      diffusion_kv_cache_skip_steps: "0,1"
-      diffusion_kv_cache_skip_layers: "0-2"
-```
-
-Legacy YAML keys `kv_cache_dtype`, `kv_cache_skip_steps`, and
-`kv_cache_skip_layers` are still accepted when constructing
-`OmniDiffusionConfig` (for example via `from_kwargs`); prefer the `diffusion_*`
-names for new configs.
-
-## Parameters
-
-| Parameter | Type | Default | Description |
-|-----------|------|---------|-------------|
-| `diffusion_kv_cache_dtype` | str \| None | `None` | Set to `"fp8"` to enable dynamic FP8 FA on supported attention backends |
-| `diffusion_kv_cache_skip_steps` | str \| None | `None` | Denoising step selector to keep in native dtype, for example `"0,1,4-6"` |
-| `diffusion_kv_cache_skip_layers` | str \| None | `None` | Transformer layer selector to keep in native dtype, for example `"0-2,10"` |
-
-Selectors use comma-separated integers and inclusive ranges. Listed steps or
-layers skip FP8 FA; all other eligible full-attention forwards use the FP8 path.
-
-## Validation and Notes
-
-1. Compare generated images or videos against a BF16 baseline with the same
-   seed, prompt, resolution, frame count, and denoising steps.
-2. Use `diffusion_kv_cache_skip_steps` for denoising steps where quality is more
-   sensitive.
-3. Use `diffusion_kv_cache_skip_layers` for transformer layers that show visible quality
-   regressions.
-4. Report both latency and quality results when enabling this option for a new
-   model. For image or video models, include visual comparison and quantitative
-   metrics when available, such as PSNR or SSIM.
@@ -195,25 +195,6 @@ stages:
    devices: "0,1"
 ```

-#### Hybrid Sharded Data Parallel (HSDP)
-
-For larger Bagel deployments on multiple GPUs, you can enable HSDP (Hybrid Sharded Data Parallel) by modifying the stage configuration (for example, [`bagel.yaml`](../../../vllm_omni/deploy/bagel.yaml)). HSDP shards transformer weights across GPUs to reduce per-GPU memory usage.
-
-1. **Enable HSDP**: Set `use_hsdp: true`.
-2. **Set shard size**: Set `hsdp_shard_size` to the number of GPUs used for sharding (for example, `4`).
-3. **Set replicate size**: Usually keep `hsdp_replicate_size: 1` unless you want replicated HSDP groups.
-4. **Set devices**: Specify the comma-separated GPU IDs used by the diffusion stage (for example, `"0,1,2,3"`).
-
-Example configuration for HSDP across 4 GPUs:
-```yaml
-  - stage_id: 1
-    devices: "0,1,2,3"
-    parallel_config:
-      use_hsdp: true
-      hsdp_shard_size: 4
-      hsdp_replicate_size: 1
-```
-
 Then pass the custom deploy YAML:

 ```bash
@@ -1,76 +0,0 @@
-# Covo-Audio-Chat (Offline Inference)
-
-## Setup
-
-Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.
-
-> **Note**
-> Covo-Audio code2wav requires `torchdiffeq`. Install it with: `pip install torchdiffeq`
-
-## Run examples
-
-Get into the example folder:
-```bash
-cd examples/offline_inference/covo_audio
-```
-
-### Audio input chat
-
-Using the default audio asset:
-```bash
-python end2end.py
-```
-
-Using a custom audio file:
-```bash
-python end2end.py --audio-path /path/to/audio.wav
-```
-
-Using a local model:
-```bash
-python end2end.py -m /path/to/Covo-Audio-Chat --output-dir ./my_output
-```
-
-### Command-line Arguments
-
-| Argument | Short | Default | Description |
-|----------|-------|---------|-------------|
-| `--model-name` | `-m` | `tencent/Covo-Audio-Chat` | Model path or HuggingFace model ID |
-| `--text` | `-t` | `请回答这段音频里的问题。` | Text prompt / question for the audio |
-| `--audio-path` | `-a` | default audio asset | Path to local audio file |
-| `--sampling-rate` | | `16000` | Sampling rate for audio loading (Hz) |
-| `--output-dir` | | `./output_audio` | Output directory for generated files |
-| `--num-prompts` | | `1` | Number of prompts to generate |
-| `--stage-configs-path` | | (auto) | Path to stage configs YAML file |
-| `--log-stats` | | `false` | Enable detailed statistics logging |
-| `--stage-init-timeout` | | `300` | Stage initialization timeout (seconds) |
-| `--batch-timeout` | | `5` | Batching timeout (seconds) |
-| `--init-timeout` | | `300` | Overall initialization timeout (seconds) |
-| `--shm-threshold-bytes` | | `65536` | Shared memory threshold (bytes) |
-
-## Pipeline
-
-Covo-Audio-Chat uses a 2-stage pipeline:
-
- **Stage 0 (fused_thinker_talker):** The 7B LLM generates interleaved text and audio tokens in a single autoregressive pass.
- **Stage 1 (code2wav):** A BigVGAN-based vocoder converts the extracted audio codes into a 24kHz WAV waveform.
-
-## Output
-
-The script generates two files per request in the output directory:
-
- `{request_id}.txt` -- prompt and generated text
- `{request_id}.wav` -- generated audio (24kHz WAV)
-
-## FAQ
-
-If you encounter `ModuleNotFoundError: No module named 'librosa'`, install it with:
-```bash
-pip install librosa
-```
-
-## Environment
-
- GPU: 1x A100 (80 GiB)
- Stage 0 (7B LLM): ~16 GiB VRAM
- Stage 1 (BigVGAN vocoder): ~2 GiB VRAM
@@ -1,222 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-This example shows how to use vLLM-Omni for running offline inference
-with the correct prompt format on Covo-Audio-Chat.
-
-Usage:
-    python end2end.py --audio-path /path/to/audio.wav
-"""
-
-import os
-
-import soundfile as sf
-from vllm.assets.audio import AudioAsset
-from vllm.multimodal.media.audio import load_audio
-from vllm.sampling_params import SamplingParams
-from vllm.utils.argparse_utils import FlexibleArgumentParser
-
-from vllm_omni.entrypoints.omni import Omni
-from vllm_omni.model_executor.models.covo_audio.prompt_utils import (
-    COVO_AUDIO_INPUT_PREFIX,
-    build_covo_audio_chat_prompt,
-)
-
-SEED = 42
-
-
-def get_audio_query(
-    question: str | None = None,
-    audio_path: str | None = None,
-    sampling_rate: int = 16000,
-) -> dict:
-    if question is None:
-        question = "请回答这段音频里的问题。"
-    user_content = COVO_AUDIO_INPUT_PREFIX + question
-    prompt = build_covo_audio_chat_prompt(user_content)
-
-    if audio_path is None:
-        audio_data = AudioAsset("mary_had_lamb").audio_and_sample_rate
-    else:
-        import numpy as np
-
-        audio_signal, sr = load_audio(audio_path, sr=sampling_rate)
-        audio_data = (audio_signal.astype(np.float32), sr)
-
-    return {
-        "prompt": prompt,
-        "multi_modal_data": {"audio": audio_data},
-        "modalities": ["audio"],
-    }
-
-
-def main(args):
-    query_result = get_audio_query(
-        question=args.text,
-        audio_path=args.audio_path,
-        sampling_rate=args.sampling_rate,
-    )
-
-    omni = Omni(
-        model=args.model_name,
-        stage_configs_path=args.stage_configs_path,
-        log_stats=args.log_stats,
-        stage_init_timeout=args.stage_init_timeout,
-        batch_timeout=args.batch_timeout,
-        init_timeout=args.init_timeout,
-        shm_threshold_bytes=args.shm_threshold_bytes,
-    )
-
-    # Stage 0: fused_thinker_talker
-    # stop_token_ids=[151645] (<|im_end|>) and ignore_eos=True are required
-    # so the model generates interleaved text+audio tokens before stopping.
-    thinker_sampling_params = SamplingParams(
-        temperature=0.0,
-        top_p=1.0,
-        top_k=-1,
-        max_tokens=2048,
-        seed=SEED,
-        detokenize=True,
-        repetition_penalty=1.05,
-        stop_token_ids=[151645],
-        ignore_eos=True,
-    )
-    # Stage 1: code2wav (audio codes, not real token IDs — skip detokenize)
-    code2wav_sampling_params = SamplingParams(
-        temperature=0.0,
-        top_p=1.0,
-        top_k=-1,
-        max_tokens=2048,
-        seed=SEED,
-        detokenize=False,
-        repetition_penalty=1.1,
-    )
-
-    sampling_params_list = [
-        thinker_sampling_params,
-        code2wav_sampling_params,
-    ]
-
-    prompts = [query_result for _ in range(args.num_prompts)]
-
-    omni_outputs = omni.generate(prompts, sampling_params_list)
-
-    output_dir = args.output_dir
-    os.makedirs(output_dir, exist_ok=True)
-
-    for stage_outputs in omni_outputs:
-        output = stage_outputs.request_output
-        if stage_outputs.final_output_type == "text":
-            request_id = output.request_id
-            text_output = output.outputs[0].text
-            prompt_text = output.prompt
-            out_txt = os.path.join(output_dir, f"{request_id}.txt")
-            lines = [
-                "Prompt:\n",
-                str(prompt_text) + "\n",
-                "vllm_text_output:\n",
-                str(text_output).strip() + "\n",
-            ]
-            try:
-                with open(out_txt, "w", encoding="utf-8") as f:
-                    f.writelines(lines)
-            except Exception as e:
-                print(f"[Warn] Failed writing text file {out_txt}: {e}")
-            print(f"Request ID: {request_id}, Text saved to {out_txt}")
-        elif stage_outputs.final_output_type == "audio":
-            request_id = output.request_id
-            audio_tensor = output.outputs[0].multimodal_output.get("audio")
-            if audio_tensor is None:
-                continue
-            output_wav = os.path.join(output_dir, f"{request_id}.wav")
-            audio_numpy = audio_tensor.float().detach().cpu().numpy()
-            if audio_numpy.ndim > 1:
-                audio_numpy = audio_numpy.flatten()
-            sf.write(output_wav, audio_numpy, samplerate=24000, format="WAV")
-            print(f"Request ID: {request_id}, Audio saved to {output_wav}")
-
-    omni.close()
-
-
-def parse_args():
-    parser = FlexibleArgumentParser(description="Offline inference demo for Covo-Audio-Chat")
-    parser.add_argument(
-        "--model-name",
-        "-m",
-        type=str,
-        default="tencent/Covo-Audio-Chat",
-        help="Model path or HuggingFace model ID.",
-    )
-    parser.add_argument(
-        "--text",
-        "-t",
-        type=str,
-        default=None,
-        help="Text prompt / question for the audio.",
-    )
-    parser.add_argument(
-        "--audio-path",
-        "-a",
-        type=str,
-        default=None,
-        help="Path to local audio file. Uses default asset if not provided.",
-    )
-    parser.add_argument(
-        "--sampling-rate",
-        type=int,
-        default=16000,
-        help="Sampling rate for audio loading (default: 16000).",
-    )
-    parser.add_argument(
-        "--stage-configs-path",
-        type=str,
-        default=None,
-        help="Path to stage configs YAML file.",
-    )
-    parser.add_argument(
-        "--log-stats",
-        action="store_true",
-        default=False,
-        help="Enable writing detailed statistics.",
-    )
-    parser.add_argument(
-        "--stage-init-timeout",
-        type=int,
-        default=300,
-        help="Timeout for initializing a single stage in seconds.",
-    )
-    parser.add_argument(
-        "--batch-timeout",
-        type=int,
-        default=5,
-        help="Timeout for batching in seconds.",
-    )
-    parser.add_argument(
-        "--init-timeout",
-        type=int,
-        default=300,
-        help="Timeout for initializing stages in seconds.",
-    )
-    parser.add_argument(
-        "--shm-threshold-bytes",
-        type=int,
-        default=65536,
-        help="Threshold for using shared memory in bytes.",
-    )
-    parser.add_argument(
-        "--output-dir",
-        default="./output_audio",
-        help="Output directory for generated files.",
-    )
-    parser.add_argument(
-        "--num-prompts",
-        type=int,
-        default=1,
-        help="Number of prompts to generate.",
-    )
-    return parser.parse_args()
-
-
-if __name__ == "__main__":
-    args = parse_args()
-    main(args)
@@ -1,155 +1,172 @@
 # HunyuanImage-3.0-Instruct

-This example runs HunyuanImage-3.0-Instruct offline with the unified deploy
-YAMLs under `vllm_omni/deploy/`.
+## Set up

-## Deploy Configs
+Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.

-| File | Topology | Default use |
-| :--- | :--- | :--- |
-| `vllm_omni/deploy/hunyuan_image3.yaml` | AR + DiT | Default for `text2img` and `img2img`. |
-| `vllm_omni/deploy/hunyuan_image3_ar.yaml` | AR only | Default for `img2text` and `text2text`. |
-| `vllm_omni/deploy/hunyuan_image3_dit.yaml` | DiT only | Standalone diffusion stage. Pass it explicitly with `--deploy-config`. |
+## Run examples

-The example chooses a deploy config automatically when `--deploy-config` and
-`--stage-configs-path` are both omitted:
+**Note**: These examples work with the default configuration on **8x NVIDIA L40S (48GB)**. For different GPU setups, modify the stage configuration to adjust device allocation and memory utilization.

-| `--modality` | `mode` passed to Omni | Default deploy |
-| :--- | :--- | :--- |
-| `text2img` | `text-to-image` | `hunyuan_image3.yaml` |
-| `img2img` | `image-editing` | `hunyuan_image3.yaml` |
-| `img2text` | `image-to-text` | `hunyuan_image3_ar.yaml` |
-| `text2text` | `text-to-text` | `hunyuan_image3_ar.yaml` |
-
-`--modality` is an offline example convenience flag. It maps to the internal
-`mode` argument passed to `Omni(...)` by this script. HunyuanImage3 uses
-separate deploy YAMLs for AR + DiT, AR-only, and DiT-only topologies, so the
-stage topology is selected by the deploy file rather than by YAML mode
-overrides.
-
-Online serving does not expose a `--modality` flag or accept `mode` as an API
-request field. Choose the deploy topology when starting the server with
-`--deploy-config`, then use the OpenAI-compatible endpoint and request shape for
-the scenario. The `modalities` request field is used by the chat completions
-path; the image endpoints infer the image task from the endpoint and payload.
-
-| Online scenario | Server deploy | Request |
-| :--- | :--- | :--- |
-| Text to image | `--deploy-config vllm_omni/deploy/hunyuan_image3.yaml` | `POST /v1/images/generations`, or `POST /v1/chat/completions` with `"modalities": ["image"]`. |
-| Image editing | `--deploy-config vllm_omni/deploy/hunyuan_image3.yaml` | `POST /v1/images/edits`. |
-| Image/text to text | `--deploy-config vllm_omni/deploy/hunyuan_image3_ar.yaml` | `POST /v1/chat/completions` for text output, for example with `"modalities": ["text"]`. |
-| DiT-only image generation | `--deploy-config vllm_omni/deploy/hunyuan_image3_dit.yaml` | `POST /v1/images/generations`. |
-
-## Run Examples
-
-Text to image, using the default AR + DiT deploy:
+Get into the hunyuan_image3 folder:

 ```bash
-python examples/offline_inference/hunyuan_image3/end2end.py \
-  --model tencent/HunyuanImage-3.0-Instruct \
-  --modality text2img \
-  --prompts "A cute cat sitting on a windowsill watching the sunset"
+cd examples/offline_inference/hunyuan_image3
 ```

-Image editing, using the default AR + DiT deploy:
+### Modality Control
+
+HunyuanImage-3.0-Instruct supports multiple modality modes. You can control the mode using the `--modality` argument:
+
+#### Text to Image (text2img)
+
+- **Pipeline**: Text → AR (CoT + latent tokens) → DiT (denoise) → VAE Decode → Image
+- **Stages Used**: Stage 0 (AR) + Stage 1 (DiT)
+- **KV Transfer**: AR sends KV cache to DiT for conditioned generation
+- **Default Config**: `hunyuan_image3_t2i.yaml`

 ```bash
-python examples/offline_inference/hunyuan_image3/end2end.py \
-  --model tencent/HunyuanImage-3.0-Instruct \
-  --modality img2img \
-  --image-path /path/to/image.png \
-  --prompts "Make the petals neon pink"
+python end2end.py --model tencent/HunyuanImage-3.0-Instruct \
+                  --modality text2img \
+                  --prompts "A cute cat sitting on a windowsill watching the sunset"
 ```

-Image to text, using the AR-only deploy:
+**With VAE tiling (required on A100 GPUs):**
+```bash
+python end2end.py --model tencent/HunyuanImage-3.0-Instruct \
+                  --modality text2img \
+                  --prompts "A cute cat sitting on a windowsill watching the sunset" \
+                  --vae-use-tiling
+```
+
+#### Image to Image (img2img)
+
+- **Pipeline**: Image + Text → AR (CoT + recaption + latent) → DiT → Edited Image
+- **Stages Used**: Stage 0 (AR) + Stage 1 (DiT)
+- **KV Transfer**: AR sends KV cache to DiT
+- **Default Config**: `hunyuan_image3_it2i.yaml`

 ```bash
-python examples/offline_inference/hunyuan_image3/end2end.py \
-  --model tencent/HunyuanImage-3.0-Instruct \
-  --modality img2text \
-  --image-path /path/to/image.jpg \
-  --prompts "Describe the content of the picture."
+python end2end.py --model tencent/HunyuanImage-3.0-Instruct \
+                  --modality img2img \
+                  --image-path /path/to/image.png \
+                  --prompts "Make the petals neon pink"
 ```

-Text to text, using the AR-only deploy:
+#### Image to Text (img2text)
+
+- **Pipeline**: Image + Question → AR → Text description
+- **Stages Used**: Stage 0 (AR) only
+- **Default Config**: `hunyuan_image3_i2t.yaml`

 ```bash
-python examples/offline_inference/hunyuan_image3/end2end.py \
-  --model tencent/HunyuanImage-3.0-Instruct \
-  --modality text2text \
-  --prompts "What is the capital of France?"
+python end2end.py --model tencent/HunyuanImage-3.0-Instruct \
+                  --modality img2text \
+                  --image-path /path/to/image.jpg \
+                  --prompts "Describe the content of the picture."
 ```

-Standalone DiT, using the DiT-only deploy explicitly:
+#### Text to Text (text2text)
+
+- **Pipeline**: Text → AR → Text
+- **Stages Used**: Stage 0 (AR) only
+- **Default Config**: `hunyuan_image3_t2t.yaml`

 ```bash
-python examples/offline_inference/hunyuan_image3/end2end.py \
-  --model tencent/HunyuanImage-3.0-Instruct \
-  --modality text2img \
-  --deploy-config vllm_omni/deploy/hunyuan_image3_dit.yaml \
-  --prompts "A cinematic portrait of an astronaut in a greenhouse"
+python end2end.py --model tencent/HunyuanImage-3.0-Instruct \
+                  --modality text2text \
+                  --prompts "What is the capital of France?"
 ```

-Override the default full AR + DiT deploy explicitly:
+### Inference Steps & Guidance

-```bash
-python examples/offline_inference/hunyuan_image3/end2end.py \
-  --model tencent/HunyuanImage-3.0-Instruct \
-  --modality text2img \
-  --deploy-config vllm_omni/deploy/hunyuan_image3.yaml \
-  --prompts "A cute cat"
-```
-
-## Additional Config
-
-You can pass diffusion worker `additional_config` from the offline example as a JSON object.
-This maps to the upstream vLLM `VllmConfig.additional_config` platform extension field:
-https://docs.vllm.ai/en/stable/api/vllm/config/#vllm.config.VllmConfig.additional_config
+Control generation quality for image modalities:

 ```bash
 python end2end.py --modality text2img \
-                  --prompts "A cute cat" \
-                  --additional-config '{"torchair_graph_config":{"enabled":true}}'
+                  --steps 50 \
+                  --guidance-scale 5.0 \
+                  --height 1024 --width 1024 \
+                  --prompts "A photo-realistic sunset over the ocean"
 ```

+### Key Arguments

-## Key Arguments
+#### 📌 Command Line Arguments (end2end.py)

-| Argument | Description |
-| :--- | :--- |
-| `--deploy-config` | Preferred config path for unified deploy YAMLs. |
-| `--stage-configs-path` | Legacy stage config path, kept only for compatibility. Prefer `--deploy-config`. |
-| `--additional-config` | JSON object forwarded to diffusion worker `additional_config`. |
-| `--modality` | Offline-only convenience flag. One of `text2img`, `img2img`, `img2text`, `text2text`. It selects prompt formatting, internal `mode`, and default deploy config for this script. Online serving uses `--deploy-config` plus the endpoint and, for chat completions, request `modalities` instead. |
-| `--steps` | Number of diffusion inference steps for image generation. |
-| `--guidance-scale` | Classifier-free guidance scale for image generation. |
-| `--height`, `--width` | Output image size for `text2img`. |
-| `--bot-task` | Override prompt mode. `none`, `think`, `recaption`, `think_recaption`, or `vanilla`. |
-| `--sys-type` | Override the system prompt type, for example `en_unified` or `en_vanilla`. |
-| `--vae-use-tiling` | Enable VAE tiling for memory reduction. |
+| Argument               | Type   | Default                              | Description                                                  |
+| :--------------------- | :----- | :----------------------------------- | :----------------------------------------------------------- |
+| `--model`              | string | `tencent/HunyuanImage-3.0-Instruct` | Model path or name                                           |
+| `--modality`           | choice | `text2img`                           | Modality: `text2img`, `img2img`, `img2text`, `text2text`     |
+| `--prompts`            | list   | `None`                               | Input text prompts                                           |
+| `--image-path`         | string | `None`                               | Input image path (for `img2img`/`img2text`)                  |
+| `--output`             | string | `.`                                  | Output directory for saved images                            |
+| `--steps`              | int    | `50`                                 | Number of inference steps                                    |
+| `--guidance-scale`     | float  | `5.0`                                | Classifier-free guidance scale                               |
+| `--seed`               | int    | `42`                                 | Random seed                                                  |
+| `--height`             | int    | `1024`                               | Output image height                                          |
+| `--width`              | int    | `1024`                               | Output image width                                           |
+| `--bot-task`           | string | auto                                 | Override prompt task (e.g. `it2i_think`, `t2i_recaption`)    |
+| `--sys-type`           | string | auto                                 | Override system prompt type (e.g. `en_unified`, `en_vanilla`) |
+| `--stage-configs-path` | string | auto                                 | Custom stage config YAML path                                |
+| `--enforce-eager`      | flag   | `False`                              | Disable torch.compile                                        |
+| `--init-timeout`       | int    | `300`                                | Initialization timeout (seconds)                             |
+| `--vae-use-tiling`     | flag   | `False`                              | Enable VAE tiling for memory optimization (required to avoid OOM on A100) |

-## Notes
+------

- `hunyuan_image3_ar.yaml` is a 4-card AR-only text/comprehension deploy.
- `hunyuan_image3_dit.yaml` is a single-stage DiT deploy with `stage_id: 0`.
- The old HunyuanImage3 YAMLs under `model_executor/stage_configs/` and `platforms/*/stage_configs/` have been folded into the deploy YAMLs.
+#### ⚙️ Stage Configurations
+
+| Config YAML                         | Modality  | Stages | GPUs   | Description                           |
+| :---------------------------------- | :-------- | :----- | :----- | :------------------------------------ |
+| `hunyuan_image3_t2i.yaml`           | text2img  | 2      | 8      | T2I with AR→DiT, 4 GPU each          |
+| `hunyuan_image3_it2i.yaml`          | img2img   | 2      | 8      | IT2I with AR→DiT, 4 GPU each         |
+| `hunyuan_image3_i2t.yaml`           | img2text  | 1      | 4      | I2T (AR only)                         |
+| `hunyuan_image3_t2t.yaml`           | text2text | 1      | 4      | T2T (AR only)                         |
+| `hunyuan_image3_t2i_2gpu.yaml`      | text2img  | 2      | 2      | T2I for 2-GPU setups                  |
+| `hunyuan_image3_moe.yaml`           | text2img  | 2      | 8      | T2I with MoE AR→DiT KV reuse          |
+| `hunyuan_image3_moe_dit_2gpu_fp8.yaml` | text2img | 2   | 2      | T2I with FP8 quantization             |
+
+------
+
+## Using MoE Config
+
+The `hunyuan_image3_moe.yaml` config enables AR→DiT KV cache reuse with 8 GPUs (4 for AR + 4 for DiT).
+
+```bash
+python end2end.py --model tencent/HunyuanImage-3.0-Instruct \
+                  --modality text2img \
+                  --stage-configs-path hunyuan_image3_moe.yaml \
+                  --prompts "A cute cat"
+```
+
+------

 ## Prompt Format

 HunyuanImage-3.0-Instruct uses an instruct chat template:

-```text
-<|startoftext|>{system_prompt}
-
-User: {<img>?}{user_prompt}
-
-Assistant: {trigger_tag?}
+```
+<|startoftext|>{system_prompt}\n\nUser: {<img>?}{user_prompt}\n\nAssistant: {trigger_tag?}
 ```

- `<img>`: Placeholder for each input image (single token; expanded by the multimodal pipeline).
- Trigger tags: `<think>` for CoT and `<recaption>` for recaptioning, placed after `Assistant: `.
- System prompt: Auto-selected from `task` and `bot_task`.
- `bot_task='vanilla'` with `task='t2i'` uses the bare pretrain template.
+- `<img>`: Placeholder for each input image (single token; expanded by the multimodal pipeline)
+- Trigger tags: `<think>` (CoT), `<recaption>` (recaptioning) — placed AFTER `Assistant: `
+- System prompt: Auto-selected based on task
+- `t2i_vanilla` is the only task that uses the bare pretrain template (no chat structure)

 The shared `vllm_omni.diffusion.models.hunyuan_image3.prompt_utils.build_prompt_tokens()`
-helper handles segment-by-segment tokenization and matches HF `apply_chat_template`.
+helper handles segment-by-segment tokenization (matches HF `apply_chat_template` byte-for-byte).
+
+------
+
+## FAQ
+
+- **OOM errors**: Decrease `gpu_memory_utilization` in the YAML stage config, use a smaller `max_num_batched_tokens`, or enable VAE tiling with `--vae-use-tiling` (required on A100 GPUs).
+- **Custom image sizes**: Use `--height` and `--width` flags (multiples of 16 recommended).
+
+| Stage             | VRAM (approx)        |
+| :---------------- | :------------------- |
+| Stage 0 (AR)      | ~15 GiB + KV Cache   |
+| Stage 1 (DiT)     | ~30 GiB              |
+| Total (8-GPU)     | ~45 GiB + KV Cache   |
@@ -1,294 +1,296 @@
 """
 HunyuanImage-3.0-Instruct unified end-to-end inference script.
+
+Supports all modalities through a single entry point:
+  - text2img:  Text → AR → DiT → Image
+  - img2img:   Text+Image → AR → DiT → Edited Image (IT2I)
+  - img2text:  Image+Text → AR → Text description (I2T)
+  - text2text: Text → AR → Text (comprehension, no image)
+
+Usage:
+    python end2end.py --modality text2img --prompts "A cute cat"
+    python end2end.py --modality img2img --image-path input.png --prompts "Make it snowy"
+    python end2end.py --modality img2text --image-path input.png --prompts "Describe this image"
 """

 import argparse
-import json
 import os
-from pathlib import Path

 from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import (
-    MAX_IMAGES_PER_REQUEST,
    build_prompt_tokens,
-    resolve_stop_token_ids,
-    resolve_sys_type,
 )
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniPromptType
+from vllm_omni.platforms import current_omni_platform

-_REPO_ROOT = Path(__file__).resolve().parents[3]
-_DEFAULT_DEPLOY_CONFIG = str(_REPO_ROOT / "vllm_omni" / "deploy" / "hunyuan_image3.yaml")
-_DEFAULT_AR_DEPLOY_CONFIG = str(_REPO_ROOT / "vllm_omni" / "deploy" / "hunyuan_image3_ar.yaml")
-
-_MODALITY_TASK_MAP: dict[str, tuple[str, str | None]] = {
-    "text2img": ("t2i", "think"),
-    "img2img": ("it2i", "think"),
-    "img2text": ("i2t", None),
-    "text2text": ("t2t", None),
+# task -> (sys_type, bot_task, trigger_tag)
+_TASK_PRESETS: dict[str, tuple[str, str | None, str | None]] = {
+    "t2t": ("en_unified", None, None),
+    "i2t": ("en_unified", None, None),
+    "it2i_think": ("en_unified", "think", "<think>"),
+    "it2i_recaption": ("en_unified", "recaption", "<recaption>"),
+    "t2i_think": ("en_unified", "think", "<think>"),
+    "t2i_recaption": ("en_unified", "recaption", "<recaption>"),
+    "t2i_vanilla": ("en_vanilla", "image", None),
 }

-_MODALITY_DEFAULT_DEPLOY_CONFIG = {
-    "text2img": _DEFAULT_DEPLOY_CONFIG,
-    "img2img": _DEFAULT_DEPLOY_CONFIG,
-    "img2text": _DEFAULT_AR_DEPLOY_CONFIG,
-    "text2text": _DEFAULT_AR_DEPLOY_CONFIG,
+# Modality → prompt_utils task mapping
+_MODALITY_TASK_MAP = {
+    "text2img": "t2i_think",
+    "img2img": "it2i_think",
+    "img2text": "i2t",
+    "text2text": "t2t",
 }

-_MODALITY_MODE = {
-    "text2img": "text-to-image",
-    "img2img": "image-editing",
-    "img2text": "image-to-text",
-    "text2text": "text-to-text",
+
+# Modality → default stage config
+_MODALITY_DEFAULT_CONFIG = {
+    "text2img": "hunyuan_image3_t2i.yaml",
+    "img2img": "hunyuan_image3_it2i.yaml",
+    "img2text": "hunyuan_image3_i2t.yaml",
+    "text2text": "hunyuan_image3_t2t.yaml",
 }


 def parse_args():
    parser = argparse.ArgumentParser(description="HunyuanImage-3.0-Instruct end-to-end inference.")
-    parser.add_argument("--model", default="tencent/HunyuanImage-3.0-Instruct", help="Model name or local path.")
+    parser.add_argument(
+        "--model",
+        default="tencent/HunyuanImage-3.0-Instruct",
+        help="Model name or local path.",
+    )
    parser.add_argument(
        "--modality",
        default="text2img",
-        choices=list(_MODALITY_TASK_MAP),
+        choices=["text2img", "img2img", "img2text", "text2text"],
+        help="Modality mode to control stage execution.",
    )
    parser.add_argument("--prompts", nargs="+", default=None, help="Input text prompts.")
    parser.add_argument(
        "--image-path",
        type=str,
        default=None,
-        help="Input image path(s) for img2img/img2text. Comma-separated for multi-image (up to 3).",
+        help="Path to input image (for img2img/img2text).",
    )
-    parser.add_argument("--output", type=str, default=".", help="Output directory to save results.")
+    parser.add_argument(
+        "--output",
+        type=str,
+        default=".",
+        help="Output directory to save results.",
+    )
+
+    # Generation parameters
    parser.add_argument("--steps", type=int, default=50, help="Number of inference steps.")
    parser.add_argument("--guidance-scale", type=float, default=5.0, help="Classifier-free guidance scale.")
    parser.add_argument("--seed", type=int, default=42, help="Random seed.")
    parser.add_argument("--height", type=int, default=1024, help="Output image height.")
    parser.add_argument("--width", type=int, default=1024, help="Output image width.")
-    parser.add_argument("--vae-use-tiling", action="store_true", help="Enable VAE tiling.")
+    parser.add_argument(
+        "--vae-use-tiling",
+        action="store_true",
+        help="Enable VAE tiling for memory optimization.",
+    )
+
+    # Prompt configuration
    parser.add_argument(
        "--bot-task",
        type=str,
        default=None,
-        choices=["none", "think", "recaption", "think_recaption", "vanilla"],
-        help="Override prompt mode. Default: auto from --modality.",
+        help="Override prompt task (e.g. it2i_think, t2i_recaption). Default: auto from modality.",
    )
-    parser.add_argument("--sys-type", type=str, default=None, help="Override system prompt type.")
-    parser.add_argument("--deploy-config", type=str, default=None, help="Custom deploy YAML path.")
-    parser.add_argument("--stage-configs-path", type=str, default=None, help="Custom legacy stage config YAML path.")
+    parser.add_argument(
+        "--sys-type",
+        type=str,
+        default=None,
+        help="Override system prompt type (e.g. en_unified, en_vanilla).",
+    )
+
+    # Omni init args
+    parser.add_argument("--stage-configs-path", type=str, default=None, help="Custom stage config YAML path.")
    parser.add_argument("--log-stats", action="store_true", default=False)
    parser.add_argument("--init-timeout", type=int, default=300, help="Initialization timeout in seconds.")
    parser.add_argument("--enforce-eager", action="store_true", help="Disable torch.compile.")
    parser.add_argument(
-        "--diffusion-kv-cache-dtype",
+        "--quantization",
        type=str,
        default=None,
-        help="Diffusion attention KV cache dtype, for example 'fp8'. Separate from vLLM --kv-cache-dtype.",
+        choices=["fp8", "int8", "gguf"],
+        help="Quantization method for the transformer. "
+        "Options: 'fp8' (FP8 W8A8 on Ada/Hopper, weight-only on older GPUs), 'int8' (Int8 W8A8), 'gguf' (GGUF quantized weights). "
+        "Default: None (no quantization, uses BF16).",
    )
    parser.add_argument(
-        "--diffusion-kv-cache-skip-steps",
+        "--gguf-model",
        type=str,
        default=None,
-        help="Denoising step selector to keep diffusion KV cache in native dtype, for example '0,1,4-6'.",
+        help=("GGUF file path or HF reference for transformer weights. Required when --quantization gguf is set."),
    )
    parser.add_argument(
-        "--diffusion-kv-cache-skip-layers",
+        "--ignored-layers",
        type=str,
        default=None,
-        help="Transformer layer selector to keep diffusion KV cache in native dtype, for example '0-2,10'.",
-    )
-    parser.add_argument(
-        "--additional-config",
-        type=str,
-        default=None,
-        help=(
-            "JSON object forwarded to Omni/additional_config, for example "
-            '\'{"torchair_graph_config":{"enabled":true}}\'. Different platforms may support different '
-            "configs. Make sure the configs are valid for the platform you are using. "
-            "Contents must be hashable."
-        ),
+        help="Comma-separated list of layer name patterns to skip quantization. "
+        "Only used when --quantization is set. "
+        "Available layers: to_qkv, to_out, add_kv_proj, to_add_out, img_mlp, txt_mlp, proj_out. "
+        "Example: --ignored-layers 'add_kv_proj,to_add_out'",
    )

+    current_omni_platform.pre_register_and_update(parser)
+
    from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults

    nullify_stage_engine_defaults(parser)
    return parser.parse_args()


-def parse_additional_config(raw_value: str | None) -> dict | None:
-    """Parse a JSON string into an additional_config mapping."""
-    if raw_value is None:
-        return None
-
-    try:
-        additional_config = json.loads(raw_value)
-    except json.JSONDecodeError as exc:
-        raise ValueError(f"Invalid --additional-config JSON: {exc}") from exc
-
-    if additional_config is None:
-        return None
-    if not isinstance(additional_config, dict):
-        raise ValueError(f"--additional-config must decode to a JSON object, got {type(additional_config).__name__}")
-    return additional_config
-
-
 def main():
    args = parse_args()
    os.makedirs(args.output, exist_ok=True)
-    additional_config = parse_additional_config(args.additional_config)

-    task, default_bot_task = _MODALITY_TASK_MAP[args.modality]
-    if args.bot_task is None:
-        bot_task: str | None = default_bot_task
-    elif args.bot_task == "none":
-        bot_task = None
-    else:
-        bot_task = args.bot_task
+    # Determine task for prompt formatting
+    task = args.bot_task or _MODALITY_TASK_MAP[args.modality]

-    if args.deploy_config is not None and args.stage_configs_path is not None:
-        raise ValueError("--deploy-config and --stage-configs-path are mutually exclusive.")
-
-    deploy_config = args.deploy_config
-    stage_configs_path = args.stage_configs_path
-    if deploy_config is None and stage_configs_path is None:
-        deploy_config = _MODALITY_DEFAULT_DEPLOY_CONFIG[args.modality]
+    # Determine stage config
+    stage_configs_path = args.stage_configs_path or _MODALITY_DEFAULT_CONFIG[args.modality]
+    # Build quantization kwargs: use quantization_config dict when
+    # ignored_layers is specified so the list flows through OmniDiffusionConfig
+    quant_kwargs: dict[str, Any] = {}
+    ignored_layers = [s.strip() for s in args.ignored_layers.split(",") if s.strip()] if args.ignored_layers else None
+    if args.quantization == "gguf":
+        if not args.gguf_model:
+            raise ValueError("--gguf-model is required when --quantization gguf is set.")
+        quant_kwargs["quantization_config"] = {
+            "method": "gguf",
+            "gguf_model": args.gguf_model,
+        }
+    elif args.quantization and ignored_layers:
+        quant_kwargs["quantization_config"] = {
+            "method": args.quantization,
+            "ignored_layers": ignored_layers,
+        }
+    elif args.quantization:
+        quant_kwargs["quantization"] = args.quantization

+    # Build Omni
    omni_kwargs = {
        "model": args.model,
        "vae_use_tiling": args.vae_use_tiling,
+        "stage_configs_path": stage_configs_path,
        "log_stats": args.log_stats,
        "init_timeout": args.init_timeout,
        "enforce_eager": args.enforce_eager,
-        "mode": _MODALITY_MODE[args.modality],
-        "diffusion_kv_cache_dtype": args.diffusion_kv_cache_dtype,
-        "diffusion_kv_cache_skip_steps": args.diffusion_kv_cache_skip_steps,
-        "diffusion_kv_cache_skip_layers": args.diffusion_kv_cache_skip_layers,
+        **quant_kwargs,
    }
-
-    if additional_config is not None:
-        omni_kwargs["additional_config"] = additional_config
-    if deploy_config is not None:
-        omni_kwargs["deploy_config"] = deploy_config
-    else:
-        omni_kwargs["stage_configs_path"] = stage_configs_path
+    if args.modality in ("text2img", "img2img"):
+        omni_kwargs["mode"] = "text-to-image"

    omni = Omni(**omni_kwargs)

+    # Prepare prompts
    prompts = args.prompts or ["A cute cat"]
-    input_images: list = []
+    if not prompts:
+        print("[Info] No prompts provided, using default.")
+        prompts = ["A cute cat"]
+
+    # Load image if needed
+    input_image = None
    if args.modality in ("img2img", "img2text"):
-        if not args.image_path:
+        if not args.image_path or not os.path.exists(args.image_path):
            raise ValueError(f"--image-path required for {args.modality}, got: {args.image_path}")
        from PIL import Image

-        image_paths = [p.strip() for p in args.image_path.split(",") if p.strip()]
-        if len(image_paths) > MAX_IMAGES_PER_REQUEST:
-            raise ValueError(
-                f"--image-path accepts at most {MAX_IMAGES_PER_REQUEST} images for "
-                f"HunyuanImage-3.0 IT2I, got {len(image_paths)}: {args.image_path}"
-            )
-        for image_path in image_paths:
-            if not os.path.exists(image_path):
-                raise ValueError(f"Image path does not exist: {image_path}")
-            input_images.append(Image.open(image_path).convert("RGB"))
-        if not input_images:
-            raise ValueError(f"--image-path produced no usable paths: {args.image_path!r}")
+        input_image = Image.open(args.image_path).convert("RGB")

+    # Load tokenizer for segment-wise prompt tokenization (matches HF
+    # apply_chat_template byte-for-byte; see build_prompt_tokens docstring).
    from transformers import AutoTokenizer

    tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
-    mm_image_payload = (input_images[0] if len(input_images) == 1 else input_images) if input_images else None

+    # Format prompts
    formatted_prompts: list[OmniPromptType] = []
-    for prompt in prompts:
-        build_kwargs: dict = {"task": task, "bot_task": bot_task, "sys_type": args.sys_type}
-        if input_images:
-            build_kwargs["num_images"] = len(input_images)
-        result = build_prompt_tokens(prompt, tokenizer, **build_kwargs)
-        token_ids = result.token_ids
-        effective_sys_type = args.sys_type or resolve_sys_type(bot_task)
+    for p in prompts:
+        token_ids = build_prompt_tokens(p, tokenizer, task=task, sys_type=args.sys_type)
+        preset_sys_type, _, _ = _TASK_PRESETS[task]
+        effective_sys_type = args.sys_type or preset_sys_type

+        # `prompt_token_ids` drives the AR stage (matches HF byte-for-byte).
+        # `prompt` and `use_system_prompt` are forwarded by ar2diffusion to
+        # the DiT stage so the diffusion pipeline can rebuild the same
+        # system prefix when constructing its model inputs.
        prompt_dict: dict = {
            "prompt_token_ids": token_ids,
-            "prompt": prompt,
+            "prompt": p,
            "use_system_prompt": effective_sys_type,
        }
+
        if args.modality == "text2img":
            prompt_dict["modalities"] = ["image"]
        elif args.modality == "img2img":
            prompt_dict["modalities"] = ["image"]
-            prompt_dict["multi_modal_data"] = {"image": mm_image_payload}
-            prompt_dict["height"] = input_images[0].height
-            prompt_dict["width"] = input_images[0].width
+            prompt_dict["multi_modal_data"] = {"image": input_image}
+            prompt_dict["height"] = input_image.height
+            prompt_dict["width"] = input_image.width
        elif args.modality == "img2text":
            prompt_dict["modalities"] = ["text"]
-            prompt_dict["multi_modal_data"] = {"image": mm_image_payload}
-        else:
+            prompt_dict["multi_modal_data"] = {"image": input_image}
+        elif args.modality == "text2text":
            prompt_dict["modalities"] = ["text"]
+
        formatted_prompts.append(prompt_dict)

+    # Build sampling params from defaults
    params_list = list(omni.default_sampling_params_list)

+    # Override diffusion params if applicable
    from vllm_omni.inputs.data import OmniDiffusionSamplingParams

-    ar_stop_token_ids = resolve_stop_token_ids(task=task, bot_task=bot_task, tokenizer=tokenizer)
-    for sp in params_list:
+    for i, sp in enumerate(params_list):
        if isinstance(sp, OmniDiffusionSamplingParams):
            sp.num_inference_steps = args.steps
            sp.guidance_scale = args.guidance_scale
-            sp.guidance_scale_provided = True
            if args.seed is not None:
                sp.seed = args.seed
-            if args.modality == "text2img":
+            if args.modality in ("text2img",):
                sp.height = args.height
                sp.width = args.width
-        elif hasattr(sp, "stop_token_ids"):
-            sp.stop_token_ids = ar_stop_token_ids

+    # Print configuration
    print(f"\n{'=' * 60}")
    print("HunyuanImage-3.0 Generation Configuration:")
    print(f"  Model: {args.model}")
    print(f"  Modality: {args.modality}")
-    print(f"  Prompt task: {task}")
-    print(f"  Bot task: {bot_task}")
-    if deploy_config is not None:
-        print(f"  Deploy config: {deploy_config}")
-    else:
-        print(f"  Stage config: {stage_configs_path}")
+    print(f"  Stage config: {stage_configs_path}")
    print(f"  Num stages: {omni.num_stages}")
    if args.modality in ("text2img", "img2img"):
        print(f"  Inference steps: {args.steps}")
        print(f"  Guidance scale: {args.guidance_scale}")
        print(f"  Seed: {args.seed}")
-        print(f"  diffusion_kv_cache_dtype: {args.diffusion_kv_cache_dtype}")
-        print(f"  diffusion_kv_cache_skip_steps: {args.diffusion_kv_cache_skip_steps}")
-        print(f"  diffusion_kv_cache_skip_layers: {args.diffusion_kv_cache_skip_layers}")
    if args.modality == "text2img":
        print(f"  Output size: {args.width}x{args.height}")
    if args.image_path:
        print(f"  Input image: {args.image_path}")
-    if additional_config is not None:
-        print(f"  Additional config: {additional_config}")
    print(f"  Prompts: {prompts}")
    print(f"{'=' * 60}\n")

+    # Generate
    omni_outputs = list(omni.generate(prompts=formatted_prompts, sampling_params_list=params_list))
+
+    # Process outputs
    img_idx = 0
    for req_output in omni_outputs:
+        # Text output (AR stage or text-only)
        ro = getattr(req_output, "request_output", None)
-        txt = ""
        if ro and getattr(ro, "outputs", None):
            txt = "".join(getattr(o, "text", "") or "" for o in ro.outputs)
-        if not txt:
-            ar_text = getattr(req_output, "custom_output", {}).get("ar_generated_text")
-            if isinstance(ar_text, list):
-                txt = "\n".join(text for text in ar_text if text)
-            else:
-                txt = ar_text or ""
-        if txt:
-            print(f"[Output] Text:\n{txt}")
+            if txt:
+                print(f"[Output] Text:\n{txt}")

+        # Image output (DiT stage)
        images = getattr(req_output, "images", None)
        if not images and ro and hasattr(ro, "images"):
            images = ro.images
+
        if images:
            for j, img in enumerate(images):
                save_path = os.path.join(args.output, f"output_{img_idx}_{j}.png")
@@ -102,24 +102,6 @@ def parse_args() -> argparse.Namespace:
        choices=["unipc", "euler"],
        help="Sampling solver for Wan2.2 pipelines. Use 'euler' for Lightning/Distill setups.",
    )
-    parser.add_argument(
-        "--diffusion-kv-cache-dtype",
-        type=str,
-        default=None,
-        help="Diffusion attention KV cache dtype (e.g. float8_e4m3fn). Separate from vLLM --kv-cache-dtype.",
-    )
-    parser.add_argument(
-        "--diffusion-kv-cache-skip-steps",
-        type=str,
-        default=None,
-        help="Diffusion KV-cache quantization skip-step selector, e.g. '0-9,20,25-30'.",
-    )
-    parser.add_argument(
-        "--diffusion-kv-cache-skip-layers",
-        type=str,
-        default=None,
-        help="Diffusion KV-cache quantization skip-layer selector, e.g. '0,1,4-8'.",
-    )
    parser.add_argument("--output", type=str, default="i2v_output.mp4", help="Path to save the video (mp4).")
    parser.add_argument("--fps", type=int, default=None, help="Frames per second for the output video.")
    parser.add_argument(
@@ -142,42 +124,6 @@ def parse_args() -> argparse.Namespace:
        action="store_true",
        help="Enable layerwise (blockwise) offloading on DiT modules.",
    )
-    parser.add_argument(
-        "--enforce-eager",
-        action="store_true",
-        help="Disable torch.compile and force eager execution.",
-    )
-    parser.add_argument(
-        "--audio-sample-rate",
-        type=int,
-        default=24000,
-        help="Sample rate for audio output when saved (default: 24000).",
-    )
-    parser.add_argument(
-        "--cache-backend",
-        type=str,
-        default=None,
-        choices=["cache_dit", "tea_cache"],
-        help=(
-            "Cache backend to use for acceleration. "
-            "Options: 'cache_dit' (DBCache + SCM + TaylorSeer), 'tea_cache' (Timestep Embedding Aware Cache). "
-            "Default: None (no cache acceleration)."
-        ),
-    )
-    parser.add_argument(
-        "--enable-diffusion-pipeline-profiler",
-        action="store_true",
-        help="Enable diffusion pipeline profiler to display stage durations.",
-    )
-    parser.add_argument(
-        "--quantization",
-        type=str,
-        default=None,
-        choices=["fp8", "mxfp8", "int8", "gguf"],
-        help="Quantization method for the transformer. mxfp8: W8A8 MXFP8 online quant (NPU). fp8: online FP8 (GPU).",
-    )
-
-    # Distributed and parallel execution
    parser.add_argument(
        "--ulysses-degree",
        type=int,
@@ -209,6 +155,28 @@ def parse_args() -> argparse.Namespace:
        default=1,
        help="Number of GPUs used for VAE patch/tile parallelism (decode).",
    )
+    parser.add_argument(
+        "--enforce-eager",
+        action="store_true",
+        help="Disable torch.compile and force eager execution.",
+    )
+    parser.add_argument(
+        "--audio-sample-rate",
+        type=int,
+        default=24000,
+        help="Sample rate for audio output when saved (default: 24000).",
+    )
+    parser.add_argument(
+        "--cache-backend",
+        type=str,
+        default=None,
+        choices=["cache_dit", "tea_cache"],
+        help=(
+            "Cache backend to use for acceleration. "
+            "Options: 'cache_dit' (DBCache + SCM + TaylorSeer), 'tea_cache' (Timestep Embedding Aware Cache). "
+            "Default: None (no cache acceleration)."
+        ),
+    )
    parser.add_argument(
        "--use-hsdp",
        action="store_true",
@@ -233,10 +201,9 @@ def parse_args() -> argparse.Namespace:
        ),
    )
    parser.add_argument(
-        "--pipeline-parallel-size",
-        type=int,
-        default=1,
-        help="Number of pipeline parallel stages.",
+        "--enable-diffusion-pipeline-profiler",
+        action="store_true",
+        help="Enable diffusion pipeline profiler to display stage durations.",
    )
    parser.add_argument(
        "--profiler-config",
@@ -334,18 +301,14 @@ def main():
        use_hsdp=args.use_hsdp,
        hsdp_shard_size=args.hsdp_shard_size,
        hsdp_replicate_size=args.hsdp_replicate_size,
-        pipeline_parallel_size=args.pipeline_parallel_size,
    )
-    omni_kwargs = dict(
+    omni = Omni(
        model=args.model,
        enable_layerwise_offload=args.enable_layerwise_offload,
        vae_use_slicing=args.vae_use_slicing,
        vae_use_tiling=args.vae_use_tiling,
        boundary_ratio=args.boundary_ratio,
        flow_shift=args.flow_shift,
-        diffusion_kv_cache_dtype=args.diffusion_kv_cache_dtype,
-        diffusion_kv_cache_skip_steps=args.diffusion_kv_cache_skip_steps,
-        diffusion_kv_cache_skip_layers=args.diffusion_kv_cache_skip_layers,
        enable_cpu_offload=args.enable_cpu_offload,
        parallel_config=parallel_config,
        enforce_eager=args.enforce_eager,
@@ -355,9 +318,6 @@ def main():
        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
        profiler_config=args.profiler_config,
    )
-    if args.quantization is not None:
-        omni_kwargs["quantization"] = args.quantization
-    omni = Omni(**omni_kwargs)

    if profiler_enabled:
        print("[Profiler] Starting profiling...")
@@ -370,13 +330,9 @@ def main():
    print(f"  Inference steps: {args.num_inference_steps}")
    print(f"  Frames: {args.num_frames}")
    print(f"  Solver: {args.sample_solver}")
-    print(f"  diffusion_kv_cache_dtype(config): {args.diffusion_kv_cache_dtype}")
-    print(f"  diffusion_kv_cache_skip_steps(config): {args.diffusion_kv_cache_skip_steps}")
-    print(f"  diffusion_kv_cache_skip_layers(config): {args.diffusion_kv_cache_skip_layers}")
    print(
        f"  Parallel configuration: cfg_parallel_size={args.cfg_parallel_size},"
-        f" tensor_parallel_size={args.tensor_parallel_size}, vae_patch_parallel_size={args.vae_patch_parallel_size},"
-        f" pipeline_parallel_size={args.pipeline_parallel_size}"
+        f" tensor_parallel_size={args.tensor_parallel_size}, vae_patch_parallel_size={args.vae_patch_parallel_size}"
    )
    print(f"  Video size: {args.width}x{args.height}")
    print(f"{'=' * 60}\n")
@@ -172,7 +172,7 @@ def run_pipeline_forward(

 def _synchronize(device: str) -> None:
    if device.startswith("cuda") and torch.cuda.is_available():
-        torch.accelerator.synchronize()
+        torch.cuda.synchronize()


 def _latency_summary(values_ms: list[float]) -> dict[str, float]:
@@ -1,208 +0,0 @@
-# SenseNova-U1-8B-MoT
-
-## Architecture
-
-SenseNova-U1 is a unified Qwen3-based LLM with Mixture-of-Tokenizers (MoT) attention. Unlike two-stage pipelines (e.g., BAGEL), it handles text encoding, optional reasoning ("think mode"), and flow-matching-based image denoising entirely within a single diffusion stage.
-
-| Feature | Description |
-| :------ | :---------- |
-| **Model type** | Unified LLM (single-stage diffusion pipeline) |
-| **Base LLM** | Qwen3 with MoT attention and 3D RoPE |
-| **Image generation** | Flow-matching Euler sampler, no separate VAE |
-| **Think mode** | Optional chain-of-thought reasoning before image generation |
-| **Parallelism** | Tensor Parallelism (TP) with fused QKV and fused gate/up projections |
-| **Modalities** | text2img, img2img, img2text (understanding), text2text |
-
-## Quick Start
-
-```bash
-cd examples/offline_inference/sensenova_u1
-
-# Text-to-image
-python end2end.py --prompt "A cute cat" --think
-
-# Image-to-image editing
-python end2end.py --prompt "Turn this into an oil painting" \
-                  --image input.png --think
-
-# Image understanding (img2text)
-python end2end.py --modality img2text \
-                  --prompt "Describe this image in detail" \
-                  --image photo.jpg
-
-# Text-to-text (pure chat)
-python end2end.py --modality text2text \
-                  --prompt "What is the capital of France?"
-
-# Custom resolution (image generation)
-python end2end.py --prompt "A futuristic cityscape at sunset" \
-                  --width 2048 --height 1024 --think
-```
-
-> **Note**: Default configuration works on a single **NVIDIA A100 (80GB)** or **H100** GPU.
-
-## Text-to-Image (t2i)
-
-Standard text-to-image with optional think mode:
-
-```bash
-python end2end.py \
-    --prompt "Close portrait of an elderly woman by a farmhouse window, textured skin, gentle smile, warm natural light" \
-    --width 1536 --height 2720 \
-    --think --print-think
-```
-
-## Image-to-Image Editing (img2img)
-
-Pass one or more `--image` paths to trigger img2img mode. The model uses the
-input image(s) as visual context for editing:
-
-```bash
-# Single input image
-python end2end.py \
-    --prompt "Add a sunset sky in the background" \
-    --image photo.jpg \
-    --width 2048 --height 2048 \
-    --think
-
-# Multiple reference images
-python end2end.py \
-    --prompt "Combine the style of Image-1 with the content of Image-2" \
-    --image style_ref.png content_ref.png \
-    --width 2048 --height 2048 \
-    --think
-```
-
-### Dual CFG for img2img
-
-img2img supports dual classifier-free guidance (text CFG + image CFG):
-
-```bash
-python end2end.py \
-    --prompt "Make the person smile" \
-    --image portrait.jpg \
-    --cfg-scale 4.0 \
-    --img-cfg-scale 2.0 \
-    --think
-```
-
-| `img_cfg_scale` | Behavior |
-| :-------------- | :------- |
-| `1.0` (default) | Text-only CFG: guidance = image-condition → full-condition |
-| `== cfg_scale` | Standard CFG: guidance = unconditional → full-condition |
-| Other value | Dual CFG: separate text and image guidance strengths |
-
-## Full Parameter Reference
-
-### General Parameters
-
-| Parameter | Default | Description |
-| :-------- | :------ | :---------- |
-| `--modality` | `auto` | Task modality: `auto`, `text2img`, `img2img`, `img2text`, `text2text` |
-| `--prompt` | "A cute cat..." | Text prompt / editing instruction / question |
-| `--image` | None | Input image path(s) for img2img or img2text |
-
-### Image Generation Parameters (text2img / img2img)
-
-| Parameter | Default | Description |
-| :-------- | :------ | :---------- |
-| `--height` | 2048 | Height of generated image (pixels) |
-| `--width` | 2048 | Width of generated image (pixels) |
-| `--seed` | 42 | Random seed for reproducibility |
-| `--num-steps` | 50 | Number of denoising steps |
-| `--cfg-scale` | 4.0 | Text classifier-free guidance scale |
-| `--img-cfg-scale` | 1.0 | Image CFG scale (img2img only, 1.0 = disabled) |
-| `--cfg-norm` | "none" | CFG normalization: none, global, channel, cfg_zero_star (t2i only) |
-| `--timestep-shift` | 3.0 | Timestep shift for flow-matching schedule |
-| `--t-eps` | 0.02 | Epsilon for timestep schedule |
-| `--think` | False | Enable think mode |
-| `--print-think` | False | Print think text to stdout |
-
-### Text Generation Parameters (text2text / img2text)
-
-| Parameter | Default | Description |
-| :-------- | :------ | :---------- |
-| `--max-tokens` | 512 | Maximum number of tokens to generate |
-| `--do-sample` | False | Use sampling instead of greedy decoding |
-| `--temperature` | 0.7 | Sampling temperature (higher = more diverse) |
-
-### Infrastructure Parameters
-
-| Parameter | Default | Description |
-| :-------- | :------ | :---------- |
-| `--model` | `SenseNova/SenseNova-U1-8B-MoT` | HuggingFace model ID or local path |
-| `--output` | `.` | Output directory for saved images |
-| `--tensor-parallel-size` | 1 | Number of GPUs for tensor parallelism |
-| `--enforce-eager` | False | Disable torch.compile |
-| `--enable-cpu-offload` | False | Enable module-wise (sequential) CPU offload to reduce peak VRAM |
-
-## Reducing GPU Memory Usage
-
-For hardware with limited VRAM, enable **module-wise CPU offload** with
-`--enable-cpu-offload`. The pipeline implements `SupportsModuleOffload`, so the
-vision encoder (`vision_model`) and the Qwen3 LLM (`language_model`) are
-swapped between CPU and GPU on demand:
-
- During text/vision encoding, the LLM is on CPU.
- During the diffusion loop, the vision encoder is on CPU.
- Lightweight FM modules stay resident on GPU.
-
-This lowers peak VRAM at the cost of extra CPU<->GPU transfers (use pinned
-memory) — useful for running SenseNova-U1 on consumer-grade GPUs.
-
-```bash
-# Text-to-image with CPU offload
-python end2end.py \
-    --prompt "A cute cat sitting on a windowsill" \
-    --width 2048 --height 2048 \
-    --enable-cpu-offload --think
-
-# Image-to-image editing with CPU offload
-python end2end.py \
-    --prompt "Turn this into an oil painting" \
-    --image input.png \
-    --width 2048 --height 2048 \
-    --enable-cpu-offload --think
-
-# Image understanding (img2text) with CPU offload
-python end2end.py \
-    --modality img2text \
-    --prompt "Describe this image in detail" \
-    --image photo.jpg \
-    --enable-cpu-offload
-```
-
-> **Notes**
-> - CPU offload is **single-GPU only** (incompatible with `--tensor-parallel-size > 1`).
-> - First-step latency is higher because of the cold-start CPU<->GPU transfers.
-> - For more details and other offloading strategies, see
->   [CPU Offloading for Diffusion Models](../../../docs/user_guide/diffusion/cpu_offload_diffusion.md).
-
-## Reproducing the E2E Test
-
-The following command reproduces the pixel-validated CI test case:
-
-```bash
-python end2end.py \
-    --prompt "Close portrait of an elderly woman by a farmhouse window, textured skin, gentle smile, warm natural light, emotional documentary look. The portrait should feel polished and natural, with sharp eyes, realistic skin texture, accurate facial anatomy, and premium lighting that keeps the face as the main focus." \
-    --width 1536 --height 2720 \
-    --seed 42 --num-steps 50 \
-    --cfg-scale 4.0 --timestep-shift 3.0 --cfg-norm none \
-    --think --print-think \
-    --output outputs
-```
-
-The corresponding pytest:
-
-```bash
-pytest -s -v tests/e2e/offline_inference/test_sensenova_u1_text2img.py \
-    -m "advanced_model" --run-level "advanced_model"
-```
-
-The img2img counterpart lives at
-`tests/e2e/offline_inference/test_sensenova_u1_img2img.py`.
-
-## Online Serving
-
-For OpenAI-compatible API serving, see
-[`examples/online_serving/sensenova_u1/`](../../examples/online_serving/sensenova_u1/).
@@ -1,281 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""
-End-to-end offline inference for SenseNova-U1-8B-MoT.
-
-Supports four modalities: text2img, img2img, img2text, text2text.
-
-Text-to-image:
-    python end2end.py --prompt "A cute cat" --think
-
-Image-to-image (editing):
-    python end2end.py --prompt "Turn this into an oil painting" \
-        --image input.png --think
-
-Image understanding:
-    python end2end.py --modality img2text \
-        --prompt "Describe this image" --image photo.jpg
-
-Text chat:
-    python end2end.py --modality text2text \
-        --prompt "What is the capital of France?"
-
-See README.md for more examples.
-"""
-
-import argparse
-import os
-
-from PIL import Image
-
-from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
-from vllm_omni.entrypoints.omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-
-def parse_args():
-    parser = argparse.ArgumentParser(
-        description="SenseNova-U1 text-to-image / image-to-image / text / understanding via vLLM-Omni.",
-    )
-    parser.add_argument(
-        "--model",
-        default="SenseNova/SenseNova-U1-8B-MoT",
-        help="HuggingFace model ID or local path.",
-    )
-    parser.add_argument(
-        "--modality",
-        default="auto",
-        choices=["auto", "text2img", "img2img", "img2text", "text2text"],
-        help="Task modality. 'auto' infers from --image (img2img/img2text if images, else t2i/t2t).",
-    )
-    parser.add_argument(
-        "--prompt",
-        default="A cute cat sitting on a windowsill, soft natural light",
-        help="Text prompt for generation or editing instruction.",
-    )
-    parser.add_argument(
-        "--image",
-        nargs="+",
-        metavar="PATH",
-        default=None,
-        help="Input image path(s) for img2img or img2text.",
-    )
-    parser.add_argument(
-        "--output",
-        type=str,
-        default=".",
-        help="Output directory to save generated images.",
-    )
-    # Image dimensions
-    parser.add_argument("--height", type=int, default=2048, help="Height of generated image.")
-    parser.add_argument("--width", type=int, default=2048, help="Width of generated image.")
-
-    # Generation parameters
-    parser.add_argument("--seed", type=int, default=42, help="Random seed for deterministic results.")
-    parser.add_argument("--num-steps", type=int, default=50, help="Number of denoising steps.")
-    parser.add_argument("--cfg-scale", type=float, default=4.0, help="Text classifier-free guidance scale.")
-    parser.add_argument(
-        "--img-cfg-scale",
-        type=float,
-        default=1.0,
-        help="Image CFG scale for img2img (1.0 = disabled). "
-        "When both --cfg-scale and --img-cfg-scale differ, dual CFG is used.",
-    )
-    parser.add_argument(
-        "--cfg-norm",
-        type=str,
-        default="none",
-        choices=["none", "global", "channel", "cfg_zero_star"],
-        help="CFG normalization mode. cfg_zero_star is only for t2i.",
-    )
-    parser.add_argument(
-        "--timestep-shift",
-        type=float,
-        default=3.0,
-        help="Timestep shift for flow-matching schedule.",
-    )
-    parser.add_argument(
-        "--t-eps",
-        type=float,
-        default=0.02,
-        help="Epsilon for flow-matching timestep schedule.",
-    )
-
-    # Text generation parameters (for text2text / img2text)
-    parser.add_argument(
-        "--max-tokens",
-        type=int,
-        default=2048,
-        help="Max tokens for text generation (text2text / img2text).",
-    )
-    parser.add_argument(
-        "--do-sample",
-        action="store_true",
-        help="Use sampling instead of greedy decoding for text generation.",
-    )
-    parser.add_argument(
-        "--temperature",
-        type=float,
-        default=0.7,
-        help="Sampling temperature for text generation.",
-    )
-
-    # Think mode
-    parser.add_argument(
-        "--think",
-        action="store_true",
-        help="Enable think mode: the model reasons about the prompt before generating the image.",
-    )
-    parser.add_argument(
-        "--print-think",
-        action="store_true",
-        help="Print the generated think text to stdout.",
-    )
-
-    # Advanced
-    parser.add_argument(
-        "--tensor-parallel-size",
-        type=int,
-        default=1,
-        help="Number of GPUs for tensor parallelism.",
-    )
-    parser.add_argument(
-        "--enforce-eager",
-        action="store_true",
-        help="Disable torch.compile and force eager execution.",
-    )
-    parser.add_argument(
-        "--enable-cpu-offload",
-        action="store_true",
-        help=(
-            "Enable model-level (sequential) CPU offloading. "
-            "Mutually exclusive between vision_model and language_model: "
-            "only one is on GPU at a time, reducing peak VRAM at the cost "
-            "of extra CPU<->GPU transfers."
-        ),
-    )
-
-    nullify_stage_engine_defaults(parser)
-    return parser.parse_args()
-
-
-def _resolve_modality(args):
-    """Determine the effective modality from CLI args."""
-    if args.modality != "auto":
-        return args.modality
-    has_images = args.image is not None
-    # Default: image-producing tasks for backward compatibility
-    return "img2img" if has_images else "text2img"
-
-
-def main():
-    args = parse_args()
-    os.makedirs(args.output, exist_ok=True)
-
-    modality = _resolve_modality(args)
-    is_text_output = modality in ("text2text", "img2text")
-    has_images = modality in ("img2img", "img2text")
-
-    omni = Omni(
-        model=args.model,
-        tensor_parallel_size=args.tensor_parallel_size,
-        enforce_eager=args.enforce_eager,
-        enable_cpu_offload=args.enable_cpu_offload,
-    )
-
-    extra_args = {
-        "cfg_scale": args.cfg_scale,
-        "cfg_norm": args.cfg_norm,
-        "timestep_shift": args.timestep_shift,
-        "cfg_interval": (0.0, 1.0),
-        "batch_size": 1,
-        "think": args.think,
-        "t_eps": args.t_eps,
-    }
-    if modality == "img2img":
-        extra_args["img_cfg_scale"] = args.img_cfg_scale
-    if is_text_output:
-        extra_args["max_tokens"] = args.max_tokens
-        extra_args["do_sample"] = args.do_sample
-        extra_args["temperature"] = args.temperature
-
-    sampling_params = OmniDiffusionSamplingParams(
-        height=args.height,
-        width=args.width,
-        seed=args.seed,
-        num_inference_steps=args.num_steps,
-        extra_args=extra_args,
-    )
-
-    print(f"\n{'=' * 60}")
-    print(f"SenseNova-U1 Configuration ({modality}):")
-    print(f"  Model          : {args.model}")
-    if not is_text_output:
-        print(f"  Image size     : {args.width}x{args.height}")
-        print(f"  Steps          : {args.num_steps}")
-        print(f"  CFG scale      : {args.cfg_scale}")
-    if modality == "img2img":
-        print(f"  Img CFG scale  : {args.img_cfg_scale}")
-    if has_images:
-        print(f"  Input images   : {args.image}")
-    if is_text_output:
-        print(f"  Max tokens     : {args.max_tokens}")
-        print(f"  Temperature    : {args.temperature}")
-    print(f"  Seed           : {args.seed}")
-    print(f"  Think mode     : {args.think}")
-    print(f"  TP size        : {args.tensor_parallel_size}")
-    print(f"{'=' * 60}\n")
-
-    # Build prompt dict
-    if has_images:
-        if not args.image:
-            raise ValueError(f"{modality} requires --image.")
-        input_images = [Image.open(p).convert("RGB") for p in args.image]
-        if is_text_output:
-            prompt_dict = {
-                "prompt": args.prompt,
-                "multi_modal_data": {"image": input_images},
-                "modalities": ["text"],
-            }
-        else:
-            prompt_dict = {
-                "prompt": args.prompt,
-                "multi_modal_data": {"image": input_images},
-                "modalities": ["img2img"],
-            }
-    elif is_text_output:
-        prompt_dict = {"prompt": args.prompt, "modalities": ["text"]}
-    else:
-        prompt_dict = {"prompt": args.prompt, "modalities": ["image"]}
-
-    outputs = list(
-        omni.generate(
-            prompts=prompt_dict,
-            sampling_params_list=sampling_params,
-        )
-    )
-
-    for req_output in outputs:
-        custom = getattr(req_output, "_custom_output", {}) or {}
-        if args.print_think and custom.get("think_text"):
-            print(f"[Think]\n{custom['think_text']}\n")
-
-        if is_text_output:
-            text = custom.get("text_output", "")
-            if not text:
-                text = getattr(req_output, "text", "") or ""
-            print(f"[Response]\n{text}")
-        else:
-            images = getattr(req_output, "images", None) or []
-            if not images:
-                print("[Warning] No images generated.")
-                continue
-            for j, img in enumerate(images):
-                save_path = os.path.join(args.output, f"sensenova_u1_output_{j}.png")
-                img.save(save_path)
-                print(f"[Output] Saved {img.size[0]}x{img.size[1]} image to {save_path}")
-
-
-if __name__ == "__main__":
-    main()
@@ -73,15 +73,13 @@ python examples/offline_inference/text_to_speech/cosyvoice3/end2end.py \
 ```

 ### Voice cloning
-If `--ref-audio` is omitted, the script downloads the upstream
-[`zero_shot_prompt.wav`](https://github.com/FunAudioLLM/CosyVoice/blob/main/asset/zero_shot_prompt.wav) from the CosyVoice repo into the current directory.
-To use your own clip, pass `--ref-audio /path/to/reference.wav`, and modify `--prompt-text` correspondingly.
+Pass a reference audio. Note that CosyVoice3's `--prompt-text` is a system-style prompt for the GPT stage, not a reference transcript:
 ```bash
 python examples/offline_inference/text_to_speech/cosyvoice3/end2end.py \
    --model pretrained_models/Fun-CosyVoice3-0.5B \
    --tokenizer pretrained_models/Fun-CosyVoice3-0.5B/CosyVoice-BlankEN \
-    --ref-audio /path/to/reference.wav \
-    --prompt-text "You are a helpful assistant.<|endofprompt|>Trascript in your ref audio clip"
+    --ref-audio prompt.wav \
+    --prompt-text "You are a helpful assistant.<|endofprompt|>Testing my voices. Why should I not?"
 ```

 ### Streaming
@@ -319,39 +317,44 @@ pip install voxcpm soundfile
 export VLLM_OMNI_VOXCPM_CODE_PATH=/path/to/VoxCPM/src
 ```

+If the native VoxCPM `config.json` does not contain HF metadata such as `model_type`, prepare a persistent HF-compatible config directory and point the stage configs to it via `VLLM_OMNI_VOXCPM_HF_CONFIG_PATH`:
+
+```bash
+export VOXCPM_MODEL=/path/to/voxcpm-model
+export VLLM_OMNI_VOXCPM_HF_CONFIG_PATH=/tmp/voxcpm_hf_config
+mkdir -p "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH"
+cp "$VOXCPM_MODEL/config.json" "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH/config.json"
+cp "$VOXCPM_MODEL/generation_config.json" "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH/generation_config.json" 2>/dev/null || true
+python3 -c 'import json, os; p=os.path.join(os.environ["VLLM_OMNI_VOXCPM_HF_CONFIG_PATH"], "config.json"); cfg=json.load(open(p, "r", encoding="utf-8")); cfg["model_type"]="voxcpm"; cfg.setdefault("architectures", ["VoxCPMForConditionalGeneration"]); json.dump(cfg, open(p, "w", encoding="utf-8"), indent=2, ensure_ascii=False)'
+```
+
 ### Quick start
 ```bash
 python examples/offline_inference/text_to_speech/voxcpm/end2end.py \
-    --model openbmb/VoxCPM-0.5B \
+    --model "$VOXCPM_MODEL" \
    --text "This is a split-stage VoxCPM synthesis example running on vLLM Omni."
 ```

 ### Voice cloning
 ```bash
 python examples/offline_inference/text_to_speech/voxcpm/end2end.py \
-    --model openbmb/VoxCPM-0.5B \
+    --model "$VOXCPM_MODEL" \
    --text "This sentence is synthesized with a cloned voice." \
    --ref-audio /path/to/reference.wav \
    --ref-text  "The exact transcript spoken in reference.wav."
 ```

 ### Streaming
-Pass `--streaming` together with the legacy async-chunk stage config:
+Pass the async-chunk stage config:
 ```bash
 python examples/offline_inference/text_to_speech/voxcpm/end2end.py \
-    --model openbmb/VoxCPM-0.5B \
-    --streaming \
+    --model "$VOXCPM_MODEL" \
    --stage-configs-path vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml \
    --text "This is a split-stage VoxCPM streaming example running on vLLM Omni."
 ```

-### Persistent HF config (optional)
-The engine auto-patches a missing `model_type` in the checkpoint's `config.json` in-memory at startup (see `vllm_omni/engine/arg_utils.py`), so no manual setup is required.
-
-If you'd rather maintain a persistent patched config dir (e.g. to share across runs or to add fields the auto-patcher doesn't set), export `VLLM_OMNI_VOXCPM_HF_CONFIG_PATH` to its location — both the bundled deploy and async-chunk yamls read it via `${oc.env:VLLM_OMNI_VOXCPM_HF_CONFIG_PATH,}`.
-
 ### Notes
- Non-streaming auto-loads `vllm_omni/deploy/voxcpm.yaml` from the model directory name / HF `model_type`. Streaming uses the legacy `voxcpm_async_chunk.yaml` passed via `--stage-configs-path`.
+- `voxcpm.yaml` is the default non-streaming stage config; `voxcpm_async_chunk.yaml` enables streaming.
 - Streaming is currently single-request oriented.
 - `--ref-text` must be the real transcript of `--ref-audio`; mismatched text degrades quality.
 - For online serving, see the [VoxCPM section in the online hub](../../online_serving/text_to_speech/README.md#voxcpm). For benchmark reporting, see [`benchmarks/voxcpm`](../../../benchmarks/voxcpm/README.md).
@@ -2,32 +2,18 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import argparse
 import os
-import urllib.request
-from pathlib import Path

 import numpy as np
 import soundfile as sf
 from vllm import SamplingParams
+from vllm.assets.audio import AudioAsset
 from vllm.multimodal.media.audio import load_audio

 from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
 from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.model_executor.models.cosyvoice3.config import CosyVoice3Config
 from vllm_omni.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokenizer
 from vllm_omni.model_executor.models.cosyvoice3.utils import extract_text_token
-from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config
-
-# Upstream zero-shot reference clip
-ZERO_SHOT_PROMPT_URL = "https://raw.githubusercontent.com/FunAudioLLM/CosyVoice/main/asset/zero_shot_prompt.wav"
-
-
-def _default_ref_audio() -> str:
-    # Download the upstream zero_shot_prompt.wav into the current dir
-    dest = Path("zero_shot_prompt.wav")
-    if not dest.exists() or dest.stat().st_size == 0:
-        print(f"Downloading default reference audio to {dest}")
-        urllib.request.urlretrieve(ZERO_SHOT_PROMPT_URL, dest)
-
-    return str(dest)


 def run_e2e():
@@ -50,15 +36,9 @@ def run_e2e():
    parser.add_argument(
        "--prompt-text",
        type=str,
-        default="You are a helpful assistant.<|endofprompt|>希望你以后，能够做的比我还好呦!",
-    )
-    parser.add_argument(
-        "--ref-audio",
-        type=str,
-        default=None,
-        help="Path to reference audio for voice cloning. "
-        "If unset, downloads the upstream CosyVoice3 zero-shot prompt audio clip",
+        default="You are a helpful assistant.<|endofprompt|>Testing my voices. Why should I not?",
    )
+    parser.add_argument("--ref-audio", type=str, default="prompt.wav")
    parser.add_argument(
        "--tokenizer",
        type=str,
@@ -86,22 +66,24 @@ def run_e2e():
    sampling_cfg = {"top_p": 0.8, "top_k": 25, "eos_token_id": 6561 + 1}

    print("Model initialized. Preparing inputs...")
-    ref_audio_path = args.ref_audio or _default_ref_audio()
-    if not os.path.exists(ref_audio_path):
-        raise FileNotFoundError(f"Audio file not found: {ref_audio_path}")
-    # Load at native sample rate
-    audio_signal, sr = load_audio(ref_audio_path, sr=None)
+    if args.ref_audio:
+        if not os.path.exists(args.ref_audio):
+            raise FileNotFoundError(f"Audio file not found: {args.ref_audio}")
+        # Load at native sample rate
+        audio_signal, sr = load_audio(args.ref_audio, sr=None)

-    # Validate sample rate before processing (similar to original CosyVoice)
-    min_sr = 16000
-    if sr < min_sr:
-        raise ValueError(
-            f"Audio sample rate {sr} Hz is too low. "
-            f"Minimum required: {min_sr} Hz. "
-            f"Please provide audio with sample rate >= {min_sr} Hz."
-        )
+        # Validate sample rate before processing (similar to original CosyVoice)
+        min_sr = 16000
+        if sr < min_sr:
+            raise ValueError(
+                f"Audio sample rate {sr} Hz is too low. "
+                f"Minimum required: {min_sr} Hz. "
+                f"Please provide audio with sample rate >= {min_sr} Hz."
+            )

-    audio_data = (audio_signal.astype(np.float32), sr)
+        audio_data = (audio_signal.astype(np.float32), sr)
+    else:
+        audio_data = AudioAsset("mary_had_lamb").audio_and_sample_rate

    prompts = {
        "prompt": args.text,
@@ -37,7 +37,7 @@ def run_e2e():
    parser.add_argument(
        "--stage-config",
        type=str,
-        default="vllm_omni/deploy/omnivoice.yaml",
+        default="vllm_omni/model_executor/stage_configs/omnivoice.yaml",
    )
    parser.add_argument(
        "--text",
@@ -14,6 +14,9 @@ from vllm.utils.argparse_utils import FlexibleArgumentParser
 from vllm_omni import AsyncOmni, Omni
 from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults

+REPO_ROOT = Path(__file__).resolve().parents[4]
+DEFAULT_SYNC_STAGE_CONFIG = REPO_ROOT / "vllm_omni" / "model_executor" / "stage_configs" / "voxcpm.yaml"
+

 def _build_prompt(args) -> dict[str, Any]:
    additional_information: dict[str, list[Any]] = {
@@ -56,6 +59,10 @@ def _extract_sample_rate(mm: dict[str, Any]) -> int:
    return int(sr_raw)


+def _is_streaming_stage_config(stage_config_path: str) -> bool:
+    return "async_chunk" in Path(stage_config_path).stem
+
+
 def _save_audio(audio: torch.Tensor, sample_rate: int, output_dir: Path, request_id: str) -> Path:
    output_dir.mkdir(parents=True, exist_ok=True)
    output_path = output_dir / f"output_{request_id}.wav"
@@ -80,7 +87,6 @@ async def _run_streaming(args) -> Path:
    omni = AsyncOmni(
        model=args.model,
        stage_configs_path=args.stage_configs_path,
-        deploy_config=args.deploy_config,
        log_stats=args.log_stats,
        stage_init_timeout=args.stage_init_timeout,
    )
@@ -127,7 +133,6 @@ def _run_sync(args) -> Path:
    omni = Omni(
        model=args.model,
        stage_configs_path=args.stage_configs_path,
-        deploy_config=args.deploy_config,
        log_stats=args.log_stats,
        stage_init_timeout=args.stage_init_timeout,
    )
@@ -159,27 +164,11 @@ def _run_sync(args) -> Path:
 def parse_args():
    parser = FlexibleArgumentParser(description="Minimal offline VoxCPM example for vLLM Omni.")
    parser.add_argument("--model", type=str, required=True, help="Local VoxCPM model directory.")
-    parser.add_argument(
-        "--deploy-config",
-        type=str,
-        default=None,
-        help=("Override the deploy config path."),
-    )
    parser.add_argument(
        "--stage-configs-path",
        type=str,
-        default=None,
-        help=(
-            "Legacy stage_args yaml path. Required for streaming "
-            "(vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml); "
-            "leave unset for non-streaming to use the auto-resolved deploy config."
-        ),
-    )
-    parser.add_argument(
-        "--streaming",
-        action="store_true",
-        default=False,
-        help="Stream audio chunks as they arrive via AsyncOmni.",
+        default=str(DEFAULT_SYNC_STAGE_CONFIG),
+        help=("Stage config path. Use voxcpm.yaml for non-streaming or voxcpm_async_chunk.yaml for streaming."),
    )
    parser.add_argument("--text", type=str, required=True, help="Input text for synthesis.")
    parser.add_argument("--ref-audio", type=str, default=None, help="Reference audio path for voice cloning.")
@@ -205,10 +194,9 @@ def parse_args():


 def main(args) -> None:
-    route = "streaming" if args.streaming else "sync"
+    route = "streaming" if _is_streaming_stage_config(args.stage_configs_path) else "sync"
    print(f"Model: {args.model}")
-    print(f"Deploy config: {args.deploy_config or '<auto from HF model_type>'}")
-    print(f"Stage configs path: {args.stage_configs_path or '<unused>'}")
+    print(f"Stage config: {args.stage_configs_path}")
    print(f"Route: {route}")
    if route == "streaming":
        asyncio.run(_run_streaming(args))
@@ -18,6 +18,7 @@ from vllm_omni import Omni
 from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults

 REPO_ROOT = Path(__file__).resolve().parents[4]
+DEFAULT_STAGE_CONFIGS_PATH = str(REPO_ROOT / "vllm_omni" / "model_executor" / "stage_configs" / "voxcpm2.yaml")
 SAMPLE_RATE = 48_000


@@ -42,11 +43,10 @@ def parse_args():
        help="Directory for output WAV files.",
    )
    parser.add_argument(
-        "--deploy-config",
+        "--stage-configs-path",
        type=str,
-        default=None,
-        help="Override the deploy config path. If unset, auto-loads "
-        "vllm_omni/deploy/voxcpm2.yaml based on the HF model_type.",
+        default=DEFAULT_STAGE_CONFIGS_PATH,
+        help="Path to the stage config YAML file.",
    )
    parser.add_argument(
        "--ref-audio",
@@ -96,7 +96,7 @@ def main():

    engine = Omni(
        model=args.model,
-        deploy_config=args.deploy_config,
+        stage_configs_path=args.stage_configs_path,
    )

    from transformers import AutoTokenizer
@@ -42,7 +42,6 @@ async def run_streaming(inputs, sampling_params_list, model_name, args, output_d
        model=model_name,
        deploy_config=args.deploy_config,
        log_stats=args.log_stats,
-        quantization=args.quantization,
    )

    # Normalize to a list so batch and single-input share the same code path
@@ -185,7 +184,7 @@ async def run_streaming(inputs, sampling_params_list, model_name, args, output_d
    )

    async_omni.shutdown()
-    torch.accelerator.empty_cache()
+    torch.cuda.empty_cache()
    gc.collect()


@@ -195,7 +194,6 @@ def run_non_streaming(inputs, sampling_params_list, model_name, args, output_dir
        model=model_name,
        log_stats=args.log_stats,
        deploy_config=args.deploy_config,
-        quantization=args.quantization,
    )

    if args.profiling_mode:
@@ -225,7 +223,7 @@ def run_non_streaming(inputs, sampling_params_list, model_name, args, output_dir
    print(f"RTF: {output_audio_dur / vllm_elapsed:.4f}")

    del llm
-    torch.accelerator.empty_cache()
+    torch.cuda.empty_cache()
    gc.collect()


@@ -307,12 +305,6 @@ def parse_args() -> Namespace:
        default=None,
        help="CFG alpha for flow-matching guidance (default: use value from stage config, typically 1.2).",
    )
-    parser.add_argument(
-        "--quantization",
-        type=str,
-        default=None,
-        help="Quantization method (e.g. 'fp8'). Applied to the language model only.",
-    )
    nullify_stage_engine_defaults(parser)
    return parser.parse_args()

@@ -387,7 +379,7 @@ def main(args: Any) -> None:
            f"--num-prompts ({args.num_prompts}) must be divisible by --concurrency ({args.concurrency})"
        )

-    torch.accelerator.empty_cache()
+    torch.cuda.empty_cache()
    gc.collect()

    if args.streaming:
@@ -142,32 +142,6 @@ def parse_args() -> argparse.Namespace:
        action="store_true",
        help="Enable layerwise (blockwise) offloading on DiT modules.",
    )
-    parser.add_argument(
-        "--audio-sample-rate",
-        type=int,
-        default=24000,
-        help="Sample rate for audio output when saved (default: 24000).",
-    )
-    parser.add_argument(
-        "--enable-diffusion-pipeline-profiler",
-        action="store_true",
-        help="Enable diffusion pipeline profiler to display stage durations.",
-    )
-    parser.add_argument(
-        "--profiler-config",
-        type=parse_profiler_config,
-        default=None,
-        help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.',
-    )
-    parser.add_argument(
-        "--quantization",
-        type=str,
-        default=None,
-        choices=["fp8", "mxfp8", "int8", "gguf"],
-        help="Quantization method for the transformer. mxfp8: W8A8 MXFP8 online quant (NPU). fp8: online FP8 (GPU).",
-    )
-
-    # Distributed and parallel execution
    parser.add_argument(
        "--ulysses-mode",
        type=str,
@@ -200,23 +174,41 @@ def parse_args() -> argparse.Namespace:
        default=1,
        help="Number of GPUs used for tensor parallelism (TP) inside the DiT.",
    )
+    parser.add_argument(
+        "--audio-sample-rate",
+        type=int,
+        default=24000,
+        help="Sample rate for audio output when saved (default: 24000).",
+    )
    parser.add_argument(
        "--vae-patch-parallel-size",
        type=int,
        default=1,
        help="Number of GPUs used for VAE patch/tile parallelism (decode).",
    )
-    parser.add_argument(
-        "--pipeline-parallel-size",
-        type=int,
-        default=1,
-        help="Number of pipeline parallel stages.",
-    )
    parser.add_argument(
        "--enable-expert-parallel",
        action="store_true",
        help="Enable expert parallelism for MoE layers.",
    )
+    parser.add_argument(
+        "--enable-diffusion-pipeline-profiler",
+        action="store_true",
+        help="Enable diffusion pipeline profiler to display stage durations.",
+    )
+    parser.add_argument(
+        "--profiler-config",
+        type=parse_profiler_config,
+        default=None,
+        help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.',
+    )
+    parser.add_argument(
+        "--quantization",
+        type=str,
+        default=None,
+        choices=["fp8", "gguf"],
+        help="Quantization method for the transformer (fp8 for online FP8 quantization).",
+    )
    parser.add_argument(
        "--use-hsdp",
        action="store_true",
@@ -270,7 +262,6 @@ def main():
        cfg_parallel_size=args.cfg_parallel_size,
        tensor_parallel_size=args.tensor_parallel_size,
        vae_patch_parallel_size=args.vae_patch_parallel_size,
-        pipeline_parallel_size=args.pipeline_parallel_size,
        enable_expert_parallel=args.enable_expert_parallel,
    )

@@ -316,8 +307,7 @@ def main():
    print(
        f"  Parallel configuration: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree},"
        f" cfg_parallel_size={args.cfg_parallel_size}, tensor_parallel_size={args.tensor_parallel_size},"
-        f" vae_patch_parallel_size={args.vae_patch_parallel_size}, pipeline_parallel_size={args.pipeline_parallel_size},"
-        f" enable_expert_parallel={args.enable_expert_parallel}"
+        f" vae_patch_parallel_size={args.vae_patch_parallel_size}, enable_expert_parallel={args.enable_expert_parallel}"
    )
    print(f"  Video size: {args.width}x{args.height}")
    print(f"{'=' * 60}\n")
@@ -68,25 +68,6 @@ vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 --tensor-parallel-size

 Or set `tensor_parallel_size` per stage in a custom deploy YAML.

-#### Hybrid Sharded Data Parallel (HSDP)
-
-For larger Bagel deployments on multiple GPUs, you can enable HSDP (Hybrid Sharded Data Parallel) by modifying the stage configuration (for example, [`bagel.yaml`](../../../vllm_omni/deploy/bagel.yaml)). HSDP shards transformer weights across GPUs to reduce per-GPU memory usage.
-
-1. **Enable HSDP**: Set `use_hsdp: true`.
-2. **Set shard size**: Set `hsdp_shard_size` to the number of GPUs used for sharding (for example, `4`).
-3. **Set replicate size**: Usually keep `hsdp_replicate_size: 1` unless you want replicated HSDP groups.
-4. **Set devices**: Specify the comma-separated GPU IDs used by the diffusion stage (for example, `"0,1,2,3"`).
-
-Example configuration for HSDP across 4 GPUs:
-```yaml
-  - stage_id: 1
-    devices: "0,1,2,3"
-    parallel_config:
-      use_hsdp: true
-      hsdp_shard_size: 4
-      hsdp_replicate_size: 1
-```
-
 ### Multi-Node Deployment

 Deploy each stage on a **separate node** for better resource utilization. Replace `<ORCHESTRATOR_IP>` with the actual IP address of your orchestrator node.
@@ -1,93 +0,0 @@
-# Covo-Audio-Chat
-
-## Installation
-
-Please refer to [README.md](../../../README.md)
-
-> **Note**
-> Covo-Audio code2wav requires `torchdiffeq`. Install it with: `pip install torchdiffeq`
-
-## Run examples (Covo-Audio-Chat)
-
-### Launch the Server
-
-```bash
-CUDA_VISIBLE_DEVICES=0 vllm serve tencent/Covo-Audio-Chat --omni \
-    --trust-remote-code --port 18091
-```
-
-> The default deploy config is located at `vllm_omni/deploy/covo_audio.yaml` and is loaded automatically by the model registry — no `--deploy-config` flag needed for default use.
-
-### Send Request
-
-Get into the example folder:
-```bash
-cd examples/online_serving/covo_audio
-```
-
-#### Audio input chat
-
-Using the default audio asset:
-```bash
-python openai_chat_completion_client.py
-```
-
-Using a custom audio file:
-```bash
-python openai_chat_completion_client.py --audio-path /path/to/audio.wav
-```
-
-#### Streaming mode
-
-```bash
-python openai_chat_completion_client.py --audio-path /path/to/audio.wav --stream
-```
-
-### Command-line Arguments
-
-The Python client supports the following arguments:
-
- `--prompt` (or `-p`): Text prompt (default: `"请回答这段音频里的问题。"`)
- `--audio-path`: Path to audio file. Supports local file paths or HTTP/HTTPS URLs. Common audio formats: WAV, MP3, OGG, FLAC, M4A. Uses the default audio asset if omitted
- `--output-audio-path` (or `-o`): Output path for generated audio (default: `./audio_0.wav`)
- `--model`: Model name for the API. Auto-detected from server if omitted
- `--stream`: Enable streaming mode
-
-## Pipeline
-
-Covo-Audio-Chat uses a 2-stage pipeline:
-
- **Stage 0 (fused_thinker_talker):** The 7B LLM generates interleaved text and audio tokens in a single autoregressive pass.
- **Stage 1 (code2wav):** A BigVGAN-based vocoder converts the extracted audio codes into a 24kHz WAV waveform.
-
-## Environment
-
- GPU: 1x H100 (80 GiB)
- Stage 0 (7B LLM): ~16 GiB VRAM
- Stage 1 (BigVGAN vocoder): ~2 GiB VRAM
-
-## Important: System Prompt
-
-Covo-Audio-Chat requires a specific system prompt to enable interleaved text+audio generation. Without it, Stage 0 produces text-only tokens, and the output audio will be **silent**.
-
-The system prompt must include the instruction:
-
-```
-请用文本和音频进行对话，交替生成5个文本token和15个音频token，音频部分使用发音人：default_female
-```
-
-The full system prompt is defined in `vllm_omni/model_executor/models/covo_audio/prompt_utils.py` (`COVO_AUDIO_SYSTEM_PROMPT`).
-
-## FAQ
-
-If you encounter `ModuleNotFoundError: No module named 'librosa'`, install it with:
-```bash
-pip install librosa
-```
-
-### Audio output is silent (0.1s, all zeros)
-
-This means Stage 0 did not generate any audio tokens. Most likely causes:
-
-1. **Missing system prompt** -- see the "Important: System Prompt" section above. The model needs the interleaved generation instruction in the system message.
-2. **Missing `ignore_eos: true` / `stop_token_ids: [151645]`** in Stage 0 sampling params -- without these, the model stops at `<|endoftext|>` before generating audio tokens.
@@ -1,216 +0,0 @@
-"""
-Client example for Covo-Audio-Chat with vllm-omni.
-
-Usage
-----
-# Start the server first:
-#   CUDA_VISIBLE_DEVICES=0 vllm serve tencent/Covo-Audio-Chat --omni --trust-remote-code
-
-# Audio input chat (uses default audio asset if --audio-path not provided):
-python openai_chat_completion_client.py
-python openai_chat_completion_client.py --audio-path /path/to/audio.wav
-
-# Streaming mode:
-python openai_chat_completion_client.py --stream
-"""
-
-import base64
-import os
-import sys
-import time
-
-from openai import OpenAI
-from vllm.assets.audio import AudioAsset
-from vllm.utils.argparse_utils import FlexibleArgumentParser
-
-from vllm_omni.model_executor.models.covo_audio.prompt_utils import (
-    COVO_AUDIO_SYSTEM_PROMPT,
-)
-
-openai_api_key = "EMPTY"
-openai_api_base = "http://localhost:18091/v1"
-
-client = OpenAI(api_key=openai_api_key, base_url=openai_api_base)
-
-SEED = 42
-
-
-def encode_base64_content_from_file(file_path: str) -> str:
-    """Encode a local file to base64 format."""
-    with open(file_path, "rb") as f:
-        return base64.b64encode(f.read()).decode("utf-8")
-
-
-def get_audio_url(audio_path: str | None) -> str:
-    """Convert an audio path to a data URL for the API."""
-    if not audio_path:
-        return AudioAsset("mary_had_lamb").url
-
-    if audio_path.startswith(("data:", "http://", "https://")):
-        return audio_path
-
-    if not os.path.exists(audio_path):
-        raise FileNotFoundError(f"Audio file not found: {audio_path}")
-
-    ext = audio_path.rsplit(".", 1)[-1].lower() if "." in audio_path else "wav"
-    mime_map = {
-        "wav": "audio/wav",
-        "mp3": "audio/mpeg",
-        "mpeg": "audio/mpeg",
-        "ogg": "audio/ogg",
-        "flac": "audio/flac",
-        "m4a": "audio/mp4",
-    }
-    mime = mime_map.get(ext, "audio/wav")
-    b64 = encode_base64_content_from_file(audio_path)
-    return f"data:{mime};base64,{b64}"
-
-
-def build_audio_messages(audio_path: str | None, prompt: str) -> list[dict]:
-    """Build messages for an audio chat request."""
-    audio_url = get_audio_url(audio_path)
-    return [
-        {"role": "system", "content": COVO_AUDIO_SYSTEM_PROMPT},
-        {
-            "role": "user",
-            "content": [
-                {"type": "audio_url", "audio_url": {"url": audio_url}},
-                {"type": "text", "text": prompt},
-            ],
-        },
-    ]
-
-
-def get_model_name(override: str | None = None) -> str:
-    """Return the model name to use in API calls."""
-    if override:
-        return override
-    # Auto-detect from the server
-    models = client.models.list()
-    return models.data[0].id
-
-
-def run(args) -> None:
-    model_name = get_model_name(args.model)
-    thinker_params = {
-        "temperature": 0.0,
-        "top_p": 1.0,
-        "top_k": -1,
-        "max_tokens": 2048,
-        "seed": SEED,
-        "detokenize": True,
-        "repetition_penalty": 1.05,
-        "stop_token_ids": [151645],
-        "ignore_eos": True,
-    }
-    code2wav_params = {
-        "temperature": 0.0,
-        "top_p": 1.0,
-        "top_k": -1,
-        "max_tokens": 2048,
-        "seed": SEED,
-        "detokenize": False,
-        "repetition_penalty": 1.05,
-    }
-
-    extra_body = {
-        "sampling_params_list": [thinker_params, code2wav_params],
-    }
-
-    messages = build_audio_messages(args.audio_path, args.prompt)
-
-    start = time.perf_counter()
-
-    if not args.stream:
-        completion = client.chat.completions.create(
-            messages=messages,
-            model=model_name,
-            extra_body=extra_body,
-            stream=False,
-        )
-        elapsed = time.perf_counter() - start
-
-        for choice in completion.choices:
-            if choice.message.audio:
-                audio_bytes = base64.b64decode(choice.message.audio.data)
-                out_path = args.output_audio_path
-                os.makedirs(os.path.dirname(out_path) or ".", exist_ok=True)
-                with open(out_path, "wb") as f:
-                    f.write(audio_bytes)
-                print(f"Audio saved to {out_path} ({len(audio_bytes)} bytes)")
-            elif choice.message.content:
-                print("Chat completion output from text:", choice.message.content)
-
-        print(f"Elapsed: {elapsed:.2f}s", file=sys.stderr)
-
-    else:
-        stream = client.chat.completions.create(
-            messages=messages,
-            model=model_name,
-            extra_body=extra_body,
-            stream=True,
-        )
-        audio_idx = 0
-        for chunk in stream:
-            modality = getattr(chunk, "modality", None)
-            elapsed = time.perf_counter() - start
-            for choice in chunk.choices:
-                content = getattr(choice.delta, "content", None) if hasattr(choice, "delta") else None
-                if not content:
-                    continue
-                if modality == "audio":
-                    audio_bytes = base64.b64decode(content)
-                    out_dir = os.path.dirname(args.output_audio_path) or "."
-                    out_path = os.path.join(out_dir, f"audio_{audio_idx}.wav")
-                    os.makedirs(out_dir, exist_ok=True)
-                    with open(out_path, "wb") as f:
-                        f.write(audio_bytes)
-                    print(f"[{elapsed:.2f}s] Audio chunk saved to {out_path}", file=sys.stderr)
-                    audio_idx += 1
-                elif modality == "text":
-                    print(content, end="", flush=True)
-
-        print()  # trailing newline
-        elapsed = time.perf_counter() - start
-        print(f"Elapsed: {elapsed:.2f}s", file=sys.stderr)
-
-
-def parse_args():
-    parser = FlexibleArgumentParser(description="Covo-Audio-Chat client example")
-    parser.add_argument(
-        "--prompt",
-        "-p",
-        type=str,
-        default="请回答这段音频里的问题。",
-        help="Text prompt.",
-    )
-    parser.add_argument(
-        "--audio-path",
-        type=str,
-        default=None,
-        help="Path to audio file. Uses default audio asset if omitted.",
-    )
-    parser.add_argument(
-        "--output-audio-path",
-        "-o",
-        type=str,
-        default="./audio_0.wav",
-        help="Output path for generated audio.",
-    )
-    parser.add_argument(
-        "--model",
-        type=str,
-        default=None,
-        help="Model name for the API. Auto-detected from server if omitted.",
-    )
-    parser.add_argument(
-        "--stream",
-        action="store_true",
-        help="Enable streaming mode.",
-    )
-    return parser.parse_args()
-
-
-if __name__ == "__main__":
-    args = parse_args()
-    run(args)