Revert "[Perf] Optimize RMSNorm in Z-Image (#3304)"

This reverts commit bb69cbc9f1.

.buildkite/bootstrap-amd-omni.sh

@@ -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
 
 # ----------------------------------------------------------------------

.buildkite/pipeline-npu-a3.yaml

@@ -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"

.buildkite/pipeline-npu.yaml

@@ -1,108 +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"
-  #   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"

.buildkite/scripts/upload_pipeline_with_skip_ci.sh

@@ -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)

.buildkite/test-amd-merge.yml

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

.buildkite/test-amd-ready.yaml

@@ -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"
       '

.buildkite/test-amd.yaml

@@ -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"

.buildkite/test-nightly.yml

@@ -13,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:
@@ -169,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
@@ -308,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:
@@ -352,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:
@@ -422,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:
@@ -515,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:
@@ -602,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
@@ -624,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"
@@ -781,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:
@@ -925,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
@@ -1078,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"
@@ -1161,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
@@ -1237,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
@@ -1266,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"
@@ -1282,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

.buildkite/test-npu-nightly.yaml

@@ -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"

.claude/skills/add-diffusion-model/SKILL.md

@@ -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.
 

.claude/skills/add-diffusion-model/references/parallelism-patterns.md

@@ -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):

.claude/skills/add-diffusion-model/references/transformer-adaptation.md

@@ -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",
         )
 ```
 

.claude/skills/add-diffusion-model/references/troubleshooting.md

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

.github/CODEOWNERS

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

.gitignore

@@ -239,7 +239,6 @@ data/
 datasets/
 *.csv
 *.json
-!tests/dfx/perf/tests/*.json
 !apps/ComfyUI-vLLM-Omni/example_workflows/*.json
 *.jsonl
 *.parquet

README.md

@@ -56,9 +56,8 @@ vLLM-Omni is flexible and easy to use with:
 
 vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, including:
 
-- **Omni-modality models** (e.g. Qwen3-Omni, Cosmos, HunyuanImage, BAGEL)
-- **TTS models** (e.g. Qwen3-TTS, VoxCPM2, Ming-Omni-TTS, CosyVoice3)
-- **Diffusion models** — image, video, and audio generation (e.g. Qwen-Image, Wan2.2, FLUX)
+- Omni-modality models (e.g. Qwen-Omni)
+- Multi-modality generation models (e.g. Qwen-Image)
 
 ## Getting Started
 

benchmarks/diffusion/README.md

@@ -32,6 +32,8 @@ python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
 
 - 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,6 +95,35 @@ 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 (the script calls `.../v1/chat/completions`).

benchmarks/diffusion/quantization_quality.py

@@ -151,7 +151,7 @@ def _generate_image(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},
@@ -163,7 +163,7 @@ 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)
 
     first = outputs[0]
     req_out = first.request_output[0] if hasattr(first, "request_output") else first
@@ -178,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": ""},
@@ -192,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):
@@ -228,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):

benchmarks/distributed/omni_connectors/cross_node_mooncake_transfer_engine.py

@@ -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)}")
 

benchmarks/tts/bench_voxcpm_offline.py

@@ -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__)
 

docker/Dockerfile.npu.ci

@@ -1,42 +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 && \
-    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" \
-      --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 []

docker/Dockerfile.npu.ci.a3

@@ -1,42 +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 && \
-    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" \
-      --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 []

docs/.nav.yml

@@ -76,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

docs/contributing/ci/CI_5levels.md

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

docs/contributing/ci/test_guide.md

@@ -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
 
     ```
 

docs/contributing/ci/tests_markers.md

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

docs/contributing/ci/tests_style.md

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

docs/contributing/model/adding_diffusion_model.md

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

docs/design/module/dit_module.md

@@ -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:

docs/models/supported_models.md

@@ -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.

docs/serving/speech_api.md

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

docs/user_guide/diffusion/attention_backends.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`:
 

docs/user_guide/diffusion/cpu_offload_diffusion.md

@@ -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.

docs/user_guide/diffusion/multi_replica_serving.md

@@ -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.

docs/user_guide/diffusion_features.md

@@ -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
 
@@ -43,6 +43,7 @@ Parallelism methods distribute computation across GPUs without quality loss (mat
 | **[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
 
@@ -107,7 +108,7 @@ The following tables show which models support each feature:
 
 | Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
 |-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|:----------------:|
-| **Bagel** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅  | ❌ | ❌ | ❌ |
+| **Bagel** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅  | ❌ | ❌ | ❌ |
 | **FLUX.1-dev** | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
 | **FLUX.1-schnell** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
 | **FLUX.2-klein** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
@@ -127,7 +128,6 @@ The following tables show which models support each feature:
 | **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** | ❌ | ✅ | ✅ | ❓ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |

docs/user_guide/quantization/fp8.md

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

docs/user_guide/quantization/modelopt.md

@@ -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.

docs/user_guide/quantization/mxfp8.md

@@ -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`.

docs/user_guide/quantization/overview.md

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

docs/user_guide/quantization/quantized_kvcache.md

@@ -1,115 +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 `kv_cache_dtype`, matching the option name
-used by vLLM's language-model KV-cache quantization. In vLLM-Omni diffusion
-pipelines, however, it 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 `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 `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 `kv_cache_dtype="fp8"`; tune `kv_cache_skip_steps` and `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,
-`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 \
-    --kv-cache-dtype fp8 \
-    --kv-cache-skip-steps "0,1" \
-    --kv-cache-skip-layers "0-2"
-```
-
-Online serving:
-
-```bash
-vllm serve <your-model> --omni --kv-cache-dtype fp8
-```
-
-Stage config:
-
-```yaml
-stage_args:
-  - stage_id: 0
-    stage_type: diffusion
-    engine_args:
-      model_stage: dit
-      kv_cache_dtype: "fp8"
-      kv_cache_skip_steps: "0,1"
-      kv_cache_skip_layers: "0-2"
-```
-
-## Parameters
-
-| Parameter | Type | Default | Description |
-|-----------|------|---------|-------------|
-| `kv_cache_dtype` | str \| None | `None` | Set to `"fp8"` to enable dynamic FP8 FA on supported attention backends |
-| `kv_cache_skip_steps` | str \| None | `None` | Denoising step selector to keep in native dtype, for example `"0,1,4-6"` |
-| `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 `kv_cache_skip_steps` for denoising steps where quality is more
-   sensitive.
-3. Use `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.

examples/offline_inference/bagel/README.md

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

examples/offline_inference/covo_audio/README.md

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

examples/offline_inference/covo_audio/end2end.py

@@ -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)

examples/offline_inference/hunyuan_image3/README.md

@@ -1,169 +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
 
-| 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` | Prompt behavior. `auto` selects the default from `--modality`; `think` adds `<think>`; `recaption` adds `<recaption>`; `vanilla` uses the text-to-image pretrain template. |
-| `--sys-type` | Override the system prompt type, for example `en_unified` or `en_vanilla`. |
-| `--vae-use-tiling` | Enable VAE tiling for memory reduction. |
+#### 📌 Command Line Arguments (end2end.py)
 
-## Notes
+| 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) |
 
-- `hunyuan_image3_ar.yaml` is a 4-card AR-only text/comprehension deploy. It sets `engine_output_type: text`, `final_output_type: text`, and text sampling defaults.
-- `hunyuan_image3_dit.yaml` is a single-stage DiT deploy with `stage_id: 0`; it does not require stage 1 or a running AR stage.
-- The old HunyuanImage3 YAMLs under `model_executor/stage_configs/` and `platforms/*/stage_configs/` have been folded into the deploy YAMLs.
-- This PR does not keep the HunyuanImage3 AR-to-DiT KV reuse wiring. The deploy YAMLs describe the topology and platform settings only.
+------
+
+#### ⚙️ 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 based on task.
-- `t2i_vanilla` is the only task that uses the bare pretrain template without chat structure.
-- The example composes the internal prompt task from `--modality` and `--bot-task`
-  before calling `prompt_utils`; for example, `img2text + think` becomes
-  `i2t_think` for prompt and stop-token lookup.
+- `<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 deploy YAML, use a smaller `max_num_batched_tokens`, or enable VAE tiling with `--vae-use-tiling`.
+- **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 |
+| Stage             | VRAM (approx)        |
+| :---------------- | :------------------- |
+| Stage 0 (AR)      | ~15 GiB + KV Cache   |
+| Stage 1 (DiT)     | ~30 GiB              |
+| Total (8-GPU)     | ~45 GiB + KV Cache   |

examples/offline_inference/hunyuan_image3/end2end.py

@@ -14,45 +14,43 @@ Usage:
 """
 
 import argparse
-import json
 import os
-from pathlib import Path
 
 from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import (
-    _TASK_PRESETS,
     build_prompt_tokens,
-    resolve_stop_token_ids,
 )
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniPromptType
 
-# Default deploy configs are absolute so this example works from any cwd.
-_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_DEFAULT_DEPLOY_CONFIG = {
-    "text2img": _DEFAULT_DEPLOY_CONFIG,
-    "img2img": _DEFAULT_DEPLOY_CONFIG,
-    "img2text": _DEFAULT_AR_DEPLOY_CONFIG,
-    "text2text": _DEFAULT_AR_DEPLOY_CONFIG,
-}
-
-_MODALITY_MODE = {
-    "text2img": "text-to-image",
-    "img2img": "image-editing",
-    "img2text": "image-to-text",
-    "text2text": "text-to-text",
+# 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 → prompt_utils task mapping
 _MODALITY_TASK_MAP = {
-    "text2img": "t2i",
-    "img2img": "it2i",
+    "text2img": "t2i_think",
+    "img2img": "it2i_think",
     "img2text": "i2t",
     "text2text": "t2t",
 }
 
 
+# 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(
@@ -96,13 +94,8 @@ def parse_args():
     parser.add_argument(
         "--bot-task",
         type=str,
-        default="auto",
-        choices=["auto", "think", "recaption", "think_recaption", "vanilla"],
-        help=(
-            "Prompt behavior. 'auto' selects the default for the modality; "
-            "'think' adds <think>; 'recaption' adds <recaption>; "
-            "'vanilla' uses the t2i pretrain template."
-        ),
+        default=None,
+        help="Override prompt task (e.g. it2i_think, t2i_recaption). Default: auto from modality.",
     )
     parser.add_argument(
         "--sys-type",
@@ -112,22 +105,10 @@ def parse_args():
     )
 
     # Omni init args
-    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("--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(
-        "--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."
-        ),
-    )
 
     from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
 
@@ -135,69 +116,27 @@ def parse_args():
     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)
 
-    # Determine task for prompt formatting from modality + bot behavior.
-    task = _MODALITY_TASK_MAP[args.modality]
-    assert task is not None
-    bot_task = args.bot_task
-    if bot_task != "auto":
-        task = task + "_" + bot_task
-    if task not in _TASK_PRESETS:
-        valid_bot_tasks = {
-            "text2img": ["think", "recaption", "vanilla"],
-            "img2img": ["think", "recaption", "think_recaption"],
-            "img2text": ["auto"],
-            "text2text": ["auto"],
-        }[args.modality]
-        raise ValueError(
-            f"--bot-task {bot_task!r} is not supported for {args.modality}. Choose from: {valid_bot_tasks}"
-        )
+    # 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 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,
     }
-
-    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
-    omni_kwargs["mode"] = _MODALITY_MODE[args.modality]
+    if args.modality in ("text2img", "img2img"):
+        omni_kwargs["mode"] = "text-to-image"
 
     omni = Omni(**omni_kwargs)
 
@@ -225,9 +164,9 @@ def main():
     # Format prompts
     formatted_prompts: list[OmniPromptType] = []
     for p in prompts:
-        result = build_prompt_tokens(p, tokenizer, task=task, sys_type=args.sys_type)
-        token_ids = result.token_ids
-        effective_sys_type = result.system_prompt_type
+        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
@@ -260,32 +199,22 @@ def main():
     # 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)
-    assert ar_stop_token_ids is not None
-    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 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}")
@@ -295,8 +224,6 @@ def main():
         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")
 
@@ -308,17 +235,10 @@ def main():
     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)

examples/offline_inference/image_to_video/image_to_video.py

@@ -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(
-        "--kv-cache-dtype",
-        type=str,
-        default=None,
-        help="Config-level KV cache dtype (e.g. float8_e4m3fn).",
-    )
-    parser.add_argument(
-        "--kv-cache-skip-steps",
-        type=str,
-        default=None,
-        help="Config-level KV-cache quantization skip-step selector, e.g. '0-9,20,25-30'.",
-    )
-    parser.add_argument(
-        "--kv-cache-skip-layers",
-        type=str,
-        default=None,
-        help="Config-level 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(
@@ -218,13 +200,6 @@ def parse_args() -> argparse.Namespace:
             "Default 1 means pure sharding (no replication). "
         ),
     )
-    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).",
-    )
     parser.add_argument(
         "--enable-diffusion-pipeline-profiler",
         action="store_true",
@@ -327,16 +302,13 @@ def main():
         hsdp_shard_size=args.hsdp_shard_size,
         hsdp_replicate_size=args.hsdp_replicate_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,
-        kv_cache_dtype=args.kv_cache_dtype,
-        kv_cache_skip_steps=args.kv_cache_skip_steps,
-        kv_cache_skip_layers=args.kv_cache_skip_layers,
         enable_cpu_offload=args.enable_cpu_offload,
         parallel_config=parallel_config,
         enforce_eager=args.enforce_eager,
@@ -346,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...")
@@ -361,9 +330,6 @@ def main():
     print(f"  Inference steps: {args.num_inference_steps}")
     print(f"  Frames: {args.num_frames}")
     print(f"  Solver: {args.sample_solver}")
-    print(f"  kv_cache_dtype(config): {args.kv_cache_dtype}")
-    print(f"  kv_cache_skip_steps(config): {args.kv_cache_skip_steps}")
-    print(f"  kv_cache_skip_layers(config): {args.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}"

examples/offline_inference/internvla_a1/end2end.py

@@ -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]:

examples/offline_inference/sensenova_u1/README.md

@@ -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/).

examples/offline_inference/sensenova_u1/end2end.py

@@ -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()

examples/offline_inference/text_to_speech/omnivoice/end2end.py

@@ -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",

examples/offline_inference/text_to_speech/voxcpm/end2end.py

@@ -11,11 +11,11 @@ import soundfile as sf
 import torch
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 
-from tests.helpers.stage_config import get_deploy_config_path
 from vllm_omni import AsyncOmni, Omni
 from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
 
-DEFAULT_SYNC_STAGE_CONFIG = get_deploy_config_path("voxcpm.yaml")
+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]:
@@ -167,7 +167,7 @@ def parse_args():
     parser.add_argument(
         "--stage-configs-path",
         type=str,
-        default=DEFAULT_SYNC_STAGE_CONFIG,
+        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.")

examples/offline_inference/text_to_speech/voxcpm2/end2end.py

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

examples/offline_inference/text_to_speech/voxtral_tts/end2end.py

@@ -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:

examples/offline_inference/text_to_video/text_to_video.py

@@ -206,8 +206,8 @@ def parse_args() -> argparse.Namespace:
         "--quantization",
         type=str,
         default=None,
-        choices=["fp8", "mxfp8", "int8", "gguf"],
-        help="Quantization method for the transformer. mxfp8: W8A8 MXFP8 online quant (NPU). fp8: online FP8 (GPU).",
+        choices=["fp8", "gguf"],
+        help="Quantization method for the transformer (fp8 for online FP8 quantization).",
     )
     parser.add_argument(
         "--use-hsdp",

examples/online_serving/bagel/README.md

@@ -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.

examples/online_serving/covo_audio/README.md

@@ -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.

examples/online_serving/covo_audio/openai_chat_completion_client.py

@@ -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)

examples/online_serving/sensenova_u1/README.md

@@ -1,163 +0,0 @@
-# SenseNova-U1 Online Serving
-
-## Launch the Server
-
-```bash
-vllm serve SenseNova/SenseNova-U1-8B-MoT --omni --port 8091
-```
-
-Or use the convenience script:
-
-```bash
-cd examples/online_serving/sensenova_u1
-bash run_server.sh
-```
-
-### Tensor Parallelism (TP)
-
-```bash
-vllm serve SenseNova/SenseNova-U1-8B-MoT --omni --port 8091 \
-    --tensor-parallel-size 2
-```
-
-## Send Requests
-
-```bash
-cd examples/online_serving/sensenova_u1
-```
-
-### Text to Image (text2img)
-
-**Python client:**
-
-```bash
-python openai_chat_client.py \
-    --prompt "A beautiful sunset over mountains" \
-    --modality text2img \
-    --height 2048 --width 2048 --num-steps 50
-```
-
-**curl:**
-
-```bash
-curl http://localhost:8091/v1/chat/completions \
-  -H "Content-Type: application/json" \
-  -d '{
-    "messages": [{"role": "user", "content": [{"type": "text", "text": "A beautiful sunset over mountains"}]}],
-    "modalities": ["image"],
-    "height": 2048,
-    "width": 2048,
-    "num_inference_steps": 50,
-    "seed": 42
-  }'
-```
-
-### Image to Image (img2img)
-
-**Python client:**
-
-```bash
-python openai_chat_client.py \
-    --prompt "Turn this into an oil painting" \
-    --modality img2img \
-    --image-url /path/to/input.jpg \
-    --height 2048 --width 2048
-```
-
-**curl:**
-
-```bash
-IMAGE_BASE64=$(base64 -w 0 input.jpg)
-
-cat <<EOF > payload.json
-{
-    "messages": [{
-      "role": "user",
-      "content": [
-        {"type": "text", "text": "Turn this into an oil painting"},
-        {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,${IMAGE_BASE64}"}}
-      ]
-    }],
-    "modalities": ["image"],
-    "height": 2048,
-    "width": 2048,
-    "num_inference_steps": 50,
-    "seed": 42
-}
-EOF
-
-curl http://localhost:8091/v1/chat/completions \
-  -H "Content-Type: application/json" \
-  -d @payload.json
-```
-
-### Image to Text (img2text)
-
-**Python client:**
-
-```bash
-python openai_chat_client.py \
-    --prompt "Describe this image in detail" \
-    --modality img2text \
-    --image-url /path/to/image.jpg
-```
-
-**curl:**
-
-```bash
-IMAGE_BASE64=$(base64 -w 0 image.jpg)
-
-cat <<EOF > payload.json
-{
-  "messages": [{
-    "role": "user",
-    "content": [
-      {"type": "text", "text": "Describe this image in detail"},
-      {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,${IMAGE_BASE64}"}}
-    ]
-  }],
-  "modalities": ["text"]
-}
-EOF
-
-curl http://localhost:8091/v1/chat/completions \
-  -H "Content-Type: application/json" \
-  -d @payload.json
-```
-
-### Text to Text (text2text)
-
-**Python client:**
-
-```bash
-python openai_chat_client.py \
-    --prompt "What is the capital of France?" \
-    --modality text2text
-```
-
-**curl:**
-
-```bash
-curl http://localhost:8091/v1/chat/completions \
-  -H "Content-Type: application/json" \
-  -d '{
-    "messages": [{"role": "user", "content": [{"type": "text", "text": "What is the capital of France?"}]}],
-    "modalities": ["text"]
-  }'
-```
-
-## Python Client Arguments
-
-| Argument | Default | Description |
-| :------- | :------ | :---------- |
-| `--prompt` / `-p` | `A cute cat` | Text prompt |
-| `--output` / `-o` | `sensenova_u1_output.png` | Output file path |
-| `--server` / `-s` | `http://localhost:8091` | Server URL |
-| `--image-url` / `-i` | `None` | Input image URL or local path (img2img/img2text) |
-| `--modality` / `-m` | `text2img` | `text2img`, `img2img`, `img2text`, `text2text` |
-| `--height` | `2048` | Image height (image generation only) |
-| `--width` | `2048` | Image width (image generation only) |
-| `--num-steps` | `50` | Number of inference steps (image generation only) |
-| `--seed` | `42` | Random seed |
-| `--cfg-scale` | `4.0` | CFG scale (image generation only) |
-| `--think` | `False` | Enable think mode |

examples/online_serving/sensenova_u1/openai_chat_client.py

@@ -1,188 +0,0 @@
-#!/usr/bin/env python3
-"""
-SenseNova-U1 OpenAI-compatible chat client for all four task modalities.
-
-Usage:
-    # Text-to-image
-    python openai_chat_client.py --prompt "A cute cat" --modality text2img
-
-    # Image-to-image editing
-    python openai_chat_client.py --prompt "Turn this into an oil painting" \
-        --modality img2img --image-url /path/to/image.jpg
-
-    # Image understanding (img2text)
-    python openai_chat_client.py --prompt "Describe this image in detail" \
-        --modality img2text --image-url /path/to/image.jpg
-
-    # Text-to-text (chat)
-    python openai_chat_client.py --prompt "What is the capital of France?" \
-        --modality text2text
-"""
-
-import argparse
-import base64
-from pathlib import Path
-
-import requests
-
-
-def _encode_image(image_url: str) -> str:
-    """Encode a local file or URL to a base64 data URI."""
-    if Path(image_url).exists():
-        with open(image_url, "rb") as f:
-            b64_data = base64.b64encode(f.read()).decode("utf-8")
-        return f"data:image/jpeg;base64,{b64_data}"
-    return image_url
-
-
-def generate(
-    prompt: str,
-    server_url: str = "http://localhost:8091",
-    image_url: str | None = None,
-    modality: str = "text2img",
-    **kwargs: object,
-) -> bytes | str | None:
-    """Send a request to the SenseNova-U1 server.
-
-    All keyword arguments (height, width, seed, cfg_scale, think, …) are
-    forwarded as top-level fields in the request payload.  The serving layer
-    maps standard keys (height, width, seed, num_inference_steps) to sampling
-    params and forwards model-specific keys (think, cfg_norm, img_cfg_scale,
-    …) to ``extra_args`` based on the pipeline's ``EXTRA_BODY_PARAMS``.
-
-    Returns:
-        Image bytes (for image outputs) or text string (for text outputs).
-    """
-    content = [{"type": "text", "text": prompt}]
-
-    if image_url:
-        content.append({"type": "image_url", "image_url": {"url": _encode_image(image_url)}})
-
-    messages = [{"role": "user", "content": content}]
-    payload: dict = {"messages": messages}
-
-    if modality in ("text2img", "img2img"):
-        payload["modalities"] = ["image"]
-    else:
-        payload["modalities"] = ["text"]
-
-    for key, val in kwargs.items():
-        if val is not None and val is not False:
-            payload[key] = val
-
-    try:
-        print(f"Sending {modality} request to {server_url}...")
-        response = requests.post(
-            f"{server_url}/v1/chat/completions",
-            headers={"Content-Type": "application/json"},
-            json=payload,
-            timeout=600,
-        )
-        response.raise_for_status()
-        data = response.json()
-
-        # Print think text if present in response metrics
-        metrics = data.get("metrics") or {}
-        think_text = metrics.get("think_text")
-        if think_text:
-            print(f"\n[Think]\n{think_text}\n")
-
-        choices = data.get("choices", [])
-
-        for choice in choices:
-            choice_content = choice.get("message", {}).get("content")
-
-            if isinstance(choice_content, list) and len(choice_content) > 0:
-                first_item = choice_content[0]
-                if isinstance(first_item, dict) and "image_url" in first_item:
-                    img_url_str = first_item["image_url"].get("url", "")
-                    if img_url_str.startswith("data:image"):
-                        _, b64_data = img_url_str.split(",", 1)
-                        return base64.b64decode(b64_data)
-
-        for choice in choices:
-            choice_content = choice.get("message", {}).get("content")
-            if isinstance(choice_content, str) and choice_content:
-                return choice_content
-
-        print(f"Unexpected response format: {choices}")
-        return None
-
-    except Exception as e:
-        print(f"Error: {e}")
-        return None
-
-
-def main():
-    parser = argparse.ArgumentParser(description="SenseNova-U1 multimodal chat client")
-    parser.add_argument("--prompt", "-p", default="A cute cat", help="Text prompt")
-    parser.add_argument("--output", "-o", default="sensenova_u1_output.png", help="Output file (for image results)")
-    parser.add_argument("--server", "-s", default="http://localhost:8091", help="Server URL")
-    parser.add_argument("--image-url", "-i", type=str, help="Input image URL or local path")
-    parser.add_argument(
-        "--modality",
-        "-m",
-        default="text2img",
-        choices=["text2img", "img2img", "img2text", "text2text"],
-        help="Task modality",
-    )
-    # Standard generation parameters
-    parser.add_argument("--height", type=int, default=2048, help="Image height")
-    parser.add_argument("--width", type=int, default=2048, help="Image width")
-    parser.add_argument("--num-steps", type=int, default=50, help="Inference steps")
-    parser.add_argument("--seed", type=int, default=42, help="Random seed")
-    # Model-specific parameters (forwarded via EXTRA_BODY_PARAMS)
-    parser.add_argument("--cfg-scale", type=float, default=4.0, help="CFG scale")
-    parser.add_argument("--img-cfg-scale", type=float, help="Image CFG scale (img2img)")
-    parser.add_argument("--cfg-norm", type=str, help="CFG normalization mode")
-    parser.add_argument("--timestep-shift", type=float, help="Timestep shift")
-    parser.add_argument("--think", action="store_true", help="Enable think mode")
-
-    args = parser.parse_args()
-    is_image = args.modality in ("text2img", "img2img")
-
-    print(f"Mode: {args.modality}")
-    if args.image_url:
-        print(f"Input Image: {args.image_url}")
-
-    extra: dict[str, object] = {
-        "seed": args.seed,
-        "think": args.think or None,
-    }
-    if is_image:
-        extra.update(
-            height=args.height,
-            width=args.width,
-            num_inference_steps=args.num_steps,
-            cfg_scale=args.cfg_scale,
-        )
-        if args.img_cfg_scale is not None:
-            extra["img_cfg_scale"] = args.img_cfg_scale
-        if args.cfg_norm is not None:
-            extra["cfg_norm"] = args.cfg_norm
-        if args.timestep_shift is not None:
-            extra["timestep_shift"] = args.timestep_shift
-
-    result = generate(
-        prompt=args.prompt,
-        server_url=args.server,
-        image_url=args.image_url,
-        modality=args.modality,
-        **extra,
-    )
-
-    if result:
-        if isinstance(result, bytes):
-            output_path = Path(args.output)
-            output_path.write_bytes(result)
-            print(f"Image saved to: {output_path}")
-            print(f"Size: {len(result) / 1024:.1f} KB")
-        elif isinstance(result, str):
-            print(f"[Response]\n{result}")
-    else:
-        print("Failed to generate response")
-        exit(1)
-
-
-if __name__ == "__main__":
-    main()

examples/online_serving/sensenova_u1/run_server.sh

@@ -1,12 +0,0 @@
-#!/bin/bash
-# SenseNova-U1 online serving startup script
-
-MODEL="${MODEL:-SenseNova/SenseNova-U1-8B-MoT}"
-PORT="${PORT:-8091}"
-
-echo "Starting SenseNova-U1 server..."
-echo "Model: $MODEL"
-echo "Port: $PORT"
-
-vllm serve "$MODEL" --omni \
-    --port "$PORT"

examples/online_serving/text_to_speech/voxcpm/run_server.sh

@@ -17,7 +17,7 @@ case "$MODE" in
         STAGE_CONFIG="vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml"
         ;;
     sync)
-        STAGE_CONFIG="vllm_omni/deploy/voxcpm.yaml"
+        STAGE_CONFIG="vllm_omni/model_executor/stage_configs/voxcpm.yaml"
         ;;
     *)
         echo "Unknown mode: $MODE"

pyproject.toml

@@ -55,16 +55,11 @@ dev = [
     "opencc>=1.2.0",
     "mistune>=3.2.0", # for example tests
     "torchmetrics>=1.4.0", # for accuracy similarity metrics
-    "ray[default]", # for RL test
     "jiwer>=3.0.0",
     "zhon>=2.0.0",
     "zhconv>=1.4.2",
     "scipy>=1.10.0",
     "funasr>=1.0.0",
-    "lpips>=0.1.4", # for quantization fp8
-    "orjson>=3.11.7", # to make v1/embeddings API fast
-    "voxcpm>=2.0.2", # for voxcpm model Test
-    "torchdiffeq>=0.2.5" # for covo_audio Test
 ]
 
 demo = [
@@ -147,12 +142,6 @@ ignore = [
 
 [tool.ruff.lint.flake8-tidy-imports.banned-api]
 "librosa".msg = "The librosa module is banned, use vllm.multimodal helpers instead"
-"torch.cuda.max_memory_allocated".msg = "Use torch.accelerator.max_memory_allocated"
-"torch.cuda.reset_peak_memory_stats".msg = "Use torch.accelerator.reset_peak_memory_stats"
-"torch.cuda.device_count".msg = "Use torch.accelerator.device_count"
-"torch.cuda.current_device".msg = "Use torch.accelerator.current_device_index"
-"torch.cuda.empty_cache".msg  = "Use torch.accelerator.empty_cache"
-"torch.cuda.synchronize".msg = "Use torch.accelerator.synchronize"
 
 [tool.mypy]
 python_version = "3.12, 3.13"
@@ -187,7 +176,6 @@ markers = [
     # function module markers
     "diffusion: Diffusion model tests",
     "omni: Omni model tests",
-    "tts: TTS model tests",
     "cache: Cache backend tests",
     "parallel: Parallelism/distributed tests",
     "example: Doc example code tests",

recipes/Qwen/Qwen3-TTS.md

@@ -1,168 +0,0 @@
-# Qwen3-TTS for text-to-speech
-
-## Summary
-
-- Vendor: Qwen
-- Model: `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` (and VoiceDesign / Base variants)
-- Task: Text-to-speech with predefined voices, voice design, or voice cloning
-- Mode: Online serving with the OpenAI-compatible `/v1/audio/speech` API
-- Maintainer: Community
-
-## When to use this recipe
-
-Use this recipe when you want a known-good starting point for serving Qwen3-TTS
-models with vLLM-Omni and validate the deployment with the existing TTS client
-examples in this repository.
-
-Qwen3-TTS supports three task types, each backed by a dedicated model checkpoint:
-
-| Task Type     | Model                                    | Description                                                  |
-| ------------- | ---------------------------------------- | ------------------------------------------------------------ |
-| `CustomVoice` | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice`  | Predefined speaker voices with optional style/emotion control |
-| `VoiceDesign` | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign`  | Generate speech from a natural language voice description     |
-| `Base`        | `Qwen/Qwen3-TTS-12Hz-1.7B-Base`         | Voice cloning from reference audio + transcript              |
-
-Smaller 0.6B variants are also available for `CustomVoice` and `Base`.
-
-## References
-
-- Upstream or canonical docs:
-  [`docs/user_guide/examples/online_serving/qwen3_tts.md`](../../docs/user_guide/examples/online_serving/qwen3_tts.md)
-- Related examples under `examples/`:
-  [`examples/online_serving/qwen3_tts/`](../../examples/online_serving/qwen3_tts/),
-  [`examples/offline_inference/qwen3_tts/`](../../examples/offline_inference/qwen3_tts/)
-- Related issue or discussion:
-  [RFC: add recipes folder](https://github.com/vllm-project/vllm-omni/issues/2645)
-
-## Environment
-
-- OS: Linux
-- Python: 3.10+
-- vLLM / vLLM-Omni: use versions from your current checkout, >=0.20.0
-
-## Command
-
-Start the server from the repository root. Pick the model that matches your
-task type:
-
-```bash
-# CustomVoice (predefined speakers with optional style control)
-vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
-    --deploy-config vllm_omni/deploy/qwen3_tts.yaml \
-    --omni --port 8091
-
-# VoiceDesign (natural language voice description)
-vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign \
-    --deploy-config vllm_omni/deploy/qwen3_tts.yaml \
-    --omni --port 8091
-
-# Base (voice cloning from reference audio)
-vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base \
-    --deploy-config vllm_omni/deploy/qwen3_tts.yaml \
-    --omni --port 8091
-```
-
-Alternatively, use the convenience script:
-
-```bash
-./examples/online_serving/qwen3_tts/run_server.sh                  # Default: CustomVoice
-./examples/online_serving/qwen3_tts/run_server.sh VoiceDesign      # VoiceDesign
-./examples/online_serving/qwen3_tts/run_server.sh Base             # Base (voice clone)
-```
-
-The bundled deploy config (`vllm_omni/deploy/qwen3_tts.yaml`) enables async
-chunking for low first-audio latency. For advanced deployment tuning, pass a
-custom deploy config:
-
-```bash
-vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
-    --deploy-config /path/to/your_qwen3_tts_overrides.yaml \
-    --omni --port 8091 --trust-remote-code
-```
-
-## Verification
-
-**Quick smoke test with curl (CustomVoice):**
-
-```bash
-curl -X POST http://localhost:8091/v1/audio/speech \
-    -H "Content-Type: application/json" \
-    -d '{
-        "input": "Hello, how are you?",
-        "voice": "vivian",
-        "language": "English"
-    }' --output output.wav
-```
-
-**CustomVoice with emotion instruction:**
-
-```bash
-curl -X POST http://localhost:8091/v1/audio/speech \
-    -H "Content-Type: application/json" \
-    -d '{
-        "input": "I am so excited!",
-        "voice": "vivian",
-        "instructions": "Speak with great enthusiasm"
-    }' --output excited.wav
-```
-
-**List available voices:**
-
-```bash
-curl http://localhost:8091/v1/audio/voices
-```
-
-**Using the Python client:**
-
-```bash
-cd examples/online_serving/qwen3_tts
-
-# CustomVoice
-python openai_speech_client.py \
-    --text "Hello, how are you?" \
-    --speaker vivian --language English
-
-# VoiceDesign (requires VoiceDesign model)
-python openai_speech_client.py \
-    --model Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign \
-    --task-type VoiceDesign \
-    --text "Hello world" \
-    --instructions "A warm, friendly female voice"
-
-# Base / voice clone (requires Base model)
-python openai_speech_client.py \
-    --model Qwen/Qwen3-TTS-12Hz-1.7B-Base \
-    --task-type Base \
-    --text "Hello, this is a cloned voice" \
-    --ref-audio /path/to/reference.wav \
-    --ref-text "Transcript of the reference audio"
-```
-
-**Streaming audio (low latency):**
-
-```bash
-curl -X POST http://localhost:8091/v1/audio/speech \
-    -H "Content-Type: application/json" \
-    -d '{
-        "input": "Hello, how are you?",
-        "voice": "vivian",
-        "language": "English",
-        "stream": true,
-        "response_format": "pcm"
-    }' --no-buffer | play -t raw -r 24000 -e signed -b 16 -c 1 -
-```
-
-**Offline inference (no server needed):**
-
-```bash
-python examples/offline_inference/qwen3_tts/end2end.py --query-type CustomVoice
-python examples/offline_inference/qwen3_tts/end2end.py --query-type CustomVoice --streaming
-```
-
-## Notes
-
-- Memory usage: The deploy config allocates `gpu_memory_utilization: 0.3` per stage (talker + code2wav share a single GPU). For the 0.6B variants or constrained GPUs, adjust via `--gpu-memory-utilization`.
-- Key flags: `--omni` is required. `--deploy-config` points to the bundled two-stage pipeline config.
-- Async chunking: Enabled by default in `qwen3_tts.yaml` for streaming-friendly first-audio latency. Streaming requires `stream=true` with `response_format="pcm"`.
-- Task/model matching: Each task type requires its matching model checkpoint. Using a CustomVoice model for a Base (voice clone) request will fail.
-- Known limitations: The server serves one model variant at a time. To switch task types (e.g., CustomVoice to Base), restart the server with the corresponding model.

recipes/README.md

@@ -29,10 +29,6 @@ recipes/
   Qwen-Image on `1x A100 80GB`, including optional step-wise continuous batching replay
 - [`Qwen/Qwen3-Omni.md`](./Qwen/Qwen3-Omni.md): online serving recipe for
   multimodal chat on `1x A100 80GB`
-- [`Tencent/Covo-Audio-Chat.md`](./Tencent/Covo-Audio-Chat.md): online
-  serving recipe for audio chat on `1x A100 80GB`
-- [`Qwen/Qwen3-TTS.md`](./Qwen/Qwen3-TTS.md): text-to-speech serving recipe
-  for Qwen3-TTS (CustomVoice / VoiceDesign / Base) on `1x H100/A100 80GB`
 - [`LTX/LTX-2.md`](./LTX/LTX-2.md): text-to-video and image-to-video serving
   recipe for LTX-2 on `1x H200 141GB`
 - [`Wan-AI/Wan2.2-I2V.md`](./Wan-AI/Wan2.2-I2V.md): image-to-video serving

recipes/SenseNova/SenseNova-U1.md

@@ -1,187 +0,0 @@
-# SenseNova-U1 for unified image generation and understanding
-
-## Summary
-
-- Vendor: SenseNova
-- Model: `SenseNova/SenseNova-U1-8B-MoT`
-- Task: text2img, img2img, img2text (visual understanding), text2text (chat)
-- Mode: Offline inference, Online serving (OpenAI-compatible API)
-- Maintainer: Community
-
-## When to use this recipe
-
-Use this recipe to run SenseNova-U1-8B-MoT via vLLM-Omni. SenseNova-U1 is a
-unified Qwen3-based LLM with Mixture-of-Tokenizers (MoT) attention that handles
-text encoding, optional chain-of-thought reasoning, flow-matching image
-denoising, and visual understanding in a single pipeline — no separate text
-encoder or VAE needed. It supports four task modalities: text-to-image,
-image-to-image editing (with dual CFG), image-to-text understanding, and
-text-to-text chat.
-
-## References
-
-- Offline example:
-  [`examples/offline_inference/sensenova_u1/end2end.py`](../../examples/offline_inference/sensenova_u1/end2end.py)
-- Online serving:
-  [`examples/online_serving/sensenova_u1/`](../../examples/online_serving/sensenova_u1/)
-- E2E tests:
-  [`tests/e2e/offline_inference/test_sensenova_u1_text2img.py`](../../tests/e2e/offline_inference/test_sensenova_u1_text2img.py),
-  [`tests/e2e/offline_inference/test_sensenova_u1_img2img.py`](../../tests/e2e/offline_inference/test_sensenova_u1_img2img.py)
-- HuggingFace model page:
-  [SenseNova/SenseNova-U1-8B-MoT](https://huggingface.co/SenseNova/SenseNova-U1-8B-MoT)
-
-## Hardware Support
-
-## GPU
-
-### 1x H200 (144GB)
-
-#### Environment
-
-- OS: Linux
-- Python: 3.12
-- Driver / runtime: NVIDIA 590.48.01, CUDA 13.1
-- vLLM-Omni version: 0.18.1.dev
-
-#### Command
-
-```bash
-python examples/offline_inference/sensenova_u1/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
-```
-
-#### Verification
-
-```bash
-pytest -s -v tests/e2e/offline_inference/test_sensenova_u1_text2img.py \
-    -m "advanced_model" --run-level "advanced_model"
-```
-
-#### Notes
-
-- E2E latency: **32.1s** (1536×2720, 50 steps, think mode, CFG scale 4.0)
-- Peak VRAM: **35.9 GB** reserved, 35.1 GB allocated
-- Model loading: 32.8 GiB, 8.7s
-- No deploy YAML needed — the engine auto-generates a single-stage diffusion config.
-- Think mode (`--think`) is recommended for higher image quality.
-
-#### Image-to-Image Editing (img2img)
-
-```bash
-python examples/offline_inference/sensenova_u1/end2end.py \
-    --prompt "Turn this into an oil painting" \
-    --image input.png \
-    --width 2048 --height 2048 \
-    --seed 42 --num-steps 50 \
-    --cfg-scale 4.0 --img-cfg-scale 1.0 --cfg-norm none \
-    --think --print-think \
-    --output outputs
-```
-
-- img2img uses dual CFG: `--cfg-scale` controls text guidance, `--img-cfg-scale`
-  controls image guidance (1.0 = image CFG disabled).
-- Pass multiple `--image` paths for multi-reference editing.
-
-#### Image Understanding (img2text)
-
-```bash
-python examples/offline_inference/sensenova_u1/end2end.py \
-    --modality img2text \
-    --prompt "Describe this image in detail" \
-    --image photo.jpg \
-    --max-tokens 512
-```
-
-#### Text-to-Text Chat (text2text)
-
-```bash
-python examples/offline_inference/sensenova_u1/end2end.py \
-    --modality text2text \
-    --prompt "Explain the theory of relativity in simple terms" \
-    --max-tokens 256
-```
-
-- For img2text and text2text, image generation parameters (height, width,
-  num-steps, cfg-scale) are ignored.
-- Use `--do-sample --temperature 0.7` for more diverse text responses.
-
-### 2x H200 (144GB) — TP=2
-
-#### Environment
-
-- OS: Linux
-- Python: 3.12
-- Driver / runtime: NVIDIA 590.48.01, CUDA 13.1
-- vLLM-Omni version: 0.18.1.dev
-
-#### Command
-
-```bash
-python examples/offline_inference/sensenova_u1/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 \
-    --tensor-parallel-size 2 \
-    --output outputs
-```
-
-#### Verification
-
-Verify the output image is generated at `outputs/sensenova_u1_output_0.png`
-with the expected 1536×2720 resolution.
-
-#### Notes
-
-- E2E latency: **28.3s** (1536×2720, 50 steps, think mode, CFG scale 4.0)
-- Peak VRAM (per GPU): **18.2 GB** reserved, 17.9 GB allocated
-- Model loading: 16.5 GiB per GPU, 7.0s
-- TP=2 provides ~12% speedup over TP=1; limited by serial CFG dual-forward
-  and communication overhead.
-- The LLM transformer uses `QKVParallelLinear` and `MergedColumnParallelLinear`
-  for fused QKV and gate/up projections with TP support.
-
-## Online Serving
-
-SenseNova-U1 supports all four modalities via the OpenAI-compatible
-`/v1/chat/completions` API.
-
-### Launch
-
-```bash
-vllm serve SenseNova/SenseNova-U1-8B-MoT --omni --port 8091
-```
-
-### Send Requests
-
-```bash
-cd examples/online_serving/sensenova_u1
-
-# Text-to-image
-python openai_chat_client.py \
-    --prompt "A beautiful sunset" --modality text2img
-
-# Image-to-image editing
-python openai_chat_client.py \
-    --prompt "Turn this into an oil painting" \
-    --modality img2img --image-url input.jpg
-
-# Image understanding
-python openai_chat_client.py \
-    --prompt "Describe this image" \
-    --modality img2text --image-url photo.jpg
-
-# Text chat
-python openai_chat_client.py \
-    --prompt "What is the capital of France?" \
-    --modality text2text
-```
-
-For full API documentation and curl examples, see
-[`examples/online_serving/sensenova_u1/README.md`](../../examples/online_serving/sensenova_u1/README.md).

recipes/Tencent-Hunyuan/HunyuanImage-3.0-Instruct.md

@@ -186,9 +186,9 @@ Validated score summary:
 
 The CI assertion threshold is `overall_mean >= 0.45`, so the smoke result is
 comfortably above the gate. The generate server and judge server run
-sequentially through the `OmniServer` fixture, with GPU memory cleanup
-between server lifetimes (for example via the `clean_gpu_memory_between_tests`
-pytest fixture in the smoke path).
+sequentially through the `OmniServer` fixture, and
+`VLLM_TEST_CLEAN_GPU_MEMORY=1` is used to wait for memory cleanup between
+server lifetimes.
 
 The lower-cost 2-GPU Instruct setup was tried for this smoke path but did not
 fit in memory. A previous 2-GPU experiment used the base HunyuanImage-3.0

recipes/Tencent/Covo-Audio-Chat.md

@@ -1,92 +0,0 @@
-# Covo-Audio-Chat for audio chat on 1x A100 80GB
-
-## Summary
-
-- Vendor: Tencent
-- Model: `tencent/Covo-Audio-Chat`
-- Task: Audio chat — audio (and optional text) input, interleaved text + audio output
-- Mode: Online serving with the OpenAI-compatible API (offline inference also supported)
-- Maintainer: Community
-
-## When to use this recipe
-
-Use this recipe when you want a known-good starting point for serving
-`tencent/Covo-Audio-Chat` with vLLM-Omni on a single 80 GB-class GPU and
-validate the deployment with the bundled client example.
-
-## References
-
-- Related example under `examples/`:
-  [`examples/online_serving/covo_audio/README.md`](../../examples/online_serving/covo_audio/README.md)
-- Offline end-to-end script:
-  [`examples/offline_inference/covo_audio/end2end.py`](../../examples/offline_inference/covo_audio/end2end.py)
-- Deploy config:
-  [`vllm_omni/deploy/covo_audio.yaml`](../../vllm_omni/deploy/covo_audio.yaml)
-- Pipeline topology:
-  [`vllm_omni/model_executor/models/covo_audio/pipeline.py`](../../vllm_omni/model_executor/models/covo_audio/pipeline.py)
-
-## Hardware Support
-
-This recipe currently documents one tested reference configuration for
-CUDA GPU serving. The fused thinker+talker stage and the BigVGAN code2wav
-stage both run on the same physical device.
-
-## GPU
-
-### 1x A100 80GB
-
-#### Environment
-
-- OS: Linux
-- Python: 3.12
-- Driver / runtime: NVIDIA CUDA 12.x with an A100 80 GB GPU
-- vLLM version: Match the repository requirements for your checkout
-- vLLM-Omni version or commit: Use the commit you are deploying from
-- Extra Python deps: `torchdiffeq`, `librosa`
-
-#### Command
-
-Install the two runtime-only deps that are not pulled in by vLLM-Omni itself:
-
-```bash
-pip install torchdiffeq librosa
-```
-
-Start the server from the repository root:
-
-```bash
-CUDA_VISIBLE_DEVICES=0 vllm serve tencent/Covo-Audio-Chat --omni \
-    --trust-remote-code --port 18091
-```
-
-The deploy config at `vllm_omni/deploy/covo_audio.yaml` is loaded
-automatically by the model registry.
-
-#### Verification
-
-With the server ready, send a test request via the bundled client:
-
-```bash
-cd examples/online_serving/covo_audio
-python openai_chat_completion_client.py
-```
-
-Expected: the client prints the generated text and writes an output WAV
-file to the current directory.
-
-#### Notes
-
-- Memory usage: Stage 0 (7B LLM) ~16 GiB + stage 1 (BigVGAN vocoder)
-  ~2 GiB, both resident on the same device.
-- Key flags: `--omni` and `--trust-remote-code` are both required.
-- System prompt: the bundled `openai_chat_completion_client.py` and
-  `end2end.py` already include `COVO_AUDIO_SYSTEM_PROMPT` (defined in
-  `vllm_omni/model_executor/models/covo_audio/prompt_utils.py`) in the
-  system message by default. When integrating your own client, make sure
-  to include it — without it, stage 0 emits text-only tokens and the
-  output audio is silent.
-- Stage 1 runs with `enforce_eager: true` and `dtype: float32` — see the
-  header of `deploy/covo_audio.yaml` for the specific reasons.
-- Known limitations: this starter recipe covers the single-GPU online-serving
-  path only; multi-GPU layouts and platform-specific overrides are not
-  documented here.

recipes/Wan-AI/Wan2.2-I2V.md

@@ -11,9 +11,7 @@
 ## When to use this recipe
 
 Use this recipe when you want to deploy the Wan2.2 14B image-to-video model
-with vLLM-Omni using multi-card parallelism. The recommended parallel strategy
-depends on the hardware backend and whether the checkpoint needs
-classifier-free guidance.
+with vLLM-Omni using multi-card parallelism. Two configurations are provided:
 
 1. **Distilled model (no negative-prompt / CFG computation)** — higher
    throughput, recommended when using a distilled checkpoint that does not
@@ -27,66 +25,6 @@ classifier-free guidance.
 
 ## Hardware Support
 
-## GPU
-
-### 8x NVIDIA H20 / H100 / A100
-
-#### Environment
-
-- OS: Linux
-- Python: 3.10+
-- Driver / runtime: NVIDIA driver with CUDA runtime supported by your PyTorch
-  build
-- Recommended operator library: Triton, installed through the vLLM/vLLM-Omni
-  Python environment
-- vLLM version: Match the repository requirements for your checkout
-- vLLM-Omni version or commit: Use the commit you are deploying from
-
-#### Command
-
-**Distilled model (no CFG, recommended for distilled checkpoints):**
-
-```bash
-CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
-vllm serve Wan-AI/Wan2.2-I2V-A14B-Diffusers \
-  --omni \
-  --use-hsdp \
-  --usp 8 \
-  --vae-patch-parallel-size 8 \
-  --vae-use-tiling
-```
-
-**Official open-source model (with CFG):**
-
-```bash
-CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
-vllm serve Wan-AI/Wan2.2-I2V-A14B-Diffusers \
-  --omni \
-  --use-hsdp \
-  --cfg-parallel-size 2 \
-  --usp 4 \
-  --vae-patch-parallel-size 8 \
-  --vae-use-tiling
-```
-
-The official model benefits from splitting the positive and negative CFG
-branches with `--cfg-parallel-size 2`, while keeping Ulysses sequence
-parallelism at `--usp 4` so the two CFG branches cover all 8 GPUs.
-
-#### Performance Notes
-
-- `--use-hsdp` is recommended for memory efficiency on the 14B DiT backbone.
-- `--vae-patch-parallel-size 8` is recommended for the standard 8-card setup;
-  disabling VAE patch parallelism can significantly increase VAE latency.
-- The CUDA Triton fused AdaLayerNorm path is not available in the current
-  release. The `VLLM_OMNI_ENABLE_TRITON_ADALN` toggle is reserved for that
-  future implementation and should not be used until the kernel lands.
-- In an 8x NVIDIA H20 test with the official model, `num_inference_steps=4`,
-  `flow_shift=12.0`, `--use-hsdp`, `--cfg-parallel-size 2`, `--usp 4`, and
-  `--vae-patch-parallel-size 8`, an internal Triton AdaLayerNorm prototype
-  reduced the measured diffusion stage by about `65 ms` at
-  `832x480, 81 frames` and about `741 ms` at `1280x720, 121 frames`.
-
 ## NPU
 
 ### 8x Ascend A2 / A3

requirements/common.txt

@@ -1,8 +1,7 @@
 # Common dependencies for all platforms
 av>=14.0.0
 omegaconf>=2.3.0
-diffusers==0.38.0
-safetensors>=0.8.0rc0
+diffusers>=0.36.0
 accelerate==1.12.0
 soundfile>=0.13.1
 cache-dit==1.3.0

tests/conftest.py

@@ -8,14 +8,6 @@ Root pytest entrypoint for the vLLM-Omni test suite.
 
 from __future__ import annotations
 
-# Before ``pytest_plugins`` and before any other test path imports vLLM, pin op
-# registration order (see :func:`tests.model_executor.helpers.bootstrap_vllm_layer_custom_op_modules`).
-# Subdir ``conftest`` hooks can run after other tests are collected/imported, which is
-# too late and can trigger duplicate ``vllm::flashinfer_rotary_embedding`` (etc.) errors.
-from tests.model_executor.helpers import bootstrap_vllm_layer_custom_op_modules
-
-bootstrap_vllm_layer_custom_op_modules()
-
 pytest_plugins = (
     "tests.helpers.fixtures.env",
     "tests.helpers.fixtures.log",

tests/core/sched/test_chunk_scheduling_coordinator.py

@@ -222,7 +222,7 @@ class TestChunkCoordinatorUpdateRequestMetadata(unittest.TestCase):
         self.assertEqual(req.prompt_token_ids, [10, 20, 30])
         self.assertEqual(req.num_computed_tokens, 0)
         self.assertIsNone(req.additional_information)
-        self.assertEqual(req._omni_initial_model_buffer, {"meta": {"left_context_size": 25}})
+        self.assertEqual(req._omni_initial_model_buffer, {"left_context_size": 25})
 
 
 class TestChunkCoordinatorPostprocess(unittest.TestCase):

tests/dfx/conftest.py

@@ -2,7 +2,6 @@ import json
 import os
 import re
 import subprocess
-import threading
 import time
 from datetime import datetime
 from pathlib import Path
@@ -378,23 +377,11 @@ def run_benchmark(
         command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, bufsize=1, universal_newlines=True
     )
 
-    if process.stdout is None or process.stderr is None:
-        raise RuntimeError("Failed to capture benchmark process output streams")
+    for line in iter(process.stdout.readline, ""):
+        print(line, end=" ")
 
-    def _forward_stream(stream) -> None:
-        try:
-            for line in iter(stream.readline, ""):
-                print(line, end="")
-        finally:
-            stream.close()
-
-    stdout_thread = threading.Thread(target=_forward_stream, args=(process.stdout,))
-    stderr_thread = threading.Thread(target=_forward_stream, args=(process.stderr,))
-    stdout_thread.start()
-    stderr_thread.start()
-    stdout_thread.join()
-    stderr_thread.join()
-    process.wait()
+    for line in iter(process.stderr.readline, ""):
+        print(line, end=" ")
 
     if "--result-dir" in command:
         index = command.index("--result-dir")

tests/dfx/perf/scripts/run_benchmark.py

@@ -17,10 +17,11 @@ from tests.dfx.conftest import (
 )
 from tests.helpers.runtime import OmniServer
 
-pytestmark = [pytest.mark.full_model]
+pytestmark = [pytest.mark.full_model, pytest.mark.omni]
 
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
 
 
 def _get_config_file_from_argv() -> str | None:

tests/dfx/perf/scripts/run_diffusion_benchmark.py

@@ -36,6 +36,7 @@ import pytest
 pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
 os.environ.setdefault("DIFFUSION_ATTENTION_BACKEND", "FLASH_ATTN")
 
 # ---------------------------------------------------------------------------

tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json

@@ -1,209 +0,0 @@
-[
-    {
-        "test_name": "test_qwen3_omni_chunk",
-        "server_params": {
-            "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
-            "extra_cli_args": ["--async-chunk"]
-        },
-        "benchmark_params": [
-            {
-                "dataset_name": "random",
-                "backend": "openai-chat-omni",
-                "endpoint": "/v1/chat/completions",
-                "num_prompts": [4, 16, 32, 64],
-                "max_concurrency": [1, 4, 8, 16],
-                "random_input_len": 2500,
-                "random_output_len": 900,
-                "ignore_eos": true,
-                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
-                "baseline": {
-                    "mean_ttft_ms": [1000, 3000, 5000, 7000],
-                    "mean_audio_ttfp_ms": [1000, 3000, 5000, 7000],
-                    "mean_audio_rtf": [0.2, 0.35, 0.6, 0.85]
-                }
-            },
-            {
-                "dataset_name": "random-mm",
-                "backend": "openai-chat-omni",
-                "endpoint": "/v1/chat/completions",
-                "num_prompts": [10],
-                "request_rate": [0.1],
-                "random_input_len": 100,
-                "random_output_len": 100,
-                "random_range_ratio": 0.0,
-                "ignore_eos": true,
-                "random_mm_base_items_per_request": 1,
-                "random_mm_num_mm_items_range_ratio": 0.5,
-                "random_mm_limit_mm_per_prompt": {
-                    "audio": 1
-                },
-                "random_mm_bucket_config": {
-                    "(0, 60, 3)": 1.0
-                },
-                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
-                "baseline": {
-                    "mean_ttft_ms": [2000],
-                    "mean_audio_ttfp_ms": [2000],
-                    "mean_audio_rtf": [0.25]
-                }
-            },
-            {
-                "dataset_name": "random-mm",
-                "backend": "openai-chat-omni",
-                "endpoint": "/v1/chat/completions",
-                "num_prompts": [40],
-                "request_rate": [0.5],
-                "random_input_len": 100,
-                "random_output_len": 100,
-                "random_range_ratio": 0.0,
-                "ignore_eos": true,
-                "random_mm_base_items_per_request": 2,
-                "random_mm_num_mm_items_range_ratio": 0.5,
-                "random_mm_limit_mm_per_prompt": {
-                    "image": 1,
-                    "video": 1
-                },
-                "random_mm_bucket_config": {
-                    "(256, 256, 1)": 0.5,
-                    "(720, 1280, 2)": 0.5
-                },
-                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
-                "baseline": {
-                    "mean_ttft_ms": [6000],
-                    "mean_audio_ttfp_ms": [6000],
-                    "mean_audio_rtf": [0.7]
-                }
-            },
-            {
-                "dataset_name": "random-mm",
-                "backend": "openai-chat-omni",
-                "endpoint": "/v1/chat/completions",
-                "num_prompts": [100],
-                "request_rate": [1.0],
-                "random_input_len": 100,
-                "random_output_len": 100,
-                "random_range_ratio": 0.0,
-                "ignore_eos": true,
-                "random_mm_base_items_per_request": 3,
-                "random_mm_num_mm_items_range_ratio": 0.5,
-                "random_mm_limit_mm_per_prompt": {
-                    "image": 1,
-                    "video": 1,
-                    "audio": 1
-                },
-                "random_mm_bucket_config": {
-                    "(256, 256, 1)": 0.34,
-                    "(720, 1280, 2)": 0.33,
-                    "(0, 60, 3)": 0.33
-                },
-                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
-                "baseline": {
-                    "mean_ttft_ms": [12000],
-                    "mean_audio_ttfp_ms": [12000],
-                    "mean_audio_rtf": [1.0]
-                }
-            },
-            {
-                "dataset_name": "random",
-                "backend": "openai-chat-omni",
-                "endpoint": "/v1/chat/completions",
-                "num_prompts": [4, 16, 32, 64, 128],
-                "max_concurrency": [1, 4, 8, 16, 32],
-                "random_input_len": 2500,
-                "random_output_len": 900,
-                "ignore_eos": true,
-                "extra_body": {
-                    "modalities": ["text"]
-                },
-                "percentile-metrics": "ttft,tpot,itl,e2el",
-                "baseline": {
-                    "mean_ttft_ms": [1000, 3000, 5000, 7000, 9000]
-                }
-            },
-            {
-                "dataset_name": "random-mm",
-                "backend": "openai-chat-omni",
-                "endpoint": "/v1/chat/completions",
-                "num_prompts": [10],
-                "request_rate": [0.1],
-                "random_input_len": 100,
-                "random_output_len": 100,
-                "random_range_ratio": 0.0,
-                "ignore_eos": true,
-                "extra_body": {
-                    "modalities": ["text"]
-                },
-                "random_mm_base_items_per_request": 1,
-                "random_mm_num_mm_items_range_ratio": 0.5,
-                "random_mm_limit_mm_per_prompt": {
-                    "audio": 1
-                },
-                "random_mm_bucket_config": {
-                    "(0, 60, 3)": 1.0
-                },
-                "percentile-metrics": "ttft,tpot,itl,e2el",
-                "baseline": {
-                    "mean_ttft_ms": [2000]
-                }
-            },
-            {
-                "dataset_name": "random-mm",
-                "backend": "openai-chat-omni",
-                "endpoint": "/v1/chat/completions",
-                "num_prompts": [40],
-                "request_rate": [0.5],
-                "random_input_len": 100,
-                "random_output_len": 100,
-                "random_range_ratio": 0.0,
-                "ignore_eos": true,
-                "extra_body": {
-                    "modalities": ["text"]
-                },
-                "random_mm_base_items_per_request": 2,
-                "random_mm_num_mm_items_range_ratio": 0.5,
-                "random_mm_limit_mm_per_prompt": {
-                    "image": 1,
-                    "video": 1
-                },
-                "random_mm_bucket_config": {
-                    "(256, 256, 1)": 0.5,
-                    "(720, 1280, 2)": 0.5
-                },
-                "percentile-metrics": "ttft,tpot,itl,e2el",
-                "baseline": {
-                    "mean_ttft_ms": [6000]
-                }
-            },
-            {
-                "dataset_name": "random-mm",
-                "backend": "openai-chat-omni",
-                "endpoint": "/v1/chat/completions",
-                "num_prompts": [100],
-                "request_rate": [1.0],
-                "random_input_len": 100,
-                "random_output_len": 100,
-                "random_range_ratio": 0.0,
-                "ignore_eos": true,
-                "extra_body": {
-                    "modalities": ["text"]
-                },
-                "random_mm_base_items_per_request": 3,
-                "random_mm_num_mm_items_range_ratio": 0.5,
-                "random_mm_limit_mm_per_prompt": {
-                    "image": 1,
-                    "video": 1,
-                    "audio": 1
-                },
-                "random_mm_bucket_config": {
-                    "(256, 256, 1)": 0.34,
-                    "(720, 1280, 2)": 0.33,
-                    "(0, 60, 3)": 0.33
-                },
-                "percentile-metrics": "ttft,tpot,itl,e2el",
-                "baseline": {
-                    "mean_ttft_ms": [6000]
-                }
-            }
-        ]
-    }
-]

tests/dfx/perf/tests/test_qwen3_omni_no_async_chunk.json

@@ -1,108 +0,0 @@
-[
-    {
-        "test_name": "test_qwen3_omni",
-        "server_params": {
-            "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
-            "extra_cli_args": ["--no-async-chunk"]
-        },
-        "benchmark_params": [
-            {
-                "dataset_name": "random",
-                "backend": "openai-chat-omni",
-                "endpoint": "/v1/chat/completions",
-                "num_prompts": [4, 16, 32, 64, 128],
-                "max_concurrency": [1, 4, 8, 16, 32],
-                "random_input_len": 2500,
-                "random_output_len": 900,
-                "ignore_eos": true,
-                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
-                "baseline": {
-                    "mean_ttft_ms": [1000, 3000, 5000, 7000, 9000],
-                    "mean_audio_ttfp_ms": [30000, 60000, 90000, 120000, 150000],
-                    "mean_audio_rtf": [0.35, 0.45, 0.55, 0.65, 0.75]
-                }
-            },
-            {
-                "dataset_name": "random-mm",
-                "backend": "openai-chat-omni",
-                "endpoint": "/v1/chat/completions",
-                "num_prompts": [10],
-                "request_rate": [0.1],
-                "random_input_len": 100,
-                "random_output_len": 100,
-                "random_range_ratio": 0.0,
-                "ignore_eos": true,
-                "random_mm_base_items_per_request": 1,
-                "random_mm_num_mm_items_range_ratio": 0.5,
-                "random_mm_limit_mm_per_prompt": {
-                    "audio": 1
-                },
-                "random_mm_bucket_config": {
-                    "(0, 60, 3)": 1.0
-                },
-                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
-                "baseline": {
-                    "mean_ttft_ms": [2000],
-                    "mean_audio_ttfp_ms": [10000],
-                    "mean_audio_rtf": [0.25]
-                }
-            },
-            {
-                "dataset_name": "random-mm",
-                "backend": "openai-chat-omni",
-                "endpoint": "/v1/chat/completions",
-                "num_prompts": [40],
-                "request_rate": [0.5],
-                "random_input_len": 100,
-                "random_output_len": 100,
-                "random_range_ratio": 0.0,
-                "ignore_eos": true,
-                "random_mm_base_items_per_request": 2,
-                "random_mm_num_mm_items_range_ratio": 0.5,
-                "random_mm_limit_mm_per_prompt": {
-                    "image": 1,
-                    "video": 1
-                },
-                "random_mm_bucket_config": {
-                    "(256, 256, 1)": 0.5,
-                    "(720, 1280, 2)": 0.5
-                },
-                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
-                "baseline": {
-                    "mean_ttft_ms": [6000],
-                    "mean_audio_ttfp_ms": [15000],
-                    "mean_audio_rtf": [0.45]
-                }
-            },
-            {
-                "dataset_name": "random-mm",
-                "backend": "openai-chat-omni",
-                "endpoint": "/v1/chat/completions",
-                "num_prompts": [100],
-                "request_rate": [1.0],
-                "random_input_len": 100,
-                "random_output_len": 100,
-                "random_range_ratio": 0.0,
-                "ignore_eos": true,
-                "random_mm_base_items_per_request": 3,
-                "random_mm_num_mm_items_range_ratio": 0.5,
-                "random_mm_limit_mm_per_prompt": {
-                    "image": 1,
-                    "video": 1,
-                    "audio": 1
-                },
-                "random_mm_bucket_config": {
-                    "(256, 256, 1)": 0.34,
-                    "(720, 1280, 2)": 0.33,
-                    "(0, 60, 3)": 0.33
-                },
-                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
-                "baseline": {
-                    "mean_ttft_ms": [12000],
-                    "mean_audio_ttfp_ms": [18000],
-                    "mean_audio_rtf": [0.9]
-                }
-            }
-        ]
-    }
-]

tests/dfx/perf/tests/test_qwen3_omni_vllm_text.json

@@ -1,137 +0,0 @@
-[
-    {
-        "test_name": "test_qwen3_omni_vllm_text",
-        "server_params": {
-            "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
-            "use_omni": false,
-            "extra_cli_args": ["--no-enable-prefix-caching"]
-        },
-        "benchmark_params": [
-            {
-                "dataset_name": "random",
-                "backend": "vllm",
-                "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
-                "endpoint": "/v1/completions",
-                "num_prompts": [4, 16, 32, 64, 128],
-                "max_concurrency": [1, 4, 8, 16, 32],
-                "random_input_len": 2500,
-                "random_output_len": 900,
-                "temperature": 0.4,
-                "top_p": 0.9,
-                "top_k": 1,
-                "seed": 42,
-                "repetition_penalty": 1.05,
-                "ignore_eos": true,
-                "extra_body": {
-                    "modalities": ["text"]
-                },
-                "percentile-metrics": "ttft,tpot,itl,e2el",
-                "baseline": {
-                    "mean_ttft_ms": [1000, 3000, 5000, 7000, 9000]
-                }
-            },
-            {
-                "dataset_name": "random-mm",
-                "backend": "vllm",
-                "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
-                "endpoint": "/v1/completions",
-                "num_prompts": [10],
-                "request_rate": [0.1],
-                "random_input_len": 100,
-                "random_output_len": 100,
-                "random_range_ratio": 0.0,
-                "temperature": 0.4,
-                "top_p": 0.9,
-                "top_k": 1,
-                "seed": 42,
-                "repetition_penalty": 1.05,
-                "ignore_eos": true,
-                "extra_body": {
-                    "modalities": ["text"]
-                },
-                "random_mm_base_items_per_request": 1,
-                "random_mm_num_mm_items_range_ratio": 0.5,
-                "random_mm_limit_mm_per_prompt": {
-                    "audio": 1
-                },
-                "random_mm_bucket_config": {
-                    "(0, 60, 3)": 1.0
-                },
-                "percentile-metrics": "ttft,tpot,itl,e2el",
-                "baseline": {
-                    "mean_ttft_ms": [2000]
-                }
-            },
-            {
-                "dataset_name": "random-mm",
-                "backend": "vllm",
-                "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
-                "endpoint": "/v1/completions",
-                "num_prompts": [40],
-                "request_rate": [0.5],
-                "random_input_len": 100,
-                "random_output_len": 100,
-                "random_range_ratio": 0.0,
-                "temperature": 0.4,
-                "top_p": 0.9,
-                "top_k": 1,
-                "seed": 42,
-                "repetition_penalty": 1.05,
-                "ignore_eos": true,
-                "extra_body": {
-                    "modalities": ["text"]
-                },
-                "random_mm_base_items_per_request": 2,
-                "random_mm_num_mm_items_range_ratio": 0.5,
-                "random_mm_limit_mm_per_prompt": {
-                    "image": 1,
-                    "video": 1
-                },
-                "random_mm_bucket_config": {
-                    "(256, 256, 1)": 0.5,
-                    "(720, 1280, 2)": 0.5
-                },
-                "percentile-metrics": "ttft,tpot,itl,e2el",
-                "baseline": {
-                    "mean_ttft_ms": [6000]
-                }
-            },
-            {
-                "dataset_name": "random-mm",
-                "backend": "vllm",
-                "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
-                "endpoint": "/v1/completions",
-                "num_prompts": [100],
-                "request_rate": [1.0],
-                "random_input_len": 100,
-                "random_output_len": 100,
-                "random_range_ratio": 0.0,
-                "temperature": 0.4,
-                "top_p": 0.9,
-                "top_k": 1,
-                "seed": 42,
-                "repetition_penalty": 1.05,
-                "ignore_eos": true,
-                "extra_body": {
-                    "modalities": ["text"]
-                },
-                "random_mm_base_items_per_request": 3,
-                "random_mm_num_mm_items_range_ratio": 0.5,
-                "random_mm_limit_mm_per_prompt": {
-                    "image": 1,
-                    "video": 1,
-                    "audio": 1
-                },
-                "random_mm_bucket_config": {
-                    "(256, 256, 1)": 0.34,
-                    "(720, 1280, 2)": 0.33,
-                    "(0, 60, 3)": 0.33
-                },
-                "percentile-metrics": "ttft,tpot,itl,e2el",
-                "baseline": {
-                    "mean_ttft_ms": [6000]
-                }
-            }
-        ]
-    }
-]

tests/dfx/perf/tests/test_qwen_omni.json

@@ -0,0 +1,450 @@
+[
+    {
+        "test_name": "test_qwen3_omni",
+        "server_params": {
+            "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+            "extra_cli_args": ["--no-async-chunk"]
+        },
+        "benchmark_params": [
+            {
+                "dataset_name": "random",
+                "backend": "openai-chat-omni",
+                "endpoint": "/v1/chat/completions",
+                "num_prompts": [4, 16, 32, 64, 128],
+                "max_concurrency": [1, 4, 8, 16, 32],
+                "random_input_len": 2500,
+                "random_output_len": 900,
+                "ignore_eos": true,
+                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
+                "baseline": {
+                    "mean_ttft_ms": [1000, 3000, 5000, 7000, 9000],
+                    "mean_audio_ttfp_ms": [30000, 60000, 90000, 120000, 150000],
+                    "mean_audio_rtf": [0.35, 0.45, 0.55, 0.65, 0.75]
+                }
+            },
+            {
+                "dataset_name": "random-mm",
+                "backend": "openai-chat-omni",
+                "endpoint": "/v1/chat/completions",
+                "num_prompts": [10],
+                "request_rate": [0.1],
+                "random_input_len": 100,
+                "random_output_len": 100,
+                "random_range_ratio": 0.0,
+                "ignore_eos": true,
+                "random_mm_base_items_per_request": 1,
+                "random_mm_num_mm_items_range_ratio": 0.5,
+                "random_mm_limit_mm_per_prompt": {
+                    "audio": 1
+                },
+                "random_mm_bucket_config": {
+                    "(0, 60, 3)": 1.0
+                },
+                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
+                "baseline": {
+                    "mean_ttft_ms": [2000],
+                    "mean_audio_ttfp_ms": [10000],
+                    "mean_audio_rtf": [0.25]
+                }
+            },
+            {
+                "dataset_name": "random-mm",
+                "backend": "openai-chat-omni",
+                "endpoint": "/v1/chat/completions",
+                "num_prompts": [40],
+                "request_rate": [0.5],
+                "random_input_len": 100,
+                "random_output_len": 100,
+                "random_range_ratio": 0.0,
+                "ignore_eos": true,
+                "random_mm_base_items_per_request": 2,
+                "random_mm_num_mm_items_range_ratio": 0.5,
+                "random_mm_limit_mm_per_prompt": {
+                    "image": 1,
+                    "video": 1
+                },
+                "random_mm_bucket_config": {
+                    "(256, 256, 1)": 0.5,
+                    "(720, 1280, 2)": 0.5
+                },
+                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
+                "baseline": {
+                    "mean_ttft_ms": [6000],
+                    "mean_audio_ttfp_ms": [15000],
+                    "mean_audio_rtf": [0.45]
+                }
+            },
+            {
+                "dataset_name": "random-mm",
+                "backend": "openai-chat-omni",
+                "endpoint": "/v1/chat/completions",
+                "num_prompts": [100],
+                "request_rate": [1.0],
+                "random_input_len": 100,
+                "random_output_len": 100,
+                "random_range_ratio": 0.0,
+                "ignore_eos": true,
+                "random_mm_base_items_per_request": 3,
+                "random_mm_num_mm_items_range_ratio": 0.5,
+                "random_mm_limit_mm_per_prompt": {
+                    "image": 1,
+                    "video": 1,
+                    "audio": 1
+                },
+                "random_mm_bucket_config": {
+                    "(256, 256, 1)": 0.34,
+                    "(720, 1280, 2)": 0.33,
+                    "(0, 60, 3)": 0.33
+                },
+                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
+                "baseline": {
+                    "mean_ttft_ms": [12000],
+                    "mean_audio_ttfp_ms": [18000],
+                    "mean_audio_rtf": [0.9]
+                }
+            }
+        ]
+    },
+    {
+        "test_name": "test_qwen3_omni_chunk",
+        "server_params": {
+            "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+            "extra_cli_args": ["--async-chunk"]
+        },
+        "benchmark_params": [
+            {
+                "dataset_name": "random",
+                "backend": "openai-chat-omni",
+                "endpoint": "/v1/chat/completions",
+                "num_prompts": [4, 16, 32, 64],
+                "max_concurrency": [1, 4, 8, 16],
+                "random_input_len": 2500,
+                "random_output_len": 900,
+                "ignore_eos": true,
+                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
+                "baseline": {
+                    "mean_ttft_ms": [1000, 3000, 5000, 7000],
+                    "mean_audio_ttfp_ms": [1000, 3000, 5000, 7000],
+                    "mean_audio_rtf": [0.2, 0.35, 0.6, 0.85]
+                }
+            },
+            {
+                "dataset_name": "random-mm",
+                "backend": "openai-chat-omni",
+                "endpoint": "/v1/chat/completions",
+                "num_prompts": [10],
+                "request_rate": [0.1],
+                "random_input_len": 100,
+                "random_output_len": 100,
+                "random_range_ratio": 0.0,
+                "ignore_eos": true,
+                "random_mm_base_items_per_request": 1,
+                "random_mm_num_mm_items_range_ratio": 0.5,
+                "random_mm_limit_mm_per_prompt": {
+                    "audio": 1
+                },
+                "random_mm_bucket_config": {
+                    "(0, 60, 3)": 1.0
+                },
+                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
+                "baseline": {
+                    "mean_ttft_ms": [2000],
+                    "mean_audio_ttfp_ms": [2000],
+                    "mean_audio_rtf": [0.25]
+                }
+            },
+            {
+                "dataset_name": "random-mm",
+                "backend": "openai-chat-omni",
+                "endpoint": "/v1/chat/completions",
+                "num_prompts": [40],
+                "request_rate": [0.5],
+                "random_input_len": 100,
+                "random_output_len": 100,
+                "random_range_ratio": 0.0,
+                "ignore_eos": true,
+                "random_mm_base_items_per_request": 2,
+                "random_mm_num_mm_items_range_ratio": 0.5,
+                "random_mm_limit_mm_per_prompt": {
+                    "image": 1,
+                    "video": 1
+                },
+                "random_mm_bucket_config": {
+                    "(256, 256, 1)": 0.5,
+                    "(720, 1280, 2)": 0.5
+                },
+                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
+                "baseline": {
+                    "mean_ttft_ms": [6000],
+                    "mean_audio_ttfp_ms": [6000],
+                    "mean_audio_rtf": [0.7]
+                }
+            },
+            {
+                "dataset_name": "random-mm",
+                "backend": "openai-chat-omni",
+                "endpoint": "/v1/chat/completions",
+                "num_prompts": [100],
+                "request_rate": [1.0],
+                "random_input_len": 100,
+                "random_output_len": 100,
+                "random_range_ratio": 0.0,
+                "ignore_eos": true,
+                "random_mm_base_items_per_request": 3,
+                "random_mm_num_mm_items_range_ratio": 0.5,
+                "random_mm_limit_mm_per_prompt": {
+                    "image": 1,
+                    "video": 1,
+                    "audio": 1
+                },
+                "random_mm_bucket_config": {
+                    "(256, 256, 1)": 0.34,
+                    "(720, 1280, 2)": 0.33,
+                    "(0, 60, 3)": 0.33
+                },
+                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
+                "baseline": {
+                    "mean_ttft_ms": [12000],
+                    "mean_audio_ttfp_ms": [12000],
+                    "mean_audio_rtf": [1.0]
+                }
+            },
+            {
+                "dataset_name": "random",
+                "backend": "openai-chat-omni",
+                "endpoint": "/v1/chat/completions",
+                "num_prompts": [4, 16, 32, 64, 128],
+                "max_concurrency": [1, 4, 8, 16, 32],
+                "random_input_len": 2500,
+                "random_output_len": 900,
+                "ignore_eos": true,
+                "extra_body": {
+                    "modalities": ["text"]
+                },
+                "percentile-metrics": "ttft,tpot,itl,e2el",
+                "baseline": {
+                    "mean_ttft_ms": [1000, 3000, 5000, 7000, 9000]
+                }
+            },
+            {
+                "dataset_name": "random-mm",
+                "backend": "openai-chat-omni",
+                "endpoint": "/v1/chat/completions",
+                "num_prompts": [10],
+                "request_rate": [0.1],
+                "random_input_len": 100,
+                "random_output_len": 100,
+                "random_range_ratio": 0.0,
+                "ignore_eos": true,
+                "extra_body": {
+                    "modalities": ["text"]
+                },
+                "random_mm_base_items_per_request": 1,
+                "random_mm_num_mm_items_range_ratio": 0.5,
+                "random_mm_limit_mm_per_prompt": {
+                    "audio": 1
+                },
+                "random_mm_bucket_config": {
+                    "(0, 60, 3)": 1.0
+                },
+                "percentile-metrics": "ttft,tpot,itl,e2el",
+                "baseline": {
+                    "mean_ttft_ms": [2000]
+                }
+            },
+            {
+                "dataset_name": "random-mm",
+                "backend": "openai-chat-omni",
+                "endpoint": "/v1/chat/completions",
+                "num_prompts": [40],
+                "request_rate": [0.5],
+                "random_input_len": 100,
+                "random_output_len": 100,
+                "random_range_ratio": 0.0,
+                "ignore_eos": true,
+                "extra_body": {
+                    "modalities": ["text"]
+                },
+                "random_mm_base_items_per_request": 2,
+                "random_mm_num_mm_items_range_ratio": 0.5,
+                "random_mm_limit_mm_per_prompt": {
+                    "image": 1,
+                    "video": 1
+                },
+                "random_mm_bucket_config": {
+                    "(256, 256, 1)": 0.5,
+                    "(720, 1280, 2)": 0.5
+                },
+                "percentile-metrics": "ttft,tpot,itl,e2el",
+                "baseline": {
+                    "mean_ttft_ms": [6000]
+                }
+            },
+            {
+                "dataset_name": "random-mm",
+                "backend": "openai-chat-omni",
+                "endpoint": "/v1/chat/completions",
+                "num_prompts": [100],
+                "request_rate": [1.0],
+                "random_input_len": 100,
+                "random_output_len": 100,
+                "random_range_ratio": 0.0,
+                "ignore_eos": true,
+                "extra_body": {
+                    "modalities": ["text"]
+                },
+                "random_mm_base_items_per_request": 3,
+                "random_mm_num_mm_items_range_ratio": 0.5,
+                "random_mm_limit_mm_per_prompt": {
+                    "image": 1,
+                    "video": 1,
+                    "audio": 1
+                },
+                "random_mm_bucket_config": {
+                    "(256, 256, 1)": 0.34,
+                    "(720, 1280, 2)": 0.33,
+                    "(0, 60, 3)": 0.33
+                },
+                "percentile-metrics": "ttft,tpot,itl,e2el",
+                "baseline": {
+                    "mean_ttft_ms": [6000]
+                }
+            }
+        ]
+    },
+    {
+        "test_name": "test_qwen3_omni_vllm_text",
+        "server_params": {
+            "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+            "use_omni": false,
+            "extra_cli_args": ["--no-enable-prefix-caching"]
+        },
+        "benchmark_params": [
+            {
+                "dataset_name": "random",
+                "backend": "vllm",
+                "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+                "endpoint": "/v1/completions",
+                "num_prompts": [4, 16, 32, 64, 128],
+                "max_concurrency": [1, 4, 8, 16, 32],
+                "random_input_len": 2500,
+                "random_output_len": 900,
+                "temperature": 0.4,
+                "top_p": 0.9,
+                "top_k": 1,
+                "seed": 42,
+                "repetition_penalty": 1.05,
+                "ignore_eos": true,
+                "extra_body": {
+                    "modalities": ["text"]
+                },
+                "percentile-metrics": "ttft,tpot,itl,e2el",
+                "baseline": {
+                    "mean_ttft_ms": [1000, 3000, 5000, 7000, 9000]
+                }
+            },
+            {
+                "dataset_name": "random-mm",
+                "backend": "vllm",
+                "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+                "endpoint": "/v1/completions",
+                "num_prompts": [10],
+                "request_rate": [0.1],
+                "random_input_len": 100,
+                "random_output_len": 100,
+                "random_range_ratio": 0.0,
+                "temperature": 0.4,
+                "top_p": 0.9,
+                "top_k": 1,
+                "seed": 42,
+                "repetition_penalty": 1.05,
+                "ignore_eos": true,
+                "extra_body": {
+                    "modalities": ["text"]
+                },
+                "random_mm_base_items_per_request": 1,
+                "random_mm_num_mm_items_range_ratio": 0.5,
+                "random_mm_limit_mm_per_prompt": {
+                    "audio": 1
+                },
+                "random_mm_bucket_config": {
+                    "(0, 60, 3)": 1.0
+                },
+                "percentile-metrics": "ttft,tpot,itl,e2el",
+                "baseline": {
+                    "mean_ttft_ms": [2000]
+                }
+            },
+            {
+                "dataset_name": "random-mm",
+                "backend": "vllm",
+                "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+                "endpoint": "/v1/completions",
+                "num_prompts": [40],
+                "request_rate": [0.5],
+                "random_input_len": 100,
+                "random_output_len": 100,
+                "random_range_ratio": 0.0,
+                "temperature": 0.4,
+                "top_p": 0.9,
+                "top_k": 1,
+                "seed": 42,
+                "repetition_penalty": 1.05,
+                "ignore_eos": true,
+                "extra_body": {
+                    "modalities": ["text"]
+                },
+                "random_mm_base_items_per_request": 2,
+                "random_mm_num_mm_items_range_ratio": 0.5,
+                "random_mm_limit_mm_per_prompt": {
+                    "image": 1,
+                    "video": 1
+                },
+                "random_mm_bucket_config": {
+                    "(256, 256, 1)": 0.5,
+                    "(720, 1280, 2)": 0.5
+                },
+                "percentile-metrics": "ttft,tpot,itl,e2el",
+                "baseline": {
+                    "mean_ttft_ms": [6000]
+                }
+            },
+            {
+                "dataset_name": "random-mm",
+                "backend": "vllm",
+                "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+                "endpoint": "/v1/completions",
+                "num_prompts": [100],
+                "request_rate": [1.0],
+                "random_input_len": 100,
+                "random_output_len": 100,
+                "random_range_ratio": 0.0,
+                "temperature": 0.4,
+                "top_p": 0.9,
+                "top_k": 1,
+                "seed": 42,
+                "repetition_penalty": 1.05,
+                "ignore_eos": true,
+                "extra_body": {
+                    "modalities": ["text"]
+                },
+                "random_mm_base_items_per_request": 3,
+                "random_mm_num_mm_items_range_ratio": 0.5,
+                "random_mm_limit_mm_per_prompt": {
+                    "image": 1,
+                    "video": 1,
+                    "audio": 1
+                },
+                "random_mm_bucket_config": {
+                    "(256, 256, 1)": 0.34,
+                    "(720, 1280, 2)": 0.33,
+                    "(0, 60, 3)": 0.33
+                },
+                "percentile-metrics": "ttft,tpot,itl,e2el",
+                "baseline": {
+                    "mean_ttft_ms": [6000]
+                }
+            }
+        ]
+    }
+]

tests/dfx/reliability/test_reliability_qwen3_omni.py

@@ -57,7 +57,7 @@ DEPLOY_CONFIGS_DIR = Path(__file__).resolve().parent.parent.parent.parent / "vll
 
 def _default_oom_device_spec() -> str:
     """Use currently visible CUDA ordinals to avoid invalid device index in sidecar."""
-    count = torch.accelerator.device_count()
+    count = torch.cuda.device_count()
     if count <= 0:
         return "0"
     return ",".join(str(i) for i in range(count))

tests/dfx/stability/scripts/test_stability_hunyuan_image.py

@@ -1,69 +0,0 @@
-"""
-Hunyuan-Image stability: OmniServer (diffusion) + ``diffusion_benchmark_serving.py``.
-
-Configuration: ``tests/dfx/stability/tests/test_hunyuan_image.json``.
-"""
-
-from __future__ import annotations
-
-from pathlib import Path
-
-import pytest
-
-from tests.dfx.conftest import (
-    create_benchmark_indices,
-    create_test_parameter_mapping,
-    create_unique_server_params,
-    load_configs,
-)
-from tests.dfx.stability.helpers import _run_one_diffusion_batch, run_stability_benchmark_loop
-
-STABILITY_DIR = Path(__file__).resolve().parent.parent
-DEPLOY_CONFIGS_DIR = STABILITY_DIR / "deploy"
-CONFIG_FILE_PATH = str(STABILITY_DIR / "tests" / "test_hunyuan_image.json")
-DEFAULT_NUM_PROMPTS_PER_BATCH = 20
-STABILITY_SERVER_TIMEOUT_ARGS = ["--stage-init-timeout", "600", "--init-timeout", "900"]
-
-try:
-    BENCHMARK_CONFIGS = load_configs(CONFIG_FILE_PATH)
-except FileNotFoundError:
-    BENCHMARK_CONFIGS = []
-
-test_params = create_unique_server_params(BENCHMARK_CONFIGS, DEPLOY_CONFIGS_DIR) if BENCHMARK_CONFIGS else []
-server_to_benchmark_mapping = create_test_parameter_mapping(BENCHMARK_CONFIGS) if BENCHMARK_CONFIGS else {}
-benchmark_indices = create_benchmark_indices(BENCHMARK_CONFIGS, server_to_benchmark_mapping)
-
-
-@pytest.mark.slow
-@pytest.mark.diffusion
-@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-@pytest.mark.parametrize("stability_benchmark_params", benchmark_indices, indirect=True)
-def test_stability_hunyuan_image(omni_server, stability_benchmark_params):
-    test_name = stability_benchmark_params["test_name"]
-    params = stability_benchmark_params["params"]
-    duration_sec = params.get("duration_sec", 300)
-    num_prompts_per_batch = params.get("num_prompts_per_batch", DEFAULT_NUM_PROMPTS_PER_BATCH)
-    request_rate = params.get("request_rate")
-    max_concurrency = params.get("max_concurrency")
-
-    bench_params = {
-        k: v
-        for k, v in params.items()
-        if k not in ("duration_sec", "request_rate", "max_concurrency", "num_prompts_per_batch")
-    }
-
-    result = run_stability_benchmark_loop(
-        host=omni_server.host,
-        port=omni_server.port,
-        model=omni_server.model,
-        duration_sec=duration_sec,
-        params=bench_params,
-        request_rate=request_rate,
-        max_concurrency=max_concurrency,
-        result_dir=str(STABILITY_DIR),
-        num_prompts_per_batch=num_prompts_per_batch,
-        run_one_batch=_run_one_diffusion_batch,
-    )
-
-    assert result.get("failed", 0) == 0, f"[{test_name}] Failed requests detected: {result.get('errors', [])}"
-    assert result.get("completed", 0) > 0, f"[{test_name}] No requests completed"

tests/dfx/stability/scripts/test_stability_qwen3_omni.py

@@ -35,7 +35,6 @@ benchmark_indices = create_benchmark_indices(BENCHMARK_CONFIGS, server_to_benchm
 
 
 @pytest.mark.slow
-@pytest.mark.omni
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
 @pytest.mark.parametrize("stability_benchmark_params", benchmark_indices, indirect=True)
 def test_stability_qwen3_omni(omni_server, stability_benchmark_params):

tests/dfx/stability/scripts/test_stability_qwen3_tts.py

@@ -35,7 +35,6 @@ benchmark_indices = create_benchmark_indices(BENCHMARK_CONFIGS, server_to_benchm
 
 
 @pytest.mark.slow
-@pytest.mark.tts
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
 @pytest.mark.parametrize("stability_benchmark_params", benchmark_indices, indirect=True)
 def test_stability_qwen3_tts(omni_server, stability_benchmark_params):

tests/dfx/stability/scripts/test_stability_qwen_image.py

@@ -35,7 +35,6 @@ benchmark_indices = create_benchmark_indices(BENCHMARK_CONFIGS, server_to_benchm
 
 
 @pytest.mark.slow
-@pytest.mark.diffusion
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
 @pytest.mark.parametrize("stability_benchmark_params", benchmark_indices, indirect=True)
 def test_stability_qwen_image(omni_server, stability_benchmark_params):

tests/dfx/stability/scripts/test_stability_wan22.py

@@ -35,7 +35,6 @@ benchmark_indices = create_benchmark_indices(BENCHMARK_CONFIGS, server_to_benchm
 
 
 @pytest.mark.slow
-@pytest.mark.diffusion
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
 @pytest.mark.parametrize("stability_benchmark_params", benchmark_indices, indirect=True)
 def test_stability_wan22(omni_server, stability_benchmark_params):

tests/dfx/stability/tests/test_hunyuan_image.json

@@ -1,36 +0,0 @@
-[
-    {
-        "test_name": "test_hunyuan_image_stability_tp2_fp8_sp2",
-        "server_params": {
-            "model": "tencent/HunyuanImage-3.0-Instruct",
-            "serve_args": {
-                "tensor-parallel-size": 2,
-                "usp": 2,
-                "quantization": "fp8",
-                "distributed-executor-backend": "mp",
-                "enforce-eager": true,
-                "enable-diffusion-pipeline-profiler": true
-            }
-        },
-        "benchmark_params": [
-            {
-                "dataset": "random",
-                "task": "t2i",
-                "backend": "vllm-omni",
-                "duration_sec": 86400,
-                "max_concurrency": 1,
-                "num_prompts_per_batch": 50,
-                "width": {
-                    "min": 512,
-                    "max": 1024
-                },
-                "height": {
-                    "min": 512,
-                    "max": 1024
-                },
-                "num_inference_steps": 50,
-                "enable_negative_prompt": true
-            }
-        ]
-    }
-]

tests/dfx/stability/tests/test_qwen3_omni.json

@@ -1,8 +1,60 @@
 [
+    {
+        "test_name": "test_qwen3_omni_stability",
+        "server_params": {
+            "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+            "stage_overrides": {
+                "2": {
+                    "max_num_batched_tokens": 1000000
+                }
+            }
+        },
+        "benchmark_params": [
+            {
+                "dataset_name": "random-mm",
+                "backend": "openai-chat-omni",
+                "endpoint": "/v1/chat/completions",
+                "duration_sec": 86400,
+                "request_rate": 0.3,
+                "num_prompts_per_batch": 10,
+                "random_input_len": {
+                    "min": 0,
+                    "max": 8000
+                },
+                "random_output_len": {
+                    "min": 0,
+                    "max": 1000
+                },
+                "random_range_ratio": 0.0,
+                "random_mm_base_items_per_request": {
+                    "min": 0,
+                    "max": 6
+                },
+                "random_mm_num_mm_items_range_ratio": 0.0,
+                "random_mm_limit_mm_per_prompt": {
+                    "image": 2,
+                    "video": 2,
+                    "audio": 2
+                },
+                "random_mm_bucket_config": {
+                    "(128-1024, 128-1024, 1)": 0.34,
+                    "(256-1080, 256-1920, 2-16)": 0.33,
+                    "(0, 1-60, 1-3)": 0.33
+                },
+                "ignore_eos": true,
+                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration"
+            }
+        ]
+    },
     {
         "test_name": "test_qwen3_omni_stability_async_chunk",
         "server_params": {
             "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+            "stage_overrides": {
+                "2": {
+                    "max_num_batched_tokens": 1000000
+                }
+            },
             "extra_cli_args": ["--async-chunk"]
         },
         "benchmark_params": [
@@ -12,7 +64,7 @@
                 "endpoint": "/v1/chat/completions",
                 "duration_sec": 86400,
                 "max_concurrency": 2,
-                "num_prompts_per_batch": 50,
+                "num_prompts_per_batch": 10,
                 "random_input_len": {
                     "min": 0,
                     "max": 8000

tests/dfx/stability/tests/test_qwen3_tts.json

@@ -5,13 +5,36 @@
             "model": "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice"
         },
         "benchmark_params": [
+            {
+                "dataset_name": "random",
+                "backend": "openai-audio-speech",
+                "endpoint": "/v1/audio/speech",
+                "duration_sec": 86400,
+                "request_rate": 0.3,
+                "num_prompts_per_batch": 10,
+                "random_input_len": {
+                    "min": 0,
+                    "max": 1000
+                },
+                "random_output_len": {
+                    "min": 0,
+                    "max": 1000
+                },
+                "random_range_ratio": 0.0,
+                "extra_body": {
+                    "voice": "Vivian",
+                    "language": "English"
+                },
+                "ignore_eos": true,
+                "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration"
+            },
             {
                 "dataset_name": "random",
                 "backend": "openai-audio-speech",
                 "endpoint": "/v1/audio/speech",
                 "duration_sec": 86400,
                 "max_concurrency": 2,
-                "num_prompts_per_batch": 50,
+                "num_prompts_per_batch": 10,
                 "random_input_len": {
                     "min": 0,
                     "max": 1000

tests/dfx/stability/tests/test_qwen_image.json

@@ -11,7 +11,7 @@
                 "backend": "vllm-omni",
                 "duration_sec": 86400,
                 "max_concurrency": 1,
-                "num_prompts_per_batch": 50,
+                "num_prompts_per_batch": 10,
                 "width": {
                     "min": 512,
                     "max": 2048

tests/dfx/stability/tests/test_wan22.json

@@ -19,7 +19,7 @@
                 "backend": "v1/videos",
                 "duration_sec": 86400,
                 "max_concurrency": 1,
-                "num_prompts_per_batch": 50,
+                "num_prompts_per_batch": 20,
                 "enable_negative_prompt": true,
                 "random_request_config": [
                     {"width": 832, "height": 480, "num_inference_steps": 4, "num_frames": 81, "fps": 16, "weight": 0.5},

tests/diffusion/attention/test_attention_config.py

@@ -1,419 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""Tests for per-role attention backend configuration (RFC: per-role-attention-backend).
-
-Tests cover:
-- AttentionSpec and AttentionConfig normalization
-- Role-aware backend resolution with category fallback
-- OmniDiffusionConfig attention shorthand handling
-- AttentionMetadata.extra field
-"""
-
-from types import SimpleNamespace
-
-import pytest
-import torch
-
-import vllm_omni.diffusion.attention.layer as layer_mod
-from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
-from vllm_omni.diffusion.attention.layer import Attention
-from vllm_omni.diffusion.config import (
-    get_current_diffusion_config,
-    get_current_diffusion_config_or_none,
-    set_current_diffusion_config,
-)
-from vllm_omni.diffusion.data import (
-    AttentionConfig,
-    AttentionSpec,
-    OmniDiffusionConfig,
-    build_attention_config,
-    parse_attention_config,
-)
-
-
-class TestAttentionSpec:
-    def test_construct_no_extra(self):
-        spec = AttentionSpec(backend="FLASH_ATTN")
-        assert spec.extra == {}
-
-    def test_mapping_extra_normalized(self):
-        spec = AttentionSpec(backend="SAGE_ATTN", extra={"quant": "int8"})
-        assert spec.backend == "SAGE_ATTN"
-        assert spec.extra == {"quant": "int8"}
-
-    def test_invalid_backend_type(self):
-        with pytest.raises(TypeError):
-            AttentionSpec(backend=123)  # type: ignore[arg-type]
-
-
-class TestAttentionConfig:
-    def test_empty_config(self):
-        config = AttentionConfig()
-        assert config.default is None
-        assert config.per_role == {}
-
-    def test_constructor_normalizes_mappings(self):
-        config = AttentionConfig(
-            default={"backend": "FLASH_ATTN"},
-            per_role={
-                "self": {"backend": "SPARSE_BLOCK", "extra": {"block_size": 128}},
-                "cross": "SAGE_ATTN",
-            },
-        )
-        assert config.default.backend == "FLASH_ATTN"
-        assert config.per_role["self"].backend == "SPARSE_BLOCK"
-        assert config.per_role["self"].extra == {"block_size": 128}
-        assert config.per_role["cross"].backend == "SAGE_ATTN"
-
-    def test_constructor_flattens_nested_per_role_tree(self):
-        config = AttentionConfig(
-            per_role={
-                "ltx2": {
-                    "audio_self": {"backend": "FLASH_ATTN"},
-                    "audio_to_video": {"backend": "SAGE_ATTN"},
-                }
-            }
-        )
-        assert config.per_role["ltx2.audio_self"].backend == "FLASH_ATTN"
-        assert config.per_role["ltx2.audio_to_video"].backend == "SAGE_ATTN"
-
-    def test_constructor_normalizes_auto_to_unset(self):
-        config = AttentionConfig(
-            default={"backend": "auto"},
-            per_role={
-                "self": "auto",
-                "cross": {"backend": "SAGE_ATTN"},
-            },
-        )
-        assert config.default is None
-        assert "self" not in config.per_role
-        assert config.per_role["cross"].backend == "SAGE_ATTN"
-
-    def test_resolve_exact_match(self):
-        config = AttentionConfig(
-            default=AttentionSpec(backend="FLASH_ATTN"),
-            per_role={
-                "self": AttentionSpec(backend="SPARSE_BLOCK"),
-                "cross": AttentionSpec(backend="SAGE_ATTN"),
-            },
-        )
-        spec, _ = config.resolve_with_source(role="self")
-        assert spec.backend == "SPARSE_BLOCK"
-
-        spec, _ = config.resolve_with_source(role="cross")
-        assert spec.backend == "SAGE_ATTN"
-
-    def test_resolve_with_source_reports_match_origin(self):
-        config = AttentionConfig(
-            default=AttentionSpec(backend="FLASH_ATTN"),
-            per_role={
-                "cross": AttentionSpec(backend="SAGE_ATTN"),
-                "ltx2.audio_to_video": AttentionSpec(backend="SPARSE_BLOCK"),
-            },
-        )
-
-        spec, source = config.resolve_with_source(role="ltx2.audio_to_video", role_category="cross")
-        assert spec is not None
-        assert spec.backend == "SPARSE_BLOCK"
-        assert source == "attention_config.per_role['ltx2.audio_to_video']"
-
-        spec, source = config.resolve_with_source(role="ltx2.video_to_audio", role_category="cross")
-        assert spec is not None
-        assert spec.backend == "SAGE_ATTN"
-        assert source == "attention_config.per_role['cross'] (role_category fallback)"
-
-        spec, source = config.resolve_with_source(role="self")
-        assert spec is not None
-        assert spec.backend == "FLASH_ATTN"
-        assert source == "attention_config.default"
-
-    def test_resolve_category_fallback(self):
-        config = AttentionConfig(
-            default=AttentionSpec(backend="FLASH_ATTN"),
-            per_role={
-                "cross": AttentionSpec(backend="SAGE_ATTN"),
-            },
-        )
-        # "ltx2.audio_to_video" falls back to category "cross"
-        spec, _ = config.resolve_with_source(role="ltx2.audio_to_video", role_category="cross")
-        assert spec.backend == "SAGE_ATTN"
-
-    def test_resolve_exact_overrides_category(self):
-        config = AttentionConfig(
-            per_role={
-                "cross": AttentionSpec(backend="SAGE_ATTN"),
-                "ltx2.audio_to_video": AttentionSpec(backend="FLASH_ATTN"),
-            },
-        )
-        # Exact match wins over category
-        spec, _ = config.resolve_with_source(role="ltx2.audio_to_video", role_category="cross")
-        assert spec.backend == "FLASH_ATTN"
-
-    def test_resolve_default_fallback(self):
-        config = AttentionConfig(
-            default=AttentionSpec(backend="FLASH_ATTN"),
-        )
-        spec, _ = config.resolve_with_source(role="self")
-        assert spec.backend == "FLASH_ATTN"
-
-        spec, _ = config.resolve_with_source(role="joint")
-        assert spec.backend == "FLASH_ATTN"
-
-    def test_resolve_returns_none_when_empty(self):
-        config = AttentionConfig()
-        spec, _ = config.resolve_with_source(role="self")
-        assert spec is None
-
-    def test_resolve_no_category_no_default(self):
-        config = AttentionConfig(
-            per_role={"self": AttentionSpec(backend="SPARSE_BLOCK")},
-        )
-        # Unknown role with no category and no default
-        spec, _ = config.resolve_with_source(role="joint")
-        assert spec is None
-
-    def test_full_ltx2_scenario(self):
-        """Test the LTX2 6-role stress test from the RFC."""
-        config = AttentionConfig(
-            default=AttentionSpec(backend="FLASH_ATTN"),
-            per_role={
-                "self": AttentionSpec(backend="SPARSE_BLOCK", extra={"block_size": 128}),
-                "cross": AttentionSpec(backend="SAGE_ATTN"),
-                "ltx2.audio_self": AttentionSpec(backend="FLASH_ATTN"),
-                "ltx2.audio_to_video": AttentionSpec(backend="FLASH_ATTN", extra={"causal_window": 64}),
-            },
-        )
-
-        # video self → exact match "self"
-        assert config.resolve_with_source("self")[0].backend == "SPARSE_BLOCK"
-
-        # audio self → exact match "ltx2.audio_self"
-        assert config.resolve_with_source("ltx2.audio_self", "self")[0].backend == "FLASH_ATTN"
-
-        # video-text cross → exact match "cross"
-        assert config.resolve_with_source("cross")[0].backend == "SAGE_ATTN"
-
-        # audio-text cross → category fallback to "cross"
-        assert config.resolve_with_source("ltx2.audio_text_cross", "cross")[0].backend == "SAGE_ATTN"
-
-        # audio-to-video → exact match
-        spec, _ = config.resolve_with_source("ltx2.audio_to_video", "cross")
-        assert spec.backend == "FLASH_ATTN"
-        assert spec.extra == {"causal_window": 64}
-
-        # video-to-audio → category fallback to "cross"
-        assert config.resolve_with_source("ltx2.video_to_audio", "cross")[0].backend == "SAGE_ATTN"
-
-
-class TestAttentionMetadataExtra:
-    def test_default_extra_is_empty(self):
-        meta = AttentionMetadata()
-        assert meta.extra == {}
-
-    def test_extra_passthrough(self):
-        block_mask = torch.ones(4, 4)
-        meta = AttentionMetadata(extra={"block_mask": block_mask, "kv_indices": [0, 1, 2]})
-        assert torch.equal(meta.extra["block_mask"], block_mask)
-        assert meta.extra["kv_indices"] == [0, 1, 2]
-
-    def test_extra_does_not_affect_existing_fields(self):
-        mask = torch.ones(2, 8)
-        meta = AttentionMetadata(attn_mask=mask, extra={"foo": "bar"})
-        assert meta.attn_mask is mask
-        assert meta.extra == {"foo": "bar"}
-
-
-class TestBuildAttentionConfig:
-    def test_env_sets_default_when_no_higher_priority_input(self, monkeypatch):
-        monkeypatch.setenv("DIFFUSION_ATTENTION_BACKEND", "TORCH_SDPA")
-
-        config = build_attention_config()
-
-        assert config.default is not None
-        assert config.default.backend == "TORCH_SDPA"
-
-    def test_attention_backend_overrides_env(self, monkeypatch):
-        monkeypatch.setenv("DIFFUSION_ATTENTION_BACKEND", "TORCH_SDPA")
-
-        config = parse_attention_config(attention_backend="SAGE_ATTN")
-
-        assert config.default is not None
-        assert config.default.backend == "SAGE_ATTN"
-
-    def test_parse_attention_config_does_not_read_env(self, monkeypatch):
-        monkeypatch.setenv("DIFFUSION_ATTENTION_BACKEND", "TORCH_SDPA")
-
-        config = parse_attention_config()
-
-        assert config.default is None
-
-    def test_attention_backend_auto_disables_env_fallback(self, monkeypatch):
-        monkeypatch.setenv("DIFFUSION_ATTENTION_BACKEND", "TORCH_SDPA")
-
-        config = parse_attention_config(attention_backend="auto")
-
-        assert config.default is None
-
-    def test_explicit_default_ignores_env(self, monkeypatch):
-        monkeypatch.setenv("DIFFUSION_ATTENTION_BACKEND", "self=FLASH_ATTN,cross=TORCH_SDPA")
-
-        config = build_attention_config(
-            AttentionConfig(default=AttentionSpec(backend="FLASH_ATTN")),
-        )
-
-        assert config.default is not None
-        assert config.default.backend == "FLASH_ATTN"
-
-    def test_env_auto_does_not_set_default(self, monkeypatch):
-        monkeypatch.setenv("DIFFUSION_ATTENTION_BACKEND", "auto")
-
-        config = build_attention_config()
-
-        assert config.default is None
-
-    def test_attention_backend_conflicts_with_explicit_default(self):
-        with pytest.raises(ValueError):
-            parse_attention_config(
-                AttentionConfig(default=AttentionSpec(backend="FLASH_ATTN")),
-                attention_backend="SAGE_ATTN",
-            )
-
-
-class TestOmniDiffusionConfigAttentionParsing:
-    """Test OmniDiffusionConfig attention shorthand and structured config."""
-
-    def test_diffusion_attention_backend_sets_default(self):
-        config = OmniDiffusionConfig.from_kwargs(diffusion_attention_backend="SAGE_ATTN")
-        assert isinstance(config.diffusion_attention_config, AttentionConfig)
-        assert config.diffusion_attention_config.default is not None
-        assert config.diffusion_attention_config.default.backend == "SAGE_ATTN"
-
-    def test_diffusion_attention_backend_auto_means_platform_default(self):
-        config = OmniDiffusionConfig.from_kwargs(diffusion_attention_backend="auto")
-        assert isinstance(config.diffusion_attention_config, AttentionConfig)
-        assert config.diffusion_attention_config.default is None
-
-    def test_diffusion_attention_backend_and_default_are_mutually_exclusive(self):
-        with pytest.raises(ValueError):
-            OmniDiffusionConfig.from_kwargs(
-                diffusion_attention_backend="SAGE_ATTN",
-                diffusion_attention_config=AttentionConfig(default=AttentionSpec(backend="FLASH_ATTN")),
-            )
-
-    def test_dict_diffusion_attention_config(self):
-        config = OmniDiffusionConfig(
-            diffusion_attention_config={
-                "default": {"backend": "FLASH_ATTN"},
-                "per_role": {"self": "SPARSE_BLOCK"},
-            }
-        )
-        assert config.diffusion_attention_config.default.backend == "FLASH_ATTN"
-        assert config.diffusion_attention_config.per_role["self"].backend == "SPARSE_BLOCK"
-
-    def test_no_diffusion_attention_config_defaults_to_empty(self):
-        config = OmniDiffusionConfig()
-        assert isinstance(config.diffusion_attention_config, AttentionConfig)
-        assert config.diffusion_attention_config.default is None
-        assert config.diffusion_attention_config.per_role == {}
-
-
-class TestCurrentDiffusionConfig:
-    def test_get_current_diffusion_config_or_none_defaults_to_none(self):
-        assert get_current_diffusion_config_or_none() is None
-
-    def test_get_current_diffusion_config_raises_when_unset(self):
-        with pytest.raises(AssertionError, match="Diffusion config is not set"):
-            get_current_diffusion_config()
-
-    def test_set_current_diffusion_config_restores_previous_value(self):
-        outer = SimpleNamespace(name="outer")
-        inner = SimpleNamespace(name="inner")
-
-        with set_current_diffusion_config(outer):
-            assert get_current_diffusion_config() is outer
-            with set_current_diffusion_config(inner):
-                assert get_current_diffusion_config() is inner
-            assert get_current_diffusion_config() is outer
-
-        assert get_current_diffusion_config_or_none() is None
-
-
-class TestAttentionInitUsesCurrentDiffusionConfig:
-    def test_attention_init_uses_current_diffusion_config_without_forward_context(self, monkeypatch):
-        class _FakeAttentionImpl:
-            def __init__(self, **kwargs):
-                self.kwargs = kwargs
-
-            def forward(self, query, key, value, attn_metadata=None):
-                return query
-
-        class _FakeBackend:
-            @staticmethod
-            def get_name() -> str:
-                return "FAKE_BACKEND"
-
-            @staticmethod
-            def get_impl_cls():
-                return _FakeAttentionImpl
-
-        captured = {}
-
-        def _fake_get_attn_backend_for_role(
-            role,
-            head_size,
-            attention_config=None,
-            role_category=None,
-        ):
-            captured["role"] = role
-            captured["head_size"] = head_size
-            captured["role_category"] = role_category
-            captured["attention_config"] = attention_config
-            return _FakeBackend, AttentionSpec(backend="TORCH_SDPA", extra={"block_size": 128})
-
-        class _FakeRingParallelAttention:
-            def __init__(self, sp_group, attn_backend_pref=None):
-                self.sp_group = sp_group
-                self.attn_backend_pref = attn_backend_pref
-
-        monkeypatch.setattr(layer_mod, "get_attn_backend_for_role", _fake_get_attn_backend_for_role)
-        monkeypatch.setattr(layer_mod.SDPABackend, "get_impl_cls", staticmethod(lambda: _FakeAttentionImpl))
-        monkeypatch.setattr(layer_mod, "build_parallel_attention_strategy", lambda **kwargs: object())
-        monkeypatch.setattr(layer_mod, "get_sp_group", lambda: SimpleNamespace(ring_group="ring-group"))
-        monkeypatch.setattr(layer_mod, "RingParallelAttention", _FakeRingParallelAttention)
-        monkeypatch.setattr(layer_mod, "is_forward_context_available", lambda: False)
-        monkeypatch.setattr(
-            layer_mod,
-            "get_forward_context",
-            lambda: (_ for _ in ()).throw(AssertionError("Attention init should not read ForwardContext")),
-        )
-
-        od_config = SimpleNamespace(
-            diffusion_attention_config=AttentionConfig(
-                default=AttentionSpec(backend="FLASH_ATTN"),
-                per_role={"cross": AttentionSpec(backend="TORCH_SDPA", extra={"block_size": 128})},
-            ),
-            parallel_config=SimpleNamespace(ring_degree=2),
-        )
-
-        with set_current_diffusion_config(od_config):
-            attn = Attention(
-                num_heads=4,
-                head_size=64,
-                causal=False,
-                softmax_scale=1.0,
-                role="cross",
-                role_category="cross",
-                qkv_layout="BSND",
-            )
-
-        assert captured["role"] == "cross"
-        assert captured["role_category"] == "cross"
-        assert captured["head_size"] == 64
-        assert captured["attention_config"] is od_config.diffusion_attention_config
-        assert attn.backend_pref == "TORCH_SDPA"
-        assert attn.attention.kwargs["backend_kwargs"] == {"block_size": 128}
-        assert attn.attention.kwargs["qkv_layout"] == "BSND"
-        assert attn.use_ring is True
-        assert attn.ring_runner is not None
-        assert attn.ring_runner.attn_backend_pref == "TORCH_SDPA"

tests/diffusion/attention/test_attention_sp.py

@@ -32,7 +32,6 @@ import pytest
 import torch
 
 from vllm_omni.diffusion.attention.layer import Attention
-from vllm_omni.diffusion.config import set_current_diffusion_config
 from vllm_omni.diffusion.data import (
     DiffusionParallelConfig,
     OmniDiffusionConfig,
@@ -404,8 +403,8 @@ def ulysses_attention_on_test_model(
         pipeline_parallel_size=1,
     )
 
-    # Set the config so Attention can access it during init and forward
-    with set_forward_context(omni_diffusion_config=od_config), set_current_diffusion_config(od_config):
+    # Set the config so Attention can access it
+    with set_forward_context(omni_diffusion_config=od_config):
         # Create model
         hidden_size = num_heads * head_size
 

tests/diffusion/attention/test_piecewise_attn.py

@@ -1,126 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""End-to-end test for ``piecewise_attn`` (CPU).
-
-Verify that running attention in segments (causal outside full-attn spans,
-bidirectional inside full-attn spans) matches running a single full SDPA call
-with the equivalent 2D attention mask.
-
-Covers:
-  * batch size = 1 and batch size > 1 (homogeneous CFG-like batch)
-  * query length == key length   (full prefill)
-  * query length <  key length   (decode-like tail slice)
-  * various full-attn-span layouts (none / start / middle / end / multi)
-"""
-
-from __future__ import annotations
-
-import pytest
-import torch
-import torch.nn.functional as F
-
-from vllm_omni.diffusion.attention.backends.utils.piecewise_attn import (
-    piecewise_attn,
-)
-
-DEVICE = torch.device("cpu")
-
-
-def _sdpa_attn_func(q, k, v, causal, softmax_scale):
-    q_ = q.transpose(1, 2)
-    k_ = k.transpose(1, 2)
-    v_ = v.transpose(1, 2)
-    attn_mask = None
-    if causal:
-        Sq, Sk = q_.shape[-2], k_.shape[-2]
-        i = torch.arange(Sq, device=q.device).unsqueeze(1)
-        j = torch.arange(Sk, device=q.device).unsqueeze(0)
-        attn_mask = j <= (i + (Sk - Sq))
-    out = F.scaled_dot_product_attention(q_, k_, v_, attn_mask=attn_mask, scale=softmax_scale)
-    return out.transpose(1, 2).contiguous()
-
-
-def _full_reference(query, key, value, global_spans, q_start, q_end, softmax_scale):
-    """Build a full 2D mask with global spans and compute reference output."""
-    Sk = key.shape[1]
-    mask = torch.tril(torch.ones(Sk, Sk, dtype=torch.bool, device=key.device))
-    for a, e in global_spans:
-        mask[a:e, :e] = True
-    mask_q = mask[q_start:q_end, :]
-    q_ = query.transpose(1, 2)
-    k_ = key.transpose(1, 2)
-    v_ = value.transpose(1, 2)
-    out = F.scaled_dot_product_attention(q_, k_, v_, attn_mask=mask_q, scale=softmax_scale)
-    return out.transpose(1, 2).contiguous()
-
-
-SPAN_CASES = [
-    pytest.param([], id="no-spans"),
-    pytest.param([(0, 10)], id="span-at-start"),
-    pytest.param([(10, 30), (54, 64)], id="multi-spans"),
-]
-
-Q_RANGE_CASES = [
-    pytest.param((0, 64), id="q_eq_k"),  # Sq == Sk (prefill)
-    pytest.param((53, 64), id="q_lt_k"),  # Sq < Sk (decode-like)
-]
-
-BATCH_CASES = [
-    pytest.param(1, id="B1"),
-    pytest.param(2, id="B2"),
-]
-
-
-@pytest.mark.parametrize("global_spans", SPAN_CASES)
-@pytest.mark.parametrize("q_range", Q_RANGE_CASES)
-@pytest.mark.parametrize("batch_size", BATCH_CASES)
-def test_piecewise_matches_full(global_spans, q_range, batch_size):
-    torch.manual_seed(0)
-    H, D, Sk = 2, 16, 64
-    q_start, q_end = q_range
-    Sq = q_end - q_start
-
-    key = torch.randn(batch_size, Sk, H, D, device=DEVICE)
-    value = torch.randn(batch_size, Sk, H, D, device=DEVICE)
-    query = torch.randn(batch_size, Sq, H, D, device=DEVICE)
-
-    full_attn_spans = [list(global_spans) for _ in range(batch_size)]
-    softmax_scale = 1.0 / (D**0.5)
-
-    got = piecewise_attn(
-        query,
-        key,
-        value,
-        full_attn_spans=full_attn_spans,
-        softmax_scale=softmax_scale,
-        attn_func=_sdpa_attn_func,
-    )
-    expected = _full_reference(query, key, value, global_spans, q_start, q_end, softmax_scale)
-    torch.testing.assert_close(got, expected, atol=1e-5, rtol=1e-5)
-
-
-def test_piecewise_span_fully_before_qstart():
-    """Spans entirely before query region produce pure causal attention."""
-    torch.manual_seed(0)
-    B, H, D, Sk = 1, 2, 16, 30
-    q_start, q_end = 15, 30
-    Sq = q_end - q_start
-
-    key = torch.randn(B, Sk, H, D, device=DEVICE)
-    value = torch.randn(B, Sk, H, D, device=DEVICE)
-    query = torch.randn(B, Sq, H, D, device=DEVICE)
-
-    global_spans = [(5, 10)]
-    full_attn_spans = [list(global_spans) for _ in range(B)]
-    softmax_scale = 1.0 / (D**0.5)
-
-    got = piecewise_attn(
-        query,
-        key,
-        value,
-        full_attn_spans=full_attn_spans,
-        softmax_scale=softmax_scale,
-        attn_func=_sdpa_attn_func,
-    )
-    expected = _full_reference(query, key, value, global_spans, q_start, q_end, softmax_scale)
-    torch.testing.assert_close(got, expected, atol=1e-5, rtol=1e-5)

tests/diffusion/attention/test_ulysses_uaa.py

@@ -12,7 +12,6 @@ import pytest
 import torch
 
 from vllm_omni.diffusion.attention.layer import Attention
-from vllm_omni.diffusion.config import set_current_diffusion_config
 from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.parallel_state import (
     destroy_distributed_env,
@@ -80,7 +79,7 @@ def _run_attention_case(
     )
     od_config = OmniDiffusionConfig(model="test", dtype=torch.float32, parallel_config=parallel_config)
 
-    with set_forward_context(omni_diffusion_config=od_config), set_current_diffusion_config(od_config):
+    with set_forward_context(omni_diffusion_config=od_config):
         attn = Attention(
             num_heads=num_heads,
             head_size=head_size,

tests/diffusion/cache/test_teacache_extractors.py

@@ -303,10 +303,7 @@ class TestFluxExtractor(BaseExtractorTest):
         """Use the SDPA backend so FLUX can be instantiated in CPU tests."""
         from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend
 
-        with patch(
-            "vllm_omni.diffusion.attention.layer.get_attn_backend_for_role",
-            return_value=(SDPABackend, None),
-        ):
+        with patch("vllm_omni.diffusion.attention.layer.get_attn_backend", return_value=SDPABackend):
             yield
 
     def get_extractor(self):

tests/diffusion/layers/test_norm.py

@@ -4,7 +4,6 @@
 
 import pytest
 import torch
-from pytest_mock import MockerFixture
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
 
@@ -299,55 +298,6 @@ def test_rmsnorm_numerical_correctness():
     torch.testing.assert_close(out, expected, atol=1e-5, rtol=1e-5)
 
 
-# ── RMSNorm compile-path regression tests ──
-
-
-def test_rmsnorm_forward_cuda_does_not_call_fused_during_compile(mocker: MockerFixture) -> None:
-    """Regression: _forward_fused must not be called during torch.compile tracing.
-
-    Under HSDP, RMSNorm.weight is a DTensor. Accessing .data on a DTensor inside
-    _forward_fused during tracing produces an orphan all-gather node outside the
-    compile boundary, causing inductor's compute_ancestors to raise KeyError.
-    The is_compiling() guard in forward_cuda/forward_hip prevents this by routing
-    to forward_native during tracing.
-
-    If someone removes the guard, this test will catch the regression by asserting
-    that _forward_fused was not called while is_compiling() returns True.
-    """
-    from vllm_omni.diffusion.layers.norm import RMSNorm
-
-    norm = RMSNorm(hidden_size=64)
-    x = torch.randn(2, 4, 64)
-
-    mock_fused = mocker.patch.object(norm, "_forward_fused", wraps=norm._forward_fused)
-    mocker.patch("torch.compiler.is_compiling", return_value=True)
-
-    out = norm.forward_cuda(x)
-
-    mock_fused.assert_not_called()
-    assert out.shape == x.shape
-
-
-def test_rmsnorm_forward_hip_does_not_call_fused_during_compile(mocker: MockerFixture) -> None:
-    """Regression: same guard must be present in forward_hip.
-
-    forward_hip is the entry point on ROCm (AMD GPU). It must behave identically
-    to forward_cuda with respect to the is_compiling() guard.
-    """
-    from vllm_omni.diffusion.layers.norm import RMSNorm
-
-    norm = RMSNorm(hidden_size=64)
-    x = torch.randn(2, 4, 64)
-
-    mock_fused = mocker.patch.object(norm, "_forward_fused", wraps=norm._forward_fused)
-    mocker.patch("torch.compiler.is_compiling", return_value=True)
-
-    out = norm.forward_hip(x)
-
-    mock_fused.assert_not_called()
-    assert out.shape == x.shape
-
-
 def test_rmsnorm_matches_reference_implementation():
     """Verify RMSNorm matches a reference implementation."""
     from vllm_omni.diffusion.layers.norm import RMSNorm
@@ -501,97 +451,3 @@ def test_rmsnorm_with_single_element_batch():
     out = norm(x)
 
     assert out.shape == (1, 1, hidden_size)
-
-
-# ── Fused vs Native correctness tests ──
-
-
-def test_rmsnorm_fused_native_parity_cuda():
-    """Verify fused and native RMSNorm produce identical results on CUDA.
-
-    The fused kernel uses vllm._custom_ops.rms_norm while native uses pure PyTorch.
-    Both should produce numerically equivalent results.
-    """
-    from vllm_omni.diffusion.layers.norm import RMSNorm
-
-    hidden_size = 128
-    eps = 1e-6
-    torch.manual_seed(42)
-
-    for dtype in [torch.float32, torch.float16, torch.bfloat16]:
-        torch.manual_seed(42)
-        norm = RMSNorm(hidden_size, eps=eps).cuda()
-        norm.weight.data = norm.weight.data.to(dtype)
-
-        x = torch.randn(4, 8, hidden_size, dtype=dtype, device="cuda")
-
-        out_fused = norm._forward_fused(x)
-        out_native = norm.forward_native(x)
-
-        torch.testing.assert_close(out_fused, out_native, atol=1e-3, rtol=1e-3)
-
-
-def test_rmsnorm_fused_native_parity_large_tensor():
-    """Verify fused and native RMSNorm parity with large tensors.
-
-    Tests realistic tensor sizes used in diffusion models.
-    """
-    from vllm_omni.diffusion.layers.norm import RMSNorm
-
-    hidden_size = 4096
-    batch_size = 16
-    seq_len = 1024
-    eps = 1e-6
-
-    for dtype in [torch.float32, torch.float16, torch.bfloat16]:
-        torch.manual_seed(123)
-        norm = RMSNorm(hidden_size, eps=eps).cuda()
-        norm.weight.data = norm.weight.data.to(dtype)
-
-        x = torch.randn(batch_size, seq_len, hidden_size, dtype=dtype, device="cuda")
-
-        out_fused = norm._forward_fused(x)
-        out_native = norm.forward_native(x)
-
-        torch.testing.assert_close(out_fused, out_native, atol=1e-2, rtol=1e-2)
-
-
-def test_rmsnorm_fused_native_parity_2d_input():
-    """Verify fused and native RMSNorm parity with 2D input tensors."""
-    from vllm_omni.diffusion.layers.norm import RMSNorm
-
-    hidden_size = 256
-    batch_size = 32
-    eps = 1e-6
-
-    for dtype in [torch.float32, torch.float16, torch.bfloat16]:
-        torch.manual_seed(456)
-        norm = RMSNorm(hidden_size, eps=eps).cuda()
-        norm.weight.data = norm.weight.data.to(dtype)
-
-        x = torch.randn(batch_size, hidden_size, dtype=dtype, device="cuda")
-
-        out_fused = norm._forward_fused(x)
-        out_native = norm.forward_native(x)
-
-        torch.testing.assert_close(out_fused, out_native, atol=1e-3, rtol=1e-3)
-
-
-def test_rmsnorm_fused_native_parity_custom_weight():
-    """Verify fused and native RMSNorm parity with non-default weights."""
-    from vllm_omni.diffusion.layers.norm import RMSNorm
-
-    hidden_size = 64
-    eps = 1e-6
-
-    for dtype in [torch.float32, torch.float16, torch.bfloat16]:
-        torch.manual_seed(789)
-        norm = RMSNorm(hidden_size, eps=eps).cuda()
-        norm.weight.data = torch.randn(hidden_size, dtype=dtype, device="cuda")
-
-        x = torch.randn(2, 4, hidden_size, dtype=dtype, device="cuda")
-
-        out_fused = norm._forward_fused(x)
-        out_native = norm.forward_native(x)
-
-        torch.testing.assert_close(out_fused, out_native, atol=1e-3, rtol=1e-3)

tests/diffusion/model_loader/test_diffusers_loader.py

@@ -1,16 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-from types import SimpleNamespace
-
 import pytest
 import torch
 import torch.nn as nn
 
-from vllm_omni.diffusion.config import get_current_diffusion_config, get_current_diffusion_config_or_none
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.model_loader.gguf_adapters import get_gguf_adapter
-from vllm_omni.diffusion.registry import initialize_model
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
 
@@ -97,74 +93,3 @@ def test_qwen_model_class_selects_qwen_gguf_adapter():
     adapter = get_gguf_adapter("dummy.gguf", object(), source, od_config)
 
     assert adapter.__class__.__name__ == "QwenImageGGUFAdapter"
-
-
-class _ConfigAwareModel(nn.Module):
-    def __init__(self, *, od_config):
-        super().__init__()
-        self.captured_config = get_current_diffusion_config()
-        self.seen_config_during_init = get_current_diffusion_config_or_none()
-        self.od_config = od_config
-
-
-def test_initialize_model_sets_current_diffusion_config_during_model_construction(monkeypatch):
-    import vllm_omni.diffusion.registry as registry_mod
-
-    od_config = SimpleNamespace(
-        model_class_name="DummyPipeline",
-        parallel_config=SimpleNamespace(vae_patch_parallel_size=1, sequence_parallel_size=1),
-        vae_use_slicing=False,
-        vae_use_tiling=False,
-    )
-
-    monkeypatch.setattr(
-        registry_mod.DiffusionModelRegistry,
-        "_try_load_model_cls",
-        staticmethod(lambda _name: _ConfigAwareModel),
-    )
-    monkeypatch.setattr(registry_mod, "_apply_sequence_parallel_if_enabled", lambda *_args, **_kwargs: None)
-
-    model = initialize_model(od_config)
-
-    assert model.captured_config is od_config
-    assert model.seen_config_during_init is od_config
-    assert get_current_diffusion_config_or_none() is None
-
-
-def test_load_model_custom_pipeline_sets_current_diffusion_config(monkeypatch):
-    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
-
-    class _DeviceContext:
-        def __init__(self, device_type: str):
-            self.type = device_type
-
-        def __enter__(self):
-            return self
-
-        def __exit__(self, exc_type, exc, tb):
-            return False
-
-    od_config = SimpleNamespace(
-        dtype=torch.float32,
-        parallel_config=SimpleNamespace(use_hsdp=False),
-        quantization_config=None,
-    )
-
-    loader = object.__new__(DiffusersPipelineLoader)
-    loader.load_weights = lambda model: None  # type: ignore[assignment]
-    loader._process_weights_after_loading = lambda model, target_device: None  # type: ignore[assignment]
-    loader._is_gguf_quantization = lambda _od_config: False  # type: ignore[assignment]
-
-    monkeypatch.setattr(loader_mod, "resolve_obj_by_qualname", lambda _name: _ConfigAwareModel)
-    monkeypatch.setattr(loader_mod.torch, "device", lambda _name: _DeviceContext("cpu"))
-
-    model = loader.load_model(
-        od_config,
-        load_device="cpu",
-        load_format="custom_pipeline",
-        custom_pipeline_name="tests.dummy.ConfigAwarePipeline",
-    )
-
-    assert model.captured_config is od_config
-    assert model.seen_config_during_init is od_config
-    assert get_current_diffusion_config_or_none() is None

tests/diffusion/model_loader/test_modelopt_fp8_adapter.py

@@ -1,94 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from types import SimpleNamespace
-
-import pytest
-import torch
-import torch.nn as nn
-
-from vllm_omni.diffusion.model_loader.checkpoint_adapters import (
-    ModelOptFp8CheckpointAdapter,
-)
-
-pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
-
-
-class _PackedModelOptModel(nn.Module):
-    def __init__(self) -> None:
-        super().__init__()
-        self.transformer = nn.Module()
-        self.transformer.block = nn.Module()
-        self.transformer.block.to_qkv = nn.Linear(2, 2, bias=False)
-
-
-class _QuantizedPackedModelOptModel(nn.Module):
-    def __init__(self) -> None:
-        super().__init__()
-        self.transformer = nn.Module()
-        self.transformer.block = nn.Module()
-        self.transformer.block.to_qkv = nn.Module()
-        self.transformer.block.to_qkv.register_parameter(
-            "weight",
-            nn.Parameter(torch.empty(2, 2, dtype=torch.float8_e4m3fn), requires_grad=False),
-        )
-        self.transformer.block.to_qkv.register_parameter(
-            "weight_scale",
-            nn.Parameter(torch.empty(1), requires_grad=False),
-        )
-        self.transformer.block.to_qkv.register_parameter(
-            "input_scale",
-            nn.Parameter(torch.empty(1), requires_grad=False),
-        )
-
-
-def _make_source() -> SimpleNamespace:
-    return SimpleNamespace(
-        subfolder="transformer",
-        prefix="transformer.",
-    )
-
-
-def test_modelopt_adapter_dequantizes_fp8_weight_for_full_precision_target():
-    model = _PackedModelOptModel()
-    adapter = ModelOptFp8CheckpointAdapter(model, _make_source())
-    fp8_weight = torch.tensor([[2.0, -4.0], [1.0, 3.0]], dtype=torch.float32).to(torch.float8_e4m3fn)
-    scale = torch.tensor([0.5], dtype=torch.float32)
-
-    adapted = list(
-        adapter.adapt(
-            iter(
-                [
-                    ("transformer.block.to_q.weight_scale", scale),
-                    ("transformer.block.to_q.input_scale", torch.tensor([1.0])),
-                    ("transformer.block.to_q.weight", fp8_weight),
-                ]
-            )
-        )
-    )
-
-    assert [name for name, _ in adapted] == ["transformer.block.to_q.weight"]
-    assert adapted[0][1].dtype == model.transformer.block.to_qkv.weight.dtype
-    assert torch.allclose(adapted[0][1], fp8_weight.to(torch.float32) * scale)
-
-
-def test_modelopt_adapter_keeps_scale_tensors_for_quantized_target():
-    model = _QuantizedPackedModelOptModel()
-    adapter = ModelOptFp8CheckpointAdapter(model, _make_source())
-    scale = torch.tensor([0.5], dtype=torch.float32)
-
-    adapted = list(
-        adapter.adapt(
-            iter(
-                [
-                    ("transformer.block.to_q.weight_scale", scale),
-                    ("transformer.block.to_q.input_scale", torch.tensor([1.0])),
-                ]
-            )
-        )
-    )
-
-    assert [name for name, _ in adapted] == [
-        "transformer.block.to_q.weight_scale",
-        "transformer.block.to_q.input_scale",
-    ]

tests/diffusion/models/dmd2/test_dmd2_request_sanitization.py

@@ -5,10 +5,8 @@ from unittest.mock import MagicMock, patch
 import pytest
 import torch
 
-from vllm_omni.diffusion.models.flux.pipeline_flux import FluxDMD2Pipeline, FluxPipeline
 from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline, LTX2T2VDMD2Pipeline
 from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_image2video import LTX2I2VDMD2Pipeline, LTX2ImageToVideoPipeline
-from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image import QwenImageDMD2Pipeline, QwenImagePipeline
 from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import Wan22Pipeline, WanT2VDMD2Pipeline
 from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v import Wan22I2VPipeline, WanI2VDMD2Pipeline
 from vllm_omni.diffusion.request import OmniDiffusionRequest, OmniDiffusionSamplingParams
@@ -21,8 +19,6 @@ _DMD2_BASE = {
     WanI2VDMD2Pipeline: Wan22I2VPipeline,
     LTX2T2VDMD2Pipeline: LTX2Pipeline,
     LTX2I2VDMD2Pipeline: LTX2ImageToVideoPipeline,
-    FluxDMD2Pipeline: FluxPipeline,
-    QwenImageDMD2Pipeline: QwenImagePipeline,
 }
 
 
@@ -53,7 +49,7 @@ def _make_request(prompts=None, **sp_kwargs) -> OmniDiffusionRequest:
 
 @pytest.fixture(
     params=list(_DMD2_BASE.keys()),
-    ids=["wan_t2v", "wan_i2v", "ltx2_t2v", "ltx2_i2v", "flux", "qwen_image"],
+    ids=["wan_t2v", "wan_i2v", "ltx2_t2v", "ltx2_i2v"],
 )
 def pipeline(request):
     return _make_pipeline(request.param)
@@ -67,13 +63,13 @@ def pipeline(request):
 def test_num_inference_steps_forced_to_dmd2_value(pipeline):
     req = _make_request(num_inference_steps=40)
     pipeline._sanitize_dmd2_request(req)
-    assert req.sampling_params.num_inference_steps == pipeline.dmd2_config.num_inference_steps
+    assert req.sampling_params.num_inference_steps == pipeline.num_inference_steps
 
 
 def test_num_inference_steps_already_correct(pipeline):
-    req = _make_request(num_inference_steps=pipeline.dmd2_config.num_inference_steps)
+    req = _make_request(num_inference_steps=pipeline.num_inference_steps)
     pipeline._sanitize_dmd2_request(req)
-    assert req.sampling_params.num_inference_steps == pipeline.dmd2_config.num_inference_steps
+    assert req.sampling_params.num_inference_steps == pipeline.num_inference_steps
 
 
 # ---------------------------------------------------------------------------
@@ -84,19 +80,19 @@ def test_num_inference_steps_already_correct(pipeline):
 def test_guidance_scale_forced_to_one(pipeline):
     req = _make_request(guidance_scale=5.0, guidance_scale_provided=True)
     pipeline._sanitize_dmd2_request(req)
-    assert req.sampling_params.guidance_scale == pipeline.dmd2_config.guidance_scale
+    assert req.sampling_params.guidance_scale == pipeline.dmd2_guidance_scale
     assert req.sampling_params.guidance_scale_provided is False
 
 
 def test_guidance_scale_already_correct(pipeline):
-    req = _make_request(guidance_scale=pipeline.dmd2_config.guidance_scale, guidance_scale_provided=False)
+    req = _make_request(guidance_scale=pipeline.dmd2_guidance_scale, guidance_scale_provided=False)
     pipeline._sanitize_dmd2_request(req)
-    assert req.sampling_params.guidance_scale == pipeline.dmd2_config.guidance_scale
+    assert req.sampling_params.guidance_scale == pipeline.dmd2_guidance_scale
 
 
 def test_guidance_scale_provided_flag_cleared(pipeline):
     """guidance_scale_provided=True must be cleared even if scale is already dmd2_guidance_scale."""
-    req = _make_request(guidance_scale=pipeline.dmd2_config.guidance_scale, guidance_scale_provided=True)
+    req = _make_request(guidance_scale=pipeline.dmd2_guidance_scale, guidance_scale_provided=True)
     pipeline._sanitize_dmd2_request(req)
     assert req.sampling_params.guidance_scale_provided is False
 
@@ -168,44 +164,15 @@ def test_multiple_prompts_all_sanitized(pipeline):
 # ---------------------------------------------------------------------------
 
 
-def test_sample_solver_stripped_from_extra_args(pipeline):
-    """[C1] defense: sample_solver must not leak into req for the base pipeline to read."""
-    req = _make_request()
-    req.sampling_params.extra_args = {"sample_solver": "euler"}
-    pipeline._sanitize_dmd2_request(req)
-    assert "sample_solver" not in req.sampling_params.extra_args
-
-
-def test_flow_shift_stripped_from_extra_args(pipeline):
-    """[C1] defense: flow_shift must not leak into req for the base pipeline to read."""
-    req = _make_request()
-    req.sampling_params.extra_args = {"flow_shift": 3.0}
-    pipeline._sanitize_dmd2_request(req)
-    assert "flow_shift" not in req.sampling_params.extra_args
-
-
-def test_unrelated_extra_args_preserved(pipeline):
-    """Sanitizer only strips sample_solver / flow_shift; other extras pass through."""
-    req = _make_request()
-    req.sampling_params.extra_args = {"sample_solver": "euler", "unrelated": 42}
-    pipeline._sanitize_dmd2_request(req)
-    assert req.sampling_params.extra_args == {"unrelated": 42}
-
-
-# ---------------------------------------------------------------------------
-# Clean request — nothing changes
-# ---------------------------------------------------------------------------
-
-
 def test_clean_request_no_changes(pipeline):
     req = _make_request(
-        guidance_scale=pipeline.dmd2_config.guidance_scale,
+        guidance_scale=pipeline.dmd2_guidance_scale,
         guidance_scale_provided=False,
         do_classifier_free_guidance=False,
         is_cfg_negative=False,
     )
     pipeline._sanitize_dmd2_request(req)
-    assert req.sampling_params.guidance_scale == pipeline.dmd2_config.guidance_scale
+    assert req.sampling_params.guidance_scale == pipeline.dmd2_guidance_scale
     assert req.sampling_params.guidance_scale_provided is False
     assert req.sampling_params.guidance_scale_2 is None
     assert req.sampling_params.true_cfg_scale is None