upd

Signed-off-by: hsliu_ustc <hsliu_ustc@noreply.gitcode.com>

.buildkite/bootstrap-amd-omni.sh

@@ -56,6 +56,88 @@ 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
@@ -91,8 +173,24 @@ upload_pipeline() {
 
     cd .buildkite
 
-    # Select test definition file: merge suite for main, ready suite for PRs
-    if [[ $BUILDKITE_BRANCH == "main" ]]; then
+    # 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
         TEST_YAML="test-amd-merge.yml"
     else
         TEST_YAML="test-amd-ready.yaml"
@@ -140,30 +238,13 @@ fi
 # Early exit start: skip pipeline if conditions are met
 # ----------------------------------------------------------------------
 
-# skip pipeline if all changed files are under docs/
+# Match CUDA Buildkite skip-ci behavior for docs/markdown-only changes.
 if [[ "${DOCS_ONLY_DISABLE}" != "1" ]]; then
-  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
+    SKIP_CI="$(resolve_skip_ci)"
+    if [[ "${SKIP_CI}" == "1" ]]; then
+        echo "[docs-only] Docs/markdown-only changes detected. Exiting before pipeline upload."
+        exit 0
     fi
-  fi
 fi
 
 # ----------------------------------------------------------------------

.buildkite/nightly-release-pipeline.yaml

@@ -0,0 +1,20 @@
+steps:
+  - label: "Build and upload wheel"
+    key: "build-wheel"
+    agents:
+      queue: cpu_queue_release
+    commands:
+      - "curl -LsSf https://astral.sh/uv/install.sh | sh"
+      - 'export PATH="$HOME/.local/bin:$PATH"'
+      - "uv venv --python=3.12 && source .venv/bin/activate"
+      - "uv pip install --upgrade build"
+      - "python3 -m build"
+      - "bash .buildkite/scripts/upload-nightly-wheels.sh"
+
+  - label: "Generate and upload wheel indices"
+    depends_on: "build-wheel"
+    allow_dependency_failure: true
+    agents:
+      queue: small_cpu_queue_release
+    commands:
+      - "bash .buildkite/scripts/generate-and-upload-nightly-index.sh"

.buildkite/pipeline-intel.yaml

@@ -10,7 +10,7 @@ steps:
           DOCKER_BUILDKIT: "1"
           # Buildkite will automatically replace this with the actual commit hash
           VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}"
-          VLLM_VERSION: "v0.18.0"
+          VLLM_VERSION: "v0.20.0"
         priority: 100
         timeout_in_minutes: 60
-        soft_fail: true
+        soft_fail: false

.buildkite/pipeline-npu-a3.yaml

@@ -0,0 +1,110 @@
+# 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

@@ -0,0 +1,108 @@
+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/pipeline.yml

@@ -1,6 +1,22 @@
+# Document 1: Buildkite loads only this block on first parse. The next step resolves docs-only skip-ci
+# from git diff, then uploads document 2. When docs-only skip applies, image-build still runs for the nightly
+# exception in upload_pipeline_with_skip_ci.sh (nightly-test label or main NIGHTLY=1); otherwise upload-nightly
+# can be skipped together with test-ready/test-merge. Weekly uploads also depend on image-build.
+
+# Document 2: appended after `---`; same file, read by upload_pipeline_with_skip_ci.sh (not evaluated as a second pipeline by Buildkite). Child uploads include test-ready.yml, test-merge.yml, test-nightly.yml, and test-weekly.yml.
+steps:
+  - label: ":github: Resolve skip-ci & upload pipeline"
+    key: upload-ci-pipeline
+    commands:
+      - "bash .buildkite/scripts/upload_pipeline_with_skip_ci.sh"
+    agents:
+      queue: "cpu_queue_premerge"
+
+---
 steps:
   - label: ":docker: Build image"
     key: image-build
+    if: __IMAGE_BUILD_IF__
     commands:
       - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
       - "docker build --progress=plain --file docker/Dockerfile.ci -t vllm-omni-ci ."
@@ -13,7 +29,7 @@ steps:
   - label: "Upload Ready Pipeline"
     depends_on: image-build
     key: upload-ready-pipeline
-    if: build.branch != "main" && build.pull_request.labels includes "ready"
+    if: __UPLOAD_READY_IF__
     commands:
       - buildkite-agent pipeline upload .buildkite/test-ready.yml
     agents:
@@ -23,18 +39,38 @@ steps:
   - label: "Upload Merge Pipeline"
     depends_on: image-build
     key: upload-merge-pipeline
-    if: build.branch == "main" && build.env("NIGHTLY") != "1"
+    if: __UPLOAD_MERGE_IF__
     commands:
       - buildkite-agent pipeline upload .buildkite/test-merge.yml
     agents:
       queue: "cpu_queue_premerge"
 
-  # L4 Test — main+NIGHTLY=1 (scheduled), or PR with label nightly-test (e.g. add label then Rebuild)
+  # 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")'
+    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-nightly.yml
     agents:
       queue: "cpu_queue_premerge"
+
+  # L5 Test — main+WEEKLY=1 (scheduled), or PR with weekly-test label
+  - label: "Upload Weekly Pipeline"
+    depends_on: image-build
+    key: upload-weekly-pipeline
+    if: >-
+      (build.branch == "main" && build.env("WEEKLY") == "1") ||
+      (build.branch != "main" && build.pull_request.labels includes "weekly-test")
+    commands:
+      - buildkite-agent pipeline upload .buildkite/test-weekly.yml
+    agents:
+      queue: "cpu_queue_premerge"

.buildkite/scripts/generate-and-upload-nightly-index.sh

@@ -0,0 +1,88 @@
+#!/usr/bin/env bash
+
+set -ex
+
+# Generate and upload wheel indices for all vllm-omni wheels in the commit directory.
+# This script should run once after all wheels have been built and uploaded.
+# All paths are under the omni/ prefix in the vllm-wheels S3 bucket.
+
+# ======== setup ========
+
+BUCKET="vllm-wheels"
+INDICES_OUTPUT_DIR="indices"
+PYTHON="${PYTHON_PROG:-python3}"
+SUBPATH="omni/$BUILDKITE_COMMIT"
+S3_COMMIT_PREFIX="s3://$BUCKET/$SUBPATH/"
+
+# detect if python3.12+ is available
+has_new_python=$($PYTHON -c "print(1 if __import__('sys').version_info >= (3,12) else 0)")
+if [[ "$has_new_python" -eq 0 ]]; then
+    # use new python from docker
+    docker pull python:3-slim
+    PYTHON="docker run --rm --user $(id -u):$(id -g) -v $(pwd):/app -w /app python:3-slim python3"
+fi
+
+echo "Using python interpreter: $PYTHON"
+echo "Python version: $($PYTHON --version)"
+
+# ======== generate and upload indices ========
+
+# list all wheels in the commit directory
+echo "Existing wheels on S3:"
+aws s3 ls "$S3_COMMIT_PREFIX" || echo "(no objects found)"
+obj_json="objects.json"
+aws s3api list-objects-v2 --bucket "$BUCKET" --prefix "$SUBPATH/" --delimiter / --output json > "$obj_json"
+mkdir -p "$INDICES_OUTPUT_DIR"
+
+# HACK: we do not need regex module here, but it is required by pre-commit hook
+# To avoid any external dependency, we simply replace it back to the stdlib re module
+sed -i.bak 's/import regex as re/import re/g' .buildkite/scripts/generate-nightly-index.py && rm -f .buildkite/scripts/generate-nightly-index.py.bak
+
+# Generate indices -- the version is just the commit hash (not omni/{commit})
+# because relative paths are computed between the index and wheel directories,
+# both of which live under the omni/ prefix in S3.
+$PYTHON .buildkite/scripts/generate-nightly-index.py \
+    --version "$BUILDKITE_COMMIT" \
+    --current-objects "$obj_json" \
+    --output-dir "$INDICES_OUTPUT_DIR" \
+    --comment "commit $BUILDKITE_COMMIT"
+
+# copy indices to /omni/{commit}/ unconditionally
+echo "Uploading indices to $S3_COMMIT_PREFIX"
+aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "$S3_COMMIT_PREFIX"
+
+# copy to /omni/nightly/ when NIGHTLY=1
+if [[ "${NIGHTLY:-}" == "1" ]]; then
+    echo "Uploading indices to overwrite /omni/nightly/"
+    aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "s3://$BUCKET/omni/nightly/"
+fi
+
+# detect version from any wheel in the commit directory
+first_wheel_key=$($PYTHON -c "import json; obj=json.load(open('$obj_json')); print(next((c['Key'] for c in obj.get('Contents', []) if c['Key'].endswith('.whl')), ''))")
+if [[ -z "$first_wheel_key" ]]; then
+    echo "Error: No wheels found in $S3_COMMIT_PREFIX"
+    exit 1
+fi
+first_wheel=$(basename "$first_wheel_key")
+aws s3 cp "s3://$BUCKET/${first_wheel_key}" "/tmp/${first_wheel}"
+version=$(unzip -p "/tmp/${first_wheel}" '**/METADATA' | grep '^Version: ' | cut -d' ' -f2)
+rm -f "/tmp/${first_wheel}"
+echo "Version in wheel: $version"
+pure_version="${version%%+*}"
+echo "Pure version (without variant): $pure_version"
+
+# re-generate and copy to /omni/{version}/ only if it does not have "dev" in the version
+if [[ "$version" != *"dev"* ]]; then
+    s3_version="v$pure_version"
+    echo "Re-generating indices for /omni/$s3_version/"
+    rm -rf "${INDICES_OUTPUT_DIR:?}"
+    mkdir -p "$INDICES_OUTPUT_DIR"
+    # wheel-dir is overridden to be the commit directory, so that the indices point to the correct wheel path
+    $PYTHON .buildkite/scripts/generate-nightly-index.py \
+        --version "$s3_version" \
+        --wheel-dir "$BUILDKITE_COMMIT" \
+        --current-objects "$obj_json" \
+        --output-dir "$INDICES_OUTPUT_DIR" \
+        --comment "version $pure_version"
+    aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "s3://$BUCKET/omni/$s3_version/"
+fi

.buildkite/scripts/generate-nightly-index.py

@@ -0,0 +1,192 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import argparse
+import json
+import re
+import sys
+from dataclasses import asdict, dataclass
+from datetime import datetime
+from pathlib import Path
+from typing import Any
+from urllib.parse import quote
+
+
+def normalize_package_name(name: str) -> str:
+    """Normalize package name per PEP 503."""
+    return re.sub(r"[-_.]+", "-", name).lower()
+
+
+if not sys.version_info >= (3, 12):
+    raise RuntimeError("This script requires Python 3.12 or higher.")
+
+INDEX_HTML_TEMPLATE = """<!DOCTYPE html>
+<html>
+  <!-- {comment} -->
+  <meta name="pypi:repository-version" content="1.0">
+  <body>
+{items}
+  </body>
+</html>
+"""
+
+
+@dataclass
+class WheelFileInfo:
+    package_name: str
+    version: str
+    build_tag: str | None
+    python_tag: str
+    abi_tag: str
+    platform_tag: str
+    filename: str
+
+
+def parse_from_filename(file: str) -> WheelFileInfo:
+    """
+    Parse wheel filename per PEP 427:
+        {package_name}-{version}(-{build_tag})?-{python_tag}-{abi_tag}-{platform_tag}.whl
+    """
+    wheel_file_re = re.compile(
+        r"^(?P<package_name>.+)-(?P<version>[^-]+?)(-(?P<build_tag>[^-]+))?-(?P<python_tag>[^-]+)-(?P<abi_tag>[^-]+)-(?P<platform_tag>[^-]+)\.whl$"
+    )
+    match = wheel_file_re.match(file)
+    if not match:
+        raise ValueError(f"Invalid wheel file name: {file}")
+
+    return WheelFileInfo(
+        package_name=match.group("package_name"),
+        version=match.group("version"),
+        build_tag=match.group("build_tag"),
+        python_tag=match.group("python_tag"),
+        abi_tag=match.group("abi_tag"),
+        platform_tag=match.group("platform_tag"),
+        filename=file,
+    )
+
+
+def generate_project_list(package_names: list[str], comment: str = "") -> str:
+    """Generate top-level PEP 503 project list HTML."""
+    href_tags = []
+    for name in sorted(package_names):
+        href_tags.append(f'    <a href="{name}/">{name}/</a><br/>')
+    return INDEX_HTML_TEMPLATE.format(items="\n".join(href_tags), comment=comment)
+
+
+def generate_package_index(
+    wheel_files: list[WheelFileInfo],
+    wheel_base_dir: Path,
+    index_base_dir: Path,
+    comment: str = "",
+) -> tuple[str, str]:
+    """Generate package index HTML and metadata JSON linking to wheel files."""
+    href_tags = []
+    metadata = []
+    for file in sorted(wheel_files, key=lambda x: x.filename):
+        relative_path = wheel_base_dir.relative_to(index_base_dir, walk_up=True) / file.filename
+        # handle '+' in URL; avoid double-encoding '/' and '%2B' (AWS S3 behavior)
+        file_path_quoted = quote(relative_path.as_posix(), safe=":%/")
+        href_tags.append(f'    <a href="{file_path_quoted}">{file.filename}</a><br/>')
+        file_meta = asdict(file)
+        file_meta["path"] = file_path_quoted
+        metadata.append(file_meta)
+    index_str = INDEX_HTML_TEMPLATE.format(items="\n".join(href_tags), comment=comment)
+    metadata_str = json.dumps(metadata, indent=2)
+    return index_str, metadata_str
+
+
+def generate_index(
+    whl_files: list[str],
+    wheel_base_dir: Path,
+    index_base_dir: Path,
+    comment: str = "",
+):
+    """
+    Generate PEP 503 index for all wheel files.
+
+    Output structure:
+        index_base_dir/
+            index.html          # project list linking to vllm-omni/
+            vllm-omni/
+                index.html      # package index linking to wheel files
+                metadata.json   # machine-readable metadata
+    """
+    parsed_files = [parse_from_filename(f) for f in whl_files]
+
+    if not parsed_files:
+        print("No wheel files found, skipping index generation.")
+        return
+
+    comment_str = f" ({comment})" if comment else ""
+    comment_tmpl = f"Generated on {datetime.now().isoformat()}{comment_str}"
+
+    # Group by normalized package name
+    packages: dict[str, list[WheelFileInfo]] = {}
+    for file in parsed_files:
+        name = normalize_package_name(file.package_name)
+        packages.setdefault(name, []).append(file)
+
+    print(f"Found packages: {list(packages.keys())}")
+
+    # Generate per-package index
+    for package, files in packages.items():
+        package_dir = index_base_dir / package
+        package_dir.mkdir(parents=True, exist_ok=True)
+        index_str, metadata_str = generate_package_index(files, wheel_base_dir, package_dir, comment)
+        with open(package_dir / "index.html", "w") as f:
+            f.write(index_str)
+        with open(package_dir / "metadata.json", "w") as f:
+            f.write(metadata_str)
+
+    # Generate top-level project list
+    project_list_str = generate_project_list(sorted(packages.keys()), comment_tmpl)
+    with open(index_base_dir / "index.html", "w") as f:
+        f.write(project_list_str)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Generate PEP 503 wheel index from S3 object listing.")
+    parser.add_argument("--version", type=str, required=True, help="Version string (e.g., commit hash)")
+    parser.add_argument("--current-objects", type=str, required=True, help="Path to JSON from S3 list-objects-v2")
+    parser.add_argument("--output-dir", type=str, required=True, help="Directory to write index files")
+    parser.add_argument("--wheel-dir", type=str, default=None, help="Wheel directory (defaults to --version)")
+    parser.add_argument("--comment", type=str, default="", help="Comment for generated HTML")
+
+    args = parser.parse_args()
+
+    version = args.version
+    if "\\" in version or "/" in version:
+        raise ValueError("Version string must not contain slashes or backslashes.")
+
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    with open(args.current_objects) as f:
+        current_objects: dict[str, list[dict[str, Any]]] = json.load(f)
+
+    wheel_files = [
+        item["Key"].split("/")[-1] for item in current_objects.get("Contents", []) if item["Key"].endswith(".whl")
+    ]
+
+    print(f"Found {len(wheel_files)} wheel files for version {version}: {wheel_files}")
+
+    # For release versions, filter to only matching non-dev wheels
+    PY_VERSION_RE = re.compile(r"^\d+\.\d+\.\d+([a-zA-Z0-9.+-]*)?$")
+    if PY_VERSION_RE.match(version):
+        wheel_files = [f for f in wheel_files if version in f and "dev" not in f]
+        print(f"Non-nightly version detected, wheel files used: {wheel_files}")
+    else:
+        print("Nightly version detected, keeping all wheel files.")
+
+    wheel_dir = (args.wheel_dir or version).strip().rstrip("/")
+    wheel_base_dir = Path(output_dir).parent / wheel_dir
+    index_base_dir = Path(output_dir)
+
+    generate_index(
+        whl_files=wheel_files,
+        wheel_base_dir=wheel_base_dir,
+        index_base_dir=index_base_dir,
+        comment=args.comment.strip(),
+    )
+    print(f"Successfully generated index in {output_dir}")

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

@@ -11,15 +11,6 @@ set -o pipefail
 export PYTHONPATH=".."
 
 # Print ROCm version
-echo "--- Confirming Clean Initial State"
-while true; do
-        sleep 3
-        if grep -q clean /opt/amdgpu/etc/gpu_state; then
-                echo "GPUs state is \"clean\""
-                break
-        fi
-done
-
 echo "--- ROCm info"
 rocminfo
 
@@ -51,25 +42,14 @@ cleanup_docker() {
 # Call the cleanup docker function
 cleanup_docker
 
-echo "--- Resetting GPUs"
-
-echo "reset" > /opt/amdgpu/etc/gpu_state
-
-while true; do
-        sleep 3
-        if grep -q clean /opt/amdgpu/etc/gpu_state; then
-                echo "GPUs state is \"clean\""
-                break
-        fi
-done
-
 echo "--- Pulling container"
-## Temporary change to use AMD Docker Hub to store the vllm-ci image
+## Temporary change to use AMD Docker Hub to store the vllm-omni image
 # to bypass the rate limit issue with ECR Public Gallery.
+# Images are now stored in a separate repository for vllm-omni, instead of vllm-ci.
 # TODO: @tjtanaa point back to ECR Public Gallery
 # once the amd agents are configured to use ECR Public Gallery.
 # image_name="public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:${BUILDKITE_COMMIT}-rocm-omni"
-image_name="rocm/vllm-ci:${BUILDKITE_COMMIT}-rocm-omni"
+image_name="rocm/vllm-omni:${BUILDKITE_COMMIT}"
 container_name="rocm_${BUILDKITE_COMMIT}_$(tr -dc A-Za-z0-9 < /dev/urandom | head -c 10; echo)"
 
 # TODO: @tjtanaa uncomment this once the amd agents are configured to use ECR Public Gallery.

.buildkite/scripts/hardware_ci/run_npu_test.sh

@@ -141,5 +141,5 @@ docker run \
     "${image_name}" \
     bash -c '
     set -e
-    VLLM_USE_MODELSCOPE=True pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
+    VLLM_USE_MODELSCOPE=True pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni_expansion.py
 '

.buildkite/scripts/upload-nightly-wheels.sh

@@ -0,0 +1,33 @@
+#!/usr/bin/env bash
+
+set -ex
+
+# Upload a single wheel to S3 under the omni/ prefix.
+# Index generation is handled separately by generate-and-upload-nightly-index.sh.
+
+BUCKET="vllm-wheels"
+SUBPATH="omni/$BUILDKITE_COMMIT"
+S3_COMMIT_PREFIX="s3://$BUCKET/$SUBPATH/"
+
+# ========= collect & upload the wheel ==========
+
+# python3 -m build outputs to dist/ by default
+wheel_files=(dist/*.whl)
+
+# Check that exactly one wheel is found
+if [[ ${#wheel_files[@]} -ne 1 ]]; then
+  echo "Error: Expected exactly one wheel file in dist/, but found ${#wheel_files[@]}"
+  exit 1
+fi
+wheel="${wheel_files[0]}"
+
+echo "Uploading wheel: $wheel"
+
+# Extract the version from the wheel
+version=$(unzip -p "$wheel" '**/METADATA' | grep '^Version: ' | cut -d' ' -f2)
+echo "Version in wheel: $version"
+
+# Upload wheel to S3
+aws s3 cp "$wheel" "$S3_COMMIT_PREFIX"
+
+echo "Wheel uploaded to $S3_COMMIT_PREFIX. Index generation is handled by a separate step."

.buildkite/scripts/upload_pipeline_with_skip_ci.sh

@@ -0,0 +1,146 @@
+#!/usr/bin/env bash
+# Evaluate docs-only skip-ci and upload continuation steps from the given pipeline YAML
+# (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
+
+# Prints a single digit to stdout: 1 = skip image CI, 0 = run. Logs go to stderr.
+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
+}
+
+SKIP_CI="$(resolve_skip_ci)"
+
+if [[ ! -f "${PIPELINE_YML}" ]]; then
+  echo "upload_pipeline_with_skip_ci: missing ${PIPELINE_YML}" >&2
+  exit 1
+fi
+
+export ROOT SKIP_CI PIPELINE_YML
+python3 <<'PY' | buildkite-agent pipeline upload
+import os
+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
+
+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
+# main NIGHTLY=1), otherwise upload-nightly (depends_on image-build) would be skipped too.
+nightly_only = (
+    '(build.pull_request.labels includes "nightly-test") '
+    '|| (build.branch == "main" && build.env("NIGHTLY") == "1")'
+)
+# Placeholder in pipeline.yml is `if: __IMAGE_BUILD_IF__` (valid YAML); replace value only.
+if skip:
+    rep = f"'{nightly_only}'"
+    ready_rep = "'false'"
+    merge_rep = "'false'"
+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\")'"
+rendered = (
+    continuation
+    .replace("__IMAGE_BUILD_IF__", rep)
+    .replace("__UPLOAD_READY_IF__", ready_rep)
+    .replace("__UPLOAD_MERGE_IF__", merge_rep)
+)
+print(rendered, end="")
+PY

.buildkite/test-amd-merge.yml

@@ -7,16 +7,9 @@ steps:
   grade: Blocking
   commands:
     - export VLLM_ROCM_USE_AITER=0
-    - "timeout 20m pytest -v -s -m 'core_model and cpu' --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"
+    # 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"
 
 - label: "Diffusion Images API LoRA E2E"
   agent_pool: mi325_1
@@ -24,7 +17,7 @@ steps:
   mirror_hardwares: [amdproduction]
   grade: Blocking
   commands:
-    - timeout 20m pytest -s -v tests/e2e/online_serving/test_images_generations_lora.py
+    - timeout 20m pytest -s -v tests/diffusion/lora/test_images_generations_lora.py
 
 - label: "Diffusion Model CPU offloading Test"
   agent_pool: mi325_1
@@ -32,15 +25,14 @@ steps:
   mirror_hardwares: [amdproduction]
   grade: Blocking
   commands:
-    - export GPU_ARCHS=gfx942
     - export VLLM_LOGGING_LEVEL=DEBUG
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
     - |
       timeout 20m bash -c '
         set +e
-        pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
+        pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py
         EXIT1=\$?
-        pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
+        pytest -s -v tests/diffusion/offloader/test_diffusion_layerwise_offload.py
         EXIT2=\$?
         exit \$((EXIT1 | EXIT2))
       '
@@ -55,7 +47,7 @@ steps:
 #     - export GPU_ARCHS=gfx942
 #     - export VLLM_LOGGING_LEVEL=DEBUG
 #     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-#     - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_model.py
+#     - 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: "Diffusion Cache Backend Test"
   agent_pool: mi325_1
@@ -63,20 +55,20 @@ steps:
   mirror_hardwares: [amdproduction]
   grade: Blocking
   commands:
-    - export GPU_ARCHS=gfx942
     - export VLLM_LOGGING_LEVEL=DEBUG
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
     - timeout 15m pytest -s -v -m "core_model and cache and diffusion and not distributed_cuda and L4"
 
-- label: "Diffusion Sequence Parallelism Test"
-  agent_pool: mi325_2
+- label: "Diffusion Sequence Parallelism Test (Need 4 GPUs)"
+  agent_pool: mi325_4
   depends_on: amd-build
   mirror_hardwares: [amdproduction]
   grade: Blocking
   commands:
     - export VLLM_LOGGING_LEVEL=DEBUG
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - timeout 20m pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py
+    - timeout 20m pytest -s -v tests/diffusion/distributed/test_sequence_parallel.py
+    - timeout 20m pytest -s -v tests/diffusion/distributed/test_ulysses_uaa_perf.py
 
 # merge-only tests
 - label: "Diffusion Tensor Parallelism Test"
@@ -85,7 +77,15 @@ steps:
   mirror_hardwares: [amdproduction]
   grade: Blocking
   commands:
-    - pytest -s -v tests/e2e/offline_inference/test_zimage_parallelism.py
+    - 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
 
 - label: "Diffusion GPU Worker Test"
   agent_pool: mi325_2
@@ -95,32 +95,26 @@ steps:
   commands:
     - timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py
 
-- label: "Benchmark & Engine Test"
-  agent_pool: mi325_2
+- label: "Engine Test"
+  agent_pool: mi325_1
   depends_on: amd-build
   mirror_hardwares: [amdproduction]
   grade: Blocking
   commands:
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - |
-      timeout 20m bash -c '
-        set +e
-        pytest -s -v tests/benchmarks/test_serve_cli.py
-        EXIT1=\$?
-        pytest -s -v tests/engine/test_async_omni_engine_abort.py
-        EXIT2=\$?
-        exit \$((EXIT1 | EXIT2))
-      '
+    - timeout 20m pytest -s -v tests/engine/test_async_omni_engine_abort.py
 
+# This is from test-nightly.yml.
+# This model seems to be important, so I am trying to run it here.
+# Since this is from test-nightly.yml,
+# I keep this non-blocking.
 - label: "Omni Model Test Qwen2-5-Omni"
-  agent_pool: mi325_2
+  agent_pool: mi325_4
   depends_on: amd-build
   mirror_hardwares: [amdproduction]
-  grade: Blocking
   commands:
-    - export VLLM_LOGGING_LEVEL=DEBUG
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - timeout 20m pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
+    - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
+    - pytest -s -v tests/e2e/*/test_qwen2_5_omni_expansion.py --run-level "core_model"
 
 - label: "Omni Model Test Qwen3-Omni"
   agent_pool: mi325_2
@@ -128,14 +122,11 @@ steps:
   mirror_hardwares: [amdproduction]
   grade: Blocking
   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 "advanced_model" --run-level "advanced_model"
+    - 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"
-  agent_pool: mi325_2
+  agent_pool: mi325_1
   depends_on: amd-build
   mirror_hardwares: [amdproduction]
   grade: Blocking
@@ -145,11 +136,77 @@ steps:
         export VLLM_LOGGING_LEVEL=DEBUG
         export VLLM_WORKER_MULTIPROC_METHOD=spawn
         export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
-        pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m "advanced_model" --run-level "advanced_model" && pytest -s -v tests/e2e/offline_inference/test_qwen3_tts_customvoice.py
+        pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py tests/e2e/offline_inference/test_qwen3_tts_customvoice.py -m "advanced_model" --run-level "advanced_model"
       '
 
 - label: "Qwen3-TTS Base E2E Test"
-  agent_pool: mi325_2
+  agent_pool: mi325_1
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - |
+      timeout 30m bash -c '
+        export VLLM_LOGGING_LEVEL=DEBUG
+        export VLLM_WORKER_MULTIPROC_METHOD=spawn
+        export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
+        pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py tests/e2e/offline_inference/test_qwen3_tts_base.py -m "advanced_model" --run-level "advanced_model"
+      '
+
+# 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
+# - label: "Bagel Text2Img Model Test (1/3)"
+#   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
+#     - 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"
+
+# - label: "Bagel Img2Img Model Test (2/3)"
+#   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
+#     - 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"
+
+# - label: "Bagel Online Serving Test (3/3)"
+#   agent_pool: mi325_1
+#   depends_on: amd-build
+#   mirror_hardwares: [amdproduction]
+#   grade: Blocking
+#   commands:
+#     - export GPU_ARCHS=gfx942
+#     - 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'
+      "
+
+- label: "Voxtral-TTS E2E Test"
+  agent_pool: mi325_1
   depends_on: amd-build
   mirror_hardwares: [amdproduction]
   grade: Blocking
@@ -158,58 +215,5 @@ steps:
       timeout 20m bash -c '
         export VLLM_LOGGING_LEVEL=DEBUG
         export VLLM_WORKER_MULTIPROC_METHOD=spawn
-        export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
-        pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py -m "advanced_model" --run-level "advanced_model" && pytest -s -v tests/e2e/offline_inference/test_qwen3_tts_base.py
+        pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py 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
-
-# split Bagel Model Test with H100 (Real Weights) into three tests
-- label: "Bagel Text2Img Model Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  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_text2img.py -m "advanced_model" --run-level "advanced_model" -k "shared_memory" -k "rocm"
-
-- label: "Bagel Img2Img Model Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  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_img2img.py -m "advanced_model" --run-level "advanced_model" -k "rocm"
-
-- label: "Bagel Online Serving Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  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_online.py -m "advanced_model" --run-level "advanced_model" -k "rocm"

.buildkite/test-amd-ready.yaml

@@ -7,15 +7,33 @@ steps:
   grade: Blocking
   commands:
     - export VLLM_ROCM_USE_AITER=0
-    - "timeout 20m pytest -v -s -m 'core_model and cpu' --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
+    # 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"
 
-- label: "Diffusion Model Test"
-  agent_pool: mi325_2
+- label: "Voxtral TTS CUDA Unit Test"
+  agent_pool: mi325_1
   depends_on: amd-build
   mirror_hardwares: [amdproduction]
   grade: Blocking
   commands:
-    - timeout 20m pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "core_model and diffusion" --run-level "core_model"
+    - timeout 10m pytest -s -v tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py
+
+- 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"
+
+- label: "Custom Pipeline Test"
+  agent_pool: mi325_1
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - timeout 20m pytest -s -v tests/e2e/offline_inference/custom_pipeline/ -m "core_model"
 
 - label: "Diffusion Model CPU offloading Test"
   agent_pool: mi325_1
@@ -23,15 +41,14 @@ steps:
   mirror_hardwares: [amdproduction]
   grade: Blocking
   commands:
-    - export GPU_ARCHS=gfx942
     - export VLLM_LOGGING_LEVEL=DEBUG
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
     - |
       timeout 20m bash -c '
         set +e
-        pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
+        pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py
         EXIT1=\$?
-        pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
+        pytest -s -v tests/diffusion/offloader/test_diffusion_layerwise_offload.py
         EXIT2=\$?
         exit \$((EXIT1 | EXIT2))
       '
@@ -46,7 +63,7 @@ steps:
 #     - export GPU_ARCHS=gfx942
 #     - export VLLM_LOGGING_LEVEL=DEBUG
 #     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-#     - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_model.py
+#     - 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: "Diffusion Cache Backend Test"
   agent_pool: mi325_1
@@ -67,7 +84,7 @@ steps:
   commands:
     - export VLLM_LOGGING_LEVEL=DEBUG
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - timeout 20m pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py -m core_model
+    - timeout 20m pytest -s -v tests/diffusion/distributed/test_sequence_parallel.py -m core_model
 
 - label: "Diffusion GPU Worker Test"
   agent_pool: mi325_2
@@ -77,47 +94,65 @@ steps:
   commands:
     - timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py
 
-- label: "Benchmark & Engine Test"
-  agent_pool: mi325_2
+- 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
   mirror_hardwares: [amdproduction]
   grade: Blocking
   commands:
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
     - |
-      timeout 30m bash -c '
-        set +e
-        pytest -s -v tests/benchmarks/test_serve_cli.py
-        EXIT1=\$?
-        pytest -s -v tests/engine/test_async_omni_engine_abort.py
-        EXIT2=\$?
-        exit \$((EXIT1 | EXIT2))
+      timeout 15m bash -c '
+              pytest -s -v tests/engine/test_async_omni_engine_abort.py
       '
 
-- label: "Omni Model Test Qwen2-5-Omni"
-  agent_pool: mi325_2
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - export VLLM_LOGGING_LEVEL=DEBUG
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - timeout 17m pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
+# - label: "Omni Model Test Qwen3-Omni"
+#   agent_pool: mi325_2
+#   depends_on: amd-build
+#   mirror_hardwares: [amdproduction]
+#   grade: Blocking
+#   commands:
+#     - export VLLM_LOGGING_LEVEL=DEBUG
+#     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+#     - timeout 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"
 
-- label: "Omni Model Test Qwen3-Omni"
-  agent_pool: mi325_2
+- label: "MiMo-Audio E2E Test with H100"
+  agent_pool: mi325_1
   depends_on: amd-build
   mirror_hardwares: [amdproduction]
   grade: Blocking
   commands:
-    - export VLLM_LOGGING_LEVEL=DEBUG
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - export VLLM_TEST_CLEAN_GPU_MEMORY=1
-    - timeout 10m pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
-    - timeout 10m pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "core_model" --run-level "core_model"
+    - |
+      timeout 30m bash -c '
+        export VLLM_LOGGING_LEVEL=DEBUG
+        export VLLM_WORKER_MULTIPROC_METHOD=spawn
+        pytest -s -v tests/e2e/online_serving/test_mimo_audio.py -m "core_model" --run-level "core_model"
+      '
 
 - label: "Qwen3-TTS E2E Test"
-  agent_pool: mi325_2
+  agent_pool: mi325_1
   depends_on: amd-build
   mirror_hardwares: [amdproduction]
   grade: Blocking
@@ -125,55 +160,68 @@ steps:
     - export VLLM_LOGGING_LEVEL=DEBUG
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
     - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
-    - timeout 20m pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m "core_model" --run-level "core_model"
+    - timeout 30m pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m "core_model" --run-level "core_model"
 
-- label: "Diffusion Image Edit Test"
+- label: "Voxtral-TTS E2E 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
+    - |
+      timeout 20m bash -c '
+        export VLLM_LOGGING_LEVEL=DEBUG
+        export VLLM_WORKER_MULTIPROC_METHOD=spawn
+        pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
+        pytest -s -v tests/e2e/offline_inference/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
+      '
 
-- label: "Bagel Text2Img Model Test"
+# TODO: Bagel test on ROCm is very unstable. @tjtanaa
+# Need to debug before reneable numerical changes across large PRs
+# - label: "Bagel Text2Img Model 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
+#     - 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"
+
+# - label: "Bagel Img2Img Model 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
+#     - 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"
+
+# - label: "Bagel Online Serving Test"
+#   agent_pool: mi325_1
+#   depends_on: amd-build
+#   mirror_hardwares: [amdproduction]
+#   grade: Blocking
+#   commands:
+#     - export GPU_ARCHS=gfx942
+#     - 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"
+
+- label: "CosyVoice3-TTS E2E Test"
   agent_pool: mi325_1
   depends_on: amd-build
   mirror_hardwares: [amdproduction]
   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_text2img.py -m "core_model" --run-level "core_model" -k "rocm"
-
-- label: "Bagel Img2Img Model Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  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_img2img.py -m "core_model" --run-level "core_model" -k "rocm"
-
-- label: "Bagel Online Serving Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  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_online.py -m "core_model" --run-level "core_model" -k "rocm"
+    - |
+      timeout 20m bash -c '
+        pytest -s -v tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py -m "full_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_t2i_model.py
+    - pytest -s -v tests/e2e/offline_inference/test_z_image.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/e2e/online_serving/test_images_generations_lora.py
+    - pytest -s -v tests/diffusion/lora/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/e2e/offline_inference/test_diffusion_cpu_offload.py
-    - pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
+    - pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py
+    - pytest -s -v tests/diffusion/offloader/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/e2e/offline_inference/test_sequence_parallel.py -m core_model
+    - pytest -s -v tests/diffusion/distributed/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/e2e/offline_inference/test_zimage_parallelism.py
+    - pytest -s -v tests/diffusion/distributed/test_zimage_parallelism.py
 
 - label: "Diffusion GPU Worker Test"
   timeout_in_minutes: 20
@@ -88,7 +88,7 @@ steps:
     - export GPU_ARCHS=gfx942
     - export VLLM_LOGGING_LEVEL=DEBUG
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
+    - pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni_expansion.py
     - pytest -s -v tests/engine/test_async_omni_engine_abort.py
 
 - label: "Omni Model Test Qwen3-Omni"
@@ -100,15 +100,14 @@ 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
+- label: "Omni Sleep Mode Test"
+  timeout_in_minutes: 40
+  agent_pool: mi325_2
   depends_on: amd-build
   mirror_hardwares: [amdproduction]
   grade: Blocking
@@ -116,4 +115,4 @@ steps:
     - 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
+    - pytest -s -v tests/entrypoints/test_omni_sleep_mode.py -m "advanced_model and omni and MI325" --run-level "advanced_model"

.buildkite/test-npu-nightly.yaml

@@ -0,0 +1,24 @@
+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"

.buildkite/test-ready.yml

@@ -1,71 +1,43 @@
+env:
+  VLLM_WORKER_MULTIPROC_METHOD: spawn
+  HF_HUB_DOWNLOAD_TIMEOUT: 300
+  HF_HUB_ETAG_TIMEOUT: 60
+
 steps:
-  - label: "Simple Unit Test"
+  - group: ":card_index_dividers: Simple Test"
     depends_on: upload-ready-pipeline
-    commands:
-      - "timeout 20m pytest -v -s -m 'core_model and cpu' --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
-    agents:
-      queue: "gpu_1_queue"
-    plugins:
-      - docker#v5.2.0:
-          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-          always-pull: true
-          propagate-environment: true
-          environment:
-            - "HF_HOME=/fsx/hf_cache"
-            - "HF_TOKEN"
-          volumes:
-            - "/fsx/hf_cache:/fsx/hf_cache"
+    steps:
+      - label: "Simple · Diffusion & Model Executor Test"
+        commands:
+          - "timeout 20m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
+        agents:
+          queue: "gpu_1_queue"
+        plugins:
+          - docker#v5.2.0:
+              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+              always-pull: true
+              propagate-environment: true
+              environment:
+                - "HF_HOME=/fsx/hf_cache"
+                - "HF_TOKEN"
+              volumes:
+                - "/fsx/hf_cache:/fsx/hf_cache"
 
-  - label: "Voxtral TTS CUDA Unit Test"
-    timeout_in_minutes: 10
-    depends_on: upload-ready-pipeline
-    commands:
-      - "timeout 10m pytest -s -v tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py"
-    agents:
-      queue: "gpu_1_queue"
-    plugins:
-      - docker#v5.2.0:
-          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-          always-pull: true
-          propagate-environment: true
-          environment:
-            - "HF_HOME=/fsx/hf_cache"
-          volumes:
-            - "/fsx/hf_cache:/fsx/hf_cache"
-
-  - label: "Diffusion Model Test"
-    depends_on: upload-ready-pipeline
-    commands:
-      - timeout 30m pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "core_model and diffusion" --run-level "core_model"
-    agents:
-      queue: "gpu_1_queue"
-    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: "Diffusion Batching Test"
-    depends_on: upload-ready-pipeline
-    commands:
-      - timeout 20m pytest -s -v tests/e2e/offline_inference/test_qwen_image_diffusion_batching.py -m "core_model and diffusion" --run-level "core_model"
-    agents:
-      queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
-    plugins:
-      - docker#v5.2.0:
-          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-          always-pull: true
-          propagate-environment: true
-          environment:
-            - "HF_HOME=/fsx/hf_cache"
-            - "HF_TOKEN"
-          volumes:
-            - "/fsx/hf_cache:/fsx/hf_cache"
+      - label: "Simple · Other Test"
+        commands:
+          - "timeout 20m pytest -sv -m 'core_model and cpu' --ignore=tests/diffusion --ignore=tests/model_executor --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
+        agents:
+          queue: "gpu_1_queue"
+        plugins:
+          - docker#v5.2.0:
+              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+              always-pull: true
+              propagate-environment: true
+              environment:
+                - "HF_HOME=/fsx/hf_cache"
+                - "HF_TOKEN"
+              volumes:
+                - "/fsx/hf_cache:/fsx/hf_cache"
 
   - label: "Custom Pipeline Test"
     depends_on: upload-ready-pipeline
@@ -83,81 +55,116 @@ steps:
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
-  - label: "Diffusion Model CPU offloading Test"
+  - group: ":card_index_dividers: Diffusion Test"
     depends_on: upload-ready-pipeline
-    commands:
-      - timeout 10m pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
-    agents:
-      queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
-    plugins:
-      - docker#v5.2.0:
-          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-          always-pull: true
-          propagate-environment: true
-          environment:
-            - "HF_HOME=/fsx/hf_cache"
-            - "HF_TOKEN"
-          volumes:
-            - "/fsx/hf_cache:/fsx/hf_cache"
+    steps:
+      - label: "Diffusion · Offloader Test"
+        commands:
+          - timeout 10m pytest -s -v tests/diffusion/offloader -m "core_model and cuda" --run-level "core_model"
+        agents:
+          queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
+        plugins:
+          - docker#v5.2.0:
+              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+              always-pull: true
+              propagate-environment: true
+              environment:
+                - "HF_HOME=/fsx/hf_cache"
+                - "HF_TOKEN"
+              volumes:
+                - "/fsx/hf_cache:/fsx/hf_cache"
 
-  - label: "Audio Generation Model Test"
-    depends_on: upload-ready-pipeline
-    commands:
-      - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_model.py
-    agents:
-      queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
-    plugins:
-      - docker#v5.2.0:
-          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-          always-pull: true
-          propagate-environment: true
-          environment:
-            - "HF_HOME=/fsx/hf_cache"
-            - "HF_TOKEN"
-          volumes:
-            - "/fsx/hf_cache:/fsx/hf_cache"
+      - label: "Diffusion · Cache Test"
+        commands:
+          - timeout 15m pytest -sv tests/diffusion/cache -m 'core_model and cuda' --run-level "core_model"
+        agents:
+          queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
+        plugins:
+          - docker#v5.2.0:
+              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+              always-pull: true
+              propagate-environment: true
+              environment:
+                - "HF_HOME=/fsx/hf_cache"
+                - "HF_TOKEN"
+              volumes:
+                - "/fsx/hf_cache:/fsx/hf_cache"
 
-  - label: "Diffusion Cache Backend Test"
-    depends_on: upload-ready-pipeline
-    commands:
-      - timeout 15m pytest -s -v -m 'core_model and cache and diffusion and not distributed_cuda and L4'
-    agents:
-      queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
-    plugins:
-      - 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: "Diffusion · Batch Test"
+        commands:
+          - timeout 15m pytest -sv tests/diffusion/batching -m 'core_model and cuda' --run-level "core_model"
+        agents:
+          queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
+        plugins:
+          - docker#v5.2.0:
+              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+              always-pull: true
+              propagate-environment: true
+              environment:
+                - "HF_HOME=/fsx/hf_cache"
+                - "HF_TOKEN"
+              volumes:
+                - "/fsx/hf_cache:/fsx/hf_cache"
 
-  - label: "Diffusion Sequence Parallelism Test"
-    depends_on: upload-ready-pipeline
-    commands:
-      - timeout 20m pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py -m core_model
-    agents:
-      queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
-    plugins:
-      - docker#v5.2.0:
-          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-          always-pull: true
-          propagate-environment: true
-          shm-size: "8gb"
-          environment:
-            - "HF_HOME=/fsx/hf_cache"
-            - "HF_TOKEN"
-          volumes:
-            - "/fsx/hf_cache:/fsx/hf_cache"
+      - label: "Diffusion · Distributed Test"
+        commands:
+          - timeout 20m pytest -sv tests/diffusion/distributed -m 'core_model and cuda' --run-level "core_model"
+        agents:
+          queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
+        plugins:
+          - docker#v5.2.0:
+              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+              always-pull: true
+              propagate-environment: true
+              shm-size: "8gb"
+              environment:
+                - "HF_HOME=/fsx/hf_cache"
+                - "HF_TOKEN"
+              volumes:
+                - "/fsx/hf_cache:/fsx/hf_cache"
 
-  - label: "Diffusion GPU Worker Test"
+      - label: "Diffusion · Model Test"
+        commands:
+          - timeout 15m pytest -sv tests/diffusion/models/ -m 'core_model and cuda' --run-level "core_model"
+        agents:
+          queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
+        plugins:
+          - docker#v5.2.0:
+              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+              always-pull: true
+              propagate-environment: true
+              environment:
+                - "HF_HOME=/fsx/hf_cache"
+                - "HF_TOKEN"
+              volumes:
+                - "/fsx/hf_cache:/fsx/hf_cache"
+
+      - label: "Diffusion · Other Test"
+        commands:
+          - timeout 15m pytest -sv tests/diffusion/*.py -m 'core_model and cuda' --run-level "core_model"
+        agents:
+          queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
+        plugins:
+          - docker#v5.2.0:
+              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+              always-pull: true
+              propagate-environment: true
+              environment:
+                - "HF_HOME=/fsx/hf_cache"
+                - "HF_TOKEN"
+              volumes:
+                - "/fsx/hf_cache:/fsx/hf_cache"
+
+
+  - label: "Engine&Model Executor Test"
     depends_on: upload-ready-pipeline
     commands:
-      - timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py
+      - |
+        timeout 15m bash -c "
+                  pytest -sv tests/engine/ tests/model_executor/ -m 'core_model and cuda' --run-level 'core_model'
+        "
     agents:
-      queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
+      queue: "gpu_1_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
     plugins:
       - docker#v5.2.0:
           image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
@@ -171,166 +178,13 @@ steps:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
 
-  - label: "Engine Test"
+  - label: "Distributed Test"
     depends_on: upload-ready-pipeline
     commands:
       - |
-        timeout 15m bash -c '
-                export VLLM_WORKER_MULTIPROC_METHOD=spawn
-                pytest -s -v tests/engine/test_async_omni_engine_abort.py
-        '
-    agents:
-      queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
-    plugins:
-      - docker#v5.2.0:
-          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-          always-pull: true
-          propagate-environment: true
-          shm-size: "8gb"
-          environment:
-            - "HF_HOME=/fsx/hf_cache"
-            - "HF_TOKEN"
-          volumes:
-            - "/fsx/hf_cache:/fsx/hf_cache"
-
-
-  - label: "Omni Model Test"
-    depends_on: upload-ready-pipeline
-    commands:
-      - |
-        timeout 17m bash -c '
-          export VLLM_LOGGING_LEVEL=DEBUG
-          export VLLM_WORKER_MULTIPROC_METHOD=spawn
-          pytest -s -v tests/e2e/online_serving/test_qwen2_5_omni.py -m "core_model" --run-level "core_model"
-        '
-    agents:
-      queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
-    plugins:
-      - docker#v5.2.0:
-          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-          always-pull: true
-          propagate-environment: true
-          environment:
-            - "HF_HOME=/fsx/hf_cache"
-            - "HF_TOKEN"
-          volumes:
-            - "/fsx/hf_cache:/fsx/hf_cache"
-
-  - label: "Omni Model Test with H100"
-    depends_on: upload-ready-pipeline
-    commands:
-      - |
-        timeout 20m bash -c '
-          export VLLM_WORKER_MULTIPROC_METHOD=spawn
-          pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "core_model" --run-level "core_model"
-        '
-    agents:
-      queue: "mithril-h100-pool"
-    plugins:
-      - kubernetes:
-          podSpec:
-            containers:
-              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-                resources:
-                  limits:
-                    nvidia.com/gpu: 2
-                volumeMounts:
-                  - name: devshm
-                    mountPath: /dev/shm
-                  - name: hf-cache
-                    mountPath: /root/.cache/huggingface
-                env:
-                  - name: HF_HOME
-                    value: /root/.cache/huggingface
-            nodeSelector:
-              node.kubernetes.io/instance-type: gpu-h100-sxm
-            volumes:
-              - name: devshm
-                emptyDir:
-                  medium: Memory
-              - name: hf-cache
-                hostPath:
-                  path: /mnt/hf-cache
-                  type: DirectoryOrCreate
-
-  - label: "MiMo-Audio E2E Test with H100"
-    depends_on: upload-ready-pipeline
-    commands:
-      - |
-        timeout 30m bash -c '
-          export VLLM_LOGGING_LEVEL=DEBUG
-          export VLLM_WORKER_MULTIPROC_METHOD=spawn
-          pytest -s -v tests/e2e/online_serving/test_mimo_audio.py -m "core_model" --run-level "core_model"
-        '
-    agents:
-      queue: "mithril-h100-pool"
-    plugins:
-      - kubernetes:
-          podSpec:
-            containers:
-              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-                resources:
-                  limits:
-                    nvidia.com/gpu: 2
-                volumeMounts:
-                  - name: devshm
-                    mountPath: /dev/shm
-                  - name: hf-cache
-                    mountPath: /root/.cache/huggingface
-                env:
-                  - name: HF_HOME
-                    value: /root/.cache/huggingface
-                  - 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: "Qwen3-TTS E2E Test"
-    depends_on: upload-ready-pipeline
-    commands:
-      - |
-        timeout 20m bash -c '
-          export VLLM_LOGGING_LEVEL=DEBUG
-          export VLLM_WORKER_MULTIPROC_METHOD=spawn
-          export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
-          pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m "core_model" --run-level "core_model"
-        '
-    agents:
-      queue: "gpu_1_queue"
-    plugins:
-      - docker#v5.2.0:
-          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-          always-pull: true
-          propagate-environment: true
-          shm-size: "8gb"
-          environment:
-            - "HF_HOME=/fsx/hf_cache"
-            - "HF_TOKEN"
-          volumes:
-            - "/fsx/hf_cache:/fsx/hf_cache"
-
-  - label: "Voxtral-TTS E2E Test"
-    timeout_in_minutes: 20
-    depends_on: upload-ready-pipeline
-    commands:
-      - |
-        timeout 20m bash -c '
-          export VLLM_LOGGING_LEVEL=DEBUG
-          export VLLM_WORKER_MULTIPROC_METHOD=spawn
-          pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
-          pytest -s -v tests/e2e/offline_inference/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
-        '
+        timeout 15m bash -c "
+                  pytest -sv tests/distributed/ -m 'core_model and cuda' --run-level 'core_model'
+        "
     agents:
       queue: "mithril-h100-pool"
     plugins:
@@ -360,169 +214,327 @@ steps:
                   path: /mnt/hf-cache
                   type: DirectoryOrCreate
 
-  # - label: "Diffusion Image Edit Test with H100 (1 GPU)"
-  #   depends_on: upload-ready-pipeline
-  #   commands:
-  #     - |
-  #       timeout 20m bash -c '
-  #         export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  #         pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py
-  #       '
-  #   agents:
-  #     queue: "mithril-h100-pool"
-  #   plugins:
-  #     - kubernetes:
-  #         podSpec:
-  #           containers:
-  #             - image: 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
-  #           nodeSelector:
-  #             node.kubernetes.io/instance-type: gpu-h100-sxm
-  #           volumes:
-  #             - name: devshm
-  #               emptyDir:
-  #                 medium: Memory
-  #             - name: hf-cache
-  #               hostPath:
-  #                 path: /mnt/hf-cache
-  #                 type: DirectoryOrCreate
-
-  - label: "Bagel Text2Img Model Test with H100"
+  - group: ":card_index_dividers: E2E Test"
     depends_on: upload-ready-pipeline
-    commands:
-      - |
-        timeout 30m bash -c '
-          export VLLM_WORKER_MULTIPROC_METHOD=spawn
-          export VLLM_TEST_CLEAN_GPU_MEMORY=1
-          pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py -m "core_model" --run-level "core_model"
-        '
-    agents:
-      queue: "mithril-h100-pool"
-    plugins:
-      - kubernetes:
-          podSpec:
-            containers:
-              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-                resources:
-                  limits:
-                    nvidia.com/gpu: 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
+    steps:
+      - label: "RLHF · VeRL-Omni E2E Test"
+        commands:
+          - |
+            timeout 20m bash -c "
+              pytest -s -v tests/e2e/offline_inference/rlhf_test/test_verl_omni_e2e.py
+            "
+        agents:
+          queue: "gpu_1_queue"
+        plugins:
+          - docker#v5.2.0:
+              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+              always-pull: true
+              propagate-environment: true
+              shm-size: "8gb"
+              environment:
+                - "HF_HOME=/fsx/hf_cache"
+                - "HF_TOKEN"
+              volumes:
+                - "/fsx/hf_cache:/fsx/hf_cache"
 
-  - label: "Bagel Img2Img Model Test with H100"
-    depends_on: upload-ready-pipeline
-    commands:
-      - |
-        timeout 30m bash -c '
-          export VLLM_WORKER_MULTIPROC_METHOD=spawn
-          export VLLM_TEST_CLEAN_GPU_MEMORY=1
-          pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py -m "core_model" --run-level "core_model"
-        '
-    agents:
-      queue: "mithril-h100-pool"
-    plugins:
-      - kubernetes:
-          podSpec:
-            containers:
-              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-                resources:
-                  limits:
-                    nvidia.com/gpu: 1
-                volumeMounts:
+      - label: "Omni · Qwen3-Omni Test"
+        commands:
+          - |
+            timeout 20m bash -c "
+              pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m 'core_model' --run-level 'core_model'
+            "
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 2
+                    volumeMounts:
+                      - name: devshm
+                        mountPath: /dev/shm
+                      - name: hf-cache
+                        mountPath: /root/.cache/huggingface
+                    env:
+                      - name: HF_HOME
+                        value: /root/.cache/huggingface
+                nodeSelector:
+                  node.kubernetes.io/instance-type: gpu-h100-sxm
+                volumes:
                   - name: devshm
-                    mountPath: /dev/shm
+                    emptyDir:
+                      medium: Memory
                   - 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
+                    hostPath:
+                      path: /mnt/hf-cache
+                      type: DirectoryOrCreate
 
-  - label: "Bagel Online Serving Test with H100"
-    depends_on: upload-ready-pipeline
-    commands:
-      - |
-        timeout 40m bash -c '
-          export VLLM_WORKER_MULTIPROC_METHOD=spawn
-          export VLLM_TEST_CLEAN_GPU_MEMORY=1
-          export VLLM_IMAGE_FETCH_TIMEOUT=60
-          pytest -s -v tests/e2e/online_serving/test_bagel_online.py -m "core_model" --run-level "core_model"
-        '
-    agents:
-      queue: "mithril-h100-pool"
-    plugins:
-      - kubernetes:
-          podSpec:
-            containers:
-              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-                resources:
-                  limits:
-                    nvidia.com/gpu: 1
-                volumeMounts:
+      - label: "Omni · VoxCPM2 Test"
+        timeout_in_minutes: 20
+        commands:
+          - |
+            timeout 20m bash -c "
+              export VLLM_LOGGING_LEVEL=DEBUG
+              export VLLM_WORKER_MULTIPROC_METHOD=spawn
+              pytest -s -v tests/e2e/offline_inference/test_voxcpm2.py -m 'core_model' --run-level 'core_model'
+            "
+        agents:
+          queue: "gpu_1_queue"
+        plugins:
+          - docker#v5.2.0:
+              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+              always-pull: true
+              propagate-environment: true
+              shm-size: "8gb"
+              environment:
+                - "HF_HOME=/fsx/hf_cache"
+                - "HF_TOKEN"
+              volumes:
+                - "/fsx/hf_cache:/fsx/hf_cache"
+
+      - label: "TTS · Qwen3-TTS CustomVoice Test"
+        commands:
+          - |
+            timeout 20m bash -c "
+              export VLLM_ALLOW_LONG_MAX_MODEL_LEN=1
+              pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m 'core_model' --run-level 'core_model'
+            "
+        agents:
+          queue: "gpu_1_queue"
+        plugins:
+          - docker#v5.2.0:
+              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+              always-pull: true
+              propagate-environment: true
+              shm-size: "8gb"
+              environment:
+                - "HF_HOME=/fsx/hf_cache"
+                - "HF_TOKEN"
+              volumes:
+                - "/fsx/hf_cache:/fsx/hf_cache"
+
+      - label: "TTS · Qwen3-TTS Base Test"
+        commands:
+          - |
+            timeout 20m bash -c "
+              export VLLM_WORKER_MULTIPROC_METHOD=spawn
+              export VLLM_ALLOW_LONG_MAX_MODEL_LEN=1
+              export VLLM_OMNI_USE_V2_RUNNER=1
+              pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py -m 'core_model' --run-level 'core_model'
+            "
+        agents:
+          queue: "gpu_1_queue"
+        plugins:
+          - docker#v5.2.0:
+              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+              always-pull: true
+              propagate-environment: true
+              shm-size: "8gb"
+              environment:
+                - "HF_HOME=/fsx/hf_cache"
+                - "HF_TOKEN"
+              volumes:
+                - "/fsx/hf_cache:/fsx/hf_cache"
+
+      - label: "Diffusion · Qwen Image Edit Test"
+        commands:
+          - |
+            timeout 20m bash -c "
+              pytest -s -v tests/e2e/online_serving/test_qwen_image_edit.py -m 'core_model' --run-level 'core_model'
+            "
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 1
+                    volumeMounts:
+                      - name: devshm
+                        mountPath: /dev/shm
+                      - name: hf-cache
+                        mountPath: /root/.cache/huggingface
+                    env:
+                      - name: HF_HOME
+                        value: /root/.cache/huggingface
+                nodeSelector:
+                  node.kubernetes.io/instance-type: gpu-h100-sxm
+                volumes:
                   - name: devshm
-                    mountPath: /dev/shm
+                    emptyDir:
+                      medium: Memory
                   - 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
+                    hostPath:
+                      path: /mnt/hf-cache
+                      type: DirectoryOrCreate
+
+      - label: "Diffusion · Bagel Test"
+        commands:
+          - |
+            timeout 40m bash -c "
+              export VLLM_IMAGE_FETCH_TIMEOUT=60
+              pytest -s -v tests/e2e/online_serving/test_bagel.py -m 'core_model' --run-level 'core_model'
+            "
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 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: "Diffusion · Wan22 Test"
+        commands:
+          - |
+            timeout 40m bash -c "
+              export VLLM_IMAGE_FETCH_TIMEOUT=60
+              pytest -s -v tests/e2e/online_serving/test_wan22_t2v.py -m 'core_model' --run-level 'core_model'
+            "
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 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: "Diffusion · Qwen Image Test"
+        commands:
+          - |
+            timeout 40m bash -c "
+              export VLLM_IMAGE_FETCH_TIMEOUT=60
+              pytest -s -v tests/e2e/online_serving/test_qwen_image.py -m 'core_model' --run-level 'core_model'
+            "
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 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: "Diffusion · Qwen Image Layered Test"
+        commands:
+          - |
+            timeout 40m bash -c "
+              export VLLM_IMAGE_FETCH_TIMEOUT=60
+              pytest -s -v tests/e2e/online_serving/test_qwen_image_layered.py -m 'core_model' --run-level 'core_model'
+            "
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 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

.buildkite/test-template-amd-omni.j2

@@ -3,7 +3,7 @@
    Last synced: 2025-12-15
    Modifications: Removed unused CUDA/NVIDIA logic, keeping only AMD tests
 #}
-{% set docker_image_amd = "rocm/vllm-ci:$BUILDKITE_COMMIT-rocm-omni" %}
+{% set docker_image_amd = "rocm/vllm-omni:$BUILDKITE_COMMIT" %}
 {% set default_working_dir = "/app/vllm-omni" %}
 
   - group: "AMD Tests"
@@ -48,6 +48,9 @@
           DOCKER_BUILDKIT: "1"
           TEST_COMMAND: |-
             (command rocm-smi || true) && cd {{ (step.working_dir or default_working_dir) | safe }}
+{% if "mi250" in step.agent_pool %}
+            python3 -m pip uninstall -y amd-aiter
+{% endif %}
             {{ indented_cmd | safe }}
         priority: 100
         {% if step.grade and step.grade == "Blocking" %}

.buildkite/test-weekly.yml

@@ -0,0 +1,85 @@
+env:
+  VLLM_WORKER_MULTIPROC_METHOD: spawn
+  HF_HUB_DOWNLOAD_TIMEOUT: 300
+  HF_HUB_ETAG_TIMEOUT: 60
+
+steps:
+  - label: "Reliability Test - qwen3-omni"
+    timeout_in_minutes: 180
+    depends_on: upload-weekly-pipeline
+    if: build.env("WEEKLY") == "1" || build.pull_request.labels includes "weekly-test"
+    commands:
+      - pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m "slow"
+    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: "Reliability Test - wan22"
+    timeout_in_minutes: 180
+    depends_on: upload-weekly-pipeline
+    if: build.env("WEEKLY") == "1" || build.pull_request.labels includes "weekly-test"
+    commands:
+      - pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m "slow"
+    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

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

@@ -0,0 +1,575 @@
+---
+name: add-diffusion-model
+description: Add a new diffusion model (text-to-image, text-to-video, image-to-video, text-to-audio, image editing) to vLLM-Omni, including Cache-DiT acceleration and parallelism support (TP, SP/USP, CFG-Parallel, HSDP). Use when integrating a new diffusion model, porting a diffusers pipeline or a custom model repo to vllm-omni, creating a new DiT transformer adapter, adding diffusion model support, or enabling multi-GPU parallelism and cache acceleration for an existing model.
+---
+
+# Adding a Diffusion Model to vLLM-Omni
+
+## Overview
+
+This skill guides you through adding a new diffusion model to vLLM-Omni. The model may come from HuggingFace Diffusers (structured pipeline) or from a private/custom repo. The workflow differs significantly depending on the source.
+
+## Prerequisites
+
+Before starting, determine:
+
+1. **Model category**: Text-to-Image, Text-to-Video, Image-to-Video, Image Editing, Text-to-Audio, or Omni
+2. **Reference source**: Diffusers pipeline, custom repo, or a combination
+3. **Model HuggingFace ID** or local checkpoint path
+4. **Architecture**: Scheduler, text encoder, VAE, transformer/backbone
+
+## Step 0: Classify the Migration Path
+
+Check the model's HF repo for `model_index.json`. This determines your path:
+
+| Scenario | How to identify | Migration path |
+|----------|----------------|----------------|
+| **Already supported** | `_class_name` in `model_index.json` matches a key in `_DIFFUSION_MODELS` in `registry.py` | Skip to Step 5 (test) and Step 7 (docs) |
+| **Diffusers-based** | Has standard `model_index.json` with `_diffusers_version`, subfolders for `transformer/`, `vae/`, etc. | Follow **Path A** below |
+| **Custom/private repo** | No diffusers `model_index.json`, weights in non-standard format, custom model code in a separate git repo | Follow **Path B** below |
+| **Hybrid** | Has some diffusers components (VAE) but custom transformer/fusion | Mix of Path A and Path B |
+
+## Path A: Diffusers-Based Model
+
+For models with a standard diffusers layout. See [references/transformer-adaptation.md](references/transformer-adaptation.md) for detailed code patterns.
+
+### A1. Analyze `model_index.json`
+
+Identify components: `transformer`, `scheduler`, `vae`, `text_encoder`, `tokenizer`.
+
+### A2. Create model directory
+
+```
+vllm_omni/diffusion/models/your_model_name/
+├── __init__.py
+├── pipeline_your_model.py
+└── your_model_transformer.py
+```
+
+### A3. Adapt transformer
+
+1. Copy from diffusers source. Remove mixins (`ModelMixin`, `ConfigMixin`, `AttentionModuleMixin`).
+2. Replace attention with `vllm_omni.diffusion.attention.layer.Attention` (QKV shape: `[B, seq, heads, head_dim]`).
+3. Add `od_config: OmniDiffusionConfig | None = None` to `__init__`.
+4. Add `load_weights()` method mapping diffusers weight names to vllm-omni names.
+5. Add class attributes: `_repeated_blocks`, `_layerwise_offload_blocks_attr`.
+
+### A4. Adapt pipeline
+
+Inherit from `nn.Module`. The key contract:
+
+```python
+class YourPipeline(nn.Module):
+    def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
+        # Load VAE, text encoder, tokenizer via from_pretrained()
+        # Instantiate transformer (weights loaded later via weights_sources)
+        self.weights_sources = [
+            DiffusersPipelineLoader.ComponentSource(
+                model_or_path=od_config.model, subfolder="transformer",
+                prefix="transformer.", fall_back_to_pt=True)]
+
+    def forward(self, req: OmniDiffusionRequest) -> DiffusionOutput:
+        # Encode prompt → prepare latents → denoise loop → VAE decode
+        return DiffusionOutput(output=output)
+
+    def load_weights(self, weights):
+        return AutoWeightsLoader(self).load_weights(weights)
+```
+
+Add post/pre-process functions in the same pipeline file. Register them in `registry.py`.
+
+### A5. Register, test, docs → continue at Step 4 below.
+
+---
+
+## Path B: Custom/Private Repo Model
+
+For models without a diffusers pipeline — weights in custom format, model code in a private repo. Real examples: DreamID-Omni, BAGEL, HunyuanImage3.
+
+### B1. Understand the reference repo
+
+Study the original model's code to identify:
+- **Model architecture files** (transformers, fusion modules, embeddings)
+- **Weight format** (safetensors, `.pth`, custom checkpoint structure)
+- **Weight loading helpers** (custom init functions, checkpoint loaders)
+- **Pre/post-processing** (image/audio transforms, tokenization, VAE encode/decode)
+- **External dependencies** (packages not on PyPI)
+- **Config format** (JSON config files, hardcoded dicts)
+
+### B2. Decide what lives WHERE
+
+This is the key design decision for custom models. Follow these placement rules:
+
+| Code type | Where to place | Example |
+|-----------|---------------|---------|
+| **Pipeline orchestration** (init, forward, denoise loop) | `vllm_omni/diffusion/models/<name>/pipeline_<name>.py` | Always required |
+| **Custom transformer/backbone** (ported and adapted to vllm-omni) | `vllm_omni/diffusion/models/<name>/<name>_transformer.py` or similar | `wan2_2.py`, `fusion.py`, `bagel_transformer.py` |
+| **Custom sub-models** (VAE, fusion, autoencoder) | `vllm_omni/diffusion/models/<name>/` as separate files | `autoencoder.py`, `fusion.py` |
+| **External dependency code** (original repo utilities) | **External repo**, installed via download script or pip | `dreamid_omni` package via git clone |
+| **Hardcoded model configs** | Module-level dicts in pipeline file | `VIDEO_CONFIG`, `AUDIO_CONFIG` dicts |
+| **Download/setup script** | `examples/offline_inference/<name>/download_<name>.py` | `download_dreamid_omni.py` |
+| **Custom `model_index.json`** | Generated by download script, placed at model root | Minimal: `{"_class_name": "YourPipeline", ...}` |
+
+### B3. Handle external dependencies
+
+If the model's code lives in a separate git repo:
+
+**Option 1: Import with graceful fallback** (recommended for models with external utils)
+
+```python
+try:
+    from external_model.utils import init_vae, load_checkpoint
+except ImportError:
+    raise ImportError(
+        "Failed to import from dependency 'external_model'. "
+        "Please run the download script first."
+    )
+```
+
+**Option 2: Port the code directly** (preferred when feasible)
+
+Copy the essential model files into `vllm_omni/diffusion/models/<name>/` and adapt them. This avoids external dependencies. BAGEL does this — `autoencoder.py` and `bagel_transformer.py` are ported directly.
+
+**Decision criteria**: Port if the code is self-contained and won't diverge. Use external deps if the model repo is actively maintained and the code is complex.
+
+### B4. Handle custom weight loading
+
+Custom models have two patterns for weight loading:
+
+**Pattern 1: Bypass standard loader** (DreamID-Omni style)
+
+When the original model has complex custom init functions that load weights in `__init__`:
+
+```python
+class CustomPipeline(nn.Module):
+    def __init__(self, *, od_config, prefix=""):
+        super().__init__()
+        model = od_config.model
+        # Load everything eagerly in __init__ using custom helpers
+        self.vae = custom_init_vae(model, device=self.device)
+        self.text_encoder = custom_init_text_encoder(model, device=self.device)
+        self.transformer = CustomFusionModel(CONFIG)
+        load_custom_checkpoint(self.transformer,
+            checkpoint_path=os.path.join(model, "model.safetensors"))
+        # NO weights_sources defined — bypasses standard loader
+
+    def load_weights(self, weights):
+        pass  # No-op — all weights loaded in __init__
+```
+
+**Pattern 2: Use standard loader with custom `load_weights`** (BAGEL style)
+
+When weights are in safetensors format but need name remapping:
+
+```python
+class CustomPipeline(nn.Module):
+    def __init__(self, *, od_config, prefix=""):
+        super().__init__()
+        # Instantiate model architecture without weights
+        self.bagel = BagelModel(config)
+        self.vae = AutoEncoder(ae_params)
+
+        # Point loader at the safetensors in the model root
+        self.weights_sources = [
+            DiffusersPipelineLoader.ComponentSource(
+                model_or_path=od_config.model,
+                subfolder=None,  # weights at root, not in subfolder
+                prefix="",
+                fall_back_to_pt=False,
+            )
+        ]
+
+    def load_weights(self, weights):
+        # Custom name remapping for non-diffusers weight names
+        params = dict(self.named_parameters())
+        loaded = set()
+        for name, tensor in weights:
+            # Remap original weight names to vllm-omni module names
+            name = self._remap_weight_name(name)
+            if name in params:
+                default_weight_loader(params[name], tensor)
+                loaded.add(name)
+        return loaded
+```
+
+### B5. Create the `model_index.json`
+
+Custom models need a `model_index.json` at the model root for vllm-omni to discover them. For custom models, this is minimal:
+
+```json
+{
+    "_class_name": "YourModelPipeline",
+    "custom_key": "path/to/custom_weights.safetensors"
+}
+```
+
+The `_class_name` must match a key in `_DIFFUSION_MODELS` in `registry.py`. Additional keys are model-specific (accessed via `od_config.model_config`).
+
+If the model's weights come from multiple HF repos, write a **download script** that:
+1. Downloads from each repo
+2. Assembles into a single directory
+3. Generates `model_index.json`
+4. Installs any external dependencies (git clone + `.pth` file)
+
+Place at: `examples/offline_inference/<name>/download_<name>.py`
+
+### B6. Handle multi-modal inputs
+
+If the model accepts images, audio, or other multi-modal inputs, implement the protocol classes from `vllm_omni/diffusion/models/interface.py`:
+
+```python
+from vllm_omni.diffusion.models.interface import SupportImageInput, SupportAudioInput
+
+class MyPipeline(nn.Module, SupportImageInput, SupportAudioInput):
+    # Protocol markers — the engine uses these to enable proper input routing
+    pass
+```
+
+Preprocessing for custom models is typically done **inside `forward()`** rather than via registered pre-process functions, since the logic is often tightly coupled to the model.
+
+### B7. Continue at Step 4 below.
+
+---
+
+## Common Steps (Both Paths)
+
+### Step 4: Register Model in registry.py
+
+Edit `vllm_omni/diffusion/registry.py`:
+
+```python
+_DIFFUSION_MODELS = {
+    "YourModelPipeline": ("your_model_name", "pipeline_your_model", "YourModelPipeline"),
+}
+_DIFFUSION_POST_PROCESS_FUNCS = {
+    "YourModelPipeline": "get_your_model_post_process_func",  # if applicable
+}
+_DIFFUSION_PRE_PROCESS_FUNCS = {
+    "YourModelPipeline": "get_your_model_pre_process_func",  # if applicable
+}
+```
+
+The registry key is the `_class_name` from `model_index.json`. The tuple is `(folder_name, module_file, class_name)`.
+
+Create `__init__.py` exporting the pipeline class and any factory functions.
+
+### Step 5: Run, Test, Debug
+
+Use the appropriate existing example script:
+
+| Category | Script |
+|----------|--------|
+| Text-to-Image | `examples/offline_inference/text_to_image/text_to_image.py` |
+| Text-to-Video | `examples/offline_inference/text_to_video/text_to_video.py` |
+| Image-to-Video | `examples/offline_inference/image_to_video/image_to_video.py` |
+| Image-to-Image | `examples/offline_inference/image_to_image/image_edit.py` |
+| Text-to-Audio | `examples/offline_inference/text_to_audio/text_to_audio.py` |
+
+For custom/Omni models that don't fit these categories, create a dedicated example script.
+
+**Validation**: No errors, output is meaningful, quality matches reference implementation.
+
+See [references/troubleshooting.md](references/troubleshooting.md) for common errors.
+
+### Step 6: Add Example Scripts
+
+For Omni or custom models, create:
+- `examples/offline_inference/your_model_name/` — offline script + README
+- `examples/online_serving/your_model_name/` — server script + client
+- Download script if weights require assembly from multiple sources
+
+### Step 7: Update Documentation
+
+Required updates:
+1. `docs/user_guide/diffusion/parallelism_acceleration.md` — parallelism support table
+2. `docs/user_guide/diffusion/cpu_offload_diffusion.md` — if CPU offload supported (add to supported models table)
+3. `docs/user_guide/diffusion/teacache.md` — if TeaCache supported
+4. `docs/user_guide/diffusion/cache_dit_acceleration.md` — if Cache-DiT supported
+5. `examples/offline_inference/xxx/README.md` — offline example docs
+6. `examples/online_serve/xxx/README.md` — online serve docs
+
+### Step 8: Add E2E Tests (Recommended)
+
+Create `tests/e2e/online_serving/test_your_model_expansion.py`.
+
+### Step 9: Add Cache-DiT Acceleration
+
+Cache-DiT accelerates inference by caching intermediate computation results across denoising steps. After your model is working correctly on a single GPU, add cache-dit support.
+
+See [references/cache-dit-patterns.md](references/cache-dit-patterns.md) for detailed code patterns.
+
+#### 9a. Determine your model type
+
+| Model Type | Description | Action |
+|------------|-------------|--------|
+| **Standard single-transformer** | One transformer with one `ModuleList` of blocks | No code needed — `CacheDiTBackend` auto-detects via `enable_cache_for_dit()` |
+| **Multi-block-list** | One transformer with multiple block lists (e.g., `transformer_blocks` + `single_transformer_blocks`) | Write custom enabler with `BlockAdapter` |
+| **Dual-transformer** | Two transformers (e.g., high-noise + low-noise) | Write custom enabler with `BlockAdapter` wrapping both |
+
+#### 9b. Standard models — verify automatic support
+
+For standard single-transformer models, test directly:
+
+```python
+omni = Omni(
+    model="your-model-name",
+    cache_backend="cache_dit",
+    cache_config={
+        "Fn_compute_blocks": 1,
+        "Bn_compute_blocks": 0,
+        "max_warmup_steps": 4,
+    }
+)
+```
+
+Check logs for "Cache-dit enabled successfully on xxx". If it works, skip to Step 9e.
+
+#### 9c. Custom architectures — write a custom enabler
+
+For multi-block-list or dual-transformer models, write a custom enabler function:
+
+```python
+from cache_dit import BlockAdapter, ForwardPattern, ParamsModifier, DBCacheConfig
+
+def enable_cache_for_your_model(pipeline, cache_config):
+    db_cache_config = DBCacheConfig(
+        num_inference_steps=None,
+        Fn_compute_blocks=cache_config.Fn_compute_blocks,
+        Bn_compute_blocks=cache_config.Bn_compute_blocks,
+        max_warmup_steps=cache_config.max_warmup_steps,
+        max_cached_steps=cache_config.max_cached_steps,
+        max_continuous_cached_steps=cache_config.max_continuous_cached_steps,
+        residual_diff_threshold=cache_config.residual_diff_threshold,
+    )
+
+    cache_dit.enable_cache(
+        BlockAdapter(
+            transformer=pipeline.transformer,
+            blocks=[
+                pipeline.transformer.transformer_blocks,
+                pipeline.transformer.single_transformer_blocks,
+            ],
+            forward_pattern=[ForwardPattern.Pattern_1, ForwardPattern.Pattern_1],
+            params_modifiers=[ParamsModifier(...)],
+        ),
+        cache_config=db_cache_config,
+    )
+
+    def refresh_cache_context(pipeline, num_inference_steps, verbose=True):
+        cache_dit.refresh_context(
+            pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose
+        )
+    return refresh_cache_context
+```
+
+#### 9d. Register the custom enabler
+
+Add your enabler to `CUSTOM_DIT_ENABLERS` in `vllm_omni/diffusion/cache/cache_dit_backend.py`:
+
+```python
+CUSTOM_DIT_ENABLERS = {
+    "Wan22Pipeline": enable_cache_for_wan22,
+    "LongCatImagePipeline": enable_cache_for_longcat_image,
+    "YourModelPipeline": enable_cache_for_your_model,  # Add here
+}
+```
+
+#### 9e. Test Cache-DiT
+
+```python
+omni = Omni(
+    model="your-model-name",
+    cache_backend="cache_dit",
+    cache_config={
+        "Fn_compute_blocks": 1, "Bn_compute_blocks": 0,
+        "max_warmup_steps": 4, "residual_diff_threshold": 0.24,
+    }
+)
+images = omni.generate("a beautiful landscape",
+    OmniDiffusionSamplingParams(num_inference_steps=50))
+```
+
+**Verify**: 1) logs show cache enabled, 2) 1.5-2x speedup, 3) output quality acceptable vs baseline.
+
+If quality degrades, lower `residual_diff_threshold` (try 0.12-0.18) or increase `max_warmup_steps` (try 6-8).
+
+---
+
+### Step 10: Add Parallelism Support
+
+After the model works on a single GPU, add multi-GPU parallelism. Add each type incrementally, testing after each addition.
+
+See [references/parallelism-patterns.md](references/parallelism-patterns.md) for detailed code patterns and API reference.
+
+**Recommended order**: TP → SP/USP → CFG Parallel → HSDP
+
+#### 10a. Tensor Parallelism (TP)
+
+Shards DiT linear layers across GPUs. Requires code changes in the transformer.
+
+**What to change in the transformer**:
+1. Replace `nn.Linear` with `ColumnParallelLinear` / `RowParallelLinear` / `QKVParallelLinear`
+2. Update `load_weights()` to handle QKV fusion with `stacked_params_mapping`
+3. Use `self.to_qkv.num_heads` (local heads) instead of total heads for split sizes
+
+```python
+from vllm.model_executor.layers.linear import (
+    QKVParallelLinear, RowParallelLinear, ColumnParallelLinear,
+)
+
+# Attention: QKV → RowParallel output
+self.to_qkv = QKVParallelLinear(dim, head_dim, num_heads, num_kv_heads)
+self.to_out = RowParallelLinear(dim, dim, input_is_parallel=True)
+
+# FFN: ColumnParallel → RowParallel
+self.w1 = ColumnParallelLinear(dim, ffn_dim)
+self.w2 = RowParallelLinear(ffn_dim, dim, input_is_parallel=True)
+```
+
+**Constraints**: `num_heads % tp_size == 0` and `num_kv_heads % tp_size == 0`.
+
+**Test**: `--tensor-parallel-size 2`
+
+#### 10b. Sequence Parallelism (SP / USP)
+
+Splits sequence tokens across GPUs. Non-intrusive via `_sp_plan` on the transformer class — no changes to `forward()`.
+
+**What to change in the transformer**:
+
+Add `_sp_plan` class attribute:
+
+```python
+from vllm_omni.diffusion.distributed.sp_plan import (
+    SequenceParallelInput, SequenceParallelOutput,
+)
+
+class YourTransformer(nn.Module):
+    _sp_plan = {
+        "blocks.0": {
+            "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3),
+        },
+        "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3),
+    }
+```
+
+If inline tensor ops (e.g., `torch.cat`) exist between shard/gather points, extract them into `nn.Module` submodules so hooks can intercept them.
+
+For RoPE that needs splitting, add an entry for the RoPE module with `split_output=True`.
+
+**Test**: `--ulysses-degree 2` (offline) or `--usp 2` (online serving)
+
+#### 10c. CFG Parallel
+
+Distributes positive/negative CFG branches across 2 GPUs. Requires the pipeline to inherit `CFGParallelMixin`.
+
+**What to change in the pipeline**:
+
+```python
+from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
+
+class YourPipeline(nn.Module, CFGParallelMixin):
+    def diffuse(self, ...) -> torch.Tensor:
+        for i, t in enumerate(timesteps):
+            positive_kwargs = {...}
+            negative_kwargs = {...} if do_true_cfg else None
+            noise_pred = self.predict_noise_maybe_with_cfg(
+                do_true_cfg=do_true_cfg, true_cfg_scale=cfg_scale,
+                positive_kwargs=positive_kwargs, negative_kwargs=negative_kwargs,
+            )
+            latents = self.scheduler_step_maybe_with_cfg(
+                noise_pred, t, latents, do_true_cfg
+            )
+        return latents
+```
+
+Override `predict_noise()` if your transformer call is non-standard. Override `combine_cfg_noise()` for multi-output models (e.g., video + audio).
+
+**Constraint**: Exactly 2 GPUs. Only for models using classifier-free guidance.
+
+**Test**: `--cfg-parallel-size 2`
+
+#### 10d. HSDP (Hybrid Sharded Data Parallel)
+
+Shards transformer weights via PyTorch FSDP2 to reduce per-GPU VRAM. No code changes to the forward pass — just add a class attribute.
+
+**What to change in the transformer**:
+
+```python
+class YourTransformer(nn.Module):
+    @staticmethod
+    def _is_transformer_block(name: str, module) -> bool:
+        return "blocks" in name and name.split(".")[-1].isdigit()
+
+    _hsdp_shard_conditions = [_is_transformer_block]
+```
+
+**Constraint**: Cannot combine with TP. For standalone HSDP, set `hsdp_shard_size` explicitly.
+
+**Test**: `--use-hsdp` or `DiffusionParallelConfig(use_hsdp=True)`
+
+#### 10e. Update parallelism documentation
+
+After adding parallelism support, update:
+1. `docs/user_guide/diffusion/parallelism_acceleration.md` — add your model to the support table
+2. Record which parallelism methods are supported (USP, Ring, CFG, TP, HSDP, VAE-Patch)
+
+### Step 11: Add CPU Offload Support
+
+Implement `SupportsComponentDiscovery` 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
+
+class YourPipeline(nn.Module, SupportsComponentDiscovery):
+    _dit_modules: ClassVar[list[str]] = ["transformer"]
+    _encoder_modules: ClassVar[list[str]] = ["text_encoder"]
+    _vae_modules: ClassVar[list[str]] = ["vae"]
+    _resident_modules: ClassVar[list[str]] = []  # optional
+```
+
+- `_dit_modules`: denoising submodules (kept on GPU during diffusion loop)
+- `_encoder_modules`: encoder/vision submodules (offloaded to CPU during diffusion loop)
+- `_vae_modules`: VAE(s) (handled by both sequential and layerwise backends)
+- `_resident_modules`: additional modules to pin on GPU during layerwise
+  offloading (e.g. embedders, connectors). Only used by the layerwise
+  backend. Optional — defaults to `[]`.
+
+All attribute names support dotted paths for nested submodules
+(e.g. `"pipe.transformer"`, `"bagel.time_embedder"`).
+
+Pipelines without `SupportsComponentDiscovery` fall back to scanning
+well-known attribute names (`transformer`, `text_encoder`, `vae`,
+etc.), which fails for non-standard names.
+
+### Step 12: Performance Profiling
+
+After verifying correctness and implementing parallelism/caching, profile the model's performance to identify bottlenecks and ensure optimal execution.
+
+See the [Profiling Single-Stage Diffusion](../../../docs/contributing/profiling.md#3-profiling-single-stage-diffusion) guide for detailed instructions on:
+1. Using the PyTorch profiler (`profiler: "torch"`) to capture detailed CPU/CUDA traces.
+2. Using Nsight Systems (`nsys`) with `profiler: "cuda"` for low-overhead CUDA traces.
+3. Controlling profiling via `omni.start_profile()` and `omni.stop_profile()`.
+
+---
+
+## Iterative Development Tips
+
+1. **Start minimal**: Basic generation first, no parallelism/caching
+2. **Use `--enforce-eager`**: Disable torch.compile during debugging
+3. **Use small models**: Test with smaller variants first
+4. **Check tensor shapes**: Most errors are reshape mismatches in attention
+5. **Add features incrementally**: Single GPU → TP → SP → CFG → HSDP → Cache-DiT
+6. **For custom models**: Get the model running with the original code first, then progressively replace components with vllm-omni equivalents
+7. **Cache-DiT before parallelism tuning**: Cache-DiT is lossy — verify quality at baseline before combining with parallelism
+8. **Combine lossless + lossy**: e.g., TP + SP + Cache-DiT for maximum throughput
+
+## Reference Files
+
+- [Transformer Adaptation](references/transformer-adaptation.md) — porting transformers from diffusers
+- [Custom Model Patterns](references/custom-model-patterns.md) — patterns for non-diffusers models
+- [Parallelism Patterns](references/parallelism-patterns.md) — TP, SP/USP, CFG parallel, HSDP implementation details
+- [Cache-DiT Patterns](references/cache-dit-patterns.md) — cache-dit acceleration for standard and custom architectures
+- [Troubleshooting](references/troubleshooting.md) — common errors and fixes

.claude/skills/add-diffusion-model/references/cache-dit-patterns.md

@@ -0,0 +1,254 @@
+# Cache-DiT Patterns Reference
+
+## Overview
+
+Cache-DiT accelerates Diffusion Transformers by caching intermediate computation results across denoising steps. Adjacent steps produce similar features, so redundant computations can be skipped.
+
+Three caching strategies:
+- **DBCache**: Dynamic block-level caching — selectively computes or caches transformer blocks based on residual differences
+- **TaylorSeer**: Calibration-based prediction using Taylor expansion to estimate block outputs
+- **SCM** (Step Computation Masking): Dynamic step skipping based on configurable policies
+
+**Typical speedup**: 1.5-2.5x depending on model and configuration.
+
+**Official docs**: https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/cache_dit
+
+## Architecture
+
+vLLM-Omni integrates cache-dit through `CacheDiTBackend`:
+
+| Component | Purpose |
+|-----------|---------|
+| `CacheDiTBackend` | Unified backend — auto-selects enabler (standard or custom) |
+| `enable_cache_for_dit()` | Default enabler for standard single-transformer models |
+| `CUSTOM_DIT_ENABLERS` dict | Registry of custom enablers keyed by pipeline class name |
+| `BlockAdapter` | Wraps complex architectures (multi-block-list or multi-transformer) |
+| `ForwardPattern` | Specifies block forward signature: `Pattern_0`, `Pattern_1`, `Pattern_2` |
+| `ParamsModifier` | Per-transformer or per-block-list config customization |
+| `DBCacheConfig` | Configuration for DBCache parameters |
+| `cache_dit.refresh_context()` | Updates cache context when `num_inference_steps` changes |
+
+**Source files**:
+- `vllm_omni/diffusion/cache/cache_dit_backend.py` — `CacheDiTBackend`, enablers, `CUSTOM_DIT_ENABLERS`
+- `vllm_omni/diffusion/cache/` — cache backend implementations
+
+## Standard Models: Automatic Support
+
+Most DiT models follow this pattern:
+- Single transformer with one `nn.ModuleList` of blocks
+- Standard forward signature
+- Compatible with cache-dit's automatic detection
+
+**Examples**: Qwen-Image, Z-Image, FLUX
+
+No code changes needed. `CacheDiTBackend` automatically uses `enable_cache_for_dit()`:
+
+```python
+from vllm_omni import Omni
+
+omni = Omni(
+    model="Qwen/Qwen-Image",
+    cache_backend="cache_dit",
+    cache_config={
+        "Fn_compute_blocks": 1,
+        "Bn_compute_blocks": 0,
+        "max_warmup_steps": 4,
+    }
+)
+```
+
+What happens automatically:
+
+```python
+def enable_cache_for_dit(pipeline, cache_config):
+    db_cache_config = DBCacheConfig(
+        num_inference_steps=None,
+        Fn_compute_blocks=cache_config.Fn_compute_blocks,
+        Bn_compute_blocks=cache_config.Bn_compute_blocks,
+        max_warmup_steps=cache_config.max_warmup_steps,
+        max_cached_steps=cache_config.max_cached_steps,
+        max_continuous_cached_steps=cache_config.max_continuous_cached_steps,
+        residual_diff_threshold=cache_config.residual_diff_threshold,
+    )
+
+    cache_dit.enable_cache(pipeline.transformer, cache_config=db_cache_config)
+
+    def refresh_cache_context(pipeline, num_inference_steps, verbose=True):
+        cache_dit.refresh_context(
+            pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose
+        )
+    return refresh_cache_context
+```
+
+## Custom Architectures: Writing Custom Enablers
+
+### When you need a custom enabler
+
+- Model has multiple block lists in one transformer (e.g., `transformer_blocks` + `single_transformer_blocks`)
+- Model has two transformers (e.g., high-noise + low-noise like Wan2.2)
+- Model uses non-standard block forward signature
+
+### Pattern 1: Multi-Block-List (LongCat-Image style)
+
+Single transformer with two block lists:
+
+```python
+import cache_dit
+from cache_dit import BlockAdapter, ForwardPattern, ParamsModifier, DBCacheConfig
+
+def enable_cache_for_your_model(pipeline, cache_config):
+    db_cache_config = DBCacheConfig(
+        num_inference_steps=None,
+        Fn_compute_blocks=cache_config.Fn_compute_blocks,
+        Bn_compute_blocks=cache_config.Bn_compute_blocks,
+        max_warmup_steps=cache_config.max_warmup_steps,
+        max_cached_steps=cache_config.max_cached_steps,
+        max_continuous_cached_steps=cache_config.max_continuous_cached_steps,
+        residual_diff_threshold=cache_config.residual_diff_threshold,
+    )
+
+    cache_dit.enable_cache(
+        BlockAdapter(
+            transformer=pipeline.transformer,
+            blocks=[
+                pipeline.transformer.transformer_blocks,
+                pipeline.transformer.single_transformer_blocks,
+            ],
+            forward_pattern=[ForwardPattern.Pattern_1, ForwardPattern.Pattern_1],
+            params_modifiers=[ParamsModifier(...)],
+        ),
+        cache_config=db_cache_config,
+    )
+
+    def refresh_cache_context(pipeline, num_inference_steps, verbose=True):
+        cache_dit.refresh_context(
+            pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose
+        )
+    return refresh_cache_context
+```
+
+For single transformer with multiple block lists, `refresh_context` works the same as standard models — call it once on the transformer.
+
+### Pattern 2: Dual-Transformer (Wan2.2 style)
+
+Two transformers with separate configs:
+
+```python
+def enable_cache_for_dual_transformer(pipeline, cache_config):
+    db_cache_config = DBCacheConfig(...)
+
+    cache_dit.enable_cache(
+        BlockAdapter(
+            transformer=[pipeline.transformer, pipeline.transformer_2],
+            blocks=[pipeline.transformer.blocks, pipeline.transformer_2.blocks],
+            forward_pattern=[ForwardPattern.Pattern_2, ForwardPattern.Pattern_2],
+            params_modifiers=[
+                ParamsModifier(...),  # Config for transformer 1
+                ParamsModifier(...),  # Config for transformer 2
+            ],
+        ),
+        cache_config=db_cache_config,
+    )
+
+    def refresh_cache_context(pipeline, num_inference_steps, verbose=True):
+        high_steps, low_steps = _split_inference_steps(num_inference_steps)
+        cache_dit.refresh_context(
+            pipeline.transformer, num_inference_steps=high_steps, verbose=verbose
+        )
+        cache_dit.refresh_context(
+            pipeline.transformer_2, num_inference_steps=low_steps, verbose=verbose
+        )
+    return refresh_cache_context
+```
+
+Key difference: `refresh_context` must be called on **each transformer separately** with its own step count.
+
+### Choosing the ForwardPattern
+
+| Pattern | Block forward signature | Example models |
+|---------|------------------------|----------------|
+| `Pattern_0` | `block(hidden_states, **kwargs)` → residual added inside block | Default |
+| `Pattern_1` | `block(hidden_states, **kwargs)` → returns `(hidden_states, ...)` tuple | FLUX-style single blocks |
+| `Pattern_2` | `block(hidden_states, **kwargs)` → `(hidden_states, ...)` with different residual pattern | Wan2.2 blocks |
+
+Inspect your block's `forward()` return type and residual connection pattern to choose the right one. See [Cache-DiT API Reference](https://cache-dit.readthedocs.io/en/latest/user_guide/CACHE_API/) for details.
+
+## Registering Custom Enablers
+
+Add your enabler to `CUSTOM_DIT_ENABLERS` in `vllm_omni/diffusion/cache/cache_dit_backend.py`:
+
+```python
+CUSTOM_DIT_ENABLERS = {
+    "Wan22Pipeline": enable_cache_for_wan22,
+    "LongCatImagePipeline": enable_cache_for_longcat_image,
+    "YourModelPipeline": enable_cache_for_your_model,
+}
+```
+
+The key must match `pipeline.__class__.__name__`.
+
+## Configuration Parameters
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `Fn_compute_blocks` | 1 | Number of blocks to always compute at the front |
+| `Bn_compute_blocks` | 0 | Number of blocks to always compute at the back |
+| `max_warmup_steps` | 4 | Steps to run without caching at the beginning |
+| `max_cached_steps` | — | Max total cached steps |
+| `max_continuous_cached_steps` | — | Max consecutive cached steps |
+| `residual_diff_threshold` | 0.24 | Threshold for deciding whether to cache a block |
+
+### Tuning for quality vs speed
+
+| Goal | Adjustments |
+|------|-------------|
+| **More speed, acceptable quality loss** | Higher `residual_diff_threshold` (0.24-0.4), lower `max_warmup_steps` (2-4) |
+| **Better quality, less speed** | Lower `residual_diff_threshold` (0.12-0.18), higher `max_warmup_steps` (6-8), lower `max_continuous_cached_steps` (2) |
+
+## Testing
+
+```python
+from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+omni = Omni(
+    model="your-model-name",
+    cache_backend="cache_dit",
+    cache_config={
+        "Fn_compute_blocks": 1,
+        "Bn_compute_blocks": 0,
+        "max_warmup_steps": 4,
+        "residual_diff_threshold": 0.24,
+    }
+)
+images = omni.generate(
+    "a beautiful landscape",
+    OmniDiffusionSamplingParams(num_inference_steps=50),
+)
+```
+
+CLI (online serving):
+
+```bash
+vllm serve your-model --omni --port 8098 \
+  --cache-backend cache_dit \
+  --cache-config '{"Fn_compute_blocks": 1, "Bn_compute_blocks": 0, "max_warmup_steps": 4}'
+```
+
+**Verification checklist**:
+1. Logs show "Cache-dit enabled successfully on xxx"
+2. Performance: 1.5-2x speedup vs no cache
+3. Quality: compare output with `cache_backend=None`
+
+## Excluded Models
+
+Models listed in `_NO_CACHE_ACCELERATION` in `vllm_omni/diffusion/registry.py` do not support cache-dit (e.g., `NextStep11Pipeline`, `StableDiffusionPipeline`). Check this set before attempting to enable cache-dit.
+
+## Reference Implementations
+
+| Model | Path | Notes |
+|-------|------|-------|
+| Standard DiT | `cache_dit_backend.py::enable_cache_for_dit` | Default enabler, automatic |
+| Wan2.2 | `cache_dit_backend.py::enable_cache_for_wan22` | Dual-transformer, auto-detects mode |
+| LongCat | `cache_dit_backend.py::enable_cache_for_longcat_image` | Multi-block-list |
+| BAGEL | `cache_dit_backend.py::enable_cache_for_bagel` | Complex omni model |

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

@@ -0,0 +1,273 @@
+# Custom Model Patterns Reference
+
+Patterns for adding models that don't come from the standard diffusers pipeline format.
+
+## Directory Structure Comparison
+
+### Diffusers-based model (e.g., Wan2.2)
+
+```
+vllm_omni/diffusion/models/wan2_2/
+├── __init__.py                    # Exports pipeline + transformer + helpers
+├── pipeline_wan2_2.py             # Pipeline: loads components via from_pretrained()
+├── pipeline_wan2_2_i2v.py         # Variant pipeline for image-to-video
+└── wan2_2_transformer.py          # Transformer: ported from diffusers, uses Attention layer
+```
+
+The transformer is loaded separately via `weights_sources` + `load_weights()`. Non-transformer components (VAE, text encoder) are loaded in `__init__` via `from_pretrained()`.
+
+### Custom model with external deps (e.g., DreamID-Omni)
+
+```
+vllm_omni/diffusion/models/dreamid_omni/
+├── __init__.py                    # Exports pipeline only
+├── pipeline_dreamid_omni.py       # Pipeline: loads ALL weights in __init__ via custom helpers
+├── fusion.py                      # Custom fusion architecture (video + audio cross-attention)
+└── wan2_2.py                      # Re-implemented Wan backbone with split API
+
+examples/offline_inference/x_to_video_audio/
+└── download_dreamid_omni.py       # Downloads weights from 3 HF repos + clones code repo
+```
+
+All weights loaded eagerly in `__init__`. `load_weights()` is a no-op. External dependency (`dreamid_omni` package) imported with try/except.
+
+### Custom model with ported code (e.g., BAGEL)
+
+```
+vllm_omni/diffusion/models/bagel/
+├── __init__.py
+├── pipeline_bagel.py              # Pipeline: instantiates models, uses weights_sources
+├── bagel_transformer.py           # Full LLM backbone (Qwen2-MoT) ported into vllm-omni
+└── autoencoder.py                 # Custom VAE ported from original repo
+```
+
+Model code is fully ported (no external dependency). Uses `weights_sources` and `load_weights()` with custom name remapping to handle non-diffusers safetensors format.
+
+## Weight Loading Patterns
+
+### Pattern 1: Standard diffusers flow (Wan2.2, Z-Image, FLUX)
+
+```
+init → create transformer (empty) → set weights_sources → [loader calls load_weights()]
+```
+
+- `weights_sources` points to safetensors in HF subfolder (e.g., `transformer/`)
+- `load_weights()` receives `(name, tensor)` pairs from the loader
+- Name remapping handles diffusers→vllm-omni differences (QKV fusion, Sequential index removal)
+
+### Pattern 2: Custom safetensors at root (BAGEL)
+
+```
+init → create all models (empty) → set weights_sources(subfolder=None) → [loader calls load_weights()]
+```
+
+- `weights_sources` points to **root** of model directory, not a subfolder
+- Weights have non-diffusers names (e.g., `bagel.language_model.model.layers.0.self_attn.q_proj.weight`)
+- `load_weights()` does heavy name normalization
+
+```python
+self.weights_sources = [
+    DiffusersPipelineLoader.ComponentSource(
+        model_or_path=od_config.model,
+        subfolder=None,      # root directory
+        prefix="",           # no prefix stripping
+        fall_back_to_pt=False,
+    )
+]
+```
+
+### Pattern 3: Fully custom loading (DreamID-Omni)
+
+```
+init → load ALL weights eagerly via custom helpers → load_weights() = no-op
+```
+
+- No `weights_sources` attribute — standard loader finds nothing to iterate
+- Custom init functions (e.g., `init_wan_vae_2_2()`, `load_fusion_checkpoint()`) handle downloading and loading
+- `load_weights()` is `pass`
+- Weights may come from multiple HF repos in different formats (`.pth`, `.safetensors`)
+
+Use this when:
+- The original model has complex, well-tested loading code you don't want to rewrite
+- Weights span multiple HF repos
+- Weight format is non-standard (e.g., a single `.pth` file, not sharded safetensors)
+
+## model_index.json for Custom Models
+
+Standard diffusers `model_index.json`:
+```json
+{
+    "_class_name": "WanPipeline",
+    "_diffusers_version": "0.35.0.dev0",
+    "scheduler": ["diffusers", "UniPCMultistepScheduler"],
+    "transformer": ["diffusers", "WanTransformer3DModel"],
+    "vae": ["diffusers", "AutoencoderKLWan"]
+}
+```
+
+Custom model `model_index.json` (minimal):
+```json
+{
+    "_class_name": "DreamIDOmniPipeline",
+    "fusion": "DreamID-Omni/dreamid_omni.safetensors"
+}
+```
+
+The only **required** field is `_class_name` — it must match a key in `_DIFFUSION_MODELS` in `registry.py`. Other fields are model-specific and accessible via `od_config.model_config` dict.
+
+## External Dependency Management
+
+### Git clone + .pth injection (DreamID-Omni pattern)
+
+```python
+def download_dependency():
+    CACHE_DIR.mkdir(parents=True, exist_ok=True)
+    with open(LOCK_FILE, "w") as f:
+        fcntl.flock(f, fcntl.LOCK_EX)
+        if not DEPENDENCY_DIR.exists():
+            subprocess.run([
+                "git", "clone", "--depth", "1",
+                REPO_URL, "--branch", BRANCH,
+                str(DEPENDENCY_DIR)
+            ], check=True)
+        fcntl.flock(f, fcntl.LOCK_UN)
+
+    # Add to Python path via .pth file
+    site_packages = Path(site.getsitepackages()[0])
+    pth_file = site_packages / "vllm_omni_dependency.pth"
+    pth_file.write_text(str(DEPENDENCY_DIR))
+```
+
+### Direct port (BAGEL pattern)
+
+Copy essential files from the original repo into `vllm_omni/diffusion/models/<name>/`. Adapt imports to use vllm-omni utilities. Benefits: no external dependency, no git clone step. Drawback: must maintain the ported code.
+
+## Multi-Modal Input/Output Protocols
+
+Custom models that handle images, audio, or video I/O should implement protocol classes:
+
+```python
+from vllm_omni.diffusion.models.interface import (
+    SupportImageInput,    # Model accepts image input
+    SupportAudioInput,    # Model accepts audio input
+    SupportAudioOutput,   # Model produces audio output
+)
+
+class MyPipeline(nn.Module, SupportImageInput, SupportAudioInput, SupportAudioOutput):
+    pass  # Protocol markers enable proper engine routing
+```
+
+The engine checks `isinstance(pipeline, SupportImageInput)` at startup to configure input validation and warmup behavior.
+
+## Hardcoded Config vs Config Files
+
+Diffusers models use `config.json` in each subfolder. Custom models often use:
+
+**Module-level config dicts** (DreamID-Omni):
+```python
+VIDEO_CONFIG = {
+    "patch_size": [1, 2, 2], "model_type": "ti2v",
+    "dim": 3072, "ffn_dim": 14336, "num_heads": 24, "num_layers": 30, ...
+}
+```
+
+**Loaded from custom JSON** (BAGEL):
+```python
+cfg_path = os.path.join(model_path, "config.json")
+with open(cfg_path) as f:
+    bagel_cfg = json.load(f)
+vae_cfg = bagel_cfg.get("vae_config", {})
+```
+
+## Custom Architecture Patterns
+
+### Split forward API (DreamID-Omni)
+
+When a fusion model needs to interleave blocks from two backbones:
+
+```python
+class WanModel(nn.Module):
+    def prepare_transformer_block_kwargs(self, x, t, context, ...):
+        # Patch embed, time embed, text embed, RoPE
+        return x, e, kwargs
+
+    def post_transformer_block_out(self, x, grid_sizes, e):
+        # Output projection, unpatchify
+        return output
+
+    def forward(self, *args, **kwargs):
+        raise NotImplementedError  # Fusion model handles block iteration
+```
+
+The `FusionModel` then iterates blocks in lock-step:
+```python
+for video_block, audio_block in zip(self.video_model.blocks, self.audio_model.blocks):
+    video_out = video_block(video_hidden, ...)
+    audio_out = audio_block(audio_hidden, ...)
+    # Cross-attend between modalities
+    video_out = cross_attention(video_out, audio_out)
+    audio_out = cross_attention(audio_out, video_out)
+```
+
+### LLM-as-denoiser (BAGEL)
+
+When the backbone is a language model that also does diffusion:
+
+```python
+class BagelModel(nn.Module):
+    def __init__(self):
+        self.language_model = Qwen2MoTForCausalLM(config)
+        self.vit_model = SiglipVisionModel(vit_config)
+```
+
+The LLM processes both text tokens and latent image tokens in a single forward pass, using KV caching for the text portion.
+
+## Pre/Post Processing for Custom Models
+
+Custom models typically handle pre/post processing **inside `forward()`** rather than via registered functions, because the logic is tightly coupled:
+
+```python
+def forward(self, req: OmniDiffusionRequest) -> DiffusionOutput:
+    # Inline preprocessing
+    image = self._load_and_resize_image(req.prompts[0].get("multi_modal_data", {}).get("image"))
+    image_latent = self._vae_encode(image)
+
+    # ... denoising loop ...
+
+    # Inline postprocessing
+    pil_image = self._decode_to_pil(latents)
+    return DiffusionOutput(output=[pil_image])
+```
+
+If pre/post functions are not registered in `_DIFFUSION_PRE_PROCESS_FUNCS` / `_DIFFUSION_POST_PROCESS_FUNCS`, the engine simply skips those steps.
+
+## Download Script Template
+
+```python
+# examples/offline_inference/<name>/download_<name>.py
+from huggingface_hub import snapshot_download
+import json, os
+
+def main(output_dir):
+    # Download model weights from HF
+    snapshot_download(repo_id="org/model-weights", local_dir=os.path.join(output_dir, "weights"))
+
+    # Download additional components if from separate repos
+    snapshot_download(repo_id="org/vae-weights", local_dir=os.path.join(output_dir, "vae"),
+        allow_patterns=["*.safetensors"])
+
+    # Generate model_index.json
+    config = {"_class_name": "YourPipeline", "custom_key": "weights/model.safetensors"}
+    with open(os.path.join(output_dir, "model_index.json"), "w") as f:
+        json.dump(config, f, indent=2)
+
+    # Install external code dependency (if needed)
+    download_dependency()
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--output-dir", default="./your_model")
+    args = parser.parse_args()
+    main(args.output_dir)
+```

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

@@ -0,0 +1,572 @@
+# Parallelism Patterns Reference
+
+## Overview
+
+vLLM-Omni supports multiple parallelism strategies for diffusion models. Each targets a different bottleneck:
+
+| Strategy | Splits | Best For | Constraint |
+|----------|--------|----------|------------|
+| Tensor Parallel (TP) | Model layers across GPUs | Latency reduction, large models | Requires fast GPU interconnect, `num_heads % tp == 0` |
+| Sequence Parallel (SP/USP) | Sequence tokens across GPUs | Long sequences (video, high-res) | Near-linear scaling |
+| CFG Parallel | Positive/negative CFG branches | Models using classifier-free guidance | Exactly 2 GPUs |
+| HSDP | Weight shards via FSDP2 | VRAM reduction | Cannot combine with TP |
+| VAE Patch Parallel | VAE decode spatial tiles | Large VAE outputs | Auto-enables tiling |
+
+**Recommended integration order**: TP → SP → CFG Parallel → HSDP
+
+**Official design docs**:
+- TP: https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/tensor_parallel
+- SP: https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/sequence_parallel
+- CFG: https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/cfg_parallel
+- HSDP: https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/hsdp
+
+---
+
+## Tensor Parallelism (TP)
+
+Replace standard `nn.Linear` with vLLM's parallel linear layers. This is the most invasive change but provides direct VRAM savings and compute speedup.
+
+### Layer replacement rules
+
+| Pattern | vLLM Layer | When to Use |
+|---------|-----------|-------------|
+| Fan-out (first in FFN) | `ColumnParallelLinear` | Projection that splits output across ranks |
+| Fan-in (second in FFN) | `RowParallelLinear` | Projection that gathers across ranks |
+| QKV projection | `QKVParallelLinear` | Fused Q/K/V for self-attention |
+| Single Q or K or V | `ColumnParallelLinear` | Separate projections (cross-attention) |
+| Attention output | `RowParallelLinear` | Output projection after attention |
+| Must not shard | `ReplicatedLinear` | Layers that must stay replicated |
+
+### MLP Block (Up-Down Pattern)
+
+```python
+from vllm.model_executor.layers.linear import (
+    ColumnParallelLinear, RowParallelLinear,
+)
+
+class TPFeedForward(nn.Module):
+    def __init__(self, dim, ffn_dim):
+        super().__init__()
+        self.fc1 = ColumnParallelLinear(dim, ffn_dim, bias=False, return_bias=False)
+        self.fc2 = RowParallelLinear(
+            ffn_dim, dim, bias=False,
+            input_is_parallel=True,  # Input already sharded from fc1
+            return_bias=False,
+        )
+
+    def forward(self, x):
+        x, _ = self.fc1(x)
+        x = torch.nn.functional.gelu(x)
+        x, _ = self.fc2(x)
+        return x
+```
+
+### Attention Block (QKV-Out Pattern)
+
+```python
+from vllm.model_executor.layers.linear import QKVParallelLinear, RowParallelLinear
+from vllm_omni.diffusion.attention.layer import Attention
+
+class TPSelfAttention(nn.Module):
+    def __init__(self, dim, num_heads, num_kv_heads=None):
+        super().__init__()
+        num_kv_heads = num_kv_heads or num_heads
+        self.head_dim = dim // num_heads
+
+        self.to_qkv = QKVParallelLinear(
+            hidden_size=dim,
+            head_size=self.head_dim,
+            total_num_heads=num_heads,
+            total_num_kv_heads=num_kv_heads,
+            bias=False,
+            return_bias=False,
+        )
+        self.to_out = RowParallelLinear(
+            dim, dim, bias=False,
+            input_is_parallel=True,
+            return_bias=False,
+        )
+        self.attn = Attention(
+            num_heads=self.to_qkv.num_heads,      # Local heads per GPU
+            head_size=self.head_dim,
+            softmax_scale=1.0 / (self.head_dim ** 0.5),
+            causal=False,
+            num_kv_heads=self.to_qkv.num_kv_heads,  # Local KV heads per GPU
+            role="self",
+        )
+
+    def forward(self, x):
+        qkv, _ = self.to_qkv(x)
+        q, k, v = qkv.split(
+            [self.to_qkv.num_heads * self.head_dim,
+             self.to_qkv.num_kv_heads * self.head_dim,
+             self.to_qkv.num_kv_heads * self.head_dim],
+            dim=-1,
+        )
+        B, S, _ = x.shape
+        q = q.view(B, S, self.to_qkv.num_heads, self.head_dim)
+        k = k.view(B, S, self.to_qkv.num_kv_heads, self.head_dim)
+        v = v.view(B, S, self.to_qkv.num_kv_heads, self.head_dim)
+        out = self.attn(q, k, v)
+        out = out.reshape(B, S, -1)
+        out, _ = self.to_out(out)
+        return out
+```
+
+### QKV Fusion in load_weights
+
+When you fuse separate Q/K/V into `QKVParallelLinear`, map diffusers' separate weight names:
+
+```python
+stacked_params_mapping = [
+    ("to_qkv", "to_q", "q"),
+    ("to_qkv", "to_k", "k"),
+    ("to_qkv", "to_v", "v"),
+]
+
+def load_weights(self, weights):
+    params = dict(self.named_parameters())
+    loaded = set()
+    for name, tensor in weights:
+        for fused_name, orig_name, shard_id in stacked_params_mapping:
+            if orig_name in name:
+                name = name.replace(orig_name, fused_name)
+                param = params[name]
+                param.weight_loader(param, tensor, shard_id)
+                loaded.add(name)
+                break
+        else:
+            if name in params:
+                param = params[name]
+                if hasattr(param, "weight_loader"):
+                    param.weight_loader(param, tensor)
+                else:
+                    default_weight_loader(param, tensor)
+                loaded.add(name)
+    return loaded
+```
+
+### RMSNorm with TP
+
+When RMSNorm sits between TP-sharded dimensions, use `DistributedRMSNorm` — it computes global RMS via all-reduce across TP ranks. See the Wan2.2 implementation for the pattern.
+
+### TP Constraints
+
+- `num_heads % tp_size == 0`
+- `num_kv_heads % tp_size == 0`
+- Use `self.to_qkv.num_heads` (local per-GPU count), not total heads, for split sizes
+
+### Testing TP
+
+```bash
+python text_to_image.py --model Your-org/your-model \
+  --tensor-parallel-size 2 --output "tp_test.png"
+```
+
+**Verify**: speedup, memory reduction proportional to TP size, quality matches single-GPU.
+
+### Reference implementations
+
+| Model | Path |
+|-------|------|
+| Z-Image | `vllm_omni/diffusion/models/z_image/z_image_transformer.py` |
+| FLUX | `vllm_omni/diffusion/models/flux/flux_transformer.py` |
+| Qwen-Image | `vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py` |
+
+---
+
+## Sequence Parallelism (SP / USP)
+
+SP splits sequence tokens across GPUs using Ulysses (all-to-all) or Ring (P2P) communication. It is applied non-intrusively via the `_sp_plan` dict — no changes to `forward()` logic.
+
+### Approach 1: Non-Intrusive `_sp_plan` (Recommended)
+
+The framework automatically registers hooks to shard inputs and gather outputs at `nn.Module` boundaries.
+
+#### Step 1: Identify module boundaries
+
+Find where tensors need sharding/gathering:
+
+```python
+class MyTransformer(nn.Module):
+    def __init__(self):
+        self.patch_embed = PatchEmbed()    # Before blocks
+        self.pos_embed = RoPE()            # RoPE may need splitting
+        self.blocks = nn.ModuleList([...]) # Blocks process sharded x
+        self.norm_out = LayerNorm()
+        self.proj_out = Linear()           # Gather after this
+
+    def forward(self, x):
+        x = self.patch_embed(x)
+        pos = self.pos_embed(x)
+        for block in self.blocks:
+            x = block(x, pos)
+        x = self.norm_out(x)
+        return self.proj_out(x)
+```
+
+#### Step 2: Handle inline operations
+
+`_sp_plan` hooks only work at `nn.Module` boundaries. Inline ops like `torch.cat()` must be extracted into submodules:
+
+```python
+# BAD: Inline — hooks can't intercept
+unified = torch.cat([x, cap_feats], dim=1)
+
+# GOOD: Extract into submodule
+class UnifiedPrepare(nn.Module):
+    def forward(self, x, cap_feats):
+        return torch.cat([x, cap_feats], dim=1)
+
+self.unified_prepare = UnifiedPrepare()
+unified = self.unified_prepare(x, cap_feats)
+```
+
+Common cases: `torch.cat()`, `pad_sequence()`, `tensor.reshape()`, complex preprocessing.
+
+#### Step 3: Write `_sp_plan`
+
+**Pattern 1: Shard at first block, gather at output** (most common)
+
+```python
+from vllm_omni.diffusion.distributed.sp_plan import (
+    SequenceParallelInput, SequenceParallelOutput,
+)
+
+class StandardTransformer(nn.Module):
+    _sp_plan = {
+        "blocks.0": {
+            "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3),
+        },
+        "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3),
+    }
+```
+
+**Pattern 2: Shard RoPE outputs separately**
+
+```python
+class TransformerWithRoPE(nn.Module):
+    _sp_plan = {
+        "rope": {
+            0: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True),
+            1: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True),
+        },
+        "blocks.0": {
+            "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3),
+        },
+        "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3),
+    }
+```
+
+**Pattern 3: Dual-stream (shard image, replicate text)**
+
+```python
+class DualStreamTransformer(nn.Module):
+    _sp_plan = {
+        "rope_preparer": {
+            2: SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True),
+            3: SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True),
+        },
+        "transformer_blocks.0": {
+            "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3),
+        },
+        "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3),
+    }
+```
+
+### API Reference
+
+**SequenceParallelInput**:
+
+| Parameter | Type | Description |
+|-----------|------|-------------|
+| `split_dim` | int | Dimension to split (usually 1 for sequence) |
+| `expected_dims` | int/None | Expected tensor rank for validation |
+| `split_output` | bool | `False`: shard input params; `True`: shard output tensors |
+| `auto_pad` | bool | Auto-pad if sequence not divisible by world_size |
+
+**SequenceParallelOutput**:
+
+| Parameter | Type | Description |
+|-----------|------|-------------|
+| `gather_dim` | int | Dimension to gather (usually 1 for sequence) |
+| `expected_dims` | int/None | Expected tensor rank for validation |
+
+**Module naming**:
+
+| Key | Meaning |
+|-----|---------|
+| `"blocks.0"` | First element of ModuleList |
+| `"blocks.*"` | All elements of ModuleList |
+| `"rope"` | Named submodule |
+
+**Dictionary value types**:
+
+| Key type | split_output | Description |
+|----------|-------------|-------------|
+| `"param_name"` (str) | False | Shard input parameter by name |
+| `0, 1, ...` (int) | True | Shard output tuple by index |
+
+### Approach 2: Intrusive Modification (Complex Cases)
+
+For dynamic sharding logic that can't be expressed via `_sp_plan`:
+
+```python
+from vllm_omni.diffusion.distributed.sp_sharding import sp_shard, sp_gather
+
+def forward(self, hidden_states, ...):
+    if self.parallel_config.sequence_parallel_size > 1:
+        hidden_states = sp_shard(hidden_states, dim=1)
+    for block in self.blocks:
+        hidden_states = block(hidden_states)
+    if self.parallel_config.sequence_parallel_size > 1:
+        hidden_states = sp_gather(hidden_states, dim=1)
+    return hidden_states
+```
+
+Use intrusive modification as a last resort — `_sp_plan` is preferred for maintainability.
+
+### UAA Mode (Experimental)
+
+`ulysses_mode="advanced_uaa"` handles arbitrary sequence lengths and head counts that aren't divisible by `ulysses_degree`. Uses variable all-to-all split sizes and temporary head padding.
+
+### Combining SP methods
+
+Ulysses and Ring can be combined: `ulysses_degree × ring_degree = total SP GPUs`.
+
+```python
+DiffusionParallelConfig(ulysses_degree=2, ring_degree=2)  # 4 GPUs total
+```
+
+### Testing SP
+
+```bash
+# Offline
+python text_to_image.py --model Your-model --ulysses-degree 2
+
+# Online serving
+vllm serve Your-model --omni --usp 2
+```
+
+### Reference implementations
+
+| Model | Path |
+|-------|------|
+| Qwen-Image | `vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py` |
+| Wan2.2 | `vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py` |
+| Z-Image | `vllm_omni/diffusion/models/z_image/z_image_transformer.py` |
+
+---
+
+## CFG Parallelism
+
+Distributes positive/negative Classifier-Free Guidance branches across 2 GPUs.
+
+### Implementation
+
+Inherit `CFGParallelMixin` and implement `diffuse()`:
+
+```python
+from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
+
+class YourPipeline(nn.Module, CFGParallelMixin):
+    def diffuse(self, latents, timesteps, prompt_embeds, negative_embeds,
+                do_true_cfg, true_cfg_scale, **kwargs):
+        for i, t in enumerate(timesteps):
+            positive_kwargs = {
+                "hidden_states": latents,
+                "encoder_hidden_states": prompt_embeds,
+                "timestep": t,
+            }
+            negative_kwargs = {
+                "hidden_states": latents,
+                "encoder_hidden_states": negative_embeds,
+                "timestep": t,
+            } if do_true_cfg else None
+
+            noise_pred = self.predict_noise_maybe_with_cfg(
+                do_true_cfg=do_true_cfg,
+                true_cfg_scale=true_cfg_scale,
+                positive_kwargs=positive_kwargs,
+                negative_kwargs=negative_kwargs,
+            )
+            latents = self.scheduler_step_maybe_with_cfg(
+                noise_pred, t, latents, do_true_cfg
+            )
+        return latents
+```
+
+### Customization hooks
+
+| Method | Override when |
+|--------|-------------|
+| `predict_noise()` | Non-standard transformer call (e.g., dual-transformer like Wan2.2) |
+| `cfg_normalize_function()` | Custom normalization (e.g., LongCat with clamping) |
+| `combine_cfg_noise()` | Multi-output models (e.g., video + audio: CFG on video, positive-only on audio) |
+
+**Custom predict_noise** (Wan2.2 — selects active transformer):
+
+```python
+def predict_noise(self, current_model=None, **kwargs):
+    if current_model is None:
+        current_model = self.transformer
+    return current_model(**kwargs)[0]
+```
+
+**Custom combine_cfg_noise** (multi-output):
+
+```python
+def combine_cfg_noise(self, positive_pred, negative_pred, scale, normalize):
+    video_pos, audio_pos = positive_pred
+    video_neg, audio_neg = negative_pred
+    video_combined = super().combine_cfg_noise(video_pos, video_neg, scale, normalize)
+    return (video_combined, audio_pos)
+```
+
+### Composite scheduler for multi-output
+
+When each output has its own schedule:
+
+```python
+class VideoAudioScheduler:
+    def __init__(self, video_scheduler, audio_scheduler):
+        self.video_scheduler = video_scheduler
+        self.audio_scheduler = audio_scheduler
+
+    def step(self, noise_pred, t, latents, return_dict=False, generator=None):
+        video_out = self.video_scheduler.step(
+            noise_pred[0], t[0], latents[0], return_dict=False, generator=generator
+        )[0]
+        audio_out = self.audio_scheduler.step(
+            noise_pred[1], t[1], latents[1], return_dict=False, generator=generator
+        )[0]
+        return ((video_out, audio_out),)
+```
+
+### Testing CFG Parallel
+
+```bash
+python text_to_image.py --model Your-model \
+  --cfg-parallel-size 2 --cfg-scale 4.0 \
+  --negative-prompt "ugly, unclear"
+```
+
+**Constraint**: `guidance_scale > 1.0` and negative prompt must be provided.
+
+### Reference implementations
+
+| Model | Path |
+|-------|------|
+| Qwen-Image | `vllm_omni/diffusion/models/qwen_image/cfg_parallel.py` |
+| Wan2.2 | `vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py` |
+| Mixin base | `vllm_omni/diffusion/distributed/cfg_parallel.py` |
+
+---
+
+## HSDP (Hybrid Sharded Data Parallel)
+
+Shards model weights across GPUs using PyTorch FSDP2. Reduces per-GPU VRAM without changing computation.
+
+### Implementation
+
+Add `_hsdp_shard_conditions` to the transformer class:
+
+```python
+class YourTransformer(nn.Module):
+    @staticmethod
+    def _is_transformer_block(name: str, module) -> bool:
+        return "blocks" in name and name.split(".")[-1].isdigit()
+
+    _hsdp_shard_conditions = [_is_transformer_block]
+```
+
+For MoE models, add additional conditions:
+
+```python
+class MoETransformer(nn.Module):
+    @staticmethod
+    def _is_transformer_block(name, module):
+        return "blocks" in name and name.split(".")[-1].isdigit()
+
+    @staticmethod
+    def _is_moe_expert(name, module):
+        return "experts" in name and name.split(".")[-1].isdigit()
+
+    _hsdp_shard_conditions = [_is_transformer_block, _is_moe_expert]
+```
+
+A module is sharded if **any** condition returns `True`.
+
+### Constraints
+
+- Cannot combine with Tensor Parallelism
+- For standalone HSDP (no other parallelism), `hsdp_shard_size` must be specified explicitly
+- Can combine with SP: HSDP reduces memory while SP distributes sequence
+
+### Testing HSDP
+
+```python
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+
+parallel_config = DiffusionParallelConfig(use_hsdp=True, hsdp_shard_size=8)
+omni = Omni(model="your-model", parallel_config=parallel_config)
+```
+
+Or CLI:
+
+```bash
+vllm serve Your-model --omni --use-hsdp
+```
+
+**Verify**: logs show "HSDP Inference: replicate_size=..., shard_size=..." and "Sharded N modules + root". Check VRAM reduction.
+
+### Reference implementations
+
+| Model | Path |
+|-------|------|
+| Wan2.2 | `vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py` |
+| HSDP Core | `vllm_omni/diffusion/distributed/hsdp.py` |
+
+---
+
+## VAE Patch Parallelism
+
+Shards VAE decode spatially across ranks using tiling:
+
+```bash
+python text_to_image.py --model Your-model --vae-patch-parallel-size 4
+```
+
+Auto-enables `--vae-use-tiling`. Uses `DistributedAutoencoderKLWan` or similar distributed VAE. Set `vae_patch_parallel_size` in `DiffusionParallelConfig`.
+
+---
+
+## Combining Parallelism Methods
+
+Common multi-GPU recipes:
+
+```bash
+# 4 GPUs: CFG (2) × Ulysses (2)
+python text_to_image.py --model Qwen/Qwen-Image \
+  --cfg-parallel-size 2 --ulysses-degree 2
+
+# 8 GPUs: Ulysses (4) × Ring (2) + VAE patch (8)
+python text_to_video.py --model Wan-AI/Wan2.2-T2V-A14B-Diffusers \
+  --ulysses-degree 4 --ring-degree 2 --vae-patch-parallel-size 8
+
+# 2 GPUs: HSDP + Ulysses (cannot combine HSDP with TP)
+vllm serve Your-model --omni --use-hsdp --usp 2
+```
+
+## Discovering Parallelism Support
+
+Check which parallelism methods a model supports:
+
+| Check | How |
+|-------|-----|
+| **Ulysses / Ring SP** | Transformer defines `_sp_plan`. Search: `grep -r '_sp_plan' vllm_omni/diffusion/models/` |
+| **CFG Parallel** | Pipeline inherits `CFGParallelMixin`. Search: `grep -r 'CFGParallelMixin' vllm_omni/diffusion/models/` |
+| **TP** | Uses `ColumnParallelLinear` / `QKVParallelLinear`. Search: `grep -r 'ParallelLinear\|QKVParallel' vllm_omni/diffusion/models/<model>/` |
+| **HSDP** | Transformer defines `_hsdp_shard_conditions`. Search: `grep -r '_hsdp_shard_conditions' vllm_omni/diffusion/models/` |
+
+The canonical per-model support table is in `docs/user_guide/diffusion/parallelism_acceleration.md`.

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

@@ -0,0 +1,238 @@
+# Transformer Adaptation Reference
+
+## Adapting a Diffusers Transformer to vLLM-Omni
+
+### Step-by-step Checklist
+
+1. Copy the transformer class from diffusers source
+2. Remove all mixin classes — inherit only from `nn.Module`
+3. Replace attention dispatch with `vllm_omni.diffusion.attention.layer.Attention`
+4. Replace logger with `vllm.logger.init_logger`
+5. Add `od_config: OmniDiffusionConfig | None = None` to `__init__`
+6. Remove training-only code (gradient checkpointing, dropout)
+7. Add `load_weights()` method for weight loading from safetensors
+8. Add class-level attributes for acceleration features
+
+### Mixin Removal
+
+Remove these diffusers mixins (and their imports):
+
+```python
+# Remove all of these:
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.attention_processor import AttentionModuleMixin
+from diffusers.loaders import PeftAdapterMixin, FromOriginalModelMixin
+
+# Replace:
+class MyTransformer(ModelMixin, ConfigMixin, AttentionModuleMixin):
+# With:
+class MyTransformer(nn.Module):
+```
+
+Also remove `@register_to_config` decorators from `__init__`.
+
+### Attention Replacement
+
+The vLLM-Omni `Attention` layer wraps backend selection (FlashAttention, SDPA, SageAttn, etc.) and supports sequence parallelism hooks.
+
+**QKV tensor shape must be `[batch, seq_len, num_heads, head_dim]`.**
+
+#### Self-Attention Pattern
+
+```python
+from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+
+class SelfAttentionBlock(nn.Module):
+    def __init__(self, dim, num_heads):
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+
+        self.to_q = nn.Linear(dim, dim)
+        self.to_k = nn.Linear(dim, dim)
+        self.to_v = nn.Linear(dim, dim)
+        self.to_out = nn.Linear(dim, dim)
+
+        self.attn = Attention(
+            num_heads=num_heads,
+            head_size=self.head_dim,
+            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):
+        B, S, _ = x.shape
+        q = self.to_q(x).view(B, S, self.num_heads, self.head_dim)
+        k = self.to_k(x).view(B, S, self.num_heads, self.head_dim)
+        v = self.to_v(x).view(B, S, self.num_heads, self.head_dim)
+
+        attn_metadata = AttentionMetadata(attn_mask=attn_mask)
+        out = self.attn(q, k, v, attn_metadata=attn_metadata)
+        out = out.reshape(B, S, -1)
+        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:
+
+```python
+from vllm.model_executor.layers.linear import (
+    QKVParallelLinear, RowParallelLinear
+)
+
+class TPSelfAttention(nn.Module):
+    def __init__(self, dim, num_heads):
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+
+        self.to_qkv = QKVParallelLinear(
+            hidden_size=dim,
+            head_size=self.head_dim,
+            total_num_heads=num_heads,
+            total_num_kv_heads=num_heads,
+        )
+        self.to_out = RowParallelLinear(dim, dim)
+
+        self.attn = Attention(
+            num_heads=num_heads,
+            head_size=self.head_dim,
+            softmax_scale=1.0 / (self.head_dim ** 0.5),
+            causal=False,
+            num_kv_heads=num_heads,
+            role="self",
+        )
+```
+
+### Logger Replacement
+
+```python
+# Replace:
+from diffusers.utils import logging
+logger = logging.get_logger(__name__)
+
+# With:
+from vllm.logger import init_logger
+logger = init_logger(__name__)
+```
+
+### Custom Layers from vLLM-Omni
+
+Available utility layers:
+
+```python
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm_omni.diffusion.layers.rope import RotaryEmbedding
+from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNorm
+```
+
+### Config Support
+
+```python
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+class MyTransformer(nn.Module):
+    def __init__(self, *, od_config=None, num_layers=28, hidden_size=3072, **kwargs):
+        super().__init__()
+        self.od_config = od_config
+        self.parallel_config = od_config.parallel_config if od_config else None
+        # ... build layers
+```
+
+The transformer config values come from `model_index.json` → `config.json` in the transformer subfolder. The pipeline uses `get_transformer_config_kwargs(od_config.tf_model_config, TransformerClass)` to filter config keys to match the `__init__` signature.
+
+### Weight Loading
+
+The `load_weights` method receives an iterable of `(name, tensor)` from safetensors files, with the prefix (e.g., `"transformer."`) already stripped by the loader.
+
+```python
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+
+class MyTransformer(nn.Module):
+    def load_weights(self, weights):
+        params = dict(self.named_parameters())
+        loaded = set()
+        for name, tensor in weights:
+            # Optional: remap names from diffusers to vllm-omni naming
+            # e.g., "ff.net.0.proj" -> "ff.net_0.proj"
+
+            if name in params:
+                param = params[name]
+                if hasattr(param, "weight_loader"):
+                    param.weight_loader(param, tensor)
+                else:
+                    default_weight_loader(param, tensor)
+                loaded.add(name)
+        return loaded
+```
+
+#### QKV Fusion in load_weights
+
+If you fused separate Q/K/V into a `QKVParallelLinear`, you need to map diffusers' separate weight names:
+
+```python
+stacked_params_mapping = [
+    ("to_qkv", "to_q", "q"),
+    ("to_qkv", "to_k", "k"),
+    ("to_qkv", "to_v", "v"),
+]
+
+def load_weights(self, weights):
+    params = dict(self.named_parameters())
+    loaded = set()
+    for name, tensor in weights:
+        for fused_name, orig_name, shard_id in stacked_params_mapping:
+            if orig_name in name:
+                name = name.replace(orig_name, fused_name)
+                param = params[name]
+                param.weight_loader(param, tensor, shard_id)
+                loaded.add(name)
+                break
+        else:
+            # Normal loading
+            ...
+    return loaded
+```
+
+### Class-Level Attributes for Features
+
+```python
+class MyTransformer(nn.Module):
+    # torch.compile: list block class names that repeat and can be compiled
+    _repeated_blocks = ["MyTransformerBlock"]
+
+    # CPU offload: attribute name of the nn.ModuleList containing blocks
+    _layerwise_offload_blocks_attr = "blocks"
+
+    # LoRA: mapping of fused param names to original param names
+    packed_modules_mapping = {"to_qkv": ["to_q", "to_k", "to_v"]}
+
+    # Sequence parallelism plan (advanced — add after basic impl works)
+    _sp_plan = {
+        "blocks.0": SequenceParallelInput(split_dim=1),
+        "proj_out": SequenceParallelOutput(gather_dim=1),
+    }
+```

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

@@ -0,0 +1,178 @@
+# Troubleshooting Reference
+
+## Common Errors When Adding a Diffusion Model
+
+### ImportError / ModuleNotFoundError
+
+**Cause**: Missing or incorrect registration.
+
+**Fix checklist**:
+1. Model registered in `vllm_omni/diffusion/registry.py` `_DIFFUSION_MODELS` dict
+2. `__init__.py` exports the pipeline class
+3. Pipeline file exists at the correct path: `vllm_omni/diffusion/models/{folder}/{file}.py`
+4. Class name in registry matches the actual class name in the file
+
+### Shape Mismatch in Attention
+
+**Symptom**: `RuntimeError: shape mismatch` or `expected 4D tensor`
+
+**Cause**: QKV tensors not reshaped to `[batch, seq_len, num_heads, head_dim]`.
+
+**Fix**: Before calling `self.attn(q, k, v, ...)`, ensure:
+```python
+q = q.view(batch, seq_len, self.num_heads, self.head_dim)
+k = k.view(batch, kv_seq_len, self.num_kv_heads, self.head_dim)
+v = v.view(batch, kv_seq_len, self.num_kv_heads, self.head_dim)
+```
+
+After attention, reshape back:
+```python
+out = out.reshape(batch, seq_len, -1)
+```
+
+### Weight Loading Failures
+
+**Symptom**: `RuntimeError: size mismatch for parameter ...` or missing keys
+
+**Debugging**:
+1. Print diffusers weight names: `safetensors.safe_open(path, "pt").keys()`
+2. Print model parameter names: `dict(model.named_parameters()).keys()`
+3. Compare and add name remappings in `load_weights()`
+
+**Common remappings needed**:
+- `ff.net.0.proj` → `ff.net_0.proj` (PyTorch Sequential indexing)
+- `.to_out.0.` → `.to_out.` (Sequential unwrapping)
+- `scale_shift_table` → moved to a wrapper module
+
+### Black/Blank/Noisy Output
+
+**Possible causes**:
+1. **Wrong latent normalization**: Check VAE expects latents scaled by `vae.config.scaling_factor`
+2. **Wrong scheduler**: Using the wrong scheduler class or wrong `flow_shift`
+3. **Missing CFG**: Some models require `guidance_scale > 1.0` with negative prompt
+4. **Wrong timestep format**: Some schedulers expect float, others expect int/long
+5. **Missing post-processing**: Raw VAE output may need denormalization
+
+**Quick test**: Run with diffusers directly using the same seed and compare latents at each step.
+
+### OOM (Out of Memory)
+
+**Solutions** (in order of preference):
+1. `--enforce-eager` to disable torch.compile (saves compile memory)
+2. `--enable-cpu-offload` for model-level offload
+3. `--enable-layerwise-offload` for block-level offload (better for large models)
+4. `--vae-use-slicing --vae-use-tiling` for VAE memory reduction
+5. Reduce resolution: `--height 480 --width 832`
+6. Use TP: `--tensor-parallel-size 2`
+
+### 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`
+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
+
+### Tensor Parallel Errors
+
+**Symptom**: `AssertionError: not divisible` or incorrect output with TP>1
+
+**Fix**:
+1. Verify `num_heads % tp_size == 0` and `num_kv_heads % tp_size == 0`
+2. Ensure `ColumnParallelLinear` / `RowParallelLinear` are used correctly
+3. Check that norms between parallel layers use distributed norm if needed
+4. Verify `load_weights` handles TP sharding for norm weights
+5. Use `self.to_qkv.num_heads` (local heads per GPU) for QKV split sizes, not total heads
+
+**Missing `input_is_parallel=True`**:
+
+`RowParallelLinear` expects sharded input from `ColumnParallelLinear`:
+```python
+self.w1 = ColumnParallelLinear(dim, hidden_dim, return_bias=False)
+self.w2 = RowParallelLinear(hidden_dim, dim, input_is_parallel=True, return_bias=False)
+```
+
+### Sequence Parallel Errors
+
+**Symptom**: Incorrect output or crashes with `--ulysses-degree N` or `--usp N`
+
+**Possible causes**:
+1. **Inline operations between shard/gather points**: `torch.cat()`, `pad_sequence()` etc. not at `nn.Module` boundaries. Fix: extract into submodule.
+2. **Wrong `split_dim`**: Check the tensor shape at the shard point. Sequence dimension is typically `dim=1` for `[B, S, D]` tensors.
+3. **RoPE not sharded**: If RoPE is computed separately, add it to `_sp_plan` with `split_output=True`.
+4. **Sequence not divisible by SP degree**: Use `auto_pad=True` in `SequenceParallelInput` or switch to `ulysses_mode="advanced_uaa"`.
+
+**Debugging**: Add `expected_dims=N` to `SequenceParallelInput`/`Output` for shape validation at runtime.
+
+### CFG Parallel Errors
+
+**Symptom**: CFG parallel not activating, no speedup
+
+**Fix checklist**:
+1. Pipeline inherits `CFGParallelMixin`
+2. `guidance_scale > 1.0`
+3. Negative prompt provided (even if empty string)
+4. `--cfg-parallel-size 2` specified
+5. `diffuse()` method calls `predict_noise_maybe_with_cfg()` and `scheduler_step_maybe_with_cfg()`
+
+**Symptom**: Different output with CFG parallel vs sequential
+
+**Possible cause**: Non-deterministic scheduler. Fix: pass `generator=torch.Generator(device).manual_seed(seed)` to `scheduler_step_maybe_with_cfg()`.
+
+### HSDP Errors
+
+**Symptom**: HSDP not activating or errors during weight loading
+
+**Fix checklist**:
+1. Transformer defines `_hsdp_shard_conditions` class attribute
+2. Shard condition functions return `True` for correct modules (test with `model.named_modules()`)
+3. Not combining with TP (HSDP and TP are incompatible)
+4. For standalone HSDP, `hsdp_shard_size` is specified explicitly
+
+**Verify**: Check logs for "HSDP Inference: replicate_size=..., shard_size=..." and "Sharded N modules + root".
+
+### Cache-DiT Not Applied
+
+**Symptom**: No speedup, no cache-related log messages
+
+**Fix checklist**:
+1. Model not in `_NO_CACHE_ACCELERATION` in `registry.py`
+2. Pipeline class name matches `CUSTOM_DIT_ENABLERS` key (if using custom enabler)
+3. `cache_backend="cache_dit"` specified
+4. Check logs for "Cache-dit enabled successfully on xxx"
+
+**Verify pipeline name**: `print(pipeline.__class__.__name__)` — must match registry key.
+
+### Cache-DiT Quality Degradation
+
+**Symptom**: Artifacts or lower quality with cache-dit
+
+**Fix**: Reduce aggressiveness:
+```python
+cache_config={
+    "residual_diff_threshold": 0.12,      # Lower from 0.24
+    "max_warmup_steps": 6,                # Increase from 4
+    "max_continuous_cached_steps": 2,      # Reduce if higher
+}
+```
+
+If quality is still poor, the model may need a custom enabler with per-block-list `ParamsModifier` tuning.
+
+### Model Not Detected / Wrong Pipeline Class
+
+**Symptom**: `ValueError: Model class ... not found in diffusion model registry`
+
+**Cause**: The model's `model_index.json` has a `_class_name` for the pipeline that doesn't match registry keys.
+
+**Fix**: The registry key must match the diffusers pipeline class name from `model_index.json`. If using a different name, map it in the registry:
+```python
+"DiffusersPipelineClassName": ("your_folder", "your_file", "YourVllmClassName"),
+```
+
+## Debugging Workflow
+
+1. **Add verbose logging**: Use `logger.info()` to print tensor shapes at each stage
+2. **Compare step-by-step**: Run diffusers and vllm-omni side by side, comparing tensors after each major operation
+3. **Use small configs**: Reduce `num_inference_steps=2`, small resolution for fast iteration
+4. **Test transformer isolation**: Feed the same input to both diffusers and vllm-omni transformers, compare outputs
+5. **Binary search for bugs**: Comment out blocks/layers to isolate where divergence starts

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

@@ -0,0 +1,504 @@
+---
+name: add-tts-model
+description: "Integrate a new text-to-speech model into vLLM-Omni from HuggingFace reference implementation through production-ready serving with streaming and CUDA graph acceleration. Use when adding a new TTS model, wiring stage separation for speech synthesis, enabling online voice generation serving, debugging TTS integration behavior, or building audio output pipelines."
+---
+
+# TTS Model Integration Workflow
+
+## Overview
+
+```
+HF Reference -> Stage Separation -> Online Serving -> Async Chunk -> CUDA Graph -> Pre-commit/DCO
+   (Phase 1)      (Phase 2)          (Phase 3)        (Phase 4)     (Phase 5)      (Phase 6)
+```
+
+Three architecture patterns are supported:
+
+- **Two-stage pipeline** (e.g. Qwen3-TTS, Fish Speech, CosyVoice3): AR
+  code-predictor → audio decoder, connected via async_chunk for low-latency
+  streaming. Use this for maximum performance.
+- **Single-stage AR via generator** (e.g. MOSS-TTS-Nano): entire model runs
+  inside one AR worker, streaming audio chunks through a per-request
+  `inference_stream()` generator. Use this when the upstream model bundles AR
+  + codec inseparably. See [references/single-stage-ar.md](references/single-stage-ar.md).
+- **Single-stage, vLLM-native base LM + side computation** (e.g. VoxCPM2):
+  the base language model runs under vLLM's PagedAttention as a normal AR
+  model; diffusion / VAE / side computations run outside vLLM and are
+  attached via the runner post-processing hook. This is a distinct pattern
+  from the generator approach above — do not confuse the two.
+
+The single-stage variants skip Phase 4 (async_chunk) but Phase 5 (CUDA graph)
+is still encouraged for the inner AR loop.
+
+## Cross-Cutting Invariants
+
+These rules apply to every TTS model regardless of architecture (AR vs AR+diffusion, single-stage vs two-stage, codec-based vs VAE-based). They surface repeatedly across PRs — check them at the end of every phase.
+
+### I1. Streaming output contract
+
+Pick exactly one per-step semantics for `forward()` and document it in the docstring:
+
+- **Delta**: yield only new audio samples produced this step. Preferred — linear cost, low memory.
+- **Cumulative**: re-decode from step 0 every call. O(N²); only acceptable if the codec has no streaming decode path.
+
+If you choose **delta**, verify the full emit→consolidate→consume chain:
+
+1. `forward()` returns `{"model_outputs": <new_chunk_only>, ...}`
+2. `_consolidate_multimodal_tensors()` in `vllm_omni/engine/output_processor.py` concatenates the audio key into one tensor at finish. If it skips the key (`continue`), offline consumers receive only the final chunk. See `output_processor.py` for the concrete list of handled modality keys.
+3. Streaming consumers (SSE, Gradio) receive per-step deltas; offline consumers (`engine.generate()`) receive a single concatenated tensor.
+
+Cumulative-vs-delta mismatch is the most common silent bug — offline RTF benchmarks pass, but users hear replays or truncation.
+
+### I2. Multimodal output consumer hygiene
+
+`outputs[0].outputs[0].multimodal_output[<key>]` can be any of `Tensor`, `list[Tensor]` (pre-consolidation snapshot), `np.ndarray`, or scalar. When writing tests, examples, and benchmarks:
+
+- **Never** use `dict.get("a") or dict.get("b")` on tensor values — Python evaluates the tensor's boolean, raising `RuntimeError: Boolean value of Tensor with more than one value is ambiguous`. Use explicit `if x is None` chains.
+- Always defensively handle the list form: `if isinstance(x, list): x = torch.cat([t.reshape(-1) for t in x], dim=0)`.
+- Assert `shape` / `dtype` / `duration` explicitly; do not rely on truthiness for presence checks.
+
+### I3. Hot-loop GPU discipline
+
+Inside any per-step model loop (AR decode, diffusion solver, CFM Euler, vocoder block loop):
+
+- No `tensor.item()`, `.cpu()`, or `.tolist()` — each triggers a GPU→CPU sync; at 10 steps × 60 frames × 4 ops that is 2400 syncs per request.
+- Prefer `dst.copy_(src)` over `dst.fill_(src.item())` when writing a scalar tensor into a buffer.
+- Prefer `torch.compile(Model.forward, fullgraph=False)` on the whole forward over per-submodule compile — fewer dispatch boundaries, larger fusion regions. Measure before choosing granularity.
+- No Python-side control flow that depends on tensor values; use `torch.where` / masking instead.
+
+Profile first, optimize second. See the profiling docs / project memory for the trace-analysis workflow.
+
+### I4. Validation pyramid
+
+Offline RTF alone is necessary but not sufficient. Every new TTS model must pass all three:
+
+| Layer | Catches | Tool |
+|-------|---------|------|
+| Offline RTF / duration check | Throughput regressions, missing audio, wrong sample rate | `end2end.py`, pytest e2e |
+| Browser streaming playback | Delta/cumulative bugs, chunk boundary glitches, TTFP regressions | Gradio demo over `/v1/audio/speech?stream=true` |
+| Concurrent requests | Per-request state leaks, codec window round-robin gaps | `max_num_seqs>1` smoke test with 4+ parallel prompts |
+
+Declaring a model "done" without all three has shipped regressions more than once.
+
+### I5. Per-request state is owned by the request, not the model
+
+If the model caches *anything* across `forward()` calls (streaming generators, codec buffers, sliding-window pads, CUDA graph state), key it by request ID:
+
+```python
+self._state: dict[str, YourState] = {}    # request_key → state
+# fetch: request_key = str(info.get("_omni_req_id", "0"))
+# free on finish: del self._state[request_key]
+```
+
+A shared buffer silently corrupts audio across concurrent requests — the symptom is crosstalk or truncation only under load.
+
+## Phase 1: HuggingFace Reference
+
+**Goal**: Understand the reference implementation and verify it produces correct audio.
+
+### Steps
+
+1. **Run the reference model** end-to-end using the official HuggingFace / GitHub code
+2. **Document the architecture**:
+   - What are the sub-models? (AR decoder, codec decoder, vocoder, etc.)
+   - What is the token vocabulary? (semantic codes, RVQ codebooks, special tokens)
+   - What is the output format? (sample rate, channels, codec type)
+3. **Capture reference outputs** for comparison during integration
+4. **Identify the config structure**: `config.json` fields, `model_type`, sub-model configs
+
+### Key Questions
+
+- How many codebooks? What are the codebook sizes?
+- What special tokens exist? (`<|voice|>`, `<|audio_start|>`, `<|im_end|>`, etc.)
+- What is the token-to-ID mapping for codec codes?
+- What is the hop length / frame rate of the codec?
+- Does the model support voice cloning? How? (reference audio encoding, speaker embeddings, etc.)
+
+### Deliverables
+
+- Working reference script that produces audio
+- Architecture diagram / notes
+- Token vocabulary mapping
+- Reference audio samples for regression testing
+
+## Phase 2: Stage Separation (Offline Inference)
+
+**Goal**: Split the model into vLLM-Omni stages and get offline inference working.
+
+### Steps
+
+1. **Register the model** in `vllm_omni/model_executor/models/registry.py`
+2. **Create config classes** (`configuration_<model>.py`) with `model_type` registration
+3. **Implement Stage 0** (AR model):
+   - Subclass appropriate base (e.g., wrap Qwen3 decoder layers)
+   - Implement `forward()` for autoregressive token generation
+   - Handle special token logic (start/stop tokens, codec token mapping)
+   - If dual-AR (like Fish Speech), implement Fast AR as a nested module
+4. **Implement Stage 1** (Decoder):
+   - Load codec weights (may need lazy loading from separate checkpoint)
+   - Implement `forward()`: codec codes -> audio waveform
+   - Return `OmniOutput` with `multimodal_outputs`
+5. **Create stage config YAML** defining both stages, memory allocation, and model paths
+6. **Create stage input processor** for prompt building
+7. **Write end2end.py** test script
+
+### Critical Parameters to Get Right
+
+| Parameter | Impact if Wrong |
+|-----------|----------------|
+| Hop length | Audio duration wrong, streaming noise |
+| Token ID mapping | Garbage codes -> noise output |
+| Codebook count/size | Shape mismatch crashes |
+| Stop token | Generation never stops or stops too early |
+| dtype / autocast | Numerical issues, silent quality degradation |
+| Repetition penalty | Must match reference (often 1.0 for TTS) |
+
+### Debugging Priority (from experience)
+
+When audio output is wrong, check in this order:
+
+1. **RoPE / attention**: Are position encodings correct? Is the attention mask right?
+2. **Normalization**: RMSNorm epsilon, layer norm placement (pre vs post)
+3. **Hop length**: Product of all upsample rates in the codec decoder
+4. **Token mapping**: Are codec IDs correctly offset from the vocabulary base?
+5. **Sampling parameters**: Temperature, top_k, top_p, repetition_penalty
+6. **Tensor layout**: Codebook-major vs frame-major ordering
+7. **dtype**: Float32 for codec decoders (autocast can corrupt audio)
+
+### Streaming Correctness Rules (single-stage and two-stage)
+
+These bugs appear in almost every new TTS PR. Check all before the first push. See also the cross-cutting invariants I1 (output contract) and I5 (per-request state) above — the rules below are the Phase 2-specific instances of those invariants:
+
+- **Accumulate codes across AR steps** — each `forward()` appends new codes; do not reset between steps or audio will be truncated (fish speech: `fix: accumulate audio_codes across steps`)
+- **Emit delta audio, not full waveform** — in streaming mode yield only the new chunk per step, not the re-decoded full waveform from step 0 (fish speech: `fix: emit delta audio not full waveform`)
+- **All return paths must emit `model_outputs`** — if any early-return branch skips setting `model_outputs`, the serving layer silently drops that step's audio (fish speech: `fix: ensure ALL return paths emit model_outputs`)
+- **Per-request state isolation** — for batched concurrent requests, key all state by request ID; a shared buffer corrupts audio across requests (fish speech: `fix: per-request vocode + delta emission`)
+- **Codec tensor device** — move codec codes to the codec decoder's device before calling decode; mismatches cause silent CPU fallback or crashes (fish speech: `fix: use model device for CUDA stream`)
+- **AR stage `max_num_seqs`** — set to at least 4 in production configs; for single-stage models this is the only stage. For two-stage models, Stage 0 (AR) needs `max_num_seqs ≥ 4` to pipeline concurrent requests; Stage 1 (codec decoder) typically uses `max_num_seqs: 1` intentionally. Default of 1 everywhere causes audio gaps under concurrency because the codec window round-robins across requests (RFC #2568)
+
+### Optional Dependency Handling
+
+Patch optional dependencies (`torchaudio` / `torchcodec` / `soundfile`) at
+the top of `load_weights()`, not at module import. Failures to do so cause
+cryptic errors only on environments missing the optional package — after
+the model is already deployed. See
+[references/optional-deps.md](references/optional-deps.md) for the full
+pattern, signature constraints, and MOSS-TTS-Nano reference.
+
+### Single-Stage AR Pattern (alternative to two-stage)
+
+When the upstream model cannot be cleanly split into an AR stage and a
+separate decoder, run the full pipeline inside a single AR worker and
+stream audio through a per-request `inference_stream()` generator keyed by
+`_omni_req_id`. Stage config must set `worker_type: ar`,
+`engine_output_type: audio`, `final_output: true`, `is_comprehension: true`,
+and `async_chunk: false` at the top level. Only extract params from
+`additional_information` that you actually forward, or pre-commit fails
+`ruff F841`.
+
+Full walkthrough with the complete `forward()` / `_create_stream_gen()`
+skeleton and stage-config fields:
+[references/single-stage-ar.md](references/single-stage-ar.md). For an
+in-tree reference, look for any single-stage AR model under
+`vllm_omni/model_executor/models/` — e.g. the MOSS-TTS-Nano integration when
+it lands.
+
+**VoxCPM2 is a different pattern** and should not reuse this skeleton — it
+runs the base LM under vLLM PagedAttention with external side-computation.
+See `plan/voxcpm2_native_ar_design.md`.
+
+### Deliverables
+
+- Model files in `vllm_omni/model_executor/models/<model_name>/`
+- Stage config YAML
+- Working `end2end.py` at `examples/offline_inference/text_to_speech/<model>/end2end.py`
+- New section in `examples/offline_inference/text_to_speech/README.md` (table row + per-model section). Do **not** create a top-level `examples/offline_inference/<model>/` dir or a per-model `README.md` inside `text_to_speech/<model>/` — the hub README is the documented surface and the mkdocs `generate_examples` hook only descends one level into `examples/<category>/`.
+
+## Phase 3: Online Serving
+
+**Goal**: Expose the model via `/v1/audio/speech` API endpoint.
+
+### Steps
+
+1. **Register in `serving_speech.py`** — add all 5 points in a **single commit**;
+   partial integration causes hard-to-debug failures. This file is modified by every
+   model PR and is the most common source of rebase conflicts — see conflict note below.
+
+   **Point 1** — stage constant (near the top, alongside the other `_*_TTS_MODEL_STAGES` sets):
+   ```python
+   _YOUR_MODEL_TTS_MODEL_STAGES = {"your_stage_key"}
+   ```
+
+   **Point 2** — union into `_TTS_MODEL_STAGES`:
+   ```python
+   _TTS_MODEL_STAGES: set[str] = (
+       ...
+       | _YOUR_MODEL_TTS_MODEL_STAGES
+   )
+   ```
+
+   **Point 3** — model type detection in `_detect_tts_model_type()`:
+   ```python
+   if model_stage in _YOUR_MODEL_TTS_MODEL_STAGES:
+       return "your_model"
+   ```
+
+   **Point 4** — validation dispatch in `_validate_tts_request()`:
+   ```python
+   if self._tts_model_type == "your_model":
+       return self._validate_your_model_request(request)
+   ```
+
+   **Point 5** — validation + parameter-builder methods:
+   ```python
+   def _validate_your_model_request(self, request) -> str | None:
+       if not request.input or not request.input.strip():
+           return "Input text cannot be empty"
+       return None
+
+   def _build_your_model_params(self, request) -> dict:
+       params = {"text": [request.input]}
+       if request.voice is not None:
+           params["voice"] = [request.voice]
+       return params
+   ```
+   Wire `_build_your_model_params` into `_create_tts_request()` alongside the other
+   model-specific param builders.
+
+   > **Two dispatch patterns coexist**: Fish Speech uses a `self._is_fish_speech` boolean
+   > instance attribute checked before `elif self._is_tts`, while all newer models
+   > (CosyVoice3, MOSS-TTS-Nano) use the `_tts_model_type` string returned by
+   > `_detect_tts_model_type()`. For new models, always use the `_tts_model_type` string
+   > pattern — do not add new `_is_*` flags.
+
+   > **Unused variable rule**: only extract fields in `_build_your_model_params` that
+   > are actually forwarded to the model. Unused extractions fail `ruff F841`.
+   > For voice-cloning fields (`ref_audio` → `prompt_audio_path`, `ref_text` →
+   > `prompt_text`), add them to the param builder and verify they reach the model call.
+
+   **Rebase conflict note**: when rebasing onto `main` after another model was merged,
+   `serving_speech.py` will conflict. Resolution: always keep *both* the upstream
+   model's additions and your own — never discard either side.
+
+2. **Handle model-specific parameters**:
+   - Voice cloning: `ref_audio` encoding and prompt injection
+   - `max_new_tokens` override in sampling params
+   - Model-specific default values
+3. **Create client scripts**: `speech_client.py`, `run_server.sh`
+4. **Test all response formats**: wav, mp3, flac, pcm
+5. **Add Gradio demo**: Interactive web UI with streaming support
+
+### Voice Cloning Pattern
+
+```python
+import base64
+from pathlib import Path
+
+def build_voice_clone_prompt(ref_audio_path: str, text: str, codec) -> list:
+    """Build prompt with reference audio for voice cloning in serving_speech.py."""
+    audio_bytes = Path(ref_audio_path).read_bytes()
+    codes = codec.encode(audio_bytes)  # Encode on CPU using model's codec (e.g., DAC)
+    token_ids = [code + codec.vocab_offset for code in codes.flatten().tolist()]
+    return [
+        {"role": "system", "content": f"<|voice|>{''.join(chr(t) for t in token_ids)}"},
+        {"role": "user", "content": text},
+    ]
+```
+
+### Deliverables
+
+- Updated `serving_speech.py` with all 5 integration points (single commit)
+- Client scripts and server launcher under `examples/online_serving/text_to_speech/<model>/`
+- Gradio demo with streaming and voice cloning UI in the same dir
+- E2E online serving test (`tests/e2e/online_serving/test_<model>.py`)
+- Buildkite CI entry in `.buildkite/test-merge.yml`
+- New section in `examples/online_serving/text_to_speech/README.md` (table row + per-model section). Do **not** create a top-level `examples/online_serving/<model>/` dir or a per-model `README.md` inside `text_to_speech/<model>/`.
+
+### E2E test pitfalls to avoid
+
+- **One `OmniServerParams` set per file.** `omni_server` is module-scoped; a second
+  id in the same file forces mid-module teardown/restart and exposes startup
+  races (`APIConnectionError` on the first request post-restart). Split variants
+  into separate files instead.
+- **No external URL fetches from the server.** CI and some dev hosts can't
+  reach `raw.githubusercontent.com` over TLS. Inline ref audio as
+  `data:audio/wav;base64,...`; the serving layer accepts both URL and data URL.
+- **Use the harness readiness gate.** The fixture waits for HTTP 200 on
+  `/health`; don't add `time.sleep` in tests. If warmup is incomplete, make
+  `/health` return non-200 until you're actually ready.
+- **Mark with `@pytest.mark.core_model` + `hardware_test(res={"cuda": "H100"})`**
+  so the test lands in `test-ready.yml` (triggered by the `ready` label) rather
+  than only nightly.
+
+## Phase 4: Async Chunk (Streaming)
+
+**Goal**: Enable inter-stage streaming so audio chunks are produced while AR generation continues.
+
+### Steps
+
+1. **Update stage config YAML**:
+   ```yaml
+   async_chunk: true
+   codec_chunk_frames: 25      # frames per chunk
+   codec_left_context_frames: 25  # overlap for smooth boundaries
+   ```
+2. **Implement chunk handling in Stage 1**:
+   - Accept partial input (chunk of codec codes)
+   - Handle left context for smooth audio boundaries
+   - Return partial audio in `OmniOutput`
+3. **Test streaming**:
+   - Verify audio quality matches non-streaming output
+   - Check for artifacts at chunk boundaries
+   - Measure TTFA (time to first audio)
+4. **Update online serving** to support `stream=true` with PCM output
+
+### Streaming Architecture
+
+```
+Stage 0 (AR)                    Stage 1 (Decoder)
+  |                                |
+  |-- chunk 0 (25 frames) ------> decode -> audio chunk 0 -> client
+  |-- chunk 1 (25 frames) ------> decode -> audio chunk 1 -> client
+  |-- chunk 2 (25 frames) ------> decode -> audio chunk 2 -> client
+  ...
+```
+
+### Key Considerations
+
+- **Left context overlap**: Prevents audible artifacts at chunk boundaries
+- **Hop length matters**: `context_audio_samples = context_frames * hop_length`
+- **First chunk latency**: Can use larger initial chunk for better quality, then smaller chunks
+
+### Deliverables
+
+- Updated stage config with async_chunk enabled
+- Smooth streaming audio without boundary artifacts
+- TTFA metrics
+
+## Phase 5: CUDA Graph Acceleration
+
+**Goal**: Capture the AR loop as a CUDA graph for significant speedup.
+
+### Steps
+
+1. **Identify the hot loop**: The AR decoding loop that runs N steps per token
+2. **Create static buffers**:
+   - KV caches with fixed max sequence length
+   - Pre-built causal masks and position tensors per step
+   - Static input/output tensors
+3. **Implement graph capture**:
+   - Warm up with real data
+   - Capture the forward pass
+   - Replay with updated inputs
+4. **Handle constraints**:
+   - Use `torch.argmax` instead of `torch.multinomial` (graph-safe)
+   - Fixed batch size (fall back to eager for other sizes)
+   - No dynamic control flow inside the graph
+
+See [references/cuda-graph-example.md](references/cuda-graph-example.md) for
+a worked skeleton (Qwen3-TTS code predictor, 16-step AR loop), performance
+expectations (3–5× on the graphed component for fixed batch_size=1), and the
+graph-safety constraints you must honor inside the captured region.
+
+### Deliverables
+
+- CUDA graph implementation for the AR hot loop
+- Benchmark script comparing eager vs graph performance
+- Documentation of constraints and fallback behavior
+
+## Phase 6: Pre-commit and DCO
+
+**Goal**: Every commit passes `pre-commit` lint and carries a DCO
+`Signed-off-by` line that matches the author email.
+
+- Install hooks once: `pre-commit install`.
+- Run `pre-commit run --files <changed-files>` before every push; accept any
+  auto-fixes, stage, re-commit.
+- Sign every commit with `git commit -s`. DCO checks that author email and
+  `Signed-off-by` email match — `git config user.email` must match your
+  GitHub account email.
+
+Common pre-commit failures, recovery commands for missing sign-off, and the
+full `pre-commit run` invocation for a TTS model:
+[references/precommit-dco.md](references/precommit-dco.md).
+
+## Integration Checklist
+
+Use this checklist when integrating a new TTS model:
+
+### Cross-Cutting Invariants (verify at end of every phase)
+- [ ] I1: `forward()` docstring states cumulative vs delta; consolidation path audited end-to-end
+- [ ] I2: Tests / examples / benchmarks never use `dict.get(a) or dict.get(b)` on tensor values; list form handled
+- [ ] I3: No `.item()` / `.cpu()` / Python branch on tensor values inside per-step loops
+- [ ] I4: Offline RTF, browser streaming playback, and concurrent-request smoke test all pass
+- [ ] I5: Any cross-step cache keyed by `_omni_req_id`; entries freed when the request finishes
+
+### Phase 1: HF Reference
+- [ ] Reference model runs and produces correct audio
+- [ ] Architecture documented (stages, codebooks, tokens, sample rate)
+- [ ] Reference audio samples saved for comparison
+
+### Phase 2: Stage Separation
+- [ ] Model registered in `registry.py`
+- [ ] Config classes created with `model_type` registration
+- [ ] Stage 0 (AR) implemented and generates correct tokens
+- [ ] Stage 1 (Decoder) produces correct audio from tokens — dtype float32 for codec decoder
+- [ ] Stage 1 `max_num_seqs` ≥ 4 in production config (default 1 causes gaps under concurrency)
+- [ ] Optional dependency fallbacks handled at `load_weights()` time (torchaudio/soundfile/etc.)
+- [ ] Streaming: codec codes accumulated across AR steps (not reset per step)
+- [ ] Streaming: delta audio emitted per chunk, not full re-decoded waveform
+- [ ] Streaming: all `forward()` return paths emit `model_outputs`
+- [ ] Streaming: per-request state keyed by request ID (not shared across requests)
+- [ ] Streaming: codec tensors moved to codec decoder device before decode
+- [ ] Stage config YAML created
+- [ ] `end2end.py` produces audio matching reference quality
+- [ ] README.md written
+
+### Phase 3: Online Serving
+- [ ] All 5 `serving_speech.py` integration points added in one commit
+- [ ] Only extract params in `_build_*_params` that are forwarded to the model call (ruff F841)
+- [ ] Prompt builder handles text input correctly
+- [ ] Voice cloning works (if supported)
+- [ ] All response formats work (wav, mp3, flac, pcm)
+- [ ] Client scripts and server launcher created
+- [ ] E2E online serving test written (`tests/e2e/online_serving/test_<model>.py`)
+- [ ] Buildkite CI entry added to `.buildkite/test-merge.yml`
+- [ ] Gradio demo working
+- [ ] Documentation added (offline + online docs, nav, supported models)
+
+### Phase 4: Async Chunk
+- [ ] Stage config updated with `async_chunk: true`
+- [ ] Stage 1 handles partial chunks correctly
+- [ ] No audio artifacts at chunk boundaries
+- [ ] Streaming via API (`stream=true`) works
+- [ ] TTFA measured and acceptable
+
+### Phase 5: CUDA Graph
+- [ ] Hot loop identified and profiled
+- [ ] Static buffers allocated
+- [ ] Graph captured and replays correctly
+- [ ] Benchmark shows meaningful speedup
+- [ ] Fallback to eager works for unsupported configs
+
+### Phase 6: Pre-commit and DCO
+- [ ] `pre-commit run --files <changed>` passes before every push
+- [ ] Every commit has `Signed-off-by` matching the author email (`git commit -s`)
+- [ ] `git config user.email` matches the email registered on your GitHub account
+- [ ] Details and failure-recovery commands: [references/precommit-dco.md](references/precommit-dco.md)
+
+## References
+
+In-skill references (details split out of the main body):
+
+- [references/single-stage-ar.md](references/single-stage-ar.md) — full `forward()` / generator skeleton for the MOSS-TTS-Nano-style pattern
+- [references/optional-deps.md](references/optional-deps.md) — torchaudio / torchcodec fallback pattern
+- [references/cuda-graph-example.md](references/cuda-graph-example.md) — Qwen3-TTS code-predictor CUDA graph skeleton
+- [references/precommit-dco.md](references/precommit-dco.md) — full pre-commit invocation, failure table, DCO recovery
+
+Project docs and adjacent skills:
+
+- [TTS audio skill](../vllm-omni-audio-tts/SKILL.md) — supported models and usage
+- [Fish Speech integration](../vllm-omni-audio-tts/references/fish-speech.md) — complete example of Phases 1–3
+- [Qwen3-TTS reference](../vllm-omni-audio-tts/references/qwen-tts.md) — complete example of all 5 phases
+- [Adding a TTS model (developer guide)](https://github.com/vllm-project/vllm-omni/blob/main/docs/contributing/model/adding_tts_model.md)
+- `plan/voxcpm2_native_ar_design.md` — VoxCPM2's vLLM-native AR + side-computation pattern (distinct from the generator-based single-stage described above)

.claude/skills/add-tts-model/references/cuda-graph-example.md

@@ -0,0 +1,42 @@
+# CUDA Graph Example: Qwen3-TTS Code Predictor
+
+Reference sketch for capturing the 16-step code-predictor AR loop as a single
+CUDA graph. Adapt the shapes, number of steps, and KV-head layout to your
+model.
+
+```python
+import torch
+
+class CodePredictorGraph:
+    """Captures the 16-step code predictor AR loop as a single CUDA graph."""
+
+    def setup_graph(self, device: torch.device, kv_heads: int = 4, head_dim: int = 64):
+        self.num_steps = 16
+        self.kv_cache = torch.zeros(1, kv_heads, self.num_steps, head_dim, device=device)
+        self.positions = torch.arange(self.num_steps, device=device)
+        self.causal_mask = torch.tril(torch.ones(self.num_steps, self.num_steps, device=device))
+        self.input_buf = torch.zeros(1, 1, kv_heads * head_dim, device=device)
+        self.output_buf = torch.zeros(1, self.num_steps, device=device, dtype=torch.long)
+        # Warm up, then: self.graph = torch.cuda.CUDAGraph(); self.graph.capture(...)
+
+    def run_graph(self, initial_input: torch.Tensor) -> torch.Tensor:
+        self.input_buf.copy_(initial_input)
+        self.graph.replay()
+        return self.output_buf.clone()
+```
+
+## Performance expectations (Qwen3-TTS code predictor)
+
+- **3–5× speedup** on the graphed component.
+- Effective only for fixed batch sizes (typically `batch_size=1`).
+- Fall back to eager for any shape/config that wasn't captured — do not try
+  to recapture per request.
+
+## Graph-safety constraints
+
+- `torch.argmax` instead of `torch.multinomial`.
+- Fixed batch size.
+- No Python control flow that branches on tensor values inside the captured
+  region (use `torch.where` / masks).
+- No `.item()`, `.cpu()`, `.tolist()` — each would break the capture or
+  cause a GPU→CPU sync during replay.

.claude/skills/add-tts-model/references/optional-deps.md

@@ -0,0 +1,47 @@
+# Optional Dependency Handling
+
+Models that rely on `torchaudio`, `torchcodec`, `soundfile`, or other optional
+packages must handle the missing-package case at import time, not at call
+time. Failing to do this causes cryptic errors only on environments without
+the optional package — after the model is already deployed.
+
+## Pattern (used in MOSS-TTS-Nano)
+
+```python
+def _patch_torchaudio_load() -> None:
+    """Fallback torchaudio.load/save to soundfile if torchcodec is unavailable."""
+    try:
+        import torchcodec  # noqa: F401
+        return  # torchcodec present, torchaudio works as-is
+    except ImportError:
+        pass
+
+    import soundfile as sf
+    import torchaudio
+
+    def _sf_load(path, **kwargs):
+        data, sr = sf.read(str(path), dtype="float32", always_2d=True)
+        return torch.from_numpy(data).T, sr
+
+    torchaudio.load = _sf_load
+    # patch .save similarly if needed
+```
+
+## Rules
+
+- Mirror the full signature of the replaced function. `torchaudio.load`
+  accepts `frame_offset`, `num_frames`, `normalize`, `channels_first`,
+  `format` — missing any of them causes `TypeError` from calling code.
+- Catch `except Exception`, not just `ImportError`. `import torchaudio`
+  itself can fail with non-`ImportError` errors on broken installs. Log the
+  exception type and message (`logger.warning("torchaudio probe failed: %s: %s",
+  type(exc).__name__, exc)`) before falling back, so unrelated errors are not
+  silently swallowed.
+- Call the patch function at the top of `load_weights()` before loading any
+  audio assets. Do not call it at module import time.
+
+## Reference implementation
+
+Any in-tree model that patches `torchaudio.load` in its `load_weights()` —
+e.g. MOSS-TTS-Nano's `modeling_moss_tts_nano.py` once that integration
+lands.

.claude/skills/add-tts-model/references/precommit-dco.md

@@ -0,0 +1,54 @@
+# Pre-commit and DCO
+
+Every commit must pass `pre-commit` lint and carry a `Signed-off-by` line
+that matches the commit author email.
+
+## Pre-commit
+
+Install hooks once:
+
+```bash
+pre-commit install
+```
+
+Run before every push on the files you changed:
+
+```bash
+pre-commit run --files \
+  vllm_omni/model_executor/models/<model_name>/*.py \
+  vllm_omni/entrypoints/openai/serving_speech.py \
+  vllm_omni/model_executor/models/registry.py \
+  tests/e2e/offline_inference/test_<model_name>.py \
+  tests/e2e/online_serving/test_<model_name>.py
+```
+
+When pre-commit **modifies files** (ruff format auto-fix), it exits non-zero
+but the changes are correct — stage the modified files and re-commit.
+
+| Failure | Root cause | Fix |
+|---------|-----------|-----|
+| `ruff F841` | Variable extracted but never forwarded to model call | Remove the extraction or wire it through |
+| `ruff E402` | Import added below function definitions | Move to top-level import block |
+| `ruff format` | Line length, spacing, quote style | Accept auto-fix, stage, re-commit |
+
+## DCO sign-off
+
+Every commit must carry `Signed-off-by: Your Name <your@email.com>`. Use
+`-s`:
+
+```bash
+git commit -s -m "feat(<model>): add <Model> TTS support"
+```
+
+Or set it permanently: `git config format.signOff true`.
+
+The DCO check verifies that the commit author email matches the
+`Signed-off-by` line. Confirm `git config user.email` matches your GitHub
+account email before committing.
+
+Fix a missing or mismatched sign-off on the latest commit:
+
+```bash
+git commit --amend -s --no-edit
+git push origin <branch> --force-with-lease
+```

.claude/skills/add-tts-model/references/single-stage-ar.md

@@ -0,0 +1,108 @@
+# Single-Stage AR Pattern
+
+When the upstream model cannot be cleanly split into an AR stage and a separate
+decoder (e.g. MOSS-TTS-Nano, or any model that bundles AR + codec via an
+`inference_stream()` generator), run the whole pipeline inside a single AR
+worker that yields audio chunks per request.
+
+This is distinct from VoxCPM2's pattern, which also runs in a single stage but
+uses vLLM's native PagedAttention on the base language model with diffusion /
+VAE side-computation outside vLLM — see
+`plan/voxcpm2_native_ar_design.md` for that variant.
+
+## Implementation
+
+1. **Single model file** — load both AR LM and codec inside
+   `modeling_<model>.py`.
+2. **Load weights in `load_weights()`**, not `__init__()` — vLLM initializes
+   distributed state before any CUDA allocations.
+3. **Stream via a per-request generator** stored in `self._stream_gens`:
+
+```python
+class YourModelForCausalLM(nn.Module):
+    def __init__(self, *, vllm_config, prefix=""):
+        super().__init__()
+        self._lm = None                   # populated in load_weights()
+        self._stream_gens: dict = {}      # request_key → generator
+
+    def load_weights(self, weights):
+        # Load self._lm here, after vLLM distributed init
+        ...
+
+    def forward(
+        self,
+        input_ids,
+        positions,
+        intermediate_tensors=None,
+        inputs_embeds=None,
+        runtime_additional_information: list[dict] | None = None,  # one dict per request
+        **kwargs,
+    ) -> OmniOutput:
+        infos = runtime_additional_information or [{}]
+        # Skip dummy/profiling calls
+        if not runtime_additional_information or all(i.get("_is_dummy") for i in infos):
+            self._ar_emit_stop_token = True
+            return OmniOutput(...)  # return empty outputs
+
+        outputs, last_flags = [], []
+        for info in infos:
+            request_key = str(info.get("_omni_req_id", "0"))  # per-request ID from vLLM
+            if request_key not in self._stream_gens:
+                self._stream_gens[request_key] = self._create_stream_gen(info)
+            try:
+                chunk, is_last = next(self._stream_gens[request_key])
+            except StopIteration:
+                chunk, is_last = torch.zeros(0), True
+            if is_last:
+                del self._stream_gens[request_key]
+            outputs.append(chunk)
+            last_flags.append(is_last)
+
+        self._ar_emit_stop_token = all(last_flags)
+        return OmniOutput(multimodal_outputs={"model_outputs": outputs, ...})
+
+    def _create_stream_gen(self, info: dict):
+        """Yield (waveform_tensor, is_last) tuples from inference_stream()."""
+        for event in self._lm.inference_stream(...):
+            if event["type"] == "audio":
+                yield event["waveform"], False
+            elif event["type"] == "result":
+                # Fallback: some models emit a single "result" event instead of
+                # incremental "audio" events — handle both paths
+                yield event.get("waveform", torch.zeros(0)), True
+                return
+        yield torch.zeros(0), True
+
+    def compute_logits(self, hidden_states, sampling_metadata):
+        # Emit EOS only after the last chunk so the AR scheduler ends the request
+        ...
+```
+
+## Key points
+
+- `runtime_additional_information` is the correct parameter name (not
+  `**kwargs`) — it carries one dict per request in the batch.
+- The request ID is `info.get("_omni_req_id")` — set by vLLM, not by user code.
+- Handle both `"audio"` (incremental) and `"result"` (final combined) event
+  types from upstream models.
+
+## Stage config
+
+Single stage with `worker_type: ar`, `engine_output_type: audio`,
+`final_output: true`, `is_comprehension: true`, and `async_chunk: false` at
+the top level. Omitting any of these causes silent misclassification in the
+serving layer.
+
+## Lint discipline
+
+Only extract variables from `additional_information` that you actually
+forward to the model call — unused extractions trip `ruff F841` in
+pre-commit.
+
+## Reference implementation
+
+Look for any single-stage AR model under
+`vllm_omni/model_executor/models/` — e.g. `moss_tts_nano/` when its
+integration lands. If none is in tree yet, follow the skeleton above and
+cross-check against the `is_comprehension: true` / `async_chunk: false`
+dispatch in `vllm_omni/entrypoints/openai/serving_speech.py`.

.claude/skills/readme.md

@@ -0,0 +1,34 @@
+# Claude Skills for vLLM-Omni
+
+This directory contains Claude Code skills maintained for the `vllm-omni`
+repository. These skills capture repeatable workflows for common contributor
+tasks such as model integration, pull request review, and release note
+generation.
+
+## Directory Structure
+
+Each skill lives in its own directory under `.claude/skills/`. A skill may
+include:
+
+- `SKILL.md`: the main workflow and operating instructions
+- `references/`: focused reference material used by the skill
+- `scripts/`: small helper scripts used by the skill
+
+## Available Skills
+
+- `add-diffusion-model`: guides integration of a new diffusion model into
+  `vllm-omni`
+- `add-omni-model`: covers addition of new omni-modality model support
+- `add-tts-model`: covers integration of new TTS models and related serving
+  workflows
+- `generate-release-note`: helps prepare release notes for repository changes
+- `review-pr`: provides a structured workflow for reviewing pull requests
+
+## Maintenance Guidelines
+
+- Keep skill names short and task-oriented.
+- Prefer repository-local paths, commands, and examples.
+- Avoid hardcoding fast-changing support matrices unless the skill is actively
+  maintained alongside those changes.
+- Treat skills as contributor tooling: optimize for clarity, actionability, and
+  low maintenance overhead.

.claude/skills/vllm-omni-npu-upgrade/SKILL.md

@@ -0,0 +1,300 @@
+---
+name: vllm-omni-npu-model-runner-upgrade
+description: "Upgrade vllm-omni NPU model runners (OmniNPUModelRunner, NPUARModelRunner, NPUGenerationModelRunner) to align with the latest vllm-ascend NPUModelRunner while preserving omni-specific logic."
+---
+
+# vLLM-Omni NPU Model Runner Upgrade Skill
+
+## Overview
+
+This skill guides the process of upgrading vllm-omni's NPU model runners to align with the latest vllm-ascend codebase while preserving omni-specific enhancements. The NPU runners are designed to run omni multimodal models (like Qwen3-Omni, Bagel, MiMoAudio) on Ascend NPUs.
+
+## File Structure
+
+### NPU Model Runner Files
+```
+vllm-omni/vllm_omni/platforms/npu/worker/
+├── __init__.py
+├── npu_model_runner.py           # OmniNPUModelRunner (base class)
+├── npu_ar_model_runner.py        # NPUARModelRunner (autoregressive)
+├── npu_ar_worker.py              # AR worker
+├── npu_generation_model_runner.py # NPUGenerationModelRunner (diffusion/non-AR)
+└── npu_generation_worker.py      # Generation worker
+```
+
+### GPU Reference Files (for omni-specific logic sync)
+```
+vllm-omni/vllm_omni/worker/
+├── __init__.py
+├── gpu_model_runner.py           # OmniGPUModelRunner
+├── gpu_ar_model_runner.py        # GPUARModelRunner
+├── gpu_ar_worker.py
+├── gpu_generation_model_runner.py
+├── gpu_generation_worker.py
+├── mixins.py
+├── base.py
+└── gpu_memory_utils.py
+```
+
+### vllm-ascend Reference Files
+```
+vllm-ascend/vllm_ascend/worker/
+├── model_runner_v1.py            # NPUModelRunner (base class to copy from)
+├── npu_input_batch.py
+├── block_table.py
+├── pcp_utils.py
+└── worker.py
+```
+
+## Inheritance Hierarchy
+
+```
+                    GPUModelRunner (vllm)
+                         |
+        +----------------+----------------+
+        |                                 |
+  OmniGPUModelRunner              NPUModelRunner (vllm-ascend)
+  (vllm_omni/worker)              (vllm_ascend/worker)
+        |                                 |
+        +----------- OmniNPUModelRunner --+
+                     (multiple inheritance)
+                            |
+            +---------------+---------------+
+            |                               |
+    NPUARModelRunner            NPUGenerationModelRunner
+    (autoregressive)            (non-autoregressive/diffusion)
+```
+
+## Omni-Specific Comment Markers
+
+Omni-specific logic is marked with comment blocks:
+```python
+# -------------------------------------- Omni-new -------------------------------------------------
+# ... omni-specific code ...
+# -------------------------------------- Omni-new -------------------------------------------------
+```
+
+Or simpler variations:
+```python
+#  -------------------------------------- Omni-new -------------------------------------------------
+#  ------------------------------------------------------------------------------------------------
+```
+
+**Important**:
+- Always preserve and add these markers when modifying code.
+- **The reference documents (`references/omni-specific-blocks.md`) may not be up-to-date.** Always grep for `Omni-new` in the GPU implementations to find the authoritative list of omni-specific blocks.
+- When you discover new omni-specific code that is not documented in the references, please update the reference files.
+
+## Key Methods Requiring Attention
+
+### OmniNPUModelRunner (npu_model_runner.py)
+
+| Method | Description | Omni-Specific Logic |
+|--------|-------------|---------------------|
+| `load_model` | Load model and initialize talker_mtp | Uses `ACLGraphWrapper` instead of `CUDAGraphWrapper`, initializes talker buffers |
+| `_dummy_run` | Warmup/profiling run | talker_mtp dummy forward, `extract_multimodal_outputs` |
+| `_model_forward` | Forward pass wrapper | Injects `model_kwargs_extra`, wraps with `OmniOutput`, NPU-specific graph updates |
+| `_talker_mtp_forward` | Talker MTP forward for Qwen3-Omni | Uses `set_ascend_forward_context` |
+
+### NPUARModelRunner (npu_ar_model_runner.py)
+
+| Method | Description | Omni-Specific Logic |
+|--------|-------------|---------------------|
+| `__init__` | Initialize with KV transfer manager | `OmniKVTransferManager` setup |
+| `execute_model` | Main inference entry | KV transfer handling, `_update_states` override, `extract_multimodal_outputs` |
+| `sample_tokens` | Token sampling | Hidden states extraction, multimodal outputs processing, `OmniModelRunnerOutput` |
+| `_resolve_global_request_id` | Request ID resolution | For disaggregated inference |
+
+### NPUGenerationModelRunner (npu_generation_model_runner.py)
+
+| Method | Description | Omni-Specific Logic |
+|--------|-------------|---------------------|
+| `_update_request_states` | Update request states for async chunk | async_chunk handling |
+| `execute_model` | Generation forward | async_chunk, `seq_token_counts`, `_run_generation_model` |
+| `sample_tokens` | Output processing | multimodal output packaging to `OmniModelRunnerOutput` |
+| `_dummy_run` | Dummy run override | model_kwargs initialization, multimodal extraction |
+| `_run_generation_model` | Run generation model | Calls `_model_forward` with sampler |
+
+## Upgrade Workflow
+
+### Step 1: Preparation
+
+1. **Identify target versions**(Use gh cli to check):
+   - We're using vllm-omni main branch
+   - Check the last release of vllm-omni
+   - Target vllm-ascend version(Just directly use the local latest vllm-ascend code)
+
+2. **Check GPU-side changes** (since last release):
+   ```bash
+   cd /root/vllm-workspace/vllm-omni
+   git log --oneline --since="<last-release-date>" -- vllm_omni/worker/
+   ```
+
+3. **Read latest vllm-ascend code**:
+   - We don't track vllm-ascend changes - just directly use the latest code from `/root/vllm-workspace/vllm-ascend/vllm_ascend/worker/model_runner_v1.py`
+   - Copy the relevant methods and re-insert omni-specific blocks
+
+### Step 2: Analyze Omni-Specific Logic
+
+For each NPU model runner file:
+
+1. **Extract existing omni-specific blocks**:
+   ```bash
+   grep -n "Omni-new" vllm_omni/platforms/npu/worker/npu_model_runner.py
+   ```
+
+2. **Document each omni block**:
+   - Which method it belongs to
+   - What functionality it provides
+   - Dependencies on other omni code
+
+### Step 3: Update Base Class (OmniNPUModelRunner)
+
+**Note**: Always check the GPU implementation `gpu_model_runner.py` for any new omni logic not yet documented in references.
+
+1. **Read the latest vllm-ascend `NPUModelRunner.load_model`**
+2. **Copy the method, keeping the structure**
+3. **Re-insert omni-specific logic** (check GPU `gpu_model_runner.py` for authoritative list):
+   - Replace `CUDAGraphWrapper` with `ACLGraphWrapper`
+   - Keep talker_mtp initialization
+   - Preserve buffer allocations for talker
+   - Check for any new omni blocks added since last sync
+
+4. **Update `_dummy_run`**:
+   - Copy from vllm-ascend
+   - Compare with GPU `_dummy_run` for omni-specific blocks
+   - Re-insert all `Omni-new` marked code from GPU version
+
+5. **Update `_model_forward`**:
+   - Keep the omni wrapper logic
+   - Update NPU-specific parts (graph params, SP all-gather)
+   - Check GPU version for any new omni logic
+
+### Step 4: Update AR Model Runner
+
+1. **Compare with GPU `gpu_ar_model_runner.py`** for any new omni features
+2. **Copy `execute_model` from vllm-ascend**
+3. **Re-insert omni blocks** (reference `references/omni-specific-blocks.md`, but note it may be incomplete):
+   - **IMPORTANT**: Always check the GPU implementation `gpu_ar_model_runner.py` for all `Omni-new` marked code blocks
+   - The reference doc may not include newly added omni logic - treat it as a starting point, not exhaustive
+   - When discovering new omni code blocks, please update `references/omni-specific-blocks.md`
+   - Common omni blocks include but are not limited to: KV transfer, multimodal outputs, sampling_metadata handling, etc.
+
+4. **Update `sample_tokens`** (also compare with GPU implementation):
+   - Compare with `gpu_ar_model_runner.py`'s `sample_tokens` method
+   - Identify all `Omni-new` marked code blocks
+   - Ensure NPU version includes all omni-specific logic
+
+### Step 5: Update Generation Model Runner
+
+**Note**: Generation model runner may have unique omni logic for diffusion/non-AR models.
+
+1. **Compare with GPU `gpu_generation_model_runner.py`** - grep for all `Omni-new` blocks
+2. **Update `execute_model`**:
+   - Check GPU version for all omni-specific blocks
+   - Keep async_chunk handling
+   - Keep `seq_token_counts` injection
+   - Update forward/context setup from vllm-ascend
+   - Look for any new omni logic not documented in references
+
+3. **Update `_dummy_run`**:
+   - Copy from vllm-ascend base
+   - Compare with GPU `_dummy_run` if exists
+   - Re-insert all omni-specific logic
+
+### Step 6: Update Imports
+
+Check and update imports at the top of each file:
+
+```python
+# Common vllm-ascend imports
+from vllm_ascend.ascend_forward_context import get_forward_context, set_ascend_forward_context
+from vllm_ascend.attention.attention_v1 import AscendAttentionState
+from vllm_ascend.attention.utils import using_paged_attention
+from vllm_ascend.compilation.acl_graph import ACLGraphWrapper, update_full_graph_params
+from vllm_ascend.ops.rotary_embedding import update_cos_sin
+from vllm_ascend.utils import enable_sp, lmhead_tp_enable
+from vllm_ascend.worker.model_runner_v1 import SEQ_LEN_WITH_MAX_PA_WORKSPACE, NPUModelRunner
+
+# Omni-specific imports
+from vllm_omni.model_executor.models.output_templates import OmniOutput
+from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
+from vllm_omni.outputs import OmniModelRunnerOutput
+from vllm_omni.distributed.omni_connectors.kv_transfer_manager import OmniKVTransferManager
+```
+
+### Step 7: Sync GPU-Side Omni Changes
+
+1. **Check recent GPU worker changes**:
+   ```bash
+   git diff <from-tag>..<to-tag> -- vllm_omni/worker/gpu_model_runner.py
+   git diff <from-tag>..<to-tag> -- vllm_omni/worker/gpu_ar_model_runner.py
+   ```
+
+2. **Identify new omni features** that need to be ported to NPU
+
+3. **Apply corresponding changes** to NPU runners
+
+### Step 8: Validation
+
+1. **Run type checking**:
+   ```bash
+   cd /root/vllm-workspace/vllm-omni
+   python -m py_compile vllm_omni/platforms/npu/worker/npu_model_runner.py
+   python -m py_compile vllm_omni/platforms/npu/worker/npu_ar_model_runner.py
+   python -m py_compile vllm_omni/platforms/npu/worker/npu_generation_model_runner.py
+   ```
+
+2. **Run import test**:
+   ```bash
+   python -c "from vllm_omni.platforms.npu.worker import *"
+   ```
+
+3. **Run model serving test** (if hardware available):
+   ```bash
+   vllm serve <model-path> --trust-remote-code
+   ```
+
+## Common Pitfalls
+
+### 1. Forward Context Differences
+- GPU uses `set_forward_context`
+- NPU uses `set_ascend_forward_context`
+- Parameters may differ slightly
+
+### 2. Graph Wrapper Differences
+- GPU: `CUDAGraphWrapper`
+- NPU: `ACLGraphWrapper`
+- Constructor parameters may differ
+
+### 3. Buffer Creation
+- GPU: `_make_buffer` returns different structure
+- NPU: May need numpy=True/False parameter
+
+### 4. Attention Metadata
+- GPU: Uses vllm attention metadata builders
+- NPU: Uses `AscendCommonAttentionMetadata`
+
+### 5. Sampling
+- GPU: Uses vllm sampler
+- NPU: Uses `AscendSampler`
+
+## Checklist Before Commit
+
+- [ ] All omni-specific comment markers preserved
+- [ ] New omni logic from GPU side synced
+- [ ] Imports updated to latest vllm-ascend
+- [ ] No `CUDAGraphWrapper` references in NPU code
+- [ ] `set_ascend_forward_context` used instead of `set_forward_context`
+- [ ] `ACLGraphWrapper` used for talker_mtp wrapping
+- [ ] Type hints match vllm-ascend signatures
+- [ ] No duplicate code blocks
+- [ ] Python syntax valid (py_compile passes)
+
+## Reference Files for Comparison
+
+When upgrading, keep these files open for reference:
+
+1. **vllm-ascend NPUModelRunner**: `/root/vllm-workspace/vllm-ascend/vllm_ascend/worker/model_runner_v1.py`
+2. **vllm GPUModelRunner**: `/root/vllm-workspace/vllm/vllm/v1/worker/gpu_model_runner.py`
+3. **vllm-omni OmniGPUModelRunner**: `/root/vllm-workspace/vllm-omni/vllm_omni/worker/gpu_model_runner.py`

.claude/skills/vllm-omni-npu-upgrade/references/gpu-to-npu-translation.md

@@ -0,0 +1,335 @@
+# GPU to NPU Translation Patterns
+
+This document provides a quick reference for translating GPU code patterns to NPU equivalents when porting omni-specific logic.
+
+## Import Translations
+
+### Forward Context
+```python
+# GPU
+from vllm.forward_context import set_forward_context
+
+# NPU
+from vllm_ascend.ascend_forward_context import set_ascend_forward_context
+```
+
+### Graph Wrapper
+```python
+# GPU
+from vllm.compilation.cuda_graph import CUDAGraphWrapper
+
+# NPU
+from vllm_ascend.compilation.acl_graph import ACLGraphWrapper
+```
+
+### Attention State
+```python
+# GPU (no equivalent - uses FlashAttention states directly)
+
+# NPU
+from vllm_ascend.attention.attention_v1 import AscendAttentionState
+```
+
+### Utilities
+```python
+# GPU
+# (directly use torch.cuda functions)
+
+# NPU
+from vllm_ascend.utils import enable_sp, lmhead_tp_enable
+from vllm_ascend.ops.rotary_embedding import update_cos_sin
+```
+
+## Context Manager Translations
+
+### Forward Context Setup
+```python
+# GPU
+with set_forward_context(
+    attn_metadata,
+    self.vllm_config,
+    num_tokens=num_tokens_padded,
+    num_tokens_across_dp=num_tokens_across_dp,
+    cudagraph_runtime_mode=cudagraph_mode,
+    batch_descriptor=batch_desc,
+):
+    # forward pass
+
+# NPU
+with set_ascend_forward_context(
+    attn_metadata,
+    self.vllm_config,
+    num_tokens=num_tokens_padded,
+    num_tokens_across_dp=num_tokens_across_dp,
+    aclgraph_runtime_mode=cudagraph_mode,  # Note: 'aclgraph' not 'cudagraph'
+    batch_descriptor=batch_desc,
+    num_actual_tokens=scheduler_output.total_num_scheduled_tokens,
+    model_instance=self.model,
+):
+    # forward pass
+```
+
+### Graph Capture Context
+```python
+# GPU
+from vllm.compilation.cuda_graph import graph_capture as cuda_graph_capture
+with cuda_graph_capture(self.device):
+    # capture
+
+# NPU
+from vllm_ascend.worker.model_runner_v1 import graph_capture
+with graph_capture(self.device):
+    # capture
+```
+
+## Graph Wrapper Usage
+
+### Creating Graph Wrapper
+```python
+# GPU
+if cudagraph_mode.has_full_cudagraphs() and has_separate_talker:
+    self.talker_mtp = CUDAGraphWrapper(
+        talker_mtp,
+        self.vllm_config,
+        runtime_mode=CUDAGraphMode.FULL
+    )
+
+# NPU
+if cudagraph_mode.has_full_cudagraphs() and has_separate_talker:
+    self.talker_mtp = ACLGraphWrapper(
+        talker_mtp,
+        self.vllm_config,
+        runtime_mode=CUDAGraphMode.FULL
+    )
+```
+
+### Checking Graph Wrapper Type
+```python
+# GPU
+if not isinstance(self.talker_mtp, CUDAGraphWrapper):
+    _cudagraph_mode = CUDAGraphMode.NONE
+
+# NPU
+if not isinstance(self.talker_mtp, ACLGraphWrapper):
+    _cudagraph_mode = CUDAGraphMode.NONE
+```
+
+## Device Operations
+
+### Synchronization
+```python
+# GPU
+torch.cuda.synchronize()
+
+# NPU
+torch.npu.synchronize()
+```
+
+### Stream Operations
+```python
+# GPU
+stream = torch.cuda.Stream(device=device)
+torch.cuda.current_stream()
+
+# NPU
+stream = torch.npu.Stream(device=device)
+torch.npu.current_stream()
+```
+
+## Attention Metadata
+
+### State Setting (NPU-specific)
+```python
+# GPU - handled internally by attention backends
+
+# NPU - explicit state setting required
+self.attn_state = AscendAttentionState.DecodeOnly
+if self.speculative_config and self.speculative_config.method == "mtp":
+    if self.vllm_config.model_config.use_mla:
+        self.attn_state = AscendAttentionState.SpecDecoding
+    else:
+        self.attn_state = AscendAttentionState.ChunkedPrefill
+```
+
+### Building Attention Metadata
+```python
+# GPU - uses vllm attention builders
+
+# NPU - may need additional parameters
+(attn_metadata, spec_decode_common_attn_metadata) = self._build_attention_metadata(
+    num_tokens=num_tokens_unpadded,
+    num_tokens_padded=num_tokens_padded,
+    num_reqs=num_reqs,
+    num_reqs_padded=num_reqs_padded,
+    max_query_len=max_num_scheduled_tokens,
+    ubatch_slices=ubatch_slices_attn,
+    logits_indices=logits_indices,
+    use_spec_decode=use_spec_decode,
+    num_scheduled_tokens=scheduler_output.num_scheduled_tokens,
+    num_scheduled_tokens_np=num_scheduled_tokens_np,
+    cascade_attn_prefix_lens=cascade_attn_prefix_lens,
+)
+```
+
+## Rotary Embedding
+
+### Update Cos/Sin Cache
+```python
+# GPU - typically handled inside attention
+
+# NPU - explicit update required before forward
+from vllm_ascend.ops.rotary_embedding import update_cos_sin
+update_cos_sin(positions)
+```
+
+## Sequence Parallelism
+
+### Enable SP Check
+```python
+# GPU - use vllm distributed utilities
+
+# NPU - use vllm-ascend wrapper
+from vllm_ascend.utils import enable_sp
+
+if enable_sp():
+    # sequence parallelism enabled
+```
+
+## Sampler
+
+### Sampler Type
+```python
+# GPU - uses vllm sampler
+self.sampler = Sampler()
+
+# NPU - uses AscendSampler
+from vllm_ascend.sample.sampler import AscendSampler
+self.sampler = AscendSampler()
+```
+
+## Input Batch
+
+### Batch Class
+```python
+# GPU
+from vllm.v1.worker.gpu_input_batch import InputBatch
+
+# NPU
+from vllm_ascend.worker.npu_input_batch import NPUInputBatch
+```
+
+## Graph Parameter Updates
+
+### Full Graph Params Update (NPU-specific)
+```python
+# GPU - not needed
+
+# NPU - required for FULL graph mode
+from vllm_ascend.compilation.acl_graph import update_full_graph_params
+
+forward_context = get_forward_context()
+if (
+    forward_context.cudagraph_runtime_mode == CUDAGraphMode.FULL
+    and not forward_context.capturing
+    and not self.use_sparse
+):
+    update_full_graph_params(
+        self.attn_backend,
+        self.update_stream,
+        forward_context,
+        num_tokens_padded,
+        self.vllm_config,
+        self.speculative_config,
+        positions.shape[0],
+    )
+```
+
+## Paged Attention Check
+
+```python
+# GPU - not typically needed
+
+# NPU
+from vllm_ascend.attention.utils import using_paged_attention
+
+if is_graph_capturing and using_paged_attention(num_tokens, self.vllm_config):
+    seq_lens = SEQ_LEN_WITH_MAX_PA_WORKSPACE
+```
+
+## Common Method Signature Differences
+
+### _dummy_run Parameters
+```python
+# GPU (v0.17.0)
+def _dummy_run(
+    self,
+    num_tokens: int,
+    cudagraph_runtime_mode: CUDAGraphMode | None = None,
+    force_attention: bool = False,
+    uniform_decode: bool = False,
+    allow_microbatching: bool = True,
+    skip_eplb: bool = False,
+    is_profile: bool = False,
+    create_mixed_batch: bool = False,
+    remove_lora: bool = True,
+    is_graph_capturing: bool = False,
+    num_active_loras: int = 0,
+) -> tuple[torch.Tensor, torch.Tensor]:
+
+# NPU (v0.17.0) - adds with_prefill, activate_lora
+def _dummy_run(
+    self,
+    num_tokens: int,
+    with_prefill: bool = False,
+    cudagraph_runtime_mode: CUDAGraphMode | None = None,
+    force_attention: bool = False,
+    uniform_decode: bool = False,
+    is_profile: bool = False,
+    create_mixed_batch: bool = False,
+    allow_microbatching: bool = True,
+    skip_eplb: bool = False,
+    remove_lora: bool = True,
+    activate_lora: bool = False,
+    is_graph_capturing: bool = False,
+    num_active_loras: int = 0,
+) -> tuple[torch.Tensor, torch.Tensor]:
+```
+
+### _model_forward Parameters
+```python
+# GPU - no num_tokens_padded
+def _model_forward(
+    self,
+    input_ids: torch.Tensor | None = None,
+    positions: torch.Tensor | None = None,
+    intermediate_tensors: IntermediateTensors | None = None,
+    inputs_embeds: torch.Tensor | None = None,
+    **model_kwargs: dict[str, Any],
+):
+
+# NPU - has num_tokens_padded as first parameter
+def _model_forward(
+    self,
+    num_tokens_padded: int,
+    input_ids: torch.Tensor | None = None,
+    positions: torch.Tensor | None = None,
+    intermediate_tensors: IntermediateTensors | None = None,
+    inputs_embeds: torch.Tensor | None = None,
+    **model_kwargs: dict[str, Any],
+):
+```
+
+## Quick Reference Table
+
+| Feature | GPU | NPU |
+|---------|-----|-----|
+| Graph wrapper | `CUDAGraphWrapper` | `ACLGraphWrapper` |
+| Forward context | `set_forward_context` | `set_ascend_forward_context` |
+| Runtime mode param | `cudagraph_runtime_mode` | `aclgraph_runtime_mode` |
+| Device sync | `torch.cuda.synchronize()` | `torch.npu.synchronize()` |
+| Stream | `torch.cuda.Stream` | `torch.npu.Stream` |
+| Current stream | `torch.cuda.current_stream()` | `torch.npu.current_stream()` |
+| Input batch | `InputBatch` | `NPUInputBatch` |
+| Sampler | `Sampler` | `AscendSampler` |
+| Attention state | N/A | `AscendAttentionState` |
+| RoPE update | N/A | `update_cos_sin()` |

.claude/skills/vllm-omni-npu-upgrade/references/omni-specific-blocks.md

@@ -0,0 +1,374 @@
+# Omni-Specific Code Blocks Reference
+
+This document catalogs omni-specific code blocks in the NPU model runners, making it easier to identify what needs to be preserved during upgrades.
+
+> **IMPORTANT**: This document may not be complete or up-to-date!
+>
+> - Always grep for `Omni-new` in the GPU implementations (`vllm_omni/worker/`) to find the authoritative list
+> - New omni features may be added that are not yet documented here
+> - When you discover new omni-specific blocks during an upgrade, please update this document
+> - Last verified: Check git history for this file
+
+## OmniNPUModelRunner (npu_model_runner.py)
+
+### load_model - Talker MTP Initialization
+
+```python
+def load_model(self, *args, **kwargs) -> None:
+    NPUModelRunner.load_model(self, *args, **kwargs)
+    # Initialize enable_sp cache to avoid get_current_vllm_config() error
+    # in _pad_for_sequence_parallelism during execute_model.
+    # This is a workaround for vllm-ascend not passing vllm_config to enable_sp().
+    enable_sp(self.vllm_config)
+    # TODO move this model specific logic to a separate class
+    # TTS model IS the talker (no .talker sub-attr); use getattr to support both Omni and TTS.
+    talker_mtp = getattr(self.model, "talker_mtp", None)
+    if talker_mtp is not None:
+        self.talker_mtp = talker_mtp  # type: ignore[assignment]
+        cudagraph_mode = self.compilation_config.cudagraph_mode
+        assert cudagraph_mode is not None
+        # Only wrap talker_mtp in CUDAGraphWrapper for Omni models that
+        # have a separate .talker sub-module.  TTS models' code predictor
+        # has internal AR loops / torch.multinomial — not graph-safe.
+        has_separate_talker = getattr(self.model, "talker", None) is not None
+        if cudagraph_mode.has_full_cudagraphs() and has_separate_talker:
+            # NOTE: Use ACLGraphWrapper on NPU, not CUDAGraphWrapper
+            self.talker_mtp = ACLGraphWrapper(talker_mtp, self.vllm_config, runtime_mode=CUDAGraphMode.FULL)
+        # TTS exposes mtp_hidden_size; Omni uses hf_text_config.hidden_size.
+        hidden_size = int(
+            getattr(self.model, "mtp_hidden_size", 0) or getattr(self.model_config.hf_text_config, "hidden_size")
+        )
+        max_batch_size = max(self.max_num_reqs, self.compilation_config.max_cudagraph_capture_size)
+        self.talker_mtp_input_ids = self._make_buffer(max_batch_size, dtype=torch.int32)
+        self.talker_mtp_inputs_embeds = self._make_buffer(
+            max_batch_size, hidden_size, dtype=self.dtype, numpy=False
+        )
+        self.last_talker_hidden = self._make_buffer(max_batch_size, hidden_size, dtype=self.dtype, numpy=False)
+        self.text_step = self._make_buffer(max_batch_size, hidden_size, dtype=self.dtype, numpy=False)
+```
+
+### _dummy_run - Talker MTP Dummy Forward
+
+Location: Inside `set_ascend_forward_context` block, before main model forward
+
+```python
+# ---------------------------------------Omni-new----------------------------------------------
+if getattr(self.model, "talker", None) is not None and hasattr(self.model, "talker_mtp"):
+    num_tokens_padded_talker_mtp = num_tokens_padded
+    if num_tokens_padded_talker_mtp == self.max_num_tokens:
+        num_tokens_padded_talker_mtp = self.talker_mtp_input_ids.gpu.shape[0]
+    outputs = self.talker_mtp(
+        self.talker_mtp_input_ids.gpu[:num_tokens_padded_talker_mtp],
+        self.talker_mtp_inputs_embeds.gpu[:num_tokens_padded_talker_mtp],
+        self.last_talker_hidden.gpu[:num_tokens_padded_talker_mtp],
+        self.text_step.gpu[:num_tokens_padded_talker_mtp],
+    )
+    self.compilation_config.cache_dir = None
+# ---------------------------------------Omni-new----------------------------------------------
+```
+
+### _dummy_run - Extract Multimodal Outputs
+
+Location: After model forward, before dummy_compute_logits
+
+```python
+# ---------------------------------------Omni-new----------------------------------------------
+hidden_states, multimodal_outputs = self.extract_multimodal_outputs(hidden_states)
+# ---------------------------------------Omni-new----------------------------------------------
+```
+
+### _model_forward - Omni Output Wrapping
+
+```python
+def _model_forward(
+    self,
+    num_tokens_padded: int,
+    input_ids: torch.Tensor | None = None,
+    positions: torch.Tensor | None = None,
+    intermediate_tensors: IntermediateTensors | None = None,
+    inputs_embeds: torch.Tensor | None = None,
+    **model_kwargs: dict[str, Any],
+):
+    """Override to combine NPUModelRunner's signature with OmniGPUModelRunner's logic."""
+    # Omni-specific: build and inject extra model kwargs
+    model_kwargs_extra = self._build_model_kwargs_extra()
+
+    # Call the model forward (same as NPUModelRunner)
+    assert self.model is not None
+    model_output = self.model(
+        input_ids=input_ids,
+        positions=positions,
+        intermediate_tensors=intermediate_tensors,
+        inputs_embeds=inputs_embeds,
+        **model_kwargs,
+        **model_kwargs_extra,
+    )
+
+    # Omni-specific: wrap output if needed
+    if not isinstance(model_output, OmniOutput) and hasattr(self.model, "make_omni_output"):
+        model_output = self.model.make_omni_output(model_output, **model_kwargs_extra)
+
+    # Omni-specific: cache model output for later sample_tokens
+    self._omni_last_model_output = model_output
+
+    # NPU-specific: update full graph params (keep from vllm-ascend)
+    forward_context = get_forward_context()
+    # ... NPU graph update logic ...
+
+    # NPU-specific: all-gather for sequence parallelism (keep from vllm-ascend)
+    if get_forward_context().sp_enabled and not isinstance(model_output, IntermediateTensors):
+        model_output = self._all_gather_hidden_states_and_aux(model_output)
+
+    return model_output
+```
+
+---
+
+## NPUARModelRunner (npu_ar_model_runner.py)
+
+### __init__ - KV Transfer Manager
+
+```python
+def __init__(self, *args, **kwargs):
+    super().__init__(*args, **kwargs)
+    self.input_ids = self._make_buffer(self.max_num_tokens, dtype=torch.int32)
+    # each model stage has their own hidden size
+    self.hidden_size = self.model_config.hf_text_config.hidden_size
+    self.inputs_embeds = self._make_buffer(self.max_num_tokens, self.hidden_size, dtype=self.dtype, numpy=False)
+    # Initialize KV cache manager (preserve vllm_config fallback behavior)
+    self.kv_transfer_manager = OmniKVTransferManager.from_vllm_config(self.vllm_config, self.model_config)
+```
+
+### execute_model - KV Transfer Before Update States
+
+Location: At the very beginning of execute_model
+
+```python
+#  -------------------------------------- Omni-new -------------------------------------------------
+# [Omni] Handle KV transfer BEFORE updating states (which removes finished requests)
+self.kv_extracted_req_ids = self.kv_transfer_manager.handle_finished_requests_kv_transfer(
+    finished_reqs=getattr(scheduler_output, "finished_requests_needing_kv_transfer", {}),
+    kv_caches=self.kv_caches,
+    block_size=self.cache_config.block_size,
+    cache_dtype=str(self.cache_config.cache_dtype),
+    request_id_resolver=self._resolve_global_request_id,
+)
+#  -------------------------------------- Omni-new -------------------------------------------------
+```
+
+### execute_model - Custom _update_states Call
+
+Location: Inside synchronize_input_prep context
+
+```python
+#  -------------------------------------- Omni-new -------------------------------------------------
+self._update_states(scheduler_output)
+#  ------------------------------------------------------------------------------------------------
+```
+
+### execute_model - Extract Multimodal Outputs
+
+Location: In post process section, after hidden_states assignment
+
+```python
+#  -------------------------------------- Omni-new -------------------------------------------------
+hidden_states, multimodal_outputs = self.extract_multimodal_outputs(hidden_states)
+
+if multimodal_outputs is not None:
+    keys_or_type = (
+        list(multimodal_outputs.keys())
+        if isinstance(multimodal_outputs, dict)
+        else type(multimodal_outputs)
+    )
+    logger.debug(f"[AR] execute_model: multimodal_outputs keys = {keys_or_type}")
+else:
+    logger.debug("[AR] execute_model: multimodal_outputs is None")
+#  -------------------------------------- Omni-new -------------------------------------------------
+```
+
+### execute_model - Compute Logits with sampling_metadata
+
+Location: In both broadcast_pp_output True and False branches
+
+```python
+#  -------------------------------------- Omni-new -------------------------------------------------
+# Try with sampling_metadata first; fall back to without for models that don't support it
+try:
+    logits = self.model.compute_logits(
+        sample_hidden_states, sampling_metadata=self.input_batch.sampling_metadata
+    )
+except TypeError:
+    logits = self.model.compute_logits(sample_hidden_states)
+#  -------------------------------------- Omni-new -------------------------------------------------
+```
+
+### sample_tokens - KV Extracted Req IDs
+
+Location: At the beginning of sample_tokens
+
+```python
+#  -------------------------------------- Omni-new -------------------------------------------------
+kv_extracted_req_ids = getattr(self, "kv_extracted_req_ids", None)
+self.kv_extracted_req_ids = None
+#  -------------------------------------- Omni-new -------------------------------------------------
+```
+
+### sample_tokens - Process Additional Information and Build Output
+
+Location: After bookkeeping sync, replacing the original output construction
+
+```python
+#  -------------------------------------- Omni-new -------------------------------------------------
+hidden_states_cpu = hidden_states.detach().to("cpu").contiguous()
+num_scheduled_tokens_np = getattr(self, "_omni_num_scheduled_tokens_np", None)
+if num_scheduled_tokens_np is None:
+    req_ids = self.input_batch.req_ids
+    num_scheduled_tokens_np = np.array(
+        [scheduler_output.num_scheduled_tokens[rid] for rid in req_ids],
+        dtype=np.int32,
+    )
+
+self._process_additional_information_updates(
+    hidden_states, multimodal_outputs, num_scheduled_tokens_np, scheduler_output
+)
+
+pooler_output: list[dict[str, object]] = []
+for rid in req_ids_output_copy:
+    idx = req_id_to_index_output_copy[rid]
+    start = int(self.query_start_loc.cpu[idx])
+    sched = int(num_scheduled_tokens_np[idx])
+    end = start + sched
+    hidden_slice = hidden_states_cpu[start:end]
+    payload: dict[str, object] = {"hidden": hidden_slice}
+    if isinstance(multimodal_outputs, dict) and multimodal_outputs:
+        # ... multimodal output slicing logic ...
+    pooler_output.append(payload)
+
+model_runner_output = OmniModelRunnerOutput(
+    req_ids=req_ids_output_copy,
+    req_id_to_index=req_id_to_index_output_copy,
+    sampled_token_ids=valid_sampled_token_ids,
+    logprobs=logprobs_lists,
+    prompt_logprobs_dict=prompt_logprobs_dict,
+    pooler_output=(pooler_output if self.vllm_config.model_config.engine_output_type != "text" else None),
+    kv_connector_output=kv_connector_output,
+)
+model_runner_output.kv_extracted_req_ids = kv_extracted_req_ids
+#  -------------------------------------- Omni-new -------------------------------------------------
+```
+
+---
+
+## NPUGenerationModelRunner (npu_generation_model_runner.py)
+
+### execute_model - Async Chunk Update
+
+Location: Inside prepare input section, before synchronize_input_prep
+
+```python
+#  -------------------------------------- Omni-new -------------------------------------------------
+if self.model_config.async_chunk and num_scheduled_tokens:
+    self._update_request_states(scheduler_output)
+#  -------------------------------------- Omni-new -------------------------------------------------
+```
+
+### execute_model - Seq Token Counts
+
+Location: After _preprocess call
+
+```python
+# [Omni] Pass token counts per request for code2wav output slicing
+model_kwargs["seq_token_counts"] = tokens
+```
+
+### execute_model - Run Generation Model
+
+Location: Inside forward context
+
+```python
+#  -------------------------------------- Omni-new -------------------------------------------------
+outputs = self._run_generation_model(
+    num_tokens_padded=num_tokens_padded,
+    input_ids=input_ids,
+    positions=positions,
+    intermediate_tensors=intermediate_tensors,
+    inputs_embeds=inputs_embeds,
+    model_kwargs=model_kwargs,
+    logits_indices=logits_indices,
+)
+_, multimodal_outputs = self.extract_multimodal_outputs(outputs)
+# -------------------------------------- Omni-new -------------------------------------------------
+```
+
+### sample_tokens - Multimodal Output Processing
+
+The entire sample_tokens method body is omni-specific for generation models:
+
+```python
+#  -------------------------------------- Omni-new -------------------------------------------------
+pooler_output: list[object] = []
+if isinstance(multimodal_outputs, torch.Tensor):
+    # ... tensor handling ...
+elif isinstance(multimodal_outputs, list):
+    # ... list handling ...
+elif isinstance(multimodal_outputs, dict):
+    # ... dict handling per request ...
+else:
+    raise RuntimeError("Unsupported diffusion output type")
+# [Omni] Copy req_id mappings to avoid async scheduling mutation.
+req_ids_output_copy = self.input_batch.req_ids.copy()
+req_id_to_index_output_copy = self.input_batch.req_id_to_index.copy()
+output = OmniModelRunnerOutput(
+    req_ids=req_ids_output_copy,
+    req_id_to_index=req_id_to_index_output_copy,
+    sampled_token_ids=[],
+    logprobs=None,
+    prompt_logprobs_dict={},
+    pooler_output=pooler_output,
+    kv_connector_output=kv_connector_output,
+    num_nans_in_logits={},
+    ec_connector_output=ec_connector_output if self.supports_mm_inputs else None,
+)
+#  -------------------------------------- Omni-new -------------------------------------------------
+```
+
+### _dummy_run - Model Kwargs Init and Multimodal Extract
+
+Location: Before model forward and after
+
+```python
+model_kwargs = self._init_model_kwargs()  # Before forward
+
+# ... forward ...
+
+# -------------------------------------- Omni-new -------------------------------------------------
+hidden_states, _ = self.extract_multimodal_outputs(hidden_states)
+# -------------------------------------------------------------------------------------------------
+```
+
+---
+
+## ExecuteModelState Extension
+
+The `ExecuteModelState` NamedTuple is extended for omni:
+
+```python
+class ExecuteModelState(NamedTuple):
+    """Ephemeral cached state transferred between execute_model() and
+    sample_tokens(), after execute_model() returns None."""
+
+    scheduler_output: SchedulerOutput
+    logits: torch.Tensor
+    spec_decode_metadata: SpecDecodeMetadata | None
+    spec_decode_common_attn_metadata: AscendCommonAttentionMetadata | None
+    hidden_states: torch.Tensor
+    sample_hidden_states: torch.Tensor
+    aux_hidden_states: list[torch.Tensor] | None
+    attn_metadata: PerLayerAttnMetadata
+    positions: torch.Tensor
+    ec_connector_output: ECConnectorOutput | None
+    cudagraph_stats: CUDAGraphStat | None
+    multimodal_outputs: Any  # <-- Omni extension
+```
+
+This extended state must be imported from `npu_ar_model_runner` in `npu_generation_model_runner`.

.claude/skills/vllm-omni-npu-upgrade/references/workflow-checklist.md

@@ -0,0 +1,222 @@
+# NPU Model Runner Upgrade Workflow Checklist
+
+> **Note**: Reference documents (`omni-specific-blocks.md`) may not be complete. Always grep for `Omni-new` in GPU implementations to find all omni-specific code blocks. Update the reference docs when discovering new blocks.
+
+## Pre-Upgrade Preparation
+
+### 1. Version Information
+- [ ] Identify current vllm-omni version: `_________`
+- [ ] Identify target vllm-ascend version: `_________`
+- [ ] Identify target vllm version: `_________`
+- [ ] Last release date for GPU worker changes: `_________`
+
+### 2. Gather Git History
+```bash
+# GPU-side omni changes since last release
+cd /root/vllm-workspace/vllm-omni
+git log --oneline --since="YYYY-MM-DD" -- vllm_omni/worker/
+
+# vllm-ascend NPUModelRunner changes
+cd /root/vllm-workspace/vllm-ascend
+git log --oneline <from-tag>..<to-tag> -- vllm_ascend/worker/model_runner_v1.py
+```
+
+### 3. Backup Current Files
+- [ ] Create backup of current NPU runners:
+  ```bash
+  cp -r vllm_omni/platforms/npu/worker vllm_omni/platforms/npu/worker.backup
+  ```
+
+---
+
+## OmniNPUModelRunner (npu_model_runner.py)
+
+### Read and Understand
+- [ ] Read current `npu_model_runner.py`
+- [ ] Read latest `vllm_ascend/worker/model_runner_v1.py`
+- [ ] Read latest `vllm_omni/worker/gpu_model_runner.py`
+
+### Method: load_model
+- [ ] Document existing omni-specific logic
+- [ ] Copy latest NPUModelRunner.load_model structure
+- [ ] Re-insert: `enable_sp(self.vllm_config)` call
+- [ ] Re-insert: talker_mtp detection and setup
+- [ ] Replace: `CUDAGraphWrapper` → `ACLGraphWrapper`
+- [ ] Re-insert: Buffer allocations (talker_mtp_input_ids, etc.)
+
+### Method: _dummy_run
+- [ ] Document existing omni-specific logic locations
+- [ ] Copy latest NPUModelRunner._dummy_run
+- [ ] Re-insert: talker_mtp dummy forward block (inside context)
+- [ ] Re-insert: `extract_multimodal_outputs` call
+- [ ] Verify: Comment markers are present
+
+### Method: _model_forward
+- [ ] Copy latest NPUModelRunner._model_forward structure
+- [ ] Re-insert: `_build_model_kwargs_extra()` call
+- [ ] Re-insert: OmniOutput wrapping logic
+- [ ] Re-insert: `_omni_last_model_output` caching
+- [ ] Keep: NPU graph params update
+- [ ] Keep: SP all-gather logic
+
+### Method: _talker_mtp_forward
+- [ ] Verify: Uses `set_ascend_forward_context`
+- [ ] Verify: Uses `ACLGraphWrapper` check
+- [ ] Sync any changes from GPU `_talker_mtp_forward`
+
+### Imports
+- [ ] Update vllm-ascend imports to latest paths
+- [ ] Verify all omni imports are present
+- [ ] Remove any deprecated imports
+
+---
+
+## NPUARModelRunner (npu_ar_model_runner.py)
+
+### Read and Understand
+- [ ] Read current `npu_ar_model_runner.py`
+- [ ] Read latest `vllm_ascend/worker/model_runner_v1.py` execute_model
+- [ ] Read latest `vllm_omni/worker/gpu_ar_model_runner.py`
+
+### Method: __init__
+- [ ] Sync any new initialization from GPU side
+- [ ] Keep: `OmniKVTransferManager` setup
+- [ ] Keep: Custom buffer allocations
+
+### Method: execute_model
+- [ ] Document all omni blocks with line numbers
+- [ ] Copy latest NPUModelRunner.execute_model structure
+- [ ] Re-insert: KV transfer handling (beginning)
+- [ ] Re-insert: Custom `_update_states` call
+- [ ] Re-insert: `extract_multimodal_outputs`
+- [ ] Re-insert: `compute_logits` with sampling_metadata try/except
+- [ ] Update: ExecuteModelState to include multimodal_outputs
+
+### Method: sample_tokens
+- [ ] Document all omni blocks
+- [ ] Copy latest NPUModelRunner.sample_tokens structure
+- [ ] Re-insert: `kv_extracted_req_ids` handling
+- [ ] Re-insert: Hidden states CPU copy
+- [ ] Re-insert: `_process_additional_information_updates`
+- [ ] Re-insert: `OmniModelRunnerOutput` construction
+
+### ExecuteModelState
+- [ ] Verify: `multimodal_outputs` field is present
+- [ ] Verify: Imported/used correctly in execute_model
+
+### Imports
+- [ ] Update all vllm-ascend imports
+- [ ] Keep omni-specific imports
+
+---
+
+## NPUGenerationModelRunner (npu_generation_model_runner.py)
+
+### Read and Understand
+- [ ] Read current `npu_generation_model_runner.py`
+- [ ] Read latest GPU `gpu_generation_model_runner.py`
+
+### Method: _update_request_states
+- [ ] Verify: async_chunk handling is correct
+- [ ] Sync any changes from GPU side
+
+### Method: execute_model
+- [ ] Document all omni blocks
+- [ ] Copy latest NPUModelRunner.execute_model base structure
+- [ ] Re-insert: async_chunk update logic
+- [ ] Re-insert: `seq_token_counts` injection
+- [ ] Re-insert: `_run_generation_model` call
+- [ ] Re-insert: `extract_multimodal_outputs`
+- [ ] Use: ExecuteModelState from npu_ar_model_runner
+
+### Method: sample_tokens
+- [ ] Keep: Entire omni multimodal output processing
+- [ ] Update: Any new output fields needed
+- [ ] Keep: `OmniModelRunnerOutput` construction
+
+### Method: _run_generation_model
+- [ ] Sync any changes from GPU side
+- [ ] Keep: `_model_forward` call with sampler
+
+### Method: _dummy_run
+- [ ] Copy latest NPUModelRunner._dummy_run
+- [ ] Re-insert: `model_kwargs = self._init_model_kwargs()`
+- [ ] Re-insert: `extract_multimodal_outputs` at end
+
+### Imports
+- [ ] Import ExecuteModelState from npu_ar_model_runner
+- [ ] Update vllm-ascend imports
+
+---
+
+## Post-Upgrade Validation
+
+### Syntax Validation
+- [ ] `python -m py_compile vllm_omni/platforms/npu/worker/npu_model_runner.py`
+- [ ] `python -m py_compile vllm_omni/platforms/npu/worker/npu_ar_model_runner.py`
+- [ ] `python -m py_compile vllm_omni/platforms/npu/worker/npu_generation_model_runner.py`
+
+### Import Validation
+- [ ] `python -c "from vllm_omni.platforms.npu.worker.npu_model_runner import OmniNPUModelRunner"`
+- [ ] `python -c "from vllm_omni.platforms.npu.worker.npu_ar_model_runner import NPUARModelRunner"`
+- [ ] `python -c "from vllm_omni.platforms.npu.worker.npu_generation_model_runner import NPUGenerationModelRunner"`
+
+### Comment Markers
+- [ ] Grep for "Omni-new" in all three files
+- [ ] Verify all omni blocks have closing markers
+
+### Code Review
+- [ ] No `CUDAGraphWrapper` references
+- [ ] All `set_forward_context` replaced with `set_ascend_forward_context`
+- [ ] Parameter names correct (`aclgraph_runtime_mode` not `cudagraph_runtime_mode`)
+- [ ] No duplicate code blocks
+- [ ] No missing imports
+
+---
+
+## Git Commit
+
+### Commit Message Template
+```
+[NPU] Upgrade model runners to align with vllm-ascend vX.Y.Z
+
+- Update OmniNPUModelRunner with latest NPUModelRunner base
+- Update NPUARModelRunner execute_model and sample_tokens
+- Update NPUGenerationModelRunner for async_chunk changes
+- Sync GPU-side omni changes from vX.Y.Z release
+- Preserve all omni-specific logic (marked with Omni-new comments)
+
+Changes from vllm-ascend:
+- <list key changes>
+
+Changes synced from GPU:
+- <list key GPU-side omni changes>
+```
+
+### Files to Stage
+- [ ] `vllm_omni/platforms/npu/worker/npu_model_runner.py`
+- [ ] `vllm_omni/platforms/npu/worker/npu_ar_model_runner.py`
+- [ ] `vllm_omni/platforms/npu/worker/npu_generation_model_runner.py`
+- [ ] Any other modified files
+
+---
+
+## Troubleshooting
+
+### Import Errors
+- Check if vllm-ascend module paths have changed
+- Verify PYTHONPATH includes both vllm-ascend and vllm-omni
+
+### Type Errors
+- Check method signatures match between GPU and NPU
+- Verify NamedTuple fields match expected structure
+
+### Runtime Errors
+- Enable debug logging: `export VLLM_LOGGING_LEVEL=DEBUG`
+- Check graph capture issues: try `--enforce-eager`
+- Check attention issues: verify AscendAttentionState usage
+
+### Performance Regression
+- Compare with previous version on same model
+- Check if graph capture is working: look for ACLGraph logs
+- Verify SP/EP configurations are correct

.gitattributes

@@ -27,6 +27,7 @@
 *.mov binary
 *.mp4 binary
 *.mp3 binary
+*.wav binary
 *.flv binary
 *.fla binary
 *.swf binary

.github/CODEOWNERS

@@ -3,15 +3,64 @@
 # 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)
-# later, we can add more reviewers here
-* @hsliuustc0106
+* @Gaohan123 @hsliuustc0106 @ywang96
 
-# Assign reviewers by directory
-/.github/ @owner1
-/tests/ @owner1
-/docs/ @owner1
+# Repository operations
+/.github/ @hsliuustc0106
 
-# Example: reviewers for specific directories
-# /examples/offline_inference/ @offline-inference-team
-# /vllm_omni/core/ @core-team
-# /vllm_omni/model_executor/ @model-team
+# 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

.github/ISSUE_TEMPLATE/750-RFC.yml

@@ -7,7 +7,7 @@ body:
 - type: markdown
   attributes:
     value: >
-      #### Please take a look at previous [RFCs](https://github.com/vllm-project/vllm-omni/issues?q=label%3ARFC+sort%3Aupdated-desc) for reference.
+      #### Please take a look at previous [RFCs](https://github.com/vllm-project/vllm-omni/issues?q=in%3Atitle%20RFC%20sort%3Aupdated-desc) for reference.
 - type: textarea
   attributes:
     label: Motivation.
@@ -21,7 +21,7 @@ body:
     description: >
       The proposed change of the RFC.
     value: |
-      Please provide the detailed design document of the RFC using the [template](https://docs.google.com/document/d/12YxSsVeD1jvL-InClkeAEnZyWFDndz_65JmXvsamuV4/edit?tab=t.0#heading=h.4ef4szrsgspp).
+      Please provide the detailed design document of the RFC using the [template](https://docs.google.com/document/d/1jcgR3cDaUQH3VczD4ZcKaJAoYWHjCmnYzHkCNyz-9fk/edit?usp=sharing).
   validations:
     required: true
 - type: textarea

.github/mergify.yml

@@ -0,0 +1,223 @@
+pull_request_rules:
+- name: label-documentation
+  description: Automatically apply documentation label
+  conditions:
+    - label != stale
+    - or:
+      - files~=^[^/]+\.md$
+      - files~=^docs/
+      - files~=^examples/
+  actions:
+    label:
+      add:
+        - documentation
+    comment:
+      message: "Documentation preview: https://vllm--{{number}}.org.readthedocs.build/en/{{number}}/"
+
+- name: comment-pre-commit-failure
+  description: Comment on PR when pre-commit check fails
+  conditions:
+    - status-failure=pre-commit
+    - -closed
+    - -draft
+  actions:
+    comment:
+      message: |
+        Hi @{{author}}, the pre-commit checks have failed. Please run:
+
+        ```bash
+        uv pip install pre-commit
+        pre-commit install
+        pre-commit run --all-files
+        ```
+
+        Then, commit the changes and push to your branch.
+
+        For future commits, `pre-commit` will run automatically on changed files before each commit.
+
+        > [!TIP]
+        > <details>
+        > <summary>Is <code>mypy</code> or <code>markdownlint</code> failing?</summary>
+        > <br/>
+        > <code>mypy</code> and <code>markdownlint</code> are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:
+        >
+        > ```bash
+        > # For mypy (substitute "3.10" with the failing version if needed)
+        > pre-commit run --hook-stage manual mypy-3.10
+        > # For markdownlint
+        > pre-commit run --hook-stage manual markdownlint
+        > ```
+        > </details>
+
+- name: comment-dco-failure
+  description: Comment on PR when DCO check fails
+  conditions:
+    - status-failure=dco
+    - -closed
+    - -draft
+  actions:
+    comment:
+      message: |
+        Hi @{{author}}, the DCO check has failed. Please click on DCO in the Checks section for instructions on how to resolve this.
+
+- name: label-ci-build
+  description: Automatically apply ci/build label
+  conditions:
+    - label != stale
+    - or:
+      - files~=^\.github/
+      - files~=\.buildkite/
+      - files~=^cmake/
+      - files=CMakeLists.txt
+      - files~=^docker/Dockerfile
+      - files~=^requirements.*\.txt
+      - files=setup.py
+  actions:
+    label:
+      add:
+        - ci/build
+
+- name: label-frontend
+  description: Automatically apply frontend label
+  conditions:
+    - label != stale
+    - files~=^vllm_omni/entrypoints/
+  actions:
+    label:
+      add:
+        - frontend
+
+- name: label-diffusion
+  description: Automatically apply diffusion label
+  conditions:
+    - label != stale
+    - files~=^vllm_omni/diffusion/
+  actions:
+    label:
+      add:
+        - diffusion
+
+- name: label-new-model
+  description: Automatically apply new-model label
+  conditions:
+    - label != stale
+    - or:
+        - files=vllm_omni/model_executor/models/registry.py
+        - files=vllm_omni/diffusion/registry.py
+  actions:
+    label:
+      add:
+        - new-model
+
+- name: label-npu
+  description: Automatically apply npu label
+  conditions:
+    - label != stale
+    - or:
+      - files~=^vllm_omni/platforms/npu/
+      - title~=(?i)NPU
+      - title~=(?i)Ascend
+  actions:
+    label:
+      add:
+        - npu
+    assign:
+      users:
+        - "gcanlin"
+
+- name: label-rocm
+  description: Automatically apply rocm label
+  conditions:
+    - label != stale
+    - or:
+      - files~=^vllm_omni/platforms/rocm/
+      - title~=(?i)AMD
+      - title~=(?i)ROCm
+  actions:
+    label:
+      add:
+        - rocm
+    assign:
+      users:
+        - "tjtanaa"
+
+- name: label-xpu
+  description: Automatically apply xpu label
+  conditions:
+    - label != stale
+    - or:
+      - files~=^vllm_omni/platforms/xpu/
+      - title~=(?i)Intel
+      - title~=(?i)XPU
+  actions:
+    label:
+      add:
+        - xpu
+    assign:
+      users:
+        - "yma11"
+
+- name: label-modelrunner
+  description: Automatically apply modelrunner label
+  conditions:
+    - label != stale
+    - or:
+      - files~=^vllm_omni/worker/
+      - title~=(?i)ModelRunner
+      - title~=(?i)worker
+  actions:
+    label:
+      add:
+        - ModelRunner
+    assign:
+      users:
+        - "gcanlin" # To check whether need to sync to NPU
+
+- name: auto-rebase if approved, ready, and 30 commits behind main
+  conditions:
+    - base = main
+    - label=ready
+    - "#approved-reviews-by >= 1"
+    - "#commits-behind >= 30"
+    - -closed
+    - -draft
+    - -conflict
+  actions:
+    rebase: {}
+
+- name: ping author on conflicts and add 'needs-rebase' label
+  conditions:
+    - label != stale
+    - conflict
+    - -closed
+  actions:
+    label:
+      add:
+        - needs-rebase
+    comment:
+      message: |
+       This pull request has merge conflicts that must be resolved before it can be
+       merged. Please rebase the PR, @{{author}}.
+
+       https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork
+
+- name: remove 'needs-rebase' label when conflict is resolved
+  conditions:
+    - -conflict
+    - -closed
+  actions:
+    label:
+      remove:
+        - needs-rebase
+
+- name: label-bug
+  description: Automatically apply bug label
+  conditions:
+    - label != stale
+    - or:
+      - title~=(?i)\bbug\b
+      - title~=(?i)\bbugfix\b
+  actions:
+    label:
+      add:
+        - bug

.github/scripts/pr_reviewer.py

@@ -1,629 +0,0 @@
-#!/usr/bin/env python3
-"""
-PR Reviewer using GLM API for vllm-omni project.
-"""
-
-import json
-import logging
-import os
-import sys
-import time
-from dataclasses import dataclass
-from typing import Any, TypedDict
-
-import requests
-
-
-# Type definitions for API responses
-class PRDetails(TypedDict):
-    """Type definition for GitHub PR details response."""
-
-    title: str
-    body: str
-    number: int
-    state: str
-    user: dict[str, Any]
-
-
-class GLMMessage(TypedDict):
-    """Type definition for GLM API message."""
-
-    role: str
-    content: str
-
-
-class GLMChoice(TypedDict):
-    """Type definition for GLM API choice."""
-
-    message: GLMMessage
-    finish_reason: str
-
-
-class GLMResponse(TypedDict):
-    """Type definition for GLM API response."""
-
-    choices: list[GLMChoice]
-    usage: dict[str, int] | None
-
-
-class GitHubComment(TypedDict):
-    """Type definition for GitHub comment."""
-
-    id: int
-    body: str
-    created_at: str
-    user: dict[str, Any]
-
-
-# Configuration
-TRIGGER_PHRASE: str = "@vllm-omni-reviewer"
-DEFAULT_GLM_API_URL: str = "https://open.bigmodel.cn/api/paas/v4/chat/completions"  # noqa: E501
-DEFAULT_GLM_MODEL: str = "glm-5"
-DEFAULT_COOLDOWN_MINUTES: int = 5
-DEFAULT_MAX_RETRIES: int = 3
-DEFAULT_RETRY_DELAY: float = 1.0
-MAX_DIFF_SIZE: int = 100_000  # Maximum diff size in characters
-
-
-@dataclass
-class Config:
-    """Configuration for the PR reviewer."""
-
-    glm_api_url: str
-    glm_model: str
-    cooldown_minutes: int
-    max_retries: int
-    retry_delay: float
-    max_diff_size: int
-
-
-# Setup logging
-logging.basicConfig(
-    level=logging.INFO,
-    format="[PR Reviewer] %(asctime)s - %(levelname)s - %(message)s",
-    datefmt="%Y-%m-%d %H:%M:%S",
-)
-logger: logging.Logger = logging.getLogger(__name__)
-
-
-def get_config() -> Config:
-    """Load configuration from environment variables with defaults."""
-    return Config(
-        glm_api_url=os.getenv("GLM_API_URL", DEFAULT_GLM_API_URL),
-        glm_model=os.getenv("GLM_MODEL", DEFAULT_GLM_MODEL),
-        cooldown_minutes=int(
-            os.getenv(
-                "PR_REVIEWER_COOLDOWN_MINUTES",
-                str(DEFAULT_COOLDOWN_MINUTES),
-            )
-        ),
-        max_retries=int(
-            os.getenv(
-                "PR_REVIEWER_MAX_RETRIES",
-                str(DEFAULT_MAX_RETRIES),
-            )
-        ),
-        retry_delay=float(os.getenv("PR_REVIEWER_RETRY_DELAY", str(DEFAULT_RETRY_DELAY))),
-        max_diff_size=int(os.getenv("PR_REVIEWER_MAX_DIFF_SIZE", str(MAX_DIFF_SIZE))),  # noqa: E501
-    )
-
-
-def get_env_var(name: str) -> str:
-    """
-    Get an environment variable or raise an error.
-
-    Args:
-        name: Name of the environment variable.
-
-    Returns:
-        The value of the environment variable.
-
-    Raises:
-        SystemExit: If the environment variable is not set.
-    """
-    value = os.environ.get(name)
-    if not value:
-        logger.error(f"Environment variable {name} is not set")
-        sys.exit(1)
-    return value
-
-
-def check_trigger(comment_body: str) -> bool:
-    """
-    Check if the comment contains the trigger phrase.
-
-    Args:
-        comment_body: The body of the comment to check.
-
-    Returns:
-        True if the trigger phrase is found, False otherwise.
-    """
-    return TRIGGER_PHRASE in comment_body
-
-
-def fetch_pr_diff(
-    repo_name: str,
-    pr_number: int,
-    token: str,
-    max_size: int = MAX_DIFF_SIZE,
-) -> str | None:
-    """
-    Fetch the diff for a pull request.
-
-    Args:
-        repo_name: The repository name in format "owner/repo".
-        pr_number: The pull request number.
-        token: GitHub authentication token.
-        max_size: Maximum diff size in characters.
-
-    Returns:
-        The diff content as a string, or None if fetching failed.
-        Returns empty string if diff is larger than max_size.
-    """
-    url: str = f"https://api.github.com/repos/{repo_name}/pulls/{pr_number}"
-    headers: dict[str, str] = {
-        "Authorization": f"Bearer {token}",
-        "Accept": "application/vnd.github.v3.diff",
-    }
-
-    logger.info(f"Fetching PR diff from {url}")
-    response = requests.get(url, headers=headers, timeout=30)
-
-    if response.status_code == 200:
-        diff: str = response.text
-        if len(diff) > max_size:
-            logger.warning(
-                f"Diff size ({len(diff)} bytes) exceeds maximum "
-                f"({max_size} bytes), truncating to first "
-                f"{max_size} bytes"
-            )
-            return diff[:max_size] + "\n\n... [Diff truncated due to size] ..."
-        logger.info(f"Successfully fetched diff ({len(diff)} bytes)")
-        return diff
-    else:
-        logger.error(f"Failed to fetch PR diff: {response.status_code}")
-        logger.error(f"Response: {response.text}")
-        return None
-
-
-def fetch_pr_details(
-    repo_name: str,
-    pr_number: int,
-    token: str,
-) -> PRDetails | None:
-    """
-    Fetch PR details including title and description.
-
-    Args:
-        repo_name: The repository name in format "owner/repo".
-        pr_number: The pull request number.
-        token: GitHub authentication token.
-
-    Returns:
-        A dictionary containing PR details, or None if fetching failed.
-    """
-    url: str = f"https://api.github.com/repos/{repo_name}/pulls/{pr_number}"
-    headers: dict[str, str] = {
-        "Authorization": f"Bearer {token}",
-        "Accept": "application/vnd.github.v3+json",
-    }
-
-    logger.info(f"Fetching PR details from {url}")
-    response = requests.get(url, headers=headers, timeout=30)
-
-    if response.status_code == 200:
-        return response.json()
-    else:
-        logger.error(f"Failed to fetch PR details: {response.status_code}")
-        return None
-
-
-def build_review_prompt(pr_title: str, pr_description: str, diff: str) -> str:
-    """
-    Build the prompt for the GLM-4.7 API.
-
-    Args:
-        pr_title: The title of the pull request.
-        pr_description: The description/body of the pull request.
-        diff: The diff content of the pull request.
-
-    Returns:
-        The formatted prompt string for the API.
-    """
-    return f"""You are an expert code reviewer for the VLLM-Omni project. \
-Please review the following pull request:
-
-## Pull Request Details
-**Title:** {pr_title}
-
-**Description:**
-{pr_description if pr_description else "No description provided."}
-
-## Code Changes (Diff)
-{diff}
-
-## Review Guidelines
-
-Please provide a comprehensive code review with the following sections:
-
-### 1. Overview
-- Brief summary of the changes
-- Overall assessment (positive, neutral, or concerns)
-
-### 2. Code Quality
-- Code style and consistency
-- Potential bugs or edge cases
-- Performance considerations
-- Error handling
-
-### 3. Architecture & Design
-- Integration with existing codebase
-- Design patterns and best practices
-- Potential improvements
-
-### 4. Security & Safety
-- Security concerns (if any)
-- Resource management
-- Input validation
-
-### 5. Testing & Documentation
-- Test coverage considerations
-- Documentation completeness
-- Examples and usage clarity
-
-### 6. Specific Suggestions
-- Line-by-line specific feedback (use `file:line` format)
-- Concrete actionable suggestions
-- Code examples for improvements (if applicable)
-
-### 7. Approval Status
-- **LGTM** (Looks Good To Me) if the PR is ready to merge
-- **LGTM with suggestions** if the PR is good but has minor suggestions
-- **Changes requested** if significant changes are needed
-
-## Important Notes
-- Be constructive and helpful
-- Focus on objective technical feedback
-- Acknowledge good practices when you see them
-- Prioritize critical issues over nitpicks
-- If the diff is empty or minimal, acknowledge this and provide
-  any relevant context-specific guidance
-
-Please format your response in Markdown with clear section headers.
-"""
-
-
-def validate_glm_response(data: dict[str, Any]) -> str | None:
-    """
-    Validate and extract content from GLM API response.
-
-    Args:
-        data: The response data from GLM API.
-
-    Returns:
-        The review content string if valid, None otherwise.
-    """
-    # Check if choices exists and is a non-empty list
-    if "choices" not in data:
-        logger.error("GLM API response missing 'choices' field")
-        logger.error(f"Response structure: {json.dumps(data, indent=2)}")
-        return None
-
-    choices = data["choices"]
-    if not isinstance(choices, list):
-        logger.error(f"GLM API 'choices' is not a list: {type(choices)}")
-        return None
-
-    if len(choices) == 0:
-        logger.error("GLM API 'choices' is an empty list")
-        return None
-
-    # Check if first choice has message
-    try:
-        first_choice = choices[0]
-        if not isinstance(first_choice, dict):
-            logger.error(f"GLM API choice is not a dict: {type(first_choice)}")
-            return None
-
-        if "message" not in first_choice:
-            logger.error("GLM API choice missing 'message' field")
-            logger.error(f"Choice structure: {json.dumps(first_choice, indent=2)}")  # noqa: E501
-            return None
-
-        message = first_choice["message"]
-        if not isinstance(message, dict):
-            logger.error(f"GLM API message is not a dict: {type(message)}")
-            return None
-
-        if "content" not in message:
-            logger.error("GLM API message missing 'content' field")
-            logger.error(f"Message structure: {json.dumps(message, indent=2)}")
-            return None
-
-        content = message["content"]
-        if not isinstance(content, str):
-            logger.error(f"GLM API content is not a string: {type(content)}")
-            return None
-
-        return content
-
-    except (KeyError, IndexError, TypeError) as e:
-        logger.error(f"Failed to parse GLM API response: {e}")
-        logger.error(f"Response: {json.dumps(data, indent=2)}")
-        return None
-
-
-def call_glm_api(prompt: str, api_key: str, config: Config) -> str | None:
-    """
-    Call the GLM-4.7 API to get code review with retry logic.
-
-    Args:
-        prompt: The prompt to send to the API.
-        api_key: The GLM API key.
-        config: Configuration object.
-
-    Returns:
-        The review content as a string, or None if all retries failed.
-    """
-    headers: dict[str, str] = {
-        "Authorization": f"Bearer {api_key}",
-        "Content-Type": "application/json",
-    }
-
-    payload: dict[str, Any] = {
-        "model": config.glm_model,
-        "messages": [{"role": "user", "content": prompt}],
-        "temperature": 0.3,
-        "max_tokens": 32000,
-        "top_p": 0.9,
-    }
-
-    last_error: str | None = None
-
-    for attempt in range(config.max_retries):
-        try:
-            logger.info(f"Calling GLM API ({config.glm_model}) - Attempt {attempt + 1}/{config.max_retries}")
-            response = requests.post(
-                config.glm_api_url,
-                headers=headers,
-                json=payload,
-                timeout=120,
-            )
-
-            if response.status_code == 200:
-                data = response.json()
-                review = validate_glm_response(data)
-                if review:
-                    logger.info(f"Successfully received review ({len(review)} chars)")  # noqa: E501
-                    return review
-                else:
-                    last_error = "Failed to validate API response structure"
-                    logger.error(last_error)
-            else:
-                last_error = f"GLM API request failed: {response.status_code} - {response.text}"
-                logger.error(last_error)
-
-        except requests.exceptions.Timeout:
-            last_error = f"GLM API request timed out (attempt {attempt + 1})"
-            logger.error(last_error)
-        except requests.exceptions.RequestException as e:
-            last_error = f"GLM API request exception: {e}"
-            logger.error(last_error)
-        except json.JSONDecodeError as e:
-            last_error = f"Failed to decode GLM API response as JSON: {e}"
-            logger.error(last_error)
-
-        # Exponential backoff before retry
-        if attempt < config.max_retries - 1:
-            wait_time: float = config.retry_delay * (2**attempt)
-            logger.info(f"Waiting {wait_time}s before retry...")  # noqa: E501
-            time.sleep(wait_time)
-
-    logger.error(
-        f"All {config.max_retries} attempts failed. Last error: {last_error}"  # noqa: E501
-    )
-    return None
-
-
-def check_cooldown(  # noqa: E501
-    repo_name: str,
-    pr_number: int,
-    token: str,
-    cooldown_minutes: int,
-) -> bool:
-    """
-    Check if the PR is within the cooldown period.
-
-    Args:
-        repo_name: The repository name in format "owner/repo".
-        pr_number: The pull request number.
-        token: GitHub authentication token.
-        cooldown_minutes: Cooldown period in minutes.
-
-    Returns:
-        True if within cooldown period (should skip), False otherwise.
-    """
-    from datetime import datetime, timedelta
-
-    url: str = (
-        f"https://api.github.com/repos/{repo_name}/issues/"
-        f"{pr_number}/comments"  # noqa: E501
-    )
-    headers: dict[str, str] = {
-        "Authorization": f"Bearer {token}",
-        "Accept": "application/vnd.github.v3+json",
-    }
-
-    logger.info(f"Checking cooldown period ({cooldown_minutes} minutes)")
-    response = requests.get(url, headers=headers, timeout=30)
-
-    if response.status_code != 200:
-        logger.warning(f"Failed to check cooldown: {response.status_code}, proceeding with review")
-        return False
-
-    comments: list[dict[str, Any]] = response.json()
-    cutoff_time: datetime = datetime.utcnow() - timedelta(minutes=cooldown_minutes)  # noqa: E501
-
-    for comment in reversed(comments):
-        # Check if this is a bot comment
-        body: str = comment.get("body", "")
-        if "VLLM-Omni PR Review" in body or "PR Reviewer Bot" in body:
-            created_at_str: str = comment.get("created_at", "")
-            try:
-                # Parse GitHub timestamp format
-                created_at = datetime.fromisoformat(created_at_str.replace("Z", "+00:00"))
-                created_at = created_at.replace(tzinfo=None)
-                if created_at > cutoff_time:
-                    logger.info(f"PR is within cooldown period (last review: {created_at_str})")
-                    return True
-            except ValueError:
-                logger.warning(f"Failed to parse comment timestamp: {created_at_str}")  # noqa: E501
-                continue
-
-    logger.info("PR is outside cooldown period, proceeding with review")
-    return False
-
-
-def post_review_comment(  # noqa: E501
-    repo_name: str,
-    pr_number: int,
-    token: str,
-    review: str,
-) -> bool:
-    """
-    Post the review as a comment on the PR.
-
-    Args:
-        repo_name: The repository name in format "owner/repo".
-        pr_number: The pull request number.
-        token: GitHub authentication token.
-        review: The review content to post.
-
-    Returns:
-        True if posting succeeded, False otherwise.
-    """
-    url: str = (
-        f"https://api.github.com/repos/{repo_name}/issues/"
-        f"{pr_number}/comments"  # noqa: E501
-    )
-    headers: dict[str, str] = {
-        "Authorization": f"Bearer {token}",
-        "Accept": "application/vnd.github.v3+json",
-    }
-
-    # Format the review comment
-    comment_body: str = f"""## 🤖 VLLM-Omni PR Review
-
-{review}
-
----
-*This review was generated automatically by the VLLM-Omni PR Reviewer Bot
-using {os.getenv("GLM_MODEL", DEFAULT_GLM_MODEL)}.*
-"""
-
-    payload: dict[str, str] = {"body": comment_body}
-
-    logger.info(f"Posting review comment to PR #{pr_number}")
-    response = requests.post(url, headers=headers, json=payload, timeout=30)
-
-    if response.status_code == 201:
-        logger.info("Successfully posted review comment")
-        return True
-    else:
-        logger.error(f"Failed to post comment: {response.status_code}")
-        logger.error(f"Response: {response.text}")
-        return False
-
-
-def main() -> int:
-    """
-    Main entry point for the PR reviewer bot.
-
-    Returns:
-        0 on success, 1 on error.
-    """
-    logger.info("VLLM-Omni PR Reviewer Bot starting...")
-
-    # Load configuration
-    config: Config = get_config()
-    logger.info(
-        f"Configuration: model={config.glm_model}, "
-        f"cooldown={config.cooldown_minutes}min, "
-        f"max_retries={config.max_retries}"
-    )
-
-    # Get environment variables
-    token: str = get_env_var("GITHUB_TOKEN")
-    api_key: str = get_env_var("GLM_API_KEY")
-    repo_name: str = get_env_var("REPO_NAME")
-    pr_number_str: str = get_env_var("PR_NUMBER")
-    comment_body: str = get_env_var("COMMENT_BODY")
-
-    try:
-        pr_number: int = int(pr_number_str)
-    except ValueError:
-        logger.error(f"Invalid PR number: {pr_number_str}")
-        return 1
-
-    logger.info(f"Repository: {repo_name}")
-    logger.info(f"PR Number: {pr_number}")
-
-    # Check if the comment contains the trigger phrase
-    if not check_trigger(comment_body):
-        logger.info(
-            f"Comment does not contain trigger phrase '{TRIGGER_PHRASE}', exiting"  # noqa: E501
-        )
-        return 0
-
-    logger.info("Trigger phrase detected! Starting review process...")
-
-    # Check cooldown period
-    if check_cooldown(repo_name, pr_number, token, config.cooldown_minutes):
-        logger.info("Skipping review due to cooldown period")
-        return 0
-
-    # Fetch PR details
-    logger.info("Step 1/4: Fetching PR details...")
-    pr_details: PRDetails | None = fetch_pr_details(repo_name, pr_number, token)  # noqa: E501
-    if not pr_details:
-        logger.error("Failed to fetch PR details")
-        return 1
-
-    pr_title: str = pr_details.get("title", "Unknown")
-    pr_description: str = pr_details.get("body", "")
-
-    logger.info(f"PR Title: {pr_title}")
-
-    # Fetch PR diff
-    logger.info("Step 2/4: Fetching PR diff...")
-    diff: str | None = fetch_pr_diff(repo_name, pr_number, token, config.max_diff_size)
-    if diff is None:
-        logger.error("Failed to fetch PR diff")
-        return 1
-
-    if not diff:
-        logger.warning("Warning: Empty diff - this might be a draft PR or no code changes")
-
-    # Build prompt
-    logger.info("Step 3/4: Building review prompt...")
-    prompt: str = build_review_prompt(pr_title, pr_description, diff)
-
-    # Call GLM API
-    logger.info("Step 4/4: Calling GLM API...")
-    review: str | None = call_glm_api(prompt, api_key, config)
-    if not review:
-        logger.error("Failed to get review from GLM API")
-        return 1
-
-    # Post review comment
-    logger.info("Posting review comment...")
-    if not post_review_comment(repo_name, pr_number, token, review):
-        logger.error("Failed to post review comment")
-        return 1
-
-    logger.info("PR review completed successfully!")
-    return 0
-
-
-if __name__ == "__main__":
-    sys.exit(main())

.github/workflows/pr-reviewer.yml

@@ -1,62 +0,0 @@
-name: VLLM-Omni PR Reviewer
-
-on:
-  issue_comment:
-    types: [created]
-
-permissions:
-  contents: read
-  pull-requests: write
-  issues: write
-
-jobs:
-  pr-reviewer:
-    name: Review Pull Request
-    runs-on: ubuntu-latest
-    timeout-minutes: 10
-    # Only run when the comment is from a collaborator/owner/member
-    if: |
-      github.event_name == 'issue_comment' &&
-      github.event.issue.pull_request != null &&
-      contains(github.event.comment.body, '@vllm-omni-reviewer') &&
-      (github.event.comment.author_association == 'MEMBER' ||
-       github.event.comment.author_association == 'COLLABORATOR' ||
-       github.event.comment.author_association == 'OWNER')
-    steps:
-      - name: Checkout repository
-        uses: actions/checkout@v6.0.2
-        with:
-          fetch-depth: 0
-
-      - name: Set up Python
-        uses: actions/setup-python@v6.2.0
-        with:
-          python-version: '3.11'
-          cache: 'pip'
-
-      - name: Install dependencies
-        run: |
-          python -m pip install --upgrade pip==24.0
-          pip install requests==2.31.0 pyyaml==6.0.1
-
-      - name: Run PR Reviewer
-        id: reviewer
-        env:
-          GLM_API_KEY: ${{ secrets.GLM_API_KEY }}
-          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-          PR_NUMBER: ${{ github.event.issue.number || github.event.pull_request.number }}
-          COMMENT_BODY: ${{ github.event.comment.body }}
-          REPO_NAME: ${{ github.repository }}
-          PR_REVIEWER_COOLDOWN_MINUTES: 5
-          PR_REVIEWER_MAX_RETRIES: 3
-        run: |
-          python .github/scripts/pr_reviewer.py 2>&1 | tee "/tmp/pr_review_${PR_NUMBER}.log"
-
-      - name: Upload review logs
-        if: always()
-        uses: actions/upload-artifact@v7.0.0
-        with:
-          name: pr-review-logs-${{ github.event.issue.number || github.event.pull_request.number }}
-          path: /tmp/pr_review_*.log
-          retention-days: 7
-          if-no-files-found: ignore

.gitignore

@@ -83,6 +83,9 @@ target/
 profile_default/
 ipython_config.py
 
+# uv
+uv.lock
+
 # pyenv
 #   For a library or package, you might want to ignore these files since the code is
 #   intended to run in multiple environments; otherwise, check them in:
@@ -155,10 +158,23 @@ cython_debug/
 
 # Claude
 CLAUDE.md
-.claude/
+/.claude/*
+!.claude/skills/
+!.claude/skills/readme.md
+!.claude/skills/add-diffusion-model/
+!.claude/skills/add-diffusion-model/SKILL.md
+!.claude/skills/add-diffusion-model/references/
+!.claude/skills/add-diffusion-model/references/*.md
+!.claude/skills/add-tts-model/
+!.claude/skills/add-tts-model/SKILL.md
+!.claude/skills/review-pr/
+!.claude/skills/review-pr/SKILL.md
+!.claude/skills/review-pr/references/
+!.claude/skills/review-pr/references/*.md
 
 # Codex
 AGENTS.md
+.codex
 .codex/
 
 # cursor
@@ -188,6 +204,7 @@ checkpoints/
 # Cache directories
 cache/
 !vllm_omni/diffusion/cache/
+!tests/diffusion/cache/
 .cache/
 diffusion_cache/
 kv_cache/
@@ -222,6 +239,7 @@ data/
 datasets/
 *.csv
 *.json
+!tests/dfx/perf/tests/*.json
 !apps/ComfyUI-vLLM-Omni/example_workflows/*.json
 *.jsonl
 *.parquet
@@ -247,3 +265,8 @@ tmp_test
 vllm_omni/_version.py
 # output files
 *.wav
+# Vendored test reference audio (small, intentionally checked in)
+!tests/assets/**/*.wav
+.worktrees/
+# CI overlay yamls materialized from tests/utils.py:_CI_OVERLAYS at test time
+tests/.ci_generated/

README.md

@@ -9,19 +9,19 @@ Easy, fast, and cheap omni-modality model serving for everyone
 </h3>
 
 <p align="center">
-| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> | <a href="docs/assets/WeChat.jpg"><b>WeChat</b></a> | <a href="https://arxiv.org/abs/2602.02204"><b>Paper</b></a> | <a href="https://docs.google.com/presentation/d/1XJWgv79lORl8rbaVvp2d5Sqs6ZEBgAgj/edit?slide=id.p1#slide=id.p1"><b>Slides</b></a> |
+| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> | <a href="docs/assets/WeChat.jpg"><b>WeChat</b></a> | <a href="https://arxiv.org/abs/2602.02204"><b>Paper</b></a> | <a href="https://docs.google.com/presentation/d/111-L8zF7A1j_YI_cR8JsblofdScdRr2f/edit?usp=sharing&ouid=110473603432222024453&rtpof=true&sd=true"><b>Slides</b></a> |
 </p>
 
 
 ---
 
 *Latest News* 🔥
-
+- [2026/05] We released [0.20.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.20.0) - refreshes the serving/runtime stack for large-scale omni workloads, and improves diffusion model performance, quantization, and hardware readiness across CUDA, ROCm, MUSA, NPU, and XPU backends.
+- [2026/03] We released [0.18.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.18.0) - strengthens the core runtime through a large entrypoint refactor and scheduler/runtime cleanups, expands unified quantization and diffusion execution, broadens multimodal model coverage, and improves production readiness across audio, omni, image, video, RL, and multi-platform deployments.
 - [2026/03] Check out our first public [project deepdive](https://youtu.be/sgwNfsNnR9I) at the vLLM Hong Kong Meetup!
 - [2026/03] **[vllm-omni-skills](https://github.com/hsliuustc0106/vllm-omni-skills)** is a community-driven collection of AI assistant skills that help developers work with vLLM-Omni more effectively. These skills can be used with popular agentic AI coding assistants like **Cursor IDE**, **Claude**, **Codex**, and more.
 - [2026/02] We released [0.16.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.16.0) - A major alignment + capability release that rebases onto **upstream vLLM v0.16.0** and significantly expands performance, distributed execution, and production readiness across **Qwen3-Omni / Qwen3-TTS**, **Bagel**, **MiMo-Audio**, **GLM-Image** and the **Diffusion (DiT) image/video stack**—while also improving platform coverage (CUDA / ROCm / NPU / XPU), CI quality, and documentation.
 - [2026/02] We released [0.14.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.14.0) - This is the first **stable release** of vLLM-Omni that expands Omni’s diffusion / image-video generation and audio / TTS stack, improves distributed execution and memory efficiency, and broadens platform/backend coverage (GPU/ROCm/NPU/XPU). It also brings meaningful upgrades to serving APIs, profiling & benchmarking, and overall stability. Please check our latest [paper](https://arxiv.org/abs/2602.02204) for architecture design and performance results.
-- [2026/01] We released [0.12.0rc1](https://github.com/vllm-project/vllm-omni/releases/tag/v0.12.0rc1) - a major RC milestone focused on maturing the diffusion stack, strengthening OpenAI-compatible serving, expanding omni-model coverage, and improving stability across platforms (GPU/NPU/ROCm).
 - [2025/11] vLLM community officially released [vllm-project/vllm-omni](https://github.com/vllm-project/vllm-omni) in order to support omni-modality models serving.
 
 ---

benchmarks/accuracy/README.md

@@ -23,5 +23,5 @@ Test guidance:
 - Local static/self-checks live in `tests/benchmarks/test_accuracy_bench_utils.py`.
 - End-to-end generation/evaluation should be validated in a remote GPU
   environment. In the current repo marker system there is `L4` but no `L5`
-  marker, so benchmark smoke tests should be wired as `advanced_model +
-  benchmark + L4` when GPU capacity is available.
+  marker, so benchmark smoke tests should be wired as `full_model +
+  benchmark + L4` for nightly when GPU capacity is available.

benchmarks/accuracy/image_to_image/README.md

@@ -99,5 +99,5 @@ Notes:
 - This flow requires the optional Hugging Face `datasets` package.
 - `generate` writes `generation_manifest.json` with local output coverage.
 - The current repo marker set exposes `L4` but not `L5`, so if you promote an
-  end-to-end smoke test into CI, use the existing `advanced_model`, `benchmark`,
-  and `L4` markers or introduce a new repo-wide marker explicitly first.
+  end-to-end smoke test into CI, use the `full_model`, `benchmark`,
+  and `L4` markers for nightly (or `advanced_model` for merge) or introduce a new repo-wide marker explicitly first.

benchmarks/build_dataset/download_process_data_seedtts.md

@@ -27,7 +27,7 @@ pip install gdown
 Download the dataset from Google Drive:
 
 ```bash
-gdown --id 1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP
+gdown 1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP
 ```
 
 ### 4. Extract the Dataset
@@ -74,7 +74,7 @@ rm meta.lst
 # Full setup and benchmark
 cd benchmarks/build_dataset
 pip install gdown
-gdown --id 1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP
+gdown  1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP
 tar -xf seedtts_testset.tar
 cp seedtts_testset/en/meta.lst meta.lst
 python extract_tts_prompts.py -i meta.lst -o top100.txt -n 100

benchmarks/build_dataset/seed_tts_design/en/meta.lst

@@ -0,0 +1,20 @@
+vd001|||The quick brown fox jumps over the lazy dog.|A warm, friendly female voice with a slight American Midwest accent, speaking at a moderate pace with natural inflection.
+vd002|||Welcome to the future of text-to-speech synthesis.|A deep, authoritative male news anchor voice, clear and professional with a measured cadence.
+vd003|||The sunset painted the sky in brilliant shades of orange and pink.|A gentle elderly female voice, soft and wise, with a slight Southern American accent.
+vd004|||Scientists have discovered a new species of deep-sea creature.|A young male voice with an Australian accent, curious and enthusiastic.
+vd005|||Breaking news: a major climate summit opens today in Geneva.|A crisp female newsreader voice, neutral accent, confident and precise.
+vd006|||In the beginning, there was darkness and silence across the universe.|A rich, dramatic bass male narrator voice, slow and deeply resonant.
+vd007|||Come closer, I have something important to tell you.|A soft, intimate female voice, slightly whispery, warm and gentle.
+vd008|||And they're off! The horses race toward the first turn at incredible speed.|An energetic male sports commentator, fast-paced and excited.
+vd009|||Once upon a time, in a land far away, lived a very clever fox.|A light, playful voice with childlike enthusiasm, bright and clear.
+vd010|||The ancient manuscript reveals secrets hidden for a thousand years.|A wise, measured elderly male voice, slow and deliberate, British English accent.
+vd011|||Good evening, ladies and gentlemen, and welcome to our show.|A sophisticated female voice with a slight French accent speaking English, elegant and refined.
+vd012|||System initialized. Running diagnostics. All systems nominal.|A clear, precise robotic-sounding voice, neutral and monotone with slight synthetic quality.
+vd013|||I hear what you are saying, and it is completely understandable to feel that way.|A warm, empathetic female therapist voice, calm and reassuring, unhurried pace.
+vd014|||Attention all units: proceed to grid reference seven-seven-alpha.|A firm, authoritative military male voice, clipped and commanding.
+vd015|||Oh my goodness, you have to try this amazing new recipe I just found!|An enthusiastic, bubbly female voice, high energy and friendly.
+vd016|||Dude, the waves were totally amazing out there today. Super happy about it!|A relaxed male voice with a California accent, casual and laid-back.
+vd017|||The quarterly results exceed expectations across all major metrics.|A sharp, businesslike female voice, confident and efficient, fast-paced delivery.
+vd018|||Chapter one. The morning sun filtered gently through the forest canopy.|A smooth, rich male audiobook narrator voice, expressive and engaging.
+vd019|||To be or not to be, that is the question.|A theatrical female voice, dramatic and expressive, stage projection quality.
+vd020|||And that is all for tonight. Stay well out there, everyone.|A warm, velvety male late-night radio DJ voice, smooth and intimate.

benchmarks/build_dataset/seed_tts_smoke/en/meta.lst

@@ -0,0 +1,20 @@
+smoke001|||The quick brown fox jumps over the lazy dog near the riverbank at sunset.
+smoke002|||Welcome to the future of text-to-speech synthesis in production systems.
+smoke003|||Yesterday the team finished rolling out the new authentication flow.
+smoke004|||She walked carefully across the wet cobblestones, careful not to slip.
+smoke005|||The conference call is scheduled for nine in the morning, Pacific time.
+smoke006|||Please remember to save your work before closing the editor.
+smoke007|||Two plus two equals four, but five hundred and forty three digits is long.
+smoke008|||I would like a coffee with oat milk and a chocolate croissant please.
+smoke009|||The library closes at eight on weekdays and six on Saturdays.
+smoke010|||During the Renaissance, art and science flourished in European cities.
+smoke011|||He whispered the secret word so quietly that no one else could hear.
+smoke012|||Our flight departs from gate twenty three at eleven fifteen.
+smoke013|||The storm knocked out power for six hours, but the backup generator kicked in.
+smoke014|||Reading a good book on a rainy afternoon is one of life's great pleasures.
+smoke015|||When the kettle whistled, she poured the hot water over the fresh tea leaves.
+smoke016|||The algorithm runs in linear time, which is a big improvement over the previous approach.
+smoke017|||In the distance, the mountains were shrouded in thick morning fog.
+smoke018|||Our company reported record revenue for the fourth quarter of the fiscal year.
+smoke019|||She explained the new policy in detail during the staff meeting this morning.
+smoke020|||The children laughed and played in the garden until the sun began to set.

benchmarks/diffusion/README.md

@@ -130,9 +130,21 @@ Warmup flags:
 
 - `--warmup-requests`: Number of warmup requests.
 - `--warmup-num-inference-steps`: Steps used during warmup.
+- `--warmup-concurrency`: Maximum concurrent warmup requests. Use this to warm
+  the same batch shape as the measured run instead of warming only batch=`1`.
 - For `--task t2v`: warmup requests are forced to use `num_frames=1` to make warmup faster and less noisy.
 
 Traffic / concurrency flags:
 
 - `--request-rate`: Target request rate (requests/second). If set to `inf`, the script sends all requests immediately.
 - `--max-concurrency`: Max number of in-flight requests (default: `1`). This can hard-cap the achieved QPS: if it is too small, requests will queue behind the semaphore, and both achieved throughput and observed SLO attainment can be skewed.
+
+### Batched warmup note
+
+For batched serving runs, warm the same in-flight shape you plan to measure.
+For example, a run with `--max-concurrency 8` should usually also use
+`--warmup-requests 8 --warmup-concurrency 8`; otherwise the first measured
+batch may still pay compile or CUDA-graph capture cost.
+
+For a Qwen-Image continuous-batching replay example, see
+[`performance_dashboard/qwen_image_serving_performance.md`](./performance_dashboard/qwen_image_serving_performance.md).

benchmarks/diffusion/backends.py

@@ -122,6 +122,18 @@ async def async_request_chat_completions(
                                     output.peak_memory_mb = first_item.get("peak_memory_mb", 0.0)
                 except (IndexError, TypeError, AttributeError):
                     pass
+
+                if (not output.stage_durations or output.peak_memory_mb == 0.0) and isinstance(
+                    resp_json.get("metrics"), dict
+                ):
+                    m = resp_json["metrics"]
+                    if not output.stage_durations and isinstance(m.get("stage_durations"), dict):
+                        output.stage_durations = m.get("stage_durations") or {}
+                    if output.peak_memory_mb == 0.0 and m.get("peak_memory_mb") is not None:
+                        try:
+                            output.peak_memory_mb = float(m.get("peak_memory_mb") or 0.0)
+                        except (TypeError, ValueError):
+                            pass
             else:
                 output.error = f"HTTP {response.status}: {await response.text()}"
                 output.success = False
@@ -306,6 +318,8 @@ async def async_request_v1_videos(
             video_bytes = await content_response.read()
             output.response_body = video_bytes
             output.success = True
+            if "stage_durations" in poll_json:
+                output.stage_durations = poll_json["stage_durations"] or {}
             if "peak_memory_mb" in poll_json:
                 output.peak_memory_mb = poll_json["peak_memory_mb"]
             elif "peak_memory_mb" in resp_json:

benchmarks/diffusion/diffusion_benchmark_serving.py

@@ -12,15 +12,15 @@ Supports multiple backends:
     - v1/videos: Use /v1/videos endpoint
 
 Usage:
-    # Video (vllm-omni backend)
+    # Video (v1/videos backend)
     t2v:
     python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
-        --backend vllm-omni --dataset vbench --task t2v --num-prompts 10 \
+        --backend v1/videos --dataset vbench --task t2v --num-prompts 10 \
         --height 480 --width 640 --fps 16 --num-frames 80
 
     i2v:
     python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
-        --backend vllm-omni --dataset vbench --task i2v --num-prompts 10
+        --backend v1/videos --dataset vbench --task i2v --num-prompts 10
 
 
     # Image (vllm-omni backend)
@@ -49,7 +49,7 @@ Usage:
         --backend openai --dataset vbench --task t2i --num-prompts 10 \
         --height 1024 --width 1024 --port 3000
 
-    # Video (v1/vedeos)
+    # Video (v1/videos)
     t2v:
     python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
         --backend v1/videos --dataset random --task t2v --num-prompts 1 \
@@ -558,6 +558,7 @@ class RandomDataset(BaseDataset):
         super().__init__(args, api_url, model)
         self.num_prompts = args.num_prompts
         self.enable_negative_prompt = enable_negative_prompt
+        self.num_input_images = max(1, args.num_input_images)
         self.random_request_config = getattr(args, "random_request_config", None)
         if self.random_request_config:
             self.random_request_config = json.loads(self.random_request_config)
@@ -580,11 +581,7 @@ class RandomDataset(BaseDataset):
 
         # Random image generate
         if self.args.task in ["i2v", "ti2v", "ti2i", "i2i"]:
-            img = Image.new("RGB", (512, 512), (255, 255, 255))
-
-            image_path = os.path.join(tempfile.gettempdir(), "diffusion_benchmark_random_image.png")
-            self._random_image_path = [image_path]
-            img.save(image_path)
+            self._random_image_path = self._generate_random_image_paths()
         else:
             self._random_image_path = None
 
@@ -619,6 +616,18 @@ class RandomDataset(BaseDataset):
     def get_requests(self) -> list[RequestFuncInput]:
         return [self[i] for i in range(len(self))]
 
+    def _generate_random_image_paths(self) -> list[str]:
+        image_paths: list[str] = []
+        for image_idx in range(self.num_input_images):
+            img = Image.new("RGB", (512, 512), (255, 255, 255))
+            image_path = os.path.join(
+                tempfile.gettempdir(),
+                f"diffusion_benchmark_random_image_{image_idx}.png",
+            )
+            img.save(image_path)
+            image_paths.append(image_path)
+        return image_paths
+
 
 def _compute_expected_latency_ms_from_base(req: RequestFuncInput, args, base_time_ms: float | None) -> float | None:
     """Compute expected execution time (ms) based on a base per-step-per-frame unit time.
@@ -742,6 +751,52 @@ async def iter_requests(
         yield req
 
 
+def _make_warmup_request(
+    requests_list: list[RequestFuncInput],
+    index: int,
+    args,
+) -> RequestFuncInput:
+    warm_req = requests_list[index % len(requests_list)]
+    if args.warmup_num_inference_steps is not None:
+        warm_req = replace(
+            warm_req,
+            num_inference_steps=args.warmup_num_inference_steps,
+        )
+    if args.task == "t2v":
+        warm_req = replace(warm_req, num_frames=1)
+    return warm_req
+
+
+async def _run_warmups(
+    requests_list: list[RequestFuncInput],
+    args,
+    session: aiohttp.ClientSession,
+    request_func,
+) -> list[tuple[RequestFuncInput, RequestFuncOutput]]:
+    if not args.warmup_requests or not requests_list:
+        return []
+
+    warmup_requests = [_make_warmup_request(requests_list, i, args) for i in range(args.warmup_requests)]
+    warmup_concurrency = min(int(args.warmup_concurrency), len(warmup_requests))
+    warmup_semaphore = asyncio.Semaphore(warmup_concurrency)
+
+    print(
+        f"Running {len(warmup_requests)} warmup request(s) "
+        f"with num_inference_steps={args.warmup_num_inference_steps} "
+        f"and warmup_concurrency={warmup_concurrency}..."
+    )
+
+    async def limited_warmup_request_func(
+        req: RequestFuncInput,
+    ) -> RequestFuncOutput:
+        async with warmup_semaphore:
+            return await request_func(req, session, None)
+
+    warmup_tasks = [asyncio.create_task(limited_warmup_request_func(req)) for req in warmup_requests]
+    warmup_outputs = await asyncio.gather(*warmup_tasks)
+    return list(zip(warmup_requests, warmup_outputs))
+
+
 def calculate_metrics(
     outputs: list[RequestFuncOutput],
     total_duration: float,
@@ -891,23 +946,12 @@ async def benchmark(args):
     pbar = tqdm(total=len(requests_list), disable=args.disable_tqdm)
 
     async with aiohttp.ClientSession() as session:
-        warmup_pairs: list[tuple[RequestFuncInput, RequestFuncOutput]] = []
-        if args.warmup_requests and requests_list:
-            print(
-                f"Running {args.warmup_requests} warmup request(s) \
-                with num_inference_steps={args.warmup_num_inference_steps}..."
-            )
-            for i in range(args.warmup_requests):
-                warm_req = requests_list[i % len(requests_list)]
-                if args.warmup_num_inference_steps is not None:
-                    warm_req = replace(
-                        warm_req,
-                        num_inference_steps=args.warmup_num_inference_steps,
-                    )
-                if args.task == "t2v":
-                    warm_req = replace(warm_req, num_frames=1)
-                warm_out = await limited_request_func(warm_req, session, None)
-                warmup_pairs.append((warm_req, warm_out))
+        warmup_pairs = await _run_warmups(
+            requests_list=requests_list,
+            args=args,
+            session=session,
+            request_func=request_func,
+        )
 
         if args.slo:
             # Prefer trace-provided per-request slo_ms. Only populate when missing.
@@ -1063,6 +1107,13 @@ if __name__ == "__main__":
         default=1,
         help="num_inference_steps used for warmup requests.",
     )
+    parser.add_argument(
+        "--warmup-concurrency",
+        type=int,
+        default=1,
+        help="Maximum number of warmup requests to run concurrently. "
+        "Set this to match the real batch shape when warming up torch.compile or CUDA graphs.",
+    )
     parser.add_argument("--width", type=int, default=None, help="Image/Video width.")
     parser.add_argument("--height", type=int, default=None, help="Image/Video height.")
     parser.add_argument("--num-frames", type=int, default=None, help="Number of frames (for video).")
@@ -1115,7 +1166,15 @@ if __name__ == "__main__":
             '{"width":768,"height":768,"num_inference_steps":20,"weight":0.85}]'
         ),
     )
+    parser.add_argument(
+        "--num-input-images",
+        type=int,
+        default=1,
+        help=(
+            "Number of synthetic input images to attach for image-conditioned tasks "
+            "(i2v, ti2v, ti2i, i2i) when using random dataset."
+        ),
+    )
 
     args = parser.parse_args()
-
     asyncio.run(benchmark(args))

benchmarks/diffusion/performance_dashboard/qwen_image_serving_performance.md

@@ -33,6 +33,12 @@ vllm serve Qwen/Qwen-Image --omni \
     --port 8091
 ```
 
+To replay step-wise continuous batching, compare
+`vllm serve Qwen/Qwen-Image --omni --port 8089 --step-execution --max-num-seqs 1`
+against the same command with `--max-num-seqs 8`. `--step-execution` is the
+feature gate, and increasing `--max-num-seqs` above `1` lets the scheduler
+keep more compatible requests active at once.
+
 ## 3.2 Key Parameters
 
 | Parameter             | Description              |
@@ -70,6 +76,7 @@ python benchmarks/diffusion/diffusion_benchmark_serving.py \
 | `--task`               | Task type (e.g., `t2i`)           |
 | `--num-prompts`        | Total number of requests          |
 | `--max-concurrency`    | Client-side concurrency           |
+| `--warmup-concurrency` | Warmup concurrency for batch-shape prewarming |
 | `--random-request-config`| JSON string defining random request |
 
 ---
@@ -133,6 +140,12 @@ python benchmarks/diffusion/diffusion_benchmark_serving.py \
     ]'
 ```
 
+For a continuous-batching A/B replay, run the same benchmark command twice
+against the two serving configs above and keep traffic plus warmup settings
+identical across both runs. If the measured run uses `--max-concurrency 8`,
+warm with `--warmup-requests 8 --warmup-concurrency 8` so the first measured
+batch does not include compile or CUDA-graph capture overhead.
+
 ---
 
 # 6. Performance Metrics
@@ -163,6 +176,7 @@ To ensure consistent and comparable benchmark results:
 * Record parallel configuration
 * Record benchmark parameters (resolution, concurrency, number of prompts)
 * Ensure no background workload on GPUs during testing
+* When comparing continuous batching, keep warmup settings the same across runs
 
 ---
 

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.cuda.reset_peak_memory_stats()
+    torch.accelerator.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.cuda.max_memory_allocated() / (1024**3)
+    peak_mem = torch.accelerator.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.cuda.reset_peak_memory_stats()
+    torch.accelerator.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.cuda.max_memory_allocated() / (1024**3)
+    peak_mem = torch.accelerator.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.cuda.empty_cache()
-        torch.cuda.synchronize()
+        torch.accelerator.empty_cache()
+        torch.accelerator.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.cuda.device_count():
-            print(f"[ERROR] GPU {args.gpu_id} not available. Found {torch.cuda.device_count()} GPUs.")
+        if args.gpu_id >= torch.accelerator.device_count():
+            print(f"[ERROR] GPU {args.gpu_id} not available. Found {torch.accelerator.device_count()} GPUs.")
             sys.exit(1)
         print(f"[INFO] Using GPU {args.gpu_id}: {torch.cuda.get_device_name(args.gpu_id)}")
 

benchmarks/fish-speech/bench_speaker_cache.py

@@ -0,0 +1,290 @@
+"""Benchmark Fish Speech voice cache: inline ref_audio vs uploaded voice.
+
+Measures TTFP improvement from DAC-code caching when using uploaded voices.
+
+Setup:
+  1. Start vllm-omni with Fish Speech S2 Pro
+  2. Provide a reference audio file for voice cloning
+
+Usage:
+    python bench_speaker_cache.py \
+        --ref-audio /path/to/reference.wav \
+        --ref-text "Transcript of the reference audio." \
+        --num-prompts 20 \
+        --port 8091
+
+The script runs two rounds:
+  A) Inline ref_audio: every request sends base64 audio (no cache)
+  B) Uploaded voice: upload once, then use voice name (cache hits after 1st)
+"""
+
+import argparse
+import asyncio
+import base64
+import json
+import os
+import sys
+import time
+from pathlib import Path
+
+import aiohttp
+
+# Allow imports from benchmarks/fish-speech/
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+
+from fish_bench_utils import (  # noqa: E402
+    BenchmarkResult,
+    RequestResult,
+    compute_stats,
+    print_benchmark_results,
+    send_streaming_request,
+)
+
+SAMPLE_RATE = 44100
+SAMPLE_WIDTH = 2
+
+PROMPTS = [
+    "Hello, welcome to the voice synthesis benchmark test.",
+    "She said she would be here by noon, but nobody showed up.",
+    "The quick brown fox jumps over the lazy dog near the riverbank.",
+    "I can't believe how beautiful the sunset looks from up here.",
+    "Please remember to bring your identification documents tomorrow morning.",
+    "Have you ever wondered what it would be like to travel through time?",
+    "The restaurant on the corner serves the best pasta I have ever tasted.",
+    "After the meeting, we should discuss the quarterly results.",
+    "Learning a new language takes patience and genuine curiosity.",
+    "The train leaves at half past seven, so we need to arrive early.",
+    "Could you please turn down the music, I'm trying to concentrate.",
+    "It was a dark and stormy night when the keeper heard a knock.",
+]
+
+
+def encode_audio_to_base64(audio_path: str) -> str:
+    """Encode a local audio file to base64 data URL."""
+    ext = audio_path.lower().rsplit(".", 1)[-1]
+    mime_map = {"wav": "audio/wav", "mp3": "audio/mpeg", "flac": "audio/flac"}
+    mime_type = mime_map.get(ext, "audio/wav")
+    with open(audio_path, "rb") as f:
+        audio_b64 = base64.b64encode(f.read()).decode("utf-8")
+    return f"data:{mime_type};base64,{audio_b64}"
+
+
+async def upload_voice(
+    host: str,
+    port: int,
+    audio_path: str,
+    ref_text: str,
+    voice_name: str = "bench_voice",
+) -> dict:
+    """Upload a voice via POST /v1/audio/voices."""
+    url = f"http://{host}:{port}/v1/audio/voices"
+    data = aiohttp.FormData()
+    data.add_field("name", voice_name)
+    data.add_field("consent", "true")
+    if ref_text:
+        data.add_field("ref_text", ref_text)
+    data.add_field(
+        "audio_sample",
+        open(audio_path, "rb"),
+        filename=os.path.basename(audio_path),
+        content_type="audio/wav",
+    )
+
+    async with aiohttp.ClientSession() as session:
+        async with session.post(url, data=data) as resp:
+            result = await resp.json()
+            print(f"  Upload response ({resp.status}): {json.dumps(result, indent=2)}")
+            return result
+
+
+async def delete_voice(host: str, port: int, voice_name: str) -> None:
+    """Delete an uploaded voice."""
+    url = f"http://{host}:{port}/v1/audio/voices/{voice_name}"
+    async with aiohttp.ClientSession() as session:
+        async with session.delete(url) as resp:
+            if resp.status == 200:
+                print(f"  Deleted voice '{voice_name}'")
+
+
+async def run_round(
+    host: str,
+    port: int,
+    num_prompts: int,
+    create_payload_fn,
+    label: str,
+    num_warmups: int = 2,
+    timeout_s: float = 120.0,
+) -> BenchmarkResult:
+    """Run one benchmark round and return results."""
+    api_url = f"http://{host}:{port}/v1/audio/speech"
+    connector = aiohttp.TCPConnector(limit=1, limit_per_host=1)
+    session = aiohttp.ClientSession(
+        connector=connector,
+        timeout=aiohttp.ClientTimeout(total=timeout_s),
+    )
+
+    try:
+        # Warmup.
+        if num_warmups > 0:
+            print(f"  [{label}] Warming up ({num_warmups} requests)...")
+            for i in range(num_warmups):
+                payload = create_payload_fn(PROMPTS[i % len(PROMPTS)])
+                r = await send_streaming_request(
+                    session,
+                    api_url,
+                    payload,
+                    SAMPLE_RATE,
+                    SAMPLE_WIDTH,
+                )
+                status = "OK" if r.success else f"FAIL: {r.error[:80]}"
+                print(f"    warmup {i + 1}: ttfp={r.ttfp * 1000:.0f}ms  {status}")
+
+        # Benchmark.
+        print(f"  [{label}] Running {num_prompts} requests (concurrency=1)...")
+        results: list[RequestResult] = []
+        start = time.perf_counter()
+        for i in range(num_prompts):
+            prompt = PROMPTS[i % len(PROMPTS)]
+            payload = create_payload_fn(prompt)
+            r = await send_streaming_request(
+                session,
+                api_url,
+                payload,
+                SAMPLE_RATE,
+                SAMPLE_WIDTH,
+            )
+            results.append(r)
+            tag = "HIT" if i > 0 and label == "uploaded_voice" else ""
+            print(
+                f"    req {i + 1:3d}: ttfp={r.ttfp * 1000:7.1f}ms  "
+                f"e2e={r.e2e * 1000:7.1f}ms  "
+                f"{'OK' if r.success else 'FAIL'} {tag}"
+            )
+        wall_time = time.perf_counter() - start
+    finally:
+        await session.close()
+
+    bench = compute_stats(results, wall_time)
+    bench.concurrency = 1
+    bench.num_prompts = num_prompts
+    bench.config_name = label
+    return bench
+
+
+async def main():
+    parser = argparse.ArgumentParser(
+        description="Benchmark Fish Speech voice cache (inline vs uploaded)",
+    )
+    parser.add_argument("--host", default="127.0.0.1")
+    parser.add_argument("--port", type=int, default=8091)
+    parser.add_argument("--ref-audio", required=True, help="Path to reference audio file")
+    parser.add_argument("--ref-text", required=True, help="Transcript of reference audio")
+    parser.add_argument("--num-prompts", type=int, default=20)
+    parser.add_argument("--num-warmups", type=int, default=2)
+    parser.add_argument("--voice-name", default="bench_voice")
+    args = parser.parse_args()
+
+    if not os.path.exists(args.ref_audio):
+        print(f"Error: ref_audio not found: {args.ref_audio}")
+        sys.exit(1)
+
+    ref_audio_b64 = encode_audio_to_base64(args.ref_audio)
+    print(f"Reference audio: {args.ref_audio} ({len(ref_audio_b64) // 1024}KB base64)")
+
+    # ---- Round A: Inline ref_audio (no cache) ----
+    print(f"\n{'=' * 60}")
+    print("Round A: INLINE ref_audio (every request sends full audio)")
+    print(f"{'=' * 60}")
+
+    def make_inline_payload(prompt: str) -> dict:
+        return {
+            "input": prompt,
+            "voice": "default",
+            "stream": True,
+            "response_format": "pcm",
+            "ref_audio": ref_audio_b64,
+            "ref_text": args.ref_text,
+            "max_new_tokens": 2048,
+        }
+
+    bench_inline = await run_round(
+        args.host,
+        args.port,
+        args.num_prompts,
+        make_inline_payload,
+        "inline_ref_audio",
+        num_warmups=args.num_warmups,
+    )
+    print_benchmark_results(bench_inline)
+
+    # ---- Upload voice ----
+    print(f"\n{'=' * 60}")
+    print("Uploading voice for cache test...")
+    print(f"{'=' * 60}")
+    await delete_voice(args.host, args.port, args.voice_name)
+    await upload_voice(
+        args.host,
+        args.port,
+        args.ref_audio,
+        args.ref_text,
+        args.voice_name,
+    )
+
+    # ---- Round B: Uploaded voice (cache hits after 1st request) ----
+    print(f"\n{'=' * 60}")
+    print("Round B: UPLOADED VOICE (cache hits after 1st request)")
+    print(f"{'=' * 60}")
+
+    def make_uploaded_payload(prompt: str) -> dict:
+        return {
+            "input": prompt,
+            "voice": args.voice_name,
+            "stream": True,
+            "response_format": "pcm",
+            "ref_text": args.ref_text,
+            "max_new_tokens": 2048,
+        }
+
+    bench_cached = await run_round(
+        args.host,
+        args.port,
+        args.num_prompts,
+        make_uploaded_payload,
+        "uploaded_voice",
+        num_warmups=args.num_warmups,
+    )
+    print_benchmark_results(bench_cached)
+
+    # ---- Comparison ----
+    print(f"\n{'=' * 60}")
+    print("COMPARISON: Inline ref_audio vs Uploaded voice (cached)")
+    print(f"{'=' * 60}")
+    print(f"{'Metric':<30} {'Inline':>12} {'Cached':>12} {'Speedup':>10}")
+    print(f"{'-' * 64}")
+
+    def fmt_speedup(inline_val: float, cached_val: float) -> str:
+        if cached_val > 0 and inline_val > 0:
+            ratio = inline_val / cached_val
+            return f"{ratio:.2f}x"
+        return "N/A"
+
+    rows = [
+        ("Mean TTFP (ms)", bench_inline.mean_ttfp_ms, bench_cached.mean_ttfp_ms),
+        ("Median TTFP (ms)", bench_inline.median_ttfp_ms, bench_cached.median_ttfp_ms),
+        ("P99 TTFP (ms)", bench_inline.p99_ttfp_ms, bench_cached.p99_ttfp_ms),
+        ("Mean E2E (ms)", bench_inline.mean_e2e_ms, bench_cached.mean_e2e_ms),
+        ("Median E2E (ms)", bench_inline.median_e2e_ms, bench_cached.median_e2e_ms),
+        ("Mean RTF", bench_inline.mean_rtf, bench_cached.mean_rtf),
+    ]
+    for label, a, b in rows:
+        print(f"{label:<30} {a:>12.1f} {b:>12.1f} {fmt_speedup(a, b):>10}")
+
+    print("\nNote: Round B request #1 is a cache MISS (cold start).")
+    print("      Requests #2+ are cache HITs (skip DAC encoding).")
+
+    # Cleanup.
+    await delete_voice(args.host, args.port, args.voice_name)
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

benchmarks/fish-speech/fish_bench_utils.py

@@ -0,0 +1,501 @@
+"""Shared benchmark infrastructure for Fish Speech serving benchmarks.
+
+Provides common dataclasses, metrics computation, streaming HTTP client,
+and result formatting used by model-specific benchmark scripts.
+
+Model-specific scripts supply a ``create_payload_fn(prompt) -> dict``
+callback and audio parameters; everything else is handled here.
+"""
+
+import asyncio
+import base64
+import json
+import time
+from collections.abc import Callable
+from dataclasses import asdict, dataclass, field
+from datetime import datetime
+from pathlib import Path
+
+import aiohttp
+import numpy as np
+from tqdm.asyncio import tqdm
+
+# ---------------------------------------------------------------------------
+# Shared test prompts (varying length for realistic workload)
+# ---------------------------------------------------------------------------
+PROMPTS = [
+    "Hello, welcome to the voice synthesis benchmark test.",
+    "She said she would be here by noon, but nobody showed up.",
+    "The quick brown fox jumps over the lazy dog near the riverbank.",
+    "I can't believe how beautiful the sunset looks from up here on the mountain.",
+    "Please remember to bring your identification documents to the appointment tomorrow morning.",
+    "Have you ever wondered what it would be like to travel through time and visit ancient civilizations?",
+    "The restaurant on the corner serves the best pasta I have ever tasted in my entire life.",
+    "After the meeting, we should discuss the quarterly results and plan for the next phase.",
+    "Learning a new language takes patience, practice, and a genuine curiosity about other cultures.",
+    "The train leaves at half past seven, so we need to arrive at the station before then.",
+    "Could you please turn down the music a little bit, I'm trying to concentrate on my work.",
+    "It was a dark and stormy night when the old lighthouse keeper heard a knock at the door.",
+]
+
+
+# ---------------------------------------------------------------------------
+# Dataclasses
+# ---------------------------------------------------------------------------
+@dataclass
+class RequestResult:
+    success: bool = False
+    ttfp: float = 0.0  # Time to first audio packet (seconds)
+    e2e: float = 0.0  # End-to-end latency (seconds)
+    audio_bytes: int = 0  # Total audio bytes received
+    audio_duration: float = 0.0  # Audio duration in seconds
+    rtf: float = 0.0  # Real-time factor = e2e / audio_duration
+    prompt: str = ""
+    error: str = ""
+
+
+@dataclass
+class BenchmarkResult:
+    config_name: str = ""
+    concurrency: int = 0
+    num_prompts: int = 0
+    completed: int = 0
+    failed: int = 0
+    duration_s: float = 0.0
+    # TTFP stats (ms)
+    mean_ttfp_ms: float = 0.0
+    median_ttfp_ms: float = 0.0
+    std_ttfp_ms: float = 0.0
+    p90_ttfp_ms: float = 0.0
+    p95_ttfp_ms: float = 0.0
+    p99_ttfp_ms: float = 0.0
+    # E2E stats (ms)
+    mean_e2e_ms: float = 0.0
+    median_e2e_ms: float = 0.0
+    std_e2e_ms: float = 0.0
+    p90_e2e_ms: float = 0.0
+    p95_e2e_ms: float = 0.0
+    p99_e2e_ms: float = 0.0
+    # RTF stats
+    mean_rtf: float = 0.0
+    median_rtf: float = 0.0
+    std_rtf: float = 0.0
+    p99_rtf: float = 0.0
+    # Audio stats
+    mean_audio_duration_s: float = 0.0
+    total_audio_duration_s: float = 0.0
+    audio_throughput: float = 0.0  # audio_duration / wall_time
+    request_throughput: float = 0.0  # requests / second
+    # Per-request details
+    per_request: list = field(default_factory=list)
+
+
+# ---------------------------------------------------------------------------
+# Audio helpers
+# ---------------------------------------------------------------------------
+def pcm_bytes_to_duration(
+    num_bytes: int,
+    sample_rate: int = 24000,
+    sample_width: int = 2,
+) -> float:
+    """Convert raw PCM byte count to duration in seconds."""
+    return num_bytes / sample_width / sample_rate
+
+
+def _is_sse_response(response: aiohttp.ClientResponse) -> bool:
+    content_type = (response.headers.get("Content-Type") or "").lower()
+    return "text/event-stream" in content_type
+
+
+async def _read_raw_audio_stream(
+    response: aiohttp.ClientResponse,
+    *,
+    start_time: float,
+) -> tuple[int, float]:
+    first_audio_at = 0.0
+    total_bytes = 0
+
+    async for chunk in response.content.iter_any():
+        if chunk and first_audio_at <= 0:
+            first_audio_at = time.perf_counter() - start_time
+        total_bytes += len(chunk)
+
+    return total_bytes, first_audio_at
+
+
+def _extract_sse_payload(raw_event: bytes) -> bytes | None:
+    data_lines: list[bytes] = []
+    for raw_line in raw_event.splitlines():
+        line = raw_line.rstrip(b"\r")
+        if line.startswith(b"data: "):
+            data_lines.append(line[6:])
+        elif line.startswith(b"data:"):
+            data_lines.append(line[5:].lstrip())
+
+    if not data_lines:
+        return None
+    return b"\n".join(data_lines).strip()
+
+
+async def _read_sse_audio_stream(
+    response: aiohttp.ClientResponse,
+    *,
+    start_time: float,
+) -> tuple[int, float]:
+    """Decode SSE events and count raw audio bytes from base64 payloads."""
+    first_audio_at = 0.0
+    total_bytes = 0
+    pending = b""
+
+    async for chunk in response.content.iter_any():
+        if not chunk:
+            continue
+        pending += chunk
+        pending = pending.replace(b"\r\n", b"\n")
+
+        while b"\n\n" in pending:
+            raw_event, pending = pending.split(b"\n\n", 1)
+            payload_bytes = _extract_sse_payload(raw_event)
+            if payload_bytes is None:
+                continue
+            if payload_bytes == b"[DONE]":
+                return total_bytes, first_audio_at
+
+            try:
+                payload = json.loads(payload_bytes)
+            except json.JSONDecodeError as exc:
+                raise ValueError(f"Invalid SSE JSON payload: {exc}") from exc
+
+            audio = payload.get("audio")
+            if not isinstance(audio, dict):
+                continue
+
+            audio_b64 = audio.get("data")
+            if not audio_b64:
+                continue
+
+            try:
+                audio_bytes = base64.b64decode(audio_b64)
+            except Exception as exc:
+                raise ValueError(f"Invalid base64 audio chunk: {exc}") from exc
+
+            if audio_bytes and first_audio_at <= 0:
+                first_audio_at = time.perf_counter() - start_time
+            total_bytes += len(audio_bytes)
+
+    return total_bytes, first_audio_at
+
+
+# ---------------------------------------------------------------------------
+# Metrics
+# ---------------------------------------------------------------------------
+def compute_stats(
+    results: list[RequestResult],
+    wall_time: float,
+) -> BenchmarkResult:
+    """Compute aggregate statistics from per-request results."""
+    successful = [r for r in results if r.success]
+    failed = [r for r in results if not r.success]
+
+    bench = BenchmarkResult(
+        completed=len(successful),
+        failed=len(failed),
+        duration_s=wall_time,
+    )
+
+    if not successful:
+        return bench
+
+    ttfps = [r.ttfp * 1000 for r in successful]
+    e2es = [r.e2e * 1000 for r in successful]
+    rtfs = [r.rtf for r in successful]
+    audio_durs = [r.audio_duration for r in successful]
+
+    bench.mean_ttfp_ms = float(np.mean(ttfps))
+    bench.median_ttfp_ms = float(np.median(ttfps))
+    bench.std_ttfp_ms = float(np.std(ttfps))
+    bench.p90_ttfp_ms = float(np.percentile(ttfps, 90))
+    bench.p95_ttfp_ms = float(np.percentile(ttfps, 95))
+    bench.p99_ttfp_ms = float(np.percentile(ttfps, 99))
+
+    bench.mean_e2e_ms = float(np.mean(e2es))
+    bench.median_e2e_ms = float(np.median(e2es))
+    bench.std_e2e_ms = float(np.std(e2es))
+    bench.p90_e2e_ms = float(np.percentile(e2es, 90))
+    bench.p95_e2e_ms = float(np.percentile(e2es, 95))
+    bench.p99_e2e_ms = float(np.percentile(e2es, 99))
+
+    bench.mean_rtf = float(np.mean(rtfs))
+    bench.median_rtf = float(np.median(rtfs))
+    bench.std_rtf = float(np.std(rtfs))
+    bench.p99_rtf = float(np.percentile(rtfs, 99))
+
+    bench.mean_audio_duration_s = float(np.mean(audio_durs))
+    bench.total_audio_duration_s = float(np.sum(audio_durs))
+    bench.audio_throughput = bench.total_audio_duration_s / wall_time
+    bench.request_throughput = len(successful) / wall_time
+
+    bench.per_request = [
+        {
+            "ttfp_ms": r.ttfp * 1000,
+            "e2e_ms": r.e2e * 1000,
+            "rtf": r.rtf,
+            "audio_duration_s": r.audio_duration,
+            "prompt": r.prompt,
+        }
+        for r in successful
+    ]
+
+    return bench
+
+
+# ---------------------------------------------------------------------------
+# Output formatting
+# ---------------------------------------------------------------------------
+def print_benchmark_results(bench: BenchmarkResult) -> None:
+    """Print benchmark results in standardized format."""
+    W = 50
+    print("")
+    print(f"{'=' * W}")
+    print(f"{'Serving Benchmark Result':^{W}}")
+    print(f"{'=' * W}")
+    print(f"{'Successful requests:':<40}{bench.completed:<10}")
+    print(f"{'Failed requests:':<40}{bench.failed:<10}")
+    print(f"{'Maximum request concurrency:':<40}{bench.concurrency:<10}")
+    print(f"{'Benchmark duration (s):':<40}{bench.duration_s:<10.2f}")
+    print(f"{'Request throughput (req/s):':<40}{bench.request_throughput:<10.2f}")
+    print(f"{'-' * W}")
+    print(f"{'End-to-end Latency':^{W}}")
+    print(f"{'-' * W}")
+    print(f"{'Mean E2EL (ms):':<40}{bench.mean_e2e_ms:<10.2f}")
+    print(f"{'Median E2EL (ms):':<40}{bench.median_e2e_ms:<10.2f}")
+    print(f"{'P99 E2EL (ms):':<40}{bench.p99_e2e_ms:<10.2f}")
+    print(f"{'=' * W}")
+    print(f"{'Audio Result':^{W}}")
+    print(f"{'=' * W}")
+    print(f"{'Total audio duration generated (s):':<40}{bench.total_audio_duration_s:<10.2f}")
+    print(f"{'Audio throughput (audio duration/s):':<40}{bench.audio_throughput:<10.2f}")
+    print(f"{'-' * W}")
+    print(f"{'Time to First Packet':^{W}}")
+    print(f"{'-' * W}")
+    print(f"{'Mean AUDIO_TTFP (ms):':<40}{bench.mean_ttfp_ms:<10.2f}")
+    print(f"{'Median AUDIO_TTFP (ms):':<40}{bench.median_ttfp_ms:<10.2f}")
+    print(f"{'P99 AUDIO_TTFP (ms):':<40}{bench.p99_ttfp_ms:<10.2f}")
+    print(f"{'-' * W}")
+    print(f"{'Real Time Factor':^{W}}")
+    print(f"{'-' * W}")
+    print(f"{'Mean AUDIO_RTF:':<40}{bench.mean_rtf:<10.3f}")
+    print(f"{'Median AUDIO_RTF:':<40}{bench.median_rtf:<10.3f}")
+    print(f"{'P99 AUDIO_RTF:':<40}{bench.p99_rtf:<10.3f}")
+    print(f"{'=' * W}")
+    print("")
+
+
+def save_results(
+    all_results: list[dict],
+    result_dir: str,
+    config_name: str,
+) -> Path:
+    """Save benchmark results as JSON and return the file path."""
+    out = Path(result_dir)
+    out.mkdir(parents=True, exist_ok=True)
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    result_file = out / f"bench_{config_name}_{timestamp}.json"
+
+    with open(result_file, "w") as f:
+        json.dump(all_results, f, indent=2)
+    print(f"Results saved to {result_file}")
+    return result_file
+
+
+# ---------------------------------------------------------------------------
+# Streaming HTTP client
+# ---------------------------------------------------------------------------
+async def send_streaming_request(
+    session: aiohttp.ClientSession,
+    api_url: str,
+    payload: dict,
+    sample_rate: int,
+    sample_width: int,
+    pbar: tqdm | None = None,
+) -> RequestResult:
+    """Send a streaming TTS request and measure latency metrics."""
+    result = RequestResult(prompt=payload.get("input", ""))
+    st = time.perf_counter()
+
+    try:
+        async with session.post(api_url, json=payload) as response:
+            if response.status != 200:
+                result.error = f"HTTP {response.status}: {await response.text()}"
+            else:
+                if _is_sse_response(response):
+                    total_bytes, result.ttfp = await _read_sse_audio_stream(
+                        response,
+                        start_time=st,
+                    )
+                else:
+                    total_bytes, result.ttfp = await _read_raw_audio_stream(
+                        response,
+                        start_time=st,
+                    )
+
+                result.e2e = time.perf_counter() - st
+                result.audio_bytes = total_bytes
+                result.audio_duration = pcm_bytes_to_duration(total_bytes, sample_rate, sample_width)
+
+                if total_bytes <= 0 or result.ttfp <= 0:
+                    result.error = "HTTP 200 but no audio bytes were received"
+                else:
+                    if result.audio_duration > 0:
+                        result.rtf = result.e2e / result.audio_duration
+                    result.success = True
+
+    except Exception as e:
+        result.error = str(e)
+        result.e2e = time.perf_counter() - st
+
+    finally:
+        if pbar:
+            pbar.update(1)
+    return result
+
+
+# ---------------------------------------------------------------------------
+# Benchmark runner
+# ---------------------------------------------------------------------------
+async def run_benchmark(
+    host: str,
+    port: int,
+    num_prompts: int,
+    max_concurrency: int,
+    create_payload_fn: Callable[[str], dict],
+    sample_rate: int,
+    sample_width: int = 2,
+    num_warmups: int = 3,
+    request_timeout_s: float = 120.0,
+) -> BenchmarkResult:
+    """Run a TTS streaming benchmark at a given concurrency level.
+
+    Args:
+        create_payload_fn: Model-specific function that takes a prompt string
+            and returns the request JSON payload dict.
+        sample_rate: PCM sample rate for audio duration calculation.
+        sample_width: PCM sample width in bytes (default 2 for 16-bit).
+    """
+    api_url = f"http://{host}:{port}/v1/audio/speech"
+
+    connector = aiohttp.TCPConnector(
+        limit=max_concurrency,
+        limit_per_host=max_concurrency,
+        keepalive_timeout=60,
+    )
+    session = aiohttp.ClientSession(
+        connector=connector,
+        timeout=aiohttp.ClientTimeout(
+            total=request_timeout_s,
+            connect=min(10.0, request_timeout_s),
+            sock_connect=min(10.0, request_timeout_s),
+            sock_read=request_timeout_s,
+        ),
+    )
+
+    try:
+        # Warmup
+        if num_warmups > 0:
+            print(f"  Warming up with {num_warmups} requests...")
+            warmup_tasks = [
+                send_streaming_request(
+                    session,
+                    api_url,
+                    create_payload_fn(PROMPTS[i % len(PROMPTS)]),
+                    sample_rate,
+                    sample_width,
+                )
+                for i in range(num_warmups)
+            ]
+            warmup_results = await asyncio.gather(*warmup_tasks)
+            warmup_ok = sum(1 for r in warmup_results if r.success)
+            if warmup_ok == 0:
+                print("  WARNING: All warmup requests failed!")
+                for r in warmup_results:
+                    if r.error:
+                        print(f"    {r.error[:200]}")
+            print(f"  Warmup done ({warmup_ok}/{num_warmups} succeeded).")
+
+        # Build request list
+        request_prompts = [PROMPTS[i % len(PROMPTS)] for i in range(num_prompts)]
+
+        # Run
+        print(f"  Running {num_prompts} requests with concurrency={max_concurrency}...")
+        semaphore = asyncio.Semaphore(max_concurrency)
+        pbar = tqdm(total=num_prompts, desc=f"  concurrency={max_concurrency}")
+
+        async def limited_request(prompt: str) -> RequestResult:
+            async with semaphore:
+                return await send_streaming_request(
+                    session,
+                    api_url,
+                    create_payload_fn(prompt),
+                    sample_rate,
+                    sample_width,
+                    pbar,
+                )
+
+        start_time = time.perf_counter()
+        tasks = [asyncio.create_task(limited_request(p)) for p in request_prompts]
+        results: list[RequestResult] = await asyncio.gather(*tasks)
+        wall_time = time.perf_counter() - start_time
+        pbar.close()
+
+    finally:
+        await session.close()
+
+    # Compute stats
+    bench = compute_stats(results, wall_time)
+    bench.concurrency = max_concurrency
+    bench.num_prompts = num_prompts
+
+    print_benchmark_results(bench)
+
+    # Print sample errors
+    failed = [r for r in results if not r.success]
+    if failed:
+        for r in failed[:3]:
+            print(f"  [ERROR] {r.error[:200]}")
+
+    return bench
+
+
+async def run_benchmark_sweep(
+    host: str,
+    port: int,
+    num_prompts: int,
+    concurrency_levels: list[int],
+    create_payload_fn: Callable[[str], dict],
+    sample_rate: int,
+    sample_width: int = 2,
+    num_warmups: int = 3,
+    request_timeout_s: float = 120.0,
+    config_name: str = "benchmark",
+    result_dir: str = "results",
+) -> list[dict]:
+    """Run benchmarks across multiple concurrency levels and save results."""
+    all_results = []
+
+    for concurrency in concurrency_levels:
+        result = await run_benchmark(
+            host=host,
+            port=port,
+            num_prompts=num_prompts,
+            max_concurrency=concurrency,
+            create_payload_fn=create_payload_fn,
+            sample_rate=sample_rate,
+            sample_width=sample_width,
+            num_warmups=num_warmups,
+            request_timeout_s=request_timeout_s,
+        )
+        result.config_name = config_name
+        all_results.append(asdict(result))
+
+    save_results(all_results, result_dir, config_name)
+    return all_results

benchmarks/glm_image/README.md

@@ -0,0 +1,157 @@
+# GLM-Image Benchmarks
+
+Benchmark GLM-Image T2I (text-to-image) and I2I (image-to-image) performance across three backends: HuggingFace baseline, vLLM-Omni offline, and vLLM-Omni online serving.
+
+## Benchmarks
+
+| Benchmark | Script | Description |
+|-----------|--------|-------------|
+| HuggingFace Baseline | `huggingface/inference.py` | Single-GPU transformers + diffusers pipeline |
+| vLLM-Omni Offline | `vllm-omni/inference.py` | Offline inference with continuous batching |
+| vLLM-Omni Online | `benchmark_glm_image.py` | Online serving via `/v1/chat/completions` |
+
+## HuggingFace Baseline
+
+Single-request sequential inference using the reference HuggingFace pipeline.
+
+```bash
+# T2I
+CUDA_VISIBLE_DEVICES=0 python benchmarks/glm_image/huggingface/inference.py \
+    --model-path /path/to/GLM-Image --mode t2i --num-prompts 10
+
+# I2I
+CUDA_VISIBLE_DEVICES=0 python benchmarks/glm_image/huggingface/inference.py \
+    --model-path /path/to/GLM-Image --mode i2i --num-prompts 10
+```
+
+### Options
+
+| Flag | Default | Description |
+|------|---------|-------------|
+| `--model-path` | `zai-org/GLM-Image` | Model path |
+| `--mode` | `t2i` | `t2i` or `i2i` |
+| `--dataset-path` | `prompt/prompt.json` | Path to prompt.json |
+| `--num-prompts` | `10` | Number of images to generate |
+| `--width` / `--height` | `1024` | Output image size |
+| `--num-inference-steps` | `50` | Diffusion denoising steps |
+| `--output-dir` | `benchmarks/glm_image/huggingface/outputs` | Output directory |
+| `--output-file` | - | JSON file for metrics |
+
+## vLLM-Omni Offline
+
+Multi-GPU offline inference with pipeline parallelism and continuous batching.
+
+```bash
+# T2I
+CUDA_VISIBLE_DEVICES=0,1 python benchmarks/glm_image/vllm-omni/inference.py \
+    --model-path /path/to/GLM-Image --mode t2i --num-prompts 10
+
+# I2I
+CUDA_VISIBLE_DEVICES=0,1 python benchmarks/glm_image/vllm-omni/inference.py \
+    --model-path /path/to/GLM-Image --mode i2i --num-prompts 10
+```
+
+### Options
+
+| Flag | Default | Description |
+|------|---------|-------------|
+| `--model-path` | `zai-org/GLM-Image` | Model path |
+| `--deploy-config` | - | Deploy config YAML |
+| `--mode` | `t2i` | `t2i` or `i2i` |
+| `--dataset-path` | `prompt/prompt.json` | Path to prompt.json |
+| `--num-prompts` | `10` | Number of images to generate |
+| `--width` / `--height` | `1024` | Output image size |
+| `--num-inference-steps` | `50` | Diffusion denoising steps |
+| `--output-dir` | `benchmarks/glm_image/vllm-omni/outputs` | Output directory |
+| `--output-file` | - | JSON file for metrics |
+| `--stage-init-timeout` | `600` | Stage initialization timeout (s) |
+
+### Latency Computation
+
+In offline mode all requests are submitted simultaneously and processed with continuous batching. The per-request latency is computed by summing the actual per-stage times (with `stage_0_gen_ms` diffed against the previous request to remove accumulated queue/scheduling wait).
+
+## vLLM-Omni Online Serving
+
+### Start the server
+
+```bash
+CUDA_VISIBLE_DEVICES=0,1 vllm serve /path/to/GLM-Image \
+    --omni --port 8091 --host 0.0.0.0 \
+    --served-model-name glm-image
+```
+
+### Run the benchmark
+
+```bash
+# T2I
+python benchmarks/glm_image/benchmark_glm_image.py \
+    --mode t2i --num-prompts 10 --model glm-image
+
+# I2I
+python benchmarks/glm_image/benchmark_glm_image.py \
+    --mode i2i --num-prompts 10 --model glm-image
+
+# Custom dataset
+python benchmarks/glm_image/benchmark_glm_image.py \
+    --mode i2i --dataset custom \
+    --dataset-path prompts.json --num-prompts 5
+```
+
+### Options
+
+| Flag | Default | Description |
+|------|---------|-------------|
+| `--mode` | `t2i` | `t2i` or `i2i` |
+| `--dataset` | `prompt` | `prompt`, `random`, or `custom` |
+| `--dataset-path` | - | JSON file path (required for `custom`) |
+| `--num-prompts` | `10` | Number of benchmark requests |
+| `--max-concurrency` | `1` | Max concurrent requests |
+| `--request-rate` | `inf` | Requests per second (Poisson arrival) |
+| `--warmup-requests` | `1` | Warmup requests before measurement |
+| `--width` / `--height` | `1024` | Output image size |
+| `--num-inference-steps` | `50` | Diffusion denoising steps |
+| `--seed` | - | Random seed |
+| `--model` | `default` | Model name (must match `--served-model-name`) |
+| `--host` | `localhost` | Server host |
+| `--port` | `8091` | Server port |
+| `--output-file` | - | JSON output file for metrics |
+| `--num-input-images` | `1` | Number of input images for random I2I |
+
+## Dataset
+
+The default dataset is hosted on [HuggingFace](https://huggingface.co/datasets/JaredforReal/glm-image-bench) (`prompt.json`). It is automatically downloaded and cached to `prompt/prompt.json` on first run. No manual setup needed.
+
+Each entry contains:
+
+- `t2i_prompt`: Text prompt for text-to-image generation
+- `i2i_prompt`: Text prompt for image-to-image editing
+- `image_url`: Source image URL for I2I (downloaded and cached on first use)
+
+Custom datasets use the same JSON format and can be provided via `--dataset-path`.
+
+## Pipeline Timings
+
+All three benchmarks report per-stage pipeline timings (in milliseconds):
+
+| Key | Description |
+|-----|-------------|
+| `preprocess_ms` | Input preprocessing (tokenization, multimodal encoding) |
+| `stage_0_gen_ms` | AR (autoregressive) model generation time |
+| `ar2diffusion_ms` | AR output to diffusion input conversion |
+| `stage_1_gen_ms` | Diffusion model denoising time |
+| `queue_wait_ms` | Queue wait time before processing |
+
+The stages are ordered by execution: `preprocess → stage_0 (AR) → ar2diffusion → stage_1 (Diffusion)`.
+
+## Sample Results
+
+Tested on 2x GPU with 10 prompts, 1024x1024, 50 denoising steps:
+
+| Backend | Mode | Latency Mean (s) | Throughput (img/s) |
+|---------|------|-------------------|--------------------|
+| HuggingFace | T2I | 72.6 | 0.014 |
+| HuggingFace | I2I | 70.9 | 0.014 |
+| vLLM-Omni Offline | T2I | 35.0 | 0.044 |
+| vLLM-Omni Offline | I2I | 31.0 | 0.053 |
+| vLLM-Omni Online | T2I | 38.8 | 0.026 |
+| vLLM-Omni Online | I2I | 34.7 | 0.029 |

benchmarks/glm_image/benchmark_glm_image.py

@@ -0,0 +1,464 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Online serving benchmark for GLM-Image (T2I and I2I modes).
+
+Sends requests to the /v1/chat/completions endpoint and reports end-to-end
+latency, throughput, and per-stage durations (when the server is started with
+--enable-diffusion-pipeline-profiler and/or --enable-ar-profiler).
+
+Supports three dataset types:
+  - prompt:   Use prompt.json (default). T2I uses t2i_prompt, I2I uses i2i_prompt
+              and sends source images from image_url.
+  - random:   Generate synthetic prompts (and random images for I2I).
+  - custom:   Load from a user-specified JSON file.
+
+Usage:
+    # T2I with prompt.json (default)
+    python benchmarks/glm_image/benchmark_glm_image.py \
+        --mode t2i --num-prompts 10
+
+    # I2I with prompt.json (downloads source images automatically)
+    python benchmarks/glm_image/benchmark_glm_image.py \
+        --mode i2i --num-prompts 10
+
+    # Random dataset
+    python benchmarks/glm_image/benchmark_glm_image.py \
+        --mode t2i --dataset random --num-prompts 20
+
+    # Custom dataset
+    python benchmarks/glm_image/benchmark_glm_image.py \
+        --mode i2i --dataset custom \
+        --dataset-path my_prompts.json --num-prompts 5
+"""
+
+import argparse
+import asyncio
+import base64
+import json
+import os
+import sys
+import tempfile
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import aiohttp
+import numpy as np
+import requests as sync_requests
+from PIL import Image
+from tqdm.asyncio import tqdm
+
+# Import backends from the diffusion benchmark (add parent dirs to path)
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "diffusion"))
+from backends import RequestFuncOutput
+
+BENCHMARK_DIR = Path(__file__).resolve().parent
+DEFAULT_PROMPT_JSON = BENCHMARK_DIR / "prompt" / "prompt.json"
+IMAGE_CACHE_DIR = BENCHMARK_DIR / "prompt" / "images"
+
+DATASET_REPO = "JaredforReal/glm-image-bench"
+DATASET_FILE = "prompt.json"
+
+
+def _ensure_prompt_json(dataset_path: str | None) -> str:
+    """Return path to prompt.json, downloading from HuggingFace if needed."""
+    if dataset_path:
+        return dataset_path
+    local = DEFAULT_PROMPT_JSON
+    if local.exists():
+        return str(local)
+    print(f"Downloading {DATASET_FILE} from {DATASET_REPO} ...")
+    try:
+        from huggingface_hub import hf_hub_download
+
+        downloaded = hf_hub_download(
+            repo_id=DATASET_REPO,
+            filename=DATASET_FILE,
+            repo_type="dataset",
+        )
+        local.parent.mkdir(parents=True, exist_ok=True)
+        import shutil
+
+        shutil.copy2(downloaded, local)
+        print(f"Saved to {local}")
+    except ImportError:
+        url = f"https://huggingface.co/datasets/{DATASET_REPO}/resolve/main/{DATASET_FILE}"
+        import urllib.request
+
+        local.parent.mkdir(parents=True, exist_ok=True)
+        urllib.request.urlretrieve(url, local)
+        print(f"Saved to {local}")
+    return str(local)
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class GLMImageRequest:
+    prompt: str
+    image_path: str | None = None  # Only for I2I mode
+
+
+def download_image(url: str) -> str:
+    """Download an image to cache and return the local path."""
+    IMAGE_CACHE_DIR.mkdir(parents=True, exist_ok=True)
+    fname = url.rsplit("/", 1)[-1]
+    local_path = IMAGE_CACHE_DIR / fname
+    if local_path.exists():
+        return str(local_path)
+    resp = sync_requests.get(url, timeout=30)
+    resp.raise_for_status()
+    local_path.write_bytes(resp.content)
+    return str(local_path)
+
+
+def encode_image_as_data_url(path: str) -> str:
+    """Encode a local image file as a base64 data URL."""
+    with open(path, "rb") as f:
+        encoded = base64.b64encode(f.read()).decode("utf-8")
+    ext = Path(path).suffix.lower()
+    mime = {"png": "image/png", ".jpg": "image/jpeg", ".jpeg": "image/jpeg"}.get(ext, "image/png")
+    return f"data:{mime};base64,{encoded}"
+
+
+# ---------------------------------------------------------------------------
+# Datasets
+# ---------------------------------------------------------------------------
+
+
+class PromptDataset:
+    """Load from prompt.json. T2I uses t2i_prompt, I2I uses i2i_prompt + image_url."""
+
+    def __init__(self, args: argparse.Namespace):
+        path = _ensure_prompt_json(args.dataset_path)
+        with open(path, encoding="utf-8") as f:
+            raw = json.load(f)
+
+        prompt_key = "t2i_prompt" if args.mode == "t2i" else "i2i_prompt"
+        self.items: list[GLMImageRequest] = []
+
+        for entry in raw:
+            prompt = entry.get(prompt_key, "").strip()
+            if not prompt:
+                continue
+            image_path = None
+            if args.mode == "i2i":
+                url = entry.get("image_url", "")
+                if url:
+                    image_path = download_image(url)
+            self.items.append(GLMImageRequest(prompt=prompt, image_path=image_path))
+
+        if args.num_prompts and len(self.items) > args.num_prompts:
+            self.items = self.items[: args.num_prompts]
+
+    def __len__(self) -> int:
+        return len(self.items)
+
+    def __getitem__(self, idx: int) -> GLMImageRequest:
+        return self.items[idx]
+
+    def get_requests(self) -> list[GLMImageRequest]:
+        return list(self.items)
+
+
+class RandomDataset:
+    """Generate synthetic prompts (and optional random images for I2I)."""
+
+    def __init__(self, args: argparse.Namespace):
+        self.args = args
+        self.num_prompts = args.num_prompts
+        self._random_image_paths: list[str] | None = None
+        if args.mode == "i2i":
+            self._random_image_paths = self._generate_random_images()
+
+    def _generate_random_images(self) -> list[str]:
+        paths: list[str] = []
+        for i in range(self.args.num_input_images):
+            img = Image.new("RGB", (512, 512), (128 + i * 30 % 128, 64, 192))
+            path = os.path.join(tempfile.gettempdir(), f"glm_image_bench_input_{i}.png")
+            img.save(path)
+            paths.append(path)
+        return paths
+
+    def __len__(self) -> int:
+        return self.num_prompts
+
+    def __getitem__(self, idx: int) -> GLMImageRequest:
+        image_path = None
+        if self._random_image_paths is not None:
+            image_path = self._random_image_paths[idx % len(self._random_image_paths)]
+        return GLMImageRequest(
+            prompt=f"A beautiful scene with vivid colors and intricate details, prompt {idx}",
+            image_path=image_path,
+        )
+
+    def get_requests(self) -> list[GLMImageRequest]:
+        return [self[i] for i in range(len(self))]
+
+
+class CustomDataset:
+    """Load from a user-specified JSON file.
+
+    Expected format:
+    [
+        {"prompt": "A cat sitting on a windowsill"},
+        {"prompt": "Make it look like winter", "image_path": "/path/to/img.png"}
+    ]
+    """
+
+    def __init__(self, args: argparse.Namespace):
+        if not args.dataset_path:
+            raise ValueError("--dataset-path is required for custom dataset")
+        with open(args.dataset_path, encoding="utf-8") as f:
+            raw = json.load(f)
+        self.items: list[GLMImageRequest] = []
+        for item in raw:
+            self.items.append(
+                GLMImageRequest(
+                    prompt=item.get("prompt", ""),
+                    image_path=item.get("image_path"),
+                )
+            )
+        if args.num_prompts and len(self.items) > args.num_prompts:
+            self.items = self.items[: args.num_prompts]
+
+    def __len__(self) -> int:
+        return len(self.items)
+
+    def __getitem__(self, idx: int) -> GLMImageRequest:
+        return self.items[idx]
+
+    def get_requests(self) -> list[GLMImageRequest]:
+        return list(self.items)
+
+
+# ---------------------------------------------------------------------------
+# Async request for GLM-Image (chat completions with image support)
+# ---------------------------------------------------------------------------
+
+
+async def async_glm_image_request(
+    req: GLMImageRequest,
+    api_url: str,
+    model: str,
+    session: aiohttp.ClientSession,
+    pbar: Any,
+    args: argparse.Namespace,
+) -> RequestFuncOutput:
+    """Send a single T2I or I2I request via chat completions endpoint."""
+    output = RequestFuncOutput()
+    output.start_time = time.perf_counter()
+
+    # Build messages
+    if req.image_path and args.mode == "i2i":
+        data_url = encode_image_as_data_url(req.image_path)
+        content = [
+            {"type": "text", "text": req.prompt},
+            {"type": "image_url", "image_url": {"url": data_url}},
+        ]
+    else:
+        content = req.prompt
+
+    messages = [{"role": "user", "content": content}]
+
+    extra_body: dict[str, Any] = {}
+    if args.height:
+        extra_body["height"] = args.height
+    if args.width:
+        extra_body["width"] = args.width
+    if args.num_inference_steps:
+        extra_body["num_inference_steps"] = args.num_inference_steps
+    if args.seed is not None:
+        extra_body["seed"] = args.seed
+
+    payload: dict[str, Any] = {
+        "model": model,
+        "messages": messages,
+    }
+    if extra_body:
+        payload["extra_body"] = extra_body
+
+    try:
+        async with session.post(api_url, json=payload) as response:
+            if response.status == 200:
+                resp_json = await response.json()
+                output.response_body = resp_json
+                output.success = True
+                try:
+                    choices = resp_json.get("choices", [])
+                    if choices and isinstance(choices, list):
+                        msg = choices[0].get("message", {})
+                        if isinstance(msg, dict):
+                            resp_content = msg.get("content", [])
+                            if resp_content and isinstance(resp_content, list) and len(resp_content) > 0:
+                                first_item = resp_content[0]
+                                if isinstance(first_item, dict):
+                                    output.stage_durations = first_item.get("stage_durations") or {}
+                                    output.peak_memory_mb = first_item.get("peak_memory_mb", 0.0)
+                except (IndexError, TypeError, AttributeError):
+                    pass
+            else:
+                output.error = f"HTTP {response.status}: {await response.text()}"
+                output.success = False
+    except Exception as e:
+        output.error = str(e)
+        output.success = False
+
+    output.latency = time.perf_counter() - output.start_time
+    if pbar:
+        pbar.update(1)
+    return output
+
+
+# ---------------------------------------------------------------------------
+# Benchmark
+# ---------------------------------------------------------------------------
+
+
+async def iter_requests(n: int, request_rate: float) -> Any:
+    import random as _random
+
+    for i in range(n):
+        if request_rate != float("inf") and i > 0:
+            await asyncio.sleep(_random.expovariate(request_rate))
+        yield i
+
+
+def calculate_metrics(outputs: list[RequestFuncOutput], total_duration: float) -> dict[str, Any]:
+    success = [o for o in outputs if o.success]
+    errors = [o for o in outputs if not o.success]
+    latencies = [o.latency for o in success]
+    peak_mems = [o.peak_memory_mb for o in success if o.peak_memory_mb > 0]
+
+    stage_duration_lists: dict[str, list[float]] = {}
+    for o in success:
+        for stage, dur in (o.stage_durations or {}).items():
+            stage_duration_lists.setdefault(stage, []).append(dur)
+
+    return {
+        "duration": total_duration,
+        "completed_requests": len(success),
+        "failed_requests": len(errors),
+        "throughput_qps": len(success) / total_duration if total_duration > 0 else 0,
+        "latency_mean": float(np.mean(latencies)) if latencies else 0,
+        "latency_median": float(np.median(latencies)) if latencies else 0,
+        "latency_p99": float(np.percentile(latencies, 99)) if latencies else 0,
+        "latency_p95": float(np.percentile(latencies, 95)) if latencies else 0,
+        "peak_memory_mb_max": max(peak_mems) if peak_mems else 0,
+        "stage_durations_mean": {s: float(np.mean(v)) for s, v in stage_duration_lists.items()},
+        "stage_durations_p50": {s: float(np.percentile(v, 50)) for s, v in stage_duration_lists.items()},
+    }
+
+
+async def benchmark(args: argparse.Namespace) -> None:
+    api_url = f"http://{args.host}:{args.port}/v1/chat/completions"
+
+    # Load dataset
+    if args.dataset == "prompt":
+        dataset = PromptDataset(args)
+    elif args.dataset == "random":
+        dataset = RandomDataset(args)
+    elif args.dataset == "custom":
+        dataset = CustomDataset(args)
+    else:
+        raise ValueError(f"Unknown dataset: {args.dataset}")
+
+    glm_requests = dataset.get_requests()
+    print(f"Prepared {len(glm_requests)} requests (mode={args.mode}, dataset={args.dataset})")
+
+    semaphore = asyncio.Semaphore(args.max_concurrency) if args.max_concurrency else None
+
+    async def limited_request(idx: int, req: GLMImageRequest, session: aiohttp.ClientSession, pbar: Any):
+        if semaphore:
+            async with semaphore:
+                return await async_glm_image_request(req, api_url, args.model, session, pbar, args)
+        return await async_glm_image_request(req, api_url, args.model, session, pbar, args)
+
+    async with aiohttp.ClientSession() as session:
+        # Warmup
+        if args.warmup_requests and glm_requests:
+            print(f"Running {args.warmup_requests} warmup request(s)...")
+            for i in range(args.warmup_requests):
+                await limited_request(i, glm_requests[i % len(glm_requests)], session, None)
+
+        # Main benchmark
+        pbar = tqdm(total=len(glm_requests), disable=args.disable_tqdm)
+        start_time = time.perf_counter()
+        tasks = []
+        async for idx in iter_requests(len(glm_requests), args.request_rate):
+            tasks.append(asyncio.create_task(limited_request(idx, glm_requests[idx], session, pbar)))
+        outputs = await asyncio.gather(*tasks)
+        total_duration = time.perf_counter() - start_time
+        pbar.close()
+
+    # Metrics
+    metrics = calculate_metrics(outputs, total_duration)
+    metrics["mode"] = args.mode
+    metrics["model"] = args.model
+    metrics["dataset"] = args.dataset
+
+    print(f"\n{' GLM-Image Online Benchmark Result ':=^60}")
+    print(f"{'Mode:':<40} {args.mode}")
+    print(f"{'Model:':<40} {args.model}")
+    print(f"{'Dataset:':<40} {args.dataset}")
+    print("-" * 50)
+    print(f"{'Benchmark duration (s):':<40} {metrics['duration']:.2f}")
+    print(f"{'Request rate:':<40} {args.request_rate}")
+    print(f"{'Max concurrency:':<40} {args.max_concurrency}")
+    print(f"{'Successful requests:':<40} {metrics['completed_requests']}/{len(glm_requests)}")
+    print("-" * 50)
+    print(f"{'Throughput (req/s):':<40} {metrics['throughput_qps']:.2f}")
+    print(f"{'Latency Mean (s):':<40} {metrics['latency_mean']:.4f}")
+    print(f"{'Latency Median (s):':<40} {metrics['latency_median']:.4f}")
+    print(f"{'Latency P95 (s):':<40} {metrics['latency_p95']:.4f}")
+    print(f"{'Latency P99 (s):':<40} {metrics['latency_p99']:.4f}")
+
+    if metrics["peak_memory_mb_max"] > 0:
+        print("-" * 50)
+        print(f"{'Peak Memory Max (MB):':<40} {metrics['peak_memory_mb_max']:.2f}")
+
+    if metrics["stage_durations_mean"]:
+        print("-" * 50)
+        print("Stage Durations Mean:")
+        for stage, val in sorted(metrics["stage_durations_mean"].items()):
+            unit = "ms" if stage.endswith("_ms") else "s"
+            print(f"  {stage + ':':<38} {val:.4f} ({unit})")
+
+    print("=" * 60)
+
+    if args.output_file:
+        with open(args.output_file, "w") as f:
+            json.dump(metrics, f, indent=2)
+        print(f"Metrics saved to {args.output_file}")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Benchmark GLM-Image T2I/I2I online serving.")
+    parser.add_argument("--mode", type=str, default="t2i", choices=["t2i", "i2i"])
+    parser.add_argument("--dataset", type=str, default="prompt", choices=["prompt", "random", "custom"])
+    parser.add_argument("--dataset-path", type=str, default=None)
+    parser.add_argument("--num-prompts", type=int, default=10)
+    parser.add_argument("--max-concurrency", type=int, default=1)
+    parser.add_argument("--request-rate", type=float, default=float("inf"))
+    parser.add_argument("--warmup-requests", type=int, default=1)
+    parser.add_argument("--width", type=int, default=1024)
+    parser.add_argument("--height", type=int, default=1024)
+    parser.add_argument("--num-inference-steps", type=int, default=50)
+    parser.add_argument("--seed", type=int, default=None)
+    parser.add_argument("--model", type=str, default="default")
+    parser.add_argument("--host", type=str, default="localhost")
+    parser.add_argument("--port", type=int, default=8091)
+    parser.add_argument("--output-file", type=str, default=None)
+    parser.add_argument("--disable-tqdm", action="store_true")
+    parser.add_argument("--num-input-images", type=int, default=1, help="For random I2I dataset.")
+    args = parser.parse_args()
+    asyncio.run(benchmark(args))
+
+
+if __name__ == "__main__":
+    main()

benchmarks/glm_image/huggingface/inference.py

@@ -0,0 +1,291 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+HuggingFace (transformers + diffusers) baseline benchmark for GLM-Image.
+
+Supports T2I and I2I modes with the prompt.json dataset.
+Downloads source images for I2I from image_url on first run and caches locally.
+
+Usage:
+    # T2I mode (text-to-image, no source images needed)
+    python benchmarks/glm_image/huggingface/inference.py \
+        --model-path zai-org/GLM-Image \
+        --mode t2i --num-prompts 10
+
+    # I2I mode (image-to-image, downloads source images)
+    python benchmarks/glm_image/huggingface/inference.py \
+        --model-path zai-org/GLM-Image \
+        --mode i2i --num-prompts 10
+
+    # With custom prompt.json
+    python benchmarks/glm_image/huggingface/inference.py \
+        --model-path zai-org/GLM-Image \
+        --mode i2i --dataset-path prompts.json --num-prompts 5
+"""
+
+import argparse
+import json
+import os
+import time
+from pathlib import Path
+
+import numpy as np
+import requests
+import torch
+from PIL import Image
+
+BENCHMARK_DIR = Path(__file__).resolve().parent.parent
+DEFAULT_PROMPT_JSON = BENCHMARK_DIR / "prompt" / "prompt.json"
+IMAGE_CACHE_DIR = BENCHMARK_DIR / "prompt" / "images"
+
+DATASET_REPO = "JaredforReal/glm-image-bench"
+DATASET_FILE = "prompt.json"
+
+
+def _ensure_prompt_json(dataset_path: str | None) -> str:
+    """Return path to prompt.json, downloading from HuggingFace if needed."""
+    if dataset_path:
+        return dataset_path
+    local = DEFAULT_PROMPT_JSON
+    if local.exists():
+        return str(local)
+    print(f"Downloading {DATASET_FILE} from {DATASET_REPO} ...")
+    try:
+        from huggingface_hub import hf_hub_download
+
+        downloaded = hf_hub_download(
+            repo_id=DATASET_REPO,
+            filename=DATASET_FILE,
+            repo_type="dataset",
+        )
+        local.parent.mkdir(parents=True, exist_ok=True)
+        import shutil
+
+        shutil.copy2(downloaded, local)
+        print(f"Saved to {local}")
+    except ImportError:
+        url = f"https://huggingface.co/datasets/{DATASET_REPO}/resolve/main/{DATASET_FILE}"
+        import urllib.request
+
+        local.parent.mkdir(parents=True, exist_ok=True)
+        urllib.request.urlretrieve(url, local)
+        print(f"Saved to {local}")
+    return str(local)
+
+
+HEIGHT = 1024
+WIDTH = 1024
+SEED = 42
+NUM_INFERENCE_STEPS = 50
+GUIDANCE_SCALE = 1.5
+
+
+# ---------------------------------------------------------------------------
+# Dataset
+# ---------------------------------------------------------------------------
+
+
+def load_dataset(
+    dataset_path: str | None,
+    mode: str,
+    num_prompts: int,
+) -> list[dict]:
+    """Load prompts from prompt.json and prepare per-request data."""
+    path = _ensure_prompt_json(dataset_path)
+    with open(path, encoding="utf-8") as f:
+        raw = json.load(f)
+
+    items = []
+    for entry in raw:
+        if mode == "t2i":
+            prompt_key = "t2i_prompt"
+        else:
+            prompt_key = "i2i_prompt"
+
+        prompt_text = entry.get(prompt_key, "").strip()
+        if not prompt_text:
+            continue
+
+        item = {"prompt": prompt_text}
+        if mode == "i2i":
+            item["image_url"] = entry.get("image_url", "")
+        items.append(item)
+
+    if num_prompts and len(items) > num_prompts:
+        items = items[:num_prompts]
+    return items
+
+
+def download_image(url: str, cache_dir: Path) -> str:
+    """Download an image to cache_dir and return the local path."""
+    cache_dir.mkdir(parents=True, exist_ok=True)
+    fname = url.rsplit("/", 1)[-1]
+    local_path = cache_dir / fname
+    if local_path.exists():
+        return str(local_path)
+    print(f"  Downloading {url} ...")
+    resp = requests.get(url, timeout=30)
+    resp.raise_for_status()
+    local_path.write_bytes(resp.content)
+    return str(local_path)
+
+
+# ---------------------------------------------------------------------------
+# Benchmark
+# ---------------------------------------------------------------------------
+
+
+def benchmark(args: argparse.Namespace) -> None:
+    from diffusers.pipelines.glm_image import GlmImagePipeline
+
+    print("=" * 60)
+    print("GLM-Image HuggingFace Baseline Benchmark")
+    print(f"Mode: {args.mode}  |  Model: {args.model_path}")
+    print(f"Size: {args.height}x{args.width}  |  Steps: {args.num_inference_steps}")
+    print("=" * 60)
+
+    # Load dataset
+    items = load_dataset(args.dataset_path, args.mode, args.num_prompts)
+    if not items:
+        print("No prompts loaded. Exiting.")
+        return
+    print(f"Loaded {len(items)} prompts for {args.mode} mode")
+
+    # Download I2I source images
+    if args.mode == "i2i":
+        print("Preparing source images...")
+        for item in items:
+            url = item.get("image_url", "")
+            if url:
+                item["image_path"] = download_image(url, IMAGE_CACHE_DIR)
+            else:
+                item["image_path"] = None
+
+    # Load pipeline
+    print(f"\nLoading pipeline from {args.model_path} ...")
+    t0 = time.perf_counter()
+    pipe = GlmImagePipeline.from_pretrained(
+        args.model_path,
+        torch_dtype=torch.bfloat16,
+        device_map="cuda",
+    )
+    init_time = time.perf_counter() - t0
+    print(f"Pipeline loaded in {init_time:.2f}s")
+
+    # Create output dir
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    # Run benchmark
+    generator = torch.Generator(device="cuda").manual_seed(args.seed)
+    latencies = []
+    success = 0
+    failed = 0
+
+    print(f"\nRunning {len(items)} requests sequentially...")
+    print("-" * 60)
+
+    for i, item in enumerate(items):
+        prompt = item["prompt"]
+        gen_kwargs: dict = {
+            "prompt": prompt,
+            "height": args.height,
+            "width": args.width,
+            "num_inference_steps": args.num_inference_steps,
+            "guidance_scale": args.guidance_scale,
+            "generator": generator,
+        }
+
+        if args.mode == "i2i":
+            img_path = item.get("image_path")
+            if img_path and os.path.exists(img_path):
+                gen_kwargs["image"] = [Image.open(img_path).convert("RGB")]
+            else:
+                print(f"  [{i + 1}] SKIP: no source image")
+                failed += 1
+                continue
+
+        t_start = time.perf_counter()
+        try:
+            result = pipe(**gen_kwargs)
+            image = result.images[0]
+            elapsed = time.perf_counter() - t_start
+            latencies.append(elapsed)
+            success += 1
+
+            out_path = os.path.join(args.output_dir, f"{i:04d}.png")
+            image.save(out_path)
+            print(f"  [{i + 1}/{len(items)}] {elapsed:.3f}s -> {out_path}")
+        except Exception as e:
+            elapsed = time.perf_counter() - t_start
+            failed += 1
+            print(f"  [{i + 1}/{len(items)}] FAILED ({elapsed:.3f}s): {e}")
+
+    # Report
+    total_gen_time = sum(latencies) if latencies else 0
+    print("\n" + "=" * 60)
+    print("HuggingFace Baseline Results")
+    print("=" * 60)
+    print(f"{'Mode:':<40} {args.mode}")
+    print(f"{'Model:':<40} {args.model_path}")
+    print(f"{'Image size:':<40} {args.height}x{args.width}")
+    print(f"{'Num inference steps:':<40} {args.num_inference_steps}")
+    print("-" * 50)
+    print(f"{'Pipeline init time (s):':<40} {init_time:.2f}")
+    print(f"{'Successful:':<40} {success}/{len(items)}")
+    print(f"{'Failed:':<40} {failed}")
+    print("-" * 50)
+    if latencies:
+        arr = np.array(latencies)
+        print(f"{'Total generation time (s):':<40} {total_gen_time:.2f}")
+        print(f"{'Throughput (img/s):':<40} {success / total_gen_time:.4f}")
+        print(f"{'Latency Mean (s):':<40} {arr.mean():.4f}")
+        print(f"{'Latency Median (s):':<40} {np.median(arr):.4f}")
+        print(f"{'Latency P95 (s):':<40} {np.percentile(arr, 95):.4f}")
+        print(f"{'Latency P99 (s):':<40} {np.percentile(arr, 99):.4f}")
+
+    print(f"\n{'Output dir:':<40} {args.output_dir}")
+    print("=" * 60)
+
+    # Save metrics JSON
+    metrics = {
+        "backend": "huggingface",
+        "mode": args.mode,
+        "model": args.model_path,
+        "height": args.height,
+        "width": args.width,
+        "num_inference_steps": args.num_inference_steps,
+        "init_time_s": init_time,
+        "completed_requests": success,
+        "failed_requests": failed,
+        "total_gen_time_s": total_gen_time,
+        "throughput_qps": success / total_gen_time if total_gen_time > 0 else 0,
+        "latency_mean": float(np.mean(latencies)) if latencies else 0,
+        "latency_median": float(np.median(latencies)) if latencies else 0,
+        "latency_p95": float(np.percentile(latencies, 95)) if latencies else 0,
+        "latency_p99": float(np.percentile(latencies, 99)) if latencies else 0,
+    }
+    if args.output_file:
+        with open(args.output_file, "w") as f:
+            json.dump(metrics, f, indent=2)
+        print(f"Metrics saved to {args.output_file}")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="GLM-Image HuggingFace baseline benchmark")
+    parser.add_argument("--model-path", type=str, default="zai-org/GLM-Image")
+    parser.add_argument("--mode", type=str, default="t2i", choices=["t2i", "i2i"])
+    parser.add_argument("--dataset-path", type=str, default=None, help="Path to prompt.json")
+    parser.add_argument("--num-prompts", type=int, default=10)
+    parser.add_argument("--height", type=int, default=HEIGHT)
+    parser.add_argument("--width", type=int, default=WIDTH)
+    parser.add_argument("--num-inference-steps", type=int, default=NUM_INFERENCE_STEPS)
+    parser.add_argument("--guidance-scale", type=float, default=GUIDANCE_SCALE)
+    parser.add_argument("--seed", type=int, default=SEED)
+    parser.add_argument("--output-dir", type=str, default="benchmarks/glm_image/huggingface/outputs")
+    parser.add_argument("--output-file", type=str, default=None, help="JSON file for metrics")
+    args = parser.parse_args()
+    benchmark(args)
+
+
+if __name__ == "__main__":
+    main()

benchmarks/glm_image/vllm-omni/inference.py

@@ -0,0 +1,505 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+vLLM-Omni offline benchmark for GLM-Image.
+
+Supports T2I and I2I modes with the prompt.json dataset.
+Downloads source images for I2I from image_url on first run and caches locally.
+
+Usage:
+    # T2I mode
+    python benchmarks/glm_image/vllm-omni/inference.py \
+        --model-path zai-org/GLM-Image \
+        --mode t2i --num-prompts 10
+
+    # I2I mode (downloads source images)
+    python benchmarks/glm_image/vllm-omni/inference.py \
+        --model-path zai-org/GLM-Image \
+        --mode i2i --num-prompts 10
+"""
+
+import argparse
+import json
+import math
+import os
+import time
+from pathlib import Path
+
+import numpy as np
+import requests
+from PIL import Image
+from vllm import SamplingParams
+
+from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+BENCHMARK_DIR = Path(__file__).resolve().parent.parent
+DEFAULT_PROMPT_JSON = BENCHMARK_DIR / "prompt" / "prompt.json"
+IMAGE_CACHE_DIR = BENCHMARK_DIR / "prompt" / "images"
+DEFAULT_DEPLOY_CONFIG = "vllm_omni/deploy/glm_image.yaml"
+
+DATASET_REPO = "JaredforReal/glm-image-bench"
+DATASET_FILE = "prompt.json"
+
+
+def _ensure_prompt_json(dataset_path: str | None) -> str:
+    """Return path to prompt.json, downloading from HuggingFace if needed."""
+    if dataset_path:
+        return dataset_path
+    local = DEFAULT_PROMPT_JSON
+    if local.exists():
+        return str(local)
+    print(f"Downloading {DATASET_FILE} from {DATASET_REPO} ...")
+    try:
+        from huggingface_hub import hf_hub_download
+
+        downloaded = hf_hub_download(
+            repo_id=DATASET_REPO,
+            filename=DATASET_FILE,
+            repo_type="dataset",
+        )
+        local.parent.mkdir(parents=True, exist_ok=True)
+        import shutil
+
+        shutil.copy2(downloaded, local)
+        print(f"Saved to {local}")
+    except ImportError:
+        url = f"https://huggingface.co/datasets/{DATASET_REPO}/resolve/main/{DATASET_FILE}"
+        import urllib.request
+
+        local.parent.mkdir(parents=True, exist_ok=True)
+        urllib.request.urlretrieve(url, local)
+        print(f"Saved to {local}")
+    return str(local)
+
+
+SEED = 42
+HEIGHT = 1024
+WIDTH = 1024
+NUM_INFERENCE_STEPS = 50
+GUIDANCE_SCALE = 1.5
+
+GLM_IMAGE_EOS_TOKEN_ID = 16385
+GLM_IMAGE_VISION_VOCAB_SIZE = 16512
+
+
+# ---------------------------------------------------------------------------
+# Dataset
+# ---------------------------------------------------------------------------
+
+
+def load_dataset(
+    dataset_path: str | None,
+    mode: str,
+    num_prompts: int,
+) -> list[dict]:
+    path = _ensure_prompt_json(dataset_path)
+    with open(path, encoding="utf-8") as f:
+        raw = json.load(f)
+
+    items = []
+    for entry in raw:
+        prompt_key = "t2i_prompt" if mode == "t2i" else "i2i_prompt"
+        prompt_text = entry.get(prompt_key, "").strip()
+        if not prompt_text:
+            continue
+
+        item = {"prompt": prompt_text}
+        if mode == "i2i":
+            item["image_url"] = entry.get("image_url", "")
+        items.append(item)
+
+    if num_prompts and len(items) > num_prompts:
+        items = items[:num_prompts]
+    return items
+
+
+def download_image(url: str, cache_dir: Path) -> str:
+    cache_dir.mkdir(parents=True, exist_ok=True)
+    fname = url.rsplit("/", 1)[-1]
+    local_path = cache_dir / fname
+    if local_path.exists():
+        return str(local_path)
+    print(f"  Downloading {url} ...")
+    resp = requests.get(url, timeout=30)
+    resp.raise_for_status()
+    local_path.write_bytes(resp.content)
+    return str(local_path)
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def compute_max_tokens(height: int, width: int, is_i2i: bool = False) -> int:
+    factor = 32
+    token_h = height // factor
+    token_w = width // factor
+    large_tokens = token_h * token_w
+
+    # Small preview tokens (half resolution in each dimension)
+
+    ratio = token_h / token_w if token_w > 0 else 1.0
+    small_token_h = max(1, int(math.sqrt(ratio) * (factor // 2)))
+    small_token_w = max(1, int(math.sqrt(1 / ratio) * (factor // 2)))
+    small_tokens = small_token_h * small_token_w
+
+    # Mode-dependent totals:
+    # - t2i: small + large + EOS
+    # - i2i: large + EOS
+    if is_i2i:
+        return large_tokens + 1
+    return small_tokens + large_tokens + 1
+
+
+def build_prompt_t2i(prompt: str, height: int, width: int, **gen_kw) -> dict:
+    return {
+        "prompt": prompt,
+        "height": height,
+        "width": width,
+        "mm_processor_kwargs": {"target_h": height, "target_w": width},
+        **gen_kw,
+    }
+
+
+def build_prompt_i2i(prompt: str, image_path: str, height: int, width: int, **gen_kw) -> dict:
+    return {
+        "prompt": prompt,
+        "height": height,
+        "width": width,
+        "mm_processor_kwargs": {"target_h": height, "target_w": width},
+        "multi_modal_data": {"image": Image.open(image_path).convert("RGB")},
+        **gen_kw,
+    }
+
+
+def resolve_deploy_config(args: argparse.Namespace) -> str:
+    if args.deploy_config:
+        return args.deploy_config
+    if os.path.exists(DEFAULT_DEPLOY_CONFIG):
+        return DEFAULT_DEPLOY_CONFIG
+    fallback = Path(__file__).resolve().parents[3] / DEFAULT_DEPLOY_CONFIG
+    if fallback.exists():
+        return str(fallback)
+    raise FileNotFoundError("Deploy config not found. Specify --deploy-config.")
+
+
+# ---------------------------------------------------------------------------
+# Benchmark
+# ---------------------------------------------------------------------------
+
+
+def benchmark(args: argparse.Namespace) -> None:
+    is_i2i = args.mode == "i2i"
+
+    print("=" * 60)
+    print("GLM-Image vLLM-Omni Benchmark")
+    print(f"Mode: {args.mode}  |  Model: {args.model_path}")
+    print(f"Size: {args.height}x{args.width}  |  Steps: {args.num_inference_steps}")
+    print("=" * 60)
+
+    # Load dataset
+    items = load_dataset(args.dataset_path, args.mode, args.num_prompts)
+    if not items:
+        print("No prompts loaded. Exiting.")
+        return
+    print(f"Loaded {len(items)} prompts for {args.mode} mode")
+
+    # Download I2I source images
+    if is_i2i:
+        print("Preparing source images...")
+        for item in items:
+            url = item.get("image_url", "")
+            if url:
+                item["image_path"] = download_image(url, IMAGE_CACHE_DIR)
+            else:
+                item["image_path"] = None
+
+    # Init Omni
+    deploy_config = resolve_deploy_config(args)
+    print(f"\nInitializing vLLM-Omni (deploy config: {deploy_config}) ...")
+    t0 = time.perf_counter()
+
+    omni = Omni(
+        model=args.model_path,
+        deploy_config=deploy_config,
+        log_stats=args.log_stats,
+        stage_init_timeout=args.stage_init_timeout,
+        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
+        enable_ar_profiler=args.enable_ar_profiler,
+    )
+
+    init_time = time.perf_counter() - t0
+    print(f"Initialized in {init_time:.2f}s")
+
+    # Sampling params
+    max_tokens = compute_max_tokens(args.height, args.width, is_i2i=is_i2i)
+    ar_params = SamplingParams(
+        temperature=0.9,
+        top_p=0.75,
+        top_k=GLM_IMAGE_VISION_VOCAB_SIZE,
+        max_tokens=max_tokens,
+        stop_token_ids=[GLM_IMAGE_EOS_TOKEN_ID],
+        seed=args.seed,
+        detokenize=False,
+        extra_args={"target_h": args.height, "target_w": args.width},
+    )
+    diff_params = OmniDiffusionSamplingParams(
+        num_inference_steps=args.num_inference_steps,
+        guidance_scale=args.guidance_scale,
+        height=args.height,
+        width=args.width,
+        seed=args.seed,
+    )
+    sampling_params_list = [ar_params, diff_params]
+
+    # Build all prompts
+    gen_kw = {
+        "seed": args.seed,
+        "num_inference_steps": args.num_inference_steps,
+        "guidance_scale": args.guidance_scale,
+    }
+    all_prompts = []
+    for item in items:
+        if is_i2i:
+            img_path = item.get("image_path")
+            if not img_path or not os.path.exists(img_path):
+                continue
+            all_prompts.append(build_prompt_i2i(item["prompt"], img_path, args.height, args.width, **gen_kw))
+        else:
+            all_prompts.append(build_prompt_t2i(item["prompt"], args.height, args.width, **gen_kw))
+
+    valid = len(all_prompts)
+    print(f"Valid prompts: {valid}")
+
+    # Create output dir
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    # Warmup: run 1 request to prime caches, CUDA graphs, etc.
+    if all_prompts:
+        print("Running warmup request...")
+        try:
+            warmup_prompt = [all_prompts[0]]
+            omni.generate(warmup_prompt, sampling_params_list, py_generator=False)
+            print("Warmup done.\n")
+        except Exception as e:
+            print(f"Warmup failed (continuing): {e}")
+
+    # Run
+    print(f"\nRunning {valid} requests...")
+    print("-" * 60)
+
+    latencies = []
+    all_stage_durations: list[dict[str, float]] = []
+    success = 0
+    failed = 0
+    wall_start = time.perf_counter()
+
+    try:
+        output_idx = 0
+        for stage_outputs in omni.generate(all_prompts, sampling_params_list, py_generator=True):
+            if stage_outputs.final_output_type == "image":
+                request_output = stage_outputs.request_output
+                request_id = getattr(request_output, "request_id", "")
+
+                images = getattr(request_output, "images", [])
+                if not images and hasattr(request_output, "multimodal_output"):
+                    mm = request_output.multimodal_output
+                    if isinstance(mm, dict):
+                        images = mm.get("images", [])
+
+                elapsed = time.perf_counter() - wall_start
+                if images:
+                    for img in images:
+                        if isinstance(img, Image.Image):
+                            out_path = os.path.join(args.output_dir, f"{output_idx:04d}.png")
+                            img.save(out_path)
+                    success += 1
+                    latencies.append(elapsed)
+                    stage_durations = getattr(stage_outputs, "stage_durations", {})
+                    if stage_durations:
+                        all_stage_durations.append(stage_durations)
+                    # Show wall-clock elapsed and pipeline breakdown if available
+                    preprocess_str = ""
+                    if "preprocess_ms" in stage_durations:
+                        preprocess_str = f" preprocess={stage_durations['preprocess_ms'] / 1000.0:.2f}s"
+                    print(f"  [{success}/{valid}] id={request_id[:8]} {elapsed:.2f}s{preprocess_str}")
+                    output_idx += 1
+                else:
+                    failed += 1
+    except Exception as e:
+        print(f"Error: {e}")
+        failed = valid - success
+
+    total_gen_time = time.perf_counter() - wall_start
+
+    # Diff stage_0_gen_ms with previous request to remove accumulated wait time.
+    # stage_0_gen_ms is measured from submit_ts (same for all requests submitted
+    # at once), so it accumulates queue/scheduling overhead across requests.
+    # Other stages and pipeline timings are per-request already.
+    _TIMING_ORDER = [
+        "preprocess_ms",
+        "stage_0_gen_ms",
+        "ar2diffusion_ms",
+        "stage_1_gen_ms",
+        "queue_wait_ms",
+    ]
+
+    per_request_actual: list[dict[str, float]] = []
+    prev_stage_0_ms = 0.0
+    for sd in all_stage_durations:
+        actual = dict(sd)
+        s0 = sd.get("stage_0_gen_ms", 0.0)
+        actual["stage_0_gen_ms"] = s0 - prev_stage_0_ms
+        prev_stage_0_ms = s0
+        per_request_actual.append(actual)
+
+    per_request_e2e_ms: list[float] = []
+    for actual in per_request_actual:
+        e2e_ms = sum(v for k, v in actual.items() if k in _TIMING_ORDER)
+        if e2e_ms > 0:
+            per_request_e2e_ms.append(e2e_ms)
+
+    # Report
+    print("\n" + "=" * 60)
+    print("vLLM-Omni Benchmark Results")
+    print("=" * 60)
+    print(f"{'Mode:':<40} {args.mode}")
+    print(f"{'Model:':<40} {args.model_path}")
+    print(f"{'Image size:':<40} {args.height}x{args.width}")
+    print(f"{'Num inference steps:':<40} {args.num_inference_steps}")
+    print("-" * 50)
+    print(f"{'Init time (s):':<40} {init_time:.2f}")
+    print(f"{'Successful:':<40} {success}/{valid}")
+    print(f"{'Failed:':<40} {failed}")
+    print("-" * 50)
+
+    if per_request_e2e_ms:
+        per_request_s = np.array(per_request_e2e_ms) / 1000.0
+        print(f"{'Total generation time (s):':<40} {total_gen_time:.2f}")
+        print(f"{'Throughput (img/s):':<40} {success / total_gen_time:.4f}")
+        print(f"{'Latency Mean (s):':<40} {per_request_s.mean():.4f}")
+        print(f"{'Latency Median (s):':<40} {np.median(per_request_s):.4f}")
+        print(f"{'Latency P95 (s):':<40} {np.percentile(per_request_s, 95):.4f}")
+        print(f"{'Latency P99 (s):':<40} {np.percentile(per_request_s, 99):.4f}")
+        print(f"{'Latency Min (s):':<40} {per_request_s.min():.4f}")
+        print(f"{'Latency Max (s):':<40} {per_request_s.max():.4f}")
+    elif latencies:
+        per_request = np.diff([0.0] + list(latencies))
+        print(f"{'Total generation time (s):':<40} {total_gen_time:.2f}")
+        print(f"{'Throughput (img/s):':<40} {success / total_gen_time:.4f}")
+        print(f"{'Latency Mean (s) [wall-clock]:':<40} {per_request.mean():.4f}")
+        print(f"{'Latency Median (s) [wall-clock]:':<40} {np.median(per_request):.4f}")
+        print(f"{'Latency P95 (s) [wall-clock]:':<40} {np.percentile(per_request, 95):.4f}")
+        print(f"{'Latency P99 (s) [wall-clock]:':<40} {np.percentile(per_request, 99):.4f}")
+        print(f"{'Latency Min (s) [wall-clock]:':<40} {per_request.min():.4f}")
+        print(f"{'Latency Max (s) [wall-clock]:':<40} {per_request.max():.4f}")
+
+    if per_request_actual:
+        print("-" * 50)
+        print("Pipeline Timings Mean:")
+        for key in _TIMING_ORDER:
+            vals = [d.get(key, 0.0) for d in per_request_actual]
+            if any(v != 0 for v in vals):
+                unit = "ms" if key.endswith("_ms") else "s"
+                print(f"  {key + ':':<38} {np.mean(vals):.4f} ({unit})")
+        # Show any extra keys not in the ordered list
+        ordered_set = set(_TIMING_ORDER)
+        extra_keys = sorted(k for k in per_request_actual[0].keys() if k not in ordered_set)
+        for key in extra_keys:
+            vals = [d.get(key, 0.0) for d in per_request_actual]
+            if any(v != 0 for v in vals):
+                unit = "ms" if key.endswith("_ms") else "s"
+                print(f"  {key + ':':<38} {np.mean(vals):.4f} ({unit})")
+
+    print(f"\n{'Output dir:':<40} {args.output_dir}")
+    print("=" * 60)
+
+    # Metrics JSON
+    metrics = {
+        "backend": "vllm-omni",
+        "mode": args.mode,
+        "model": args.model_path,
+        "height": args.height,
+        "width": args.width,
+        "num_inference_steps": args.num_inference_steps,
+        "init_time_s": init_time,
+        "completed_requests": success,
+        "failed_requests": failed,
+        "total_gen_time_s": total_gen_time,
+        "throughput_qps": success / total_gen_time if total_gen_time > 0 else 0,
+    }
+    if per_request_e2e_ms:
+        per_request_s = np.array(per_request_e2e_ms) / 1000.0
+        metrics["latency_mean"] = float(per_request_s.mean())
+        metrics["latency_median"] = float(np.median(per_request_s))
+        metrics["latency_p95"] = float(np.percentile(per_request_s, 95))
+        metrics["latency_p99"] = float(np.percentile(per_request_s, 99))
+    elif latencies:
+        per_request = np.diff([0.0] + list(latencies))
+        metrics["latency_mean"] = float(per_request.mean())
+        metrics["latency_median"] = float(np.median(per_request))
+        metrics["latency_p95"] = float(np.percentile(per_request, 95))
+        metrics["latency_p99"] = float(np.percentile(per_request, 99))
+    else:
+        metrics["latency_mean"] = 0
+        metrics["latency_median"] = 0
+        metrics["latency_p95"] = 0
+        metrics["latency_p99"] = 0
+    if per_request_actual:
+        all_keys = list(_TIMING_ORDER) + sorted(k for k in per_request_actual[0].keys() if k not in set(_TIMING_ORDER))
+        stage_metrics = {}
+        for key in all_keys:
+            vals = [d.get(key, 0.0) for d in per_request_actual]
+            stage_metrics[key] = {
+                "mean": float(np.mean(vals)),
+                "median": float(np.median(vals)),
+                "p95": float(np.percentile(vals, 95)),
+            }
+        metrics["stage_durations"] = stage_metrics
+    if args.output_file:
+        with open(args.output_file, "w") as f:
+            json.dump(metrics, f, indent=2)
+        print(f"Metrics saved to {args.output_file}")
+
+    omni.close()
+    print("Done!")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="GLM-Image vLLM-Omni offline benchmark")
+    parser.add_argument("--model-path", type=str, default="zai-org/GLM-Image")
+    parser.add_argument("--deploy-config", type=str, default=None, help="Deploy config YAML")
+    parser.add_argument("--mode", type=str, default="t2i", choices=["t2i", "i2i"])
+    parser.add_argument("--dataset-path", type=str, default=None, help="Path to prompt.json")
+    parser.add_argument("--num-prompts", type=int, default=10)
+    parser.add_argument("--height", type=int, default=HEIGHT)
+    parser.add_argument("--width", type=int, default=WIDTH)
+    parser.add_argument("--num-inference-steps", type=int, default=NUM_INFERENCE_STEPS)
+    parser.add_argument("--guidance-scale", type=float, default=GUIDANCE_SCALE)
+    parser.add_argument("--seed", type=int, default=SEED)
+    parser.add_argument("--output-dir", type=str, default="benchmarks/glm_image/vllm-omni/outputs")
+    parser.add_argument("--output-file", type=str, default=None, help="JSON file for metrics")
+    parser.add_argument("--stage-init-timeout", type=int, default=600)
+    parser.add_argument(
+        "--enable-diffusion-pipeline-profiler",
+        action="store_true",
+        help="Enable diffusion pipeline profiler for stage-level timing",
+    )
+    parser.add_argument(
+        "--enable-ar-profiler",
+        action="store_true",
+        help="Enable AR stage profiler to include AR timing in stage_durations",
+    )
+    parser.add_argument(
+        "--log-stats",
+        action="store_true",
+        help="Enable detailed per-request pipeline stats logging",
+    )
+    args = parser.parse_args()
+    benchmark(args)
+
+
+if __name__ == "__main__":
+    main()

benchmarks/qwen3-omni/README.md

@@ -1,86 +0,0 @@
-# Benchmarks Guide
-
-This README explains how to (1) prepare benchmark datasets and (2) run the provided Qwen3-Omni benchmarks.
-
-## 1) Prepare the dataset (SeedTTS top100)
-
-```bash
-cd benchmarks/build_dataset
-pip install gdown
-
-# Download SeedTTS test set from Google Drive
-gdown --id 1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP
-
-# Extract
-tar -xf seedtts_testset.tar
-
-# Copy metadata and extract top-100 prompts
-cp seedtts_testset/en/meta.lst meta.lst
-python extract_prompts.py -i meta.lst -o top100.txt -n 100
-
-# (Optional) clean up to save space
-rm -rf seedtts_testset seedtts_testset.tar meta.lst
-```
-
-Artifacts:
-- `benchmarks/build_dataset/top100.txt` — 100 text prompts (one per line).
-
-## 2) Run benchmarks
-
-All commands assume repo root (`vllm-omni`).
-
-### A. Transformers benchmark (offline, HF Transformers)
-
-```
-bash benchmarks/qwen3-omni/transformers/eval_qwen3_moe_omni_transformers.sh
-```
-
-What it does:
-- Runs `qwen3_omni_moe_transformers.py` over `top100.txt` with `--num-prompts 100`.
-- Outputs to `benchmarks/qwen3-omni/transformers/benchmark_results/`:
-  - `perf_stats.json` — aggregated & per-prompt TPS/latency (thinker/talker/code2wav/overall).
-  - `results.json` — per-prompt outputs and audio paths.
-  - `audio/` — ~100 generated `.wav` files.
-
-Key checks:
-- `overall_tps` and `*_tps_avg` should be non-zero and reasonably stable.
-- Investigate any 0/NaN or unusually low TPS / long-tail latency.
-
-### B. vLLM Omni end-to-end benchmark (pipeline)
-
-```
-bash benchmarks/qwen3-omni/vllm_omni/eval_qwen3_moe_omni.sh
-```
-
-What it does:
-- Runs `examples/offline_inference/qwen3_omni/end2end.py` with `--log-stats`.
-- Uses `benchmarks/build_dataset/top100.txt` and writes to:
-  - Logs: `benchmarks/qwen3-omni/vllm_omni/logs/`
-    - `omni_pipeline_text.orchestrator.stats.jsonl` — per-stage latency stats.
-    - `omni_pipeline_text.overall.stats.jsonl` — end-to-end latency/TPS.
-    - `omni_pipeline_text.stage{0,1,2}.log` — per-stage detailed logs/errors.
-  - Outputs: `benchmarks/qwen3-omni/vllm_omni/outputs/` — ~100 text and `.wav` files.
-
-Key checks:
-- Overall stats: end-to-end latency/TPS should be reasonable.
-- Orchestrator stats: per-stage latency should be stable; investigate long tails.
-- Stage logs: ensure no errors and no unusually slow stages.
-
-
-## Performance snapshot
-
-The chart below summarizes our measured Qwen3-Omni MoE end-to-end benchmark, comparing vLLM-Omni against HF Transformers. It shows the overall throughput advantage for vLLM-Omni. These are actual experiment results—please refer to this performance when evaluating or reproducing the benchmark.
-
-![vLLM-Omni vs HF](./vllm-omni-vs-hf.png)
-
-## Directory layout
-- `benchmarks/build_dataset/` — dataset prep utilities (e.g., SeedTTS top100).
-- `benchmarks/<model>/vllm_omni/` — vLLM-Omni pipeline benchmarks, logs, outputs.
-- Add new tasks under `benchmarks/<model>/...` with the same pattern: `transformers/`, `vllm_omni/`, task-specific README, and (optionally) dataset prep notes.
-- `benchmarks/<model>/vllm-omni-vs-hf.png` — current performance snapshot (overall throughput comparison).
-- `benchmarks/<model>/transformers/` — HF Transformers benchmarks (offline reference).
-
-## Troubleshooting
-- Make sure GPU/driver/FlashAttention2 requirements are met for the chosen model.
-- If downloads fail, confirm network access to Google Drive (`gdown`) and Hugging Face.
-- If audio files are missing, check for errors in stage logs or model generation.***

benchmarks/qwen3-omni/transformers/eval_qwen3_moe_omni_transformers.sh

@@ -1,26 +0,0 @@
-#!/bin/bash
-# Qwen3-Omni Transformers Benchmark Evaluation Script
-# This script must be run from the vllm-omni root directory
-
-# Get the directory where this script is located
-SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
-
-# Navigate to vllm-omni root directory (4 levels up from script location)
-VLLM_OMNI_ROOT="$(cd "$SCRIPT_DIR/../../.." && pwd)"
-cd "$VLLM_OMNI_ROOT" || { echo "Error: Failed to navigate to vllm-omni directory"; exit 1; }
-
-echo "Working directory: $(pwd)"
-# Verify we're in the correct directory and run benchmark
-if [[ ! -f "benchmarks/qwen3-omni/transformers/qwen3_omni_moe_transformers.py" ]]; then
-    echo "Error: Not in vllm-omni root directory. Please run from vllm-omni folder."
-else
-    cd benchmarks/qwen3-omni/transformers
-
-    python qwen3_omni_moe_transformers.py --prompts-file ../../build_dataset/top100.txt --num-prompts 100
-
-    echo "Logs and outputs are saved to $(pwd)/benchmark_results:"
-    echo "  - perf_stats.json    Aggregated/per-prompt TPS and latency (thinker/talker/code2wav/overall)"
-    echo "  - results.json       Per-prompt outputs and audio paths"
-    echo "  - audio/             Generated wav files, there should be 100 wav file generated"
-    echo "Key checks: overall_tps and *_tps_avg should be non-zero and stable; investigate 0/NaN or unusually low TPS/long-tail latency."
-fi

benchmarks/qwen3-omni/transformers/qwen3_omni_moe_model.py

@@ -1,265 +0,0 @@
-import time
-
-import torch
-from transformers import Qwen3OmniMoeForConditionalGeneration
-
-
-class Qwen3OmniMoeForConditionalGenerationWithLogging(Qwen3OmniMoeForConditionalGeneration):
-    @torch.no_grad()
-    def generate(
-        self,
-        input_ids: torch.Tensor | None = None,
-        speaker: str = "Ethan",
-        use_audio_in_video: bool = False,
-        return_audio: bool | None = None,
-        thinker_max_new_tokens: int = 1024,
-        thinker_eos_token_id: int = 151645,
-        talker_max_new_tokens: int = 4096,
-        talker_do_sample: bool = True,
-        talker_top_k: int = 50,
-        talker_top_p: float = 1.0,
-        talker_temperature: float = 0.9,
-        talker_repetition_penalty: float = 1.05,
-        **kwargs,
-    ):
-        total_t0 = time.time()
-        perf_stats = {
-            "thinker_tokens": 0,
-            "thinker_time_s": 0.0,
-            "thinker_tps": 0.0,
-            "talker_tokens": 0,
-            "talker_time_s": 0.0,
-            "talker_tps": 0.0,
-            "code2wav_tokens": 0,
-            "code2wav_time_s": 0.0,
-            "code2wav_tps": 0.0,
-            "total_tokens": 0,
-            "total_time_s": 0.0,
-            "total_tps": 0.0,
-        }
-        if return_audio and not self.has_talker:
-            raise ValueError(
-                "Cannot use talker when talker module not initialized. "
-                "Use `enable_talker` method or set enable_talker in config "
-                "to enable talker."
-            )
-        if return_audio is None:
-            return_audio = self.has_talker
-
-        shared_kwargs = {"use_audio_in_video": use_audio_in_video}
-        thinker_kwargs = {
-            "max_new_tokens": thinker_max_new_tokens,
-            "eos_token_id": thinker_eos_token_id,
-        }
-
-        talker_kwargs = {}
-        token2wav_kwargs = {}
-        if return_audio:
-            speaker_id = self.config.talker_config.speaker_id.get(speaker.lower())
-            if speaker_id is None:
-                raise NotImplementedError(f"Speaker {speaker} not implemented")
-            if input_ids.shape[0] != 1:
-                raise NotImplementedError("Qwen3-Omni currently does not support batched inference with audio output")
-            talker_suppressed_tokens = [
-                i
-                for i in range(
-                    self.config.talker_config.text_config.vocab_size - 1024,
-                    self.config.talker_config.text_config.vocab_size,
-                )
-                if i != self.config.talker_config.codec_eos_token_id
-            ]  # Suppress additional special tokens, should not be predicted
-            talker_kwargs = {
-                "max_new_tokens": talker_max_new_tokens,
-                "do_sample": talker_do_sample,
-                "top_k": talker_top_k,
-                "top_p": talker_top_p,
-                "temperature": talker_temperature,
-                "eos_token_id": self.config.talker_config.codec_eos_token_id,
-                "repetition_penalty": talker_repetition_penalty,
-                "suppress_tokens": talker_suppressed_tokens,
-                "output_hidden_states": True,
-                "return_dict_in_generate": True,
-            }
-            token2wav_kwargs = {}
-
-        for key, value in kwargs.items():
-            if key.startswith("thinker_"):
-                thinker_kwargs[key[len("thinker_") :]] = value
-            elif key.startswith("talker_"):
-                talker_kwargs[key[len("talker_") :]] = value
-            elif key.startswith("token2wav_"):
-                token2wav_kwargs[key[len("token2wav_") :]] = value
-            # Process special input values
-            elif key == "feature_attention_mask":
-                thinker_kwargs[key] = value
-                talker_kwargs["audio_feature_lengths"] = torch.sum(value, dim=1)
-            elif key in ("input_features", "attention_mask"):
-                thinker_kwargs[key] = value
-            # Put other key to shared kwargs
-            else:
-                shared_kwargs[key] = value
-
-        # Merge kwargs
-        for key, value in shared_kwargs.items():
-            if key not in thinker_kwargs:
-                thinker_kwargs[key] = value
-            if key not in talker_kwargs and key in ["image_grid_thw", "video_grid_thw", "video_second_per_grid"]:
-                talker_kwargs[key] = value
-            if key not in token2wav_kwargs:
-                token2wav_kwargs[key] = value
-
-        # 1. Generate from thinker module
-        generate_audio = return_audio and self.has_talker
-        if generate_audio:
-            thinker_kwargs["output_hidden_states"] = True
-            thinker_kwargs["return_dict_in_generate"] = True
-
-        t0 = time.time()
-        thinker_result = self.thinker.generate(input_ids=input_ids, **thinker_kwargs)
-        t1 = time.time()
-        perf_stats["thinker_time_s"] = max(0.0, t1 - t0)
-        try:
-            prompt_len = int(input_ids.shape[1]) if input_ids is not None else 0
-            total_len = int(thinker_result.sequences.shape[-1])
-            thinker_out_len = max(0, total_len - prompt_len)
-        except Exception:
-            thinker_out_len = 0
-        perf_stats["thinker_tokens"] = thinker_out_len
-        perf_stats["thinker_tps"] = (
-            (thinker_out_len / perf_stats["thinker_time_s"]) if perf_stats["thinker_time_s"] > 0 else 0.0
-        )
-
-        if not generate_audio:
-            perf_stats["total_tokens"] = perf_stats["thinker_tokens"]
-            perf_stats["total_time_s"] = time.time() - total_t0
-            perf_stats["total_tps"] = (
-                (perf_stats["total_tokens"] / perf_stats["total_time_s"]) if perf_stats["total_time_s"] > 0 else 0.0
-            )
-            # attach stats to self
-            setattr(self, "_perf_stats_last", perf_stats)
-            if not hasattr(self, "_perf_stats_history"):
-                setattr(self, "_perf_stats_history", [])
-            self._perf_stats_history.append(perf_stats)
-            return thinker_result, None
-
-        # 2. Prepare talker input
-        thinker_embed = torch.cat([hidden_states[0] for hidden_states in thinker_result.hidden_states], dim=1).to(
-            self.talker.device
-        )  # [1 t d]
-        thinker_hidden = torch.cat(
-            [
-                hidden_states[self.config.talker_config.accept_hidden_layer]
-                for hidden_states in thinker_result.hidden_states
-            ],
-            dim=1,
-        ).to(self.talker.device)  # [1 t d]
-
-        im_start_indexes = torch.cat(
-            (
-                torch.nonzero(input_ids[0] == self.config.im_start_token_id).squeeze(),
-                torch.tensor([thinker_result.sequences.shape[-1]], device=input_ids.device, dtype=input_ids.dtype),
-            ),
-            dim=-1,
-        ).to(self.talker.device)  # Shape [n_starts + 1]; Take batch 0 since batched inference is not supported here.
-        multimodal_mask = (
-            (thinker_result.sequences == self.config.thinker_config.audio_token_id) |
-            (thinker_result.sequences == self.config.thinker_config.image_token_id) |
-            (thinker_result.sequences == self.config.thinker_config.video_token_id)
-        ).to(self.talker.device)  # [1 t] # fmt: skip
-
-        talker_special_tokens = torch.tensor(
-            [[self.config.tts_bos_token_id, self.config.tts_eos_token_id, self.config.tts_pad_token_id]],
-            device=self.thinker.device,
-            dtype=input_ids.dtype,
-        )
-        tts_bos_embed, tts_eos_embed, tts_pad_embed = (
-            self.talker.text_projection(self.thinker.get_input_embeddings()(talker_special_tokens))
-            .to(self.talker.device)
-            .chunk(3, dim=1)
-        )  # 3 * [1 1 d]
-
-        talker_input_embeds = []  # [1 t d]
-        talker_input_ids = []
-        # For every chatml parts
-        for i in range(len(im_start_indexes) - 1):
-            im_start_index = im_start_indexes[i]
-            segment_end_index = im_start_indexes[i + 1]
-            role_token = input_ids[0][im_start_index + 1]
-            # Talker should ignore thinker system prompt
-            if role_token == self.config.system_token_id:
-                continue
-            # Talker takes word embeddings for tokens and hidden state from `accept_hidden_layer` for multimodal inputs
-            elif role_token == self.config.user_token_id:
-                talker_user_part = self._get_talker_user_parts(
-                    im_start_index, segment_end_index, multimodal_mask, thinker_hidden, thinker_embed
-                )
-                talker_input_embeds.append(talker_user_part)
-                talker_input_ids.append(thinker_result.sequences[:, im_start_index:segment_end_index])
-            # Take assistant output (for now)
-            elif role_token == self.config.assistant_token_id and i == len(im_start_indexes) - 2:
-                talker_assistant_embeds, talker_assistant_ids, trailing_text_hidden = self._get_talker_assistant_parts(
-                    im_start_index,
-                    segment_end_index,
-                    speaker_id,
-                    thinker_embed,
-                    tts_pad_embed,
-                    tts_bos_embed,
-                    tts_eos_embed,
-                )
-                talker_input_embeds.append(talker_assistant_embeds)
-                talker_input_ids.append(talker_assistant_ids)
-            # History assistant output (ignore for now)
-            elif role_token == self.config.assistant_token_id and i != len(im_start_indexes) - 2:
-                continue
-            else:
-                raise AssertionError("Expect role id after <|im_start|> (assistant, user, system)")
-        talker_input_embed = torch.cat([embed.to(self.talker.device) for embed in talker_input_embeds], dim=1)
-        talker_input_id = torch.cat([embed.to(self.talker.device) for embed in talker_input_ids], dim=1)
-        t2 = time.time()
-        talker_result = self.talker.generate(
-            inputs_embeds=talker_input_embed,
-            trailing_text_hidden=trailing_text_hidden,
-            tts_pad_embed=tts_pad_embed,
-            talker_input_ids=talker_input_id,  # Not use input_ids to prevent repetition penalty out of bound
-            **talker_kwargs,
-        )
-        t3 = time.time()
-        perf_stats["talker_time_s"] = max(0.0, t3 - t2)
-        talker_codes = (
-            torch.stack([hid[-1] for hid in talker_result.hidden_states if hid[-1] is not None], dim=1)
-            .transpose(1, 2)
-            .to(self.code2wav.device)
-        )
-        try:
-            # codes shape: (B, num_quantizers, T). We log T as token length.
-            perf_stats["talker_tokens"] = int(talker_codes.shape[-1])
-        except Exception:
-            perf_stats["talker_tokens"] = 0
-        perf_stats["talker_tps"] = (
-            (perf_stats["talker_tokens"] / perf_stats["talker_time_s"]) if perf_stats["talker_time_s"] > 0 else 0.0
-        )
-        t4 = time.time()
-        talker_wavs = self.code2wav.chunked_decode(talker_codes, chunk_size=300, left_context_size=25).float()
-        t5 = time.time()
-        perf_stats["code2wav_time_s"] = max(0.0, t5 - t4)
-        perf_stats["code2wav_tokens"] = perf_stats["talker_tokens"]  # same T, not times 16
-        perf_stats["code2wav_tps"] = (
-            (perf_stats["code2wav_tokens"] / perf_stats["code2wav_time_s"])
-            if perf_stats["code2wav_time_s"] > 0
-            else 0.0
-        )
-        perf_stats["total_tokens"] = perf_stats["thinker_tokens"] + perf_stats["talker_tokens"]
-        perf_stats["total_time_s"] = time.time() - total_t0
-        perf_stats["total_tps"] = (
-            (perf_stats["total_tokens"] / perf_stats["total_time_s"]) if perf_stats["total_time_s"] > 0 else 0.0
-        )
-        setattr(self, "_perf_stats_last", perf_stats)
-        if not hasattr(self, "_perf_stats_history"):
-            setattr(self, "_perf_stats_history", [])
-        self._perf_stats_history.append(perf_stats)
-        return thinker_result, talker_wavs.float()
-
-
-__all__ = [
-    "Qwen3OmniMoeForConditionalGenerationWithLogging",
-]

benchmarks/qwen3-omni/transformers/qwen3_omni_moe_transformers.py

@@ -1,275 +0,0 @@
-import argparse
-import json
-import os
-
-import soundfile as sf
-from qwen3_omni_moe_model import Qwen3OmniMoeForConditionalGenerationWithLogging
-from qwen_omni_utils import process_mm_info
-from tqdm import tqdm
-from transformers import Qwen3OmniMoeProcessor
-
-MODEL_PATH = "Qwen/Qwen3-Omni-30B-A3B-Instruct"
-# MODEL_PATH = "Qwen/Qwen3-Omni-30B-A3B-Thinking"
-
-
-def load_prompts(prompts_file: str) -> list[str]:
-    """Load prompts from a text file, one prompt per line."""
-    prompts = []
-    with open(prompts_file, encoding="utf-8") as f:
-        for line in f:
-            line = line.strip()
-            if line:
-                prompts.append(line)
-    return prompts
-
-
-def run_benchmark(
-    model,
-    processor,
-    prompts: list[str],
-    output_dir: str = "benchmark_results",
-    speaker: str = "Ethan",
-    use_audio_in_video: bool = True,
-):
-    """
-    Run benchmark on a list of prompts and collect performance stats.
-
-    Args:
-        model: The Qwen3OmniMoe model
-        processor: The Qwen3OmniMoe processor
-        prompts: List of text prompts to process
-        output_dir: Directory to save results
-        speaker: Speaker voice for audio output
-        use_audio_in_video: Whether to use audio in video
-
-    Returns:
-        tuple: (aggregated_stats, results, audio_outputs)
-            - aggregated_stats: dict with aggregated performance statistics
-            - results: list of dicts with per-prompt results
-            - audio_outputs: list of audio tensors/arrays (or None if no audio)
-    """
-    os.makedirs(output_dir, exist_ok=True)
-    audio_dir = os.path.join(output_dir, "audio")
-    os.makedirs(audio_dir, exist_ok=True)
-
-    all_stats = []
-    results = []
-    audio_outputs = []
-
-    for idx, prompt in enumerate(tqdm(prompts, desc="Processing prompts")):
-        conversation = [
-            {
-                "role": "user",
-                "content": [{"type": "text", "text": prompt}],
-            },
-        ]
-
-        # Preparation for inference
-        text = processor.apply_chat_template(conversation, add_generation_prompt=True, tokenize=False)
-        audios, images, videos = process_mm_info(conversation, use_audio_in_video=use_audio_in_video)
-        inputs = processor(
-            text=text,
-            audio=audios,
-            images=images,
-            videos=videos,
-            return_tensors="pt",
-            padding=True,
-            use_audio_in_video=use_audio_in_video,
-        )
-        inputs = inputs.to(model.device).to(model.dtype)
-
-        # Inference: Generation of the output text and audio
-        text_ids, audio = model.generate(
-            **inputs, speaker=speaker, thinker_return_dict_in_generate=True, use_audio_in_video=use_audio_in_video
-        )
-
-        # Decode output text
-        output_text = processor.batch_decode(
-            text_ids.sequences[:, inputs["input_ids"].shape[1] :],
-            skip_special_tokens=True,
-            clean_up_tokenization_spaces=False,
-        )[0]
-
-        # Collect performance stats
-        perf_stats = None
-        if hasattr(model, "_perf_stats_last"):
-            perf_stats = model._perf_stats_last.copy()
-            perf_stats["prompt_idx"] = idx
-            perf_stats["prompt"] = prompt
-            all_stats.append(perf_stats)
-
-        # Save audio and collect audio output
-        audio_path = None
-        audio_data = None
-        if audio is not None:
-            audio_data = audio.reshape(-1).detach().cpu().numpy()
-            audio_path = os.path.join(audio_dir, f"output_{idx:04d}.wav")
-            sf.write(
-                audio_path,
-                audio_data,
-                samplerate=24000,
-            )
-            audio_outputs.append(audio_data)
-        else:
-            audio_outputs.append(None)
-
-        # Save result
-        result = {
-            "idx": idx,
-            "prompt": prompt,
-            "output": output_text,
-            "audio_path": audio_path,
-            "perf_stats": perf_stats,
-        }
-        results.append(result)
-
-    # Aggregate statistics
-    aggregated_stats = aggregate_stats(all_stats)
-
-    # Save all results
-    results_path = os.path.join(output_dir, "results.json")
-    with open(results_path, "w", encoding="utf-8") as f:
-        json.dump(results, f, ensure_ascii=False, indent=2)
-
-    # Save aggregated stats
-    stats_path = os.path.join(output_dir, "perf_stats.json")
-    with open(stats_path, "w", encoding="utf-8") as f:
-        json.dump({"aggregated": aggregated_stats, "per_prompt": all_stats}, f, ensure_ascii=False, indent=2)
-
-    # Count saved audio files
-    num_audio_saved = sum(1 for a in audio_outputs if a is not None)
-    print(f"\nSaved {num_audio_saved} audio files to {audio_dir}/")
-
-    return aggregated_stats, results, audio_outputs
-
-
-def aggregate_stats(all_stats: list[dict]) -> dict:
-    """Aggregate performance statistics from multiple runs."""
-    if not all_stats:
-        return {}
-
-    keys = [
-        "thinker_tokens",
-        "thinker_time_s",
-        "thinker_tps",
-        "talker_tokens",
-        "talker_time_s",
-        "talker_tps",
-        "code2wav_tokens",
-        "code2wav_time_s",
-        "code2wav_tps",
-        "total_tokens",
-        "total_time_s",
-        "total_tps",
-    ]
-
-    aggregated = {
-        "num_samples": len(all_stats),
-    }
-
-    for key in keys:
-        values = [s.get(key, 0) for s in all_stats if key in s]
-        if values:
-            aggregated[f"{key}_sum"] = sum(values)
-            aggregated[f"{key}_avg"] = sum(values) / len(values)
-            aggregated[f"{key}_min"] = min(values)
-            aggregated[f"{key}_max"] = max(values)
-
-    # Calculate overall throughput
-    total_tokens = aggregated.get("total_tokens_sum", 0)
-    total_time = aggregated.get("total_time_s_sum", 0)
-    if total_time > 0:
-        aggregated["overall_tps"] = total_tokens / total_time
-
-    return aggregated
-
-
-def print_stats(stats: dict):
-    """Print performance statistics in a formatted way."""
-    print("\n" + "=" * 60)
-    print("Performance Statistics Summary")
-    print("=" * 60)
-
-    print(f"\nNumber of samples: {stats.get('num_samples', 0)}")
-
-    print("\n--- Thinker ---")
-    print(f"  Total tokens:  {stats.get('thinker_tokens_sum', 0):.0f}")
-    print(f"  Total time:    {stats.get('thinker_time_s_sum', 0):.2f}s")
-    print(f"  Avg TPS:       {stats.get('thinker_tps_avg', 0):.2f}")
-    print(f"  Min TPS:       {stats.get('thinker_tps_min', 0):.2f}")
-    print(f"  Max TPS:       {stats.get('thinker_tps_max', 0):.2f}")
-
-    print("\n--- Talker ---")
-    print(f"  Total tokens:  {stats.get('talker_tokens_sum', 0):.0f}")
-    print(f"  Total time:    {stats.get('talker_time_s_sum', 0):.2f}s")
-    print(f"  Avg TPS:       {stats.get('talker_tps_avg', 0):.2f}")
-    print(f"  Min TPS:       {stats.get('talker_tps_min', 0):.2f}")
-    print(f"  Max TPS:       {stats.get('talker_tps_max', 0):.2f}")
-
-    print("\n--- Code2Wav ---")
-    print(f"  Total tokens:  {stats.get('code2wav_tokens_sum', 0):.0f}")
-    print(f"  Total time:    {stats.get('code2wav_time_s_sum', 0):.2f}s")
-    print(f"  Avg TPS:       {stats.get('code2wav_tps_avg', 0):.2f}")
-    print(f"  Min TPS:       {stats.get('code2wav_tps_min', 0):.2f}")
-    print(f"  Max TPS:       {stats.get('code2wav_tps_max', 0):.2f}")
-
-    print("\n--- Overall ---")
-    print(f"  Total tokens:  {stats.get('total_tokens_sum', 0):.0f}")
-    print(f"  Total time:    {stats.get('total_time_s_sum', 0):.2f}s")
-    print(f"  Overall TPS:   {stats.get('overall_tps', 0):.2f}")
-    print(f"  Avg TPS:       {stats.get('total_tps_avg', 0):.2f}")
-    print(f"  Min TPS:       {stats.get('total_tps_min', 0):.2f}")
-    print(f"  Max TPS:       {stats.get('total_tps_max', 0):.2f}")
-
-    print("=" * 60 + "\n")
-
-
-def main():
-    parser = argparse.ArgumentParser(description="Qwen3-Omni Benchmark Script")
-    parser.add_argument(
-        "--prompts-file",
-        type=str,
-        default="benchmark/build_dataset/top100.txt",
-        help="Path to the prompts file (one prompt per line)",
-    )
-    parser.add_argument(
-        "--output-dir", type=str, default="benchmark_results", help="Directory to save benchmark results"
-    )
-    parser.add_argument("--model-path", type=str, default=MODEL_PATH, help="Path to the model")
-    parser.add_argument("--speaker", type=str, default="Ethan", help="Speaker voice for audio output")
-    parser.add_argument("--num-prompts", type=int, default=None, help="Number of prompts to process (default: all)")
-    args = parser.parse_args()
-
-    # Load model and processor
-    print(f"Loading model from {args.model_path}...")
-    model = Qwen3OmniMoeForConditionalGenerationWithLogging.from_pretrained(
-        args.model_path,
-        dtype="auto",
-        device_map="auto",
-        attn_implementation="flash_attention_2",
-    )
-    processor = Qwen3OmniMoeProcessor.from_pretrained(args.model_path)
-
-    # Benchmark mode
-    print(f"Loading prompts from {args.prompts_file}...")
-    prompts = load_prompts(args.prompts_file)
-
-    if args.num_prompts:
-        prompts = prompts[: args.num_prompts]
-
-    print(f"Running benchmark on {len(prompts)} prompts...")
-
-    aggregated_stats, results, audio_outputs = run_benchmark(
-        model=model,
-        processor=processor,
-        prompts=prompts,
-        output_dir=args.output_dir,
-        speaker=args.speaker,
-    )
-
-    print_stats(aggregated_stats)
-    print(f"\nResults saved to {args.output_dir}/")
-
-
-if __name__ == "__main__":
-    main()

benchmarks/qwen3-omni/vllm_omni/eval_qwen3_moe_omni.sh

@@ -1,37 +0,0 @@
-#!/bin/bash
-# Qwen3-Omni Benchmark Evaluation Script
-# This script must be run from the vllm-omni root directory
-
-# Get the directory where this script is located
-SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
-
-# Navigate to vllm-omni root directory (4 levels up from script location)
-VLLM_OMNI_ROOT="$(cd "$SCRIPT_DIR/../../.." && pwd)"
-cd "$VLLM_OMNI_ROOT" || { echo "Error: Failed to navigate to vllm-omni directory"; exit 1; }
-
-echo "Working directory: $(pwd)"
-
-# Verify we're in the correct directory and run benchmark
-if [[ ! -d "benchmarks/qwen3-omni/vllm_omni" ]]; then
-    echo "Error: Not in vllm-omni root directory. Please run from vllm-omni folder."
-else
-    log_dir=benchmarks/qwen3-omni/vllm_omni/logs
-    outputs_dir=benchmarks/qwen3-omni/vllm_omni/outputs
-    end2end_script_path=examples/offline_inference/qwen3_omni/end2end.py
-    build_dataset_path=benchmarks/build_dataset/top100.txt
-
-    python $end2end_script_path --output-wav $outputs_dir \
-                      --query-type text \
-                      --txt-prompts $build_dataset_path \
-                      --log-stats \
-                      --log-dir $log_dir
-    echo "Logs and outputs are saved in ${log_dir} and ${outputs_dir} respectively:"
-    echo "  - omni_pipeline_text                       run dir/base name"
-    echo "  - omni_pipeline_text.orchestrator.stats.jsonl  orchestrator-stage latency stats"
-    echo "  - omni_pipeline_text.overall.stats.jsonl       overall latency/TPS stats"
-    echo "  - omni_pipeline_text.stage0.log                per-stage detailed logs"
-    echo "  - omni_pipeline_text.stage1.log"
-    echo "  - omni_pipeline_text.stage2.log"
-    echo "Key checks: overall.stats.jsonl for end-to-end latency/TPS; orchestrator.stats.jsonl for stable per-stage latency; stage*.log for errors or long tails."
-    echo "  - outputs/             Generated txt and wav files, there should be 100 text and wav files generated respectively"
-fi

benchmarks/qwen3-tts/README.md

@@ -1,100 +0,0 @@
-# Qwen3-TTS Benchmark
-
-Benchmarks for Qwen3-TTS text-to-speech models, comparing vLLM-Omni streaming serving against HuggingFace Transformers offline inference.
-
-## Prerequisites
-
-```bash
-pip install matplotlib aiohttp soundfile numpy tqdm
-pip install qwen_tts  # for HF baseline
-```
-
-## Quick Start
-
-Run the full benchmark (vllm-omni + HF baseline) with a single command:
-
-```bash
-cd benchmarks/qwen3-tts
-bash run_benchmark.sh
-```
-
-Results (JSON + PNG plots) are saved to `results/`.
-
-### Common options
-
-```bash
-# Only vllm-omni (skip HF baseline)
-bash run_benchmark.sh --async-only
-
-# Only HF baseline
-bash run_benchmark.sh --hf-only
-
-# Use a different model (e.g. 1.7B)
-MODEL=Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice bash run_benchmark.sh --async-only
-
-# Use bs16 config for higher throughput
-STAGE_CONFIG=vllm_omni/configs/qwen3_tts_bs16.yaml bash run_benchmark.sh --async-only
-
-# Custom GPU, prompt count, concurrency levels
-GPU_DEVICE=1 NUM_PROMPTS=20 CONCURRENCY="1 4" bash run_benchmark.sh
-```
-
-## Manual Steps
-
-### 1) Start the vLLM-Omni server
-
-```bash
-CUDA_VISIBLE_DEVICES=0 python -m vllm_omni.entrypoints.cli.main serve \
-    "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice" \
-    --omni --host 127.0.0.1 --port 8000 \
-    --stage-configs-path benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs1.yaml \
-    --trust-remote-code
-```
-
-### 2) Run online serving benchmark
-
-```bash
-python benchmarks/qwen3-tts/vllm_omni/bench_tts_serve.py \
-    --port 8000 \
-    --num-prompts 50 \
-    --max-concurrency 1 4 10 \
-    --config-name "async_chunk" \
-    --result-dir results/
-```
-
-### 3) Run HuggingFace baseline
-
-```bash
-python benchmarks/qwen3-tts/transformers/bench_tts_hf.py \
-    --model "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice" \
-    --num-prompts 50 \
-    --gpu-device 0 \
-    --result-dir results/
-```
-
-### 4) Generate comparison plots
-
-```bash
-python benchmarks/qwen3-tts/plot_results.py \
-    --results results/bench_async_chunk_*.json results/bench_hf_transformers_*.json \
-    --labels "vllm-omni" "hf_transformers" \
-    --output results/comparison.png
-```
-
-## Stage Configs
-
-| Config | max_num_seqs | Description |
-|--------|:------------:|-------------|
-| `vllm_omni/configs/qwen3_tts_bs1.yaml` | 1 | Single-request processing (lowest latency) |
-| `vllm_omni/configs/qwen3_tts_bs16.yaml` | 16 | High-throughput concurrent processing |
-
-All configs use a 2-stage pipeline (Talker -> Code2Wav) with `async_chunk` streaming enabled. The `SharedMemoryConnector` streams codec frames (25-frame chunks with 25-frame context overlap) between stages.
-
-The model is specified via the CLI `--model` flag (or `MODEL` env var), so the same configs work for both the 0.6B and 1.7B model variants.
-
-## Metrics
-
-- **TTFP (Time to First Audio Packet)**: Time from request to first audio chunk (streaming latency)
-- **E2E (End-to-End Latency)**: Total time from request to complete audio response
-- **RTF (Real-Time Factor)**: E2E latency / audio duration. RTF < 1.0 means faster-than-real-time synthesis
-- **Throughput**: Total audio seconds generated per wall-clock second

benchmarks/qwen3-tts/plot_results.py

@@ -1,254 +0,0 @@
-"""Plot Qwen3-TTS benchmark results.
-
-Generates comparison bar charts similar to the async_chunk design doc:
-- TTFP (Time-to-First-Packet) across concurrency levels
-- E2E latency across concurrency levels
-- RTF (Real-Time Factor) across concurrency levels
-
-Usage:
-    # Compare two configs (async_chunk vs no_async_chunk):
-    python plot_results.py \
-        --results results/bench_async_chunk_*.json results/bench_no_async_chunk_*.json \
-        --labels "async_chunk" "no_async_chunk" \
-        --output results/qwen3_tts_benchmark.png
-
-    # Single config:
-    python plot_results.py \
-        --results results/bench_async_chunk_*.json \
-        --labels "async_chunk" \
-        --output results/qwen3_tts_benchmark.png
-"""
-
-import argparse
-import json
-from pathlib import Path
-
-import matplotlib.pyplot as plt
-import numpy as np
-
-
-def load_results(result_files: list[str]) -> list[list[dict]]:
-    """Load benchmark results from JSON files."""
-    all_results = []
-    for f in result_files:
-        with open(f) as fh:
-            data = json.load(fh)
-        all_results.append(data)
-    return all_results
-
-
-def plot_comparison(
-    all_results: list[list[dict]],
-    labels: list[str],
-    output_path: str,
-    title_prefix: str = "Qwen3-TTS",
-):
-    """Generate comparison bar charts."""
-    n_configs = len(all_results)
-
-    # Collect concurrency levels present in ALL configs (skip missing data)
-    all_concurrencies = [set(r["concurrency"] for r in results) for results in all_results]
-    concurrencies = sorted(set.union(*all_concurrencies))
-
-    # Build data arrays, using None for missing concurrency levels
-    ttfp_data = {label: [] for label in labels}
-    e2e_data = {label: [] for label in labels}
-    rtf_data = {label: [] for label in labels}
-    throughput_data = {label: [] for label in labels}
-
-    for results, label in zip(all_results, labels):
-        conc_map = {r["concurrency"]: r for r in results}
-        for c in concurrencies:
-            r = conc_map.get(c)
-            ttfp_data[label].append(r["mean_ttfp_ms"] if r else None)
-            e2e_data[label].append(r["mean_e2e_ms"] if r else None)
-            rtf_data[label].append(r["mean_rtf"] if r else None)
-            throughput_data[label].append(r["audio_throughput"] if r else None)
-
-    fig, axes = plt.subplots(2, 2, figsize=(14, 10))
-    fig.suptitle(f"{title_prefix} Performance Benchmark", fontsize=16, fontweight="bold")
-
-    x = np.arange(len(concurrencies))
-    width = 0.35 if n_configs == 2 else 0.5
-    if n_configs > 1:
-        offsets = np.linspace(-width / 2 * (n_configs - 1), width / 2 * (n_configs - 1), n_configs)
-    else:
-        offsets = [0]
-
-    colors = ["#2196F3", "#FF5722", "#4CAF50", "#FFC107"]
-
-    def plot_metric(ax, data_dict, ylabel, title, fmt=".1f"):
-        bars = []
-        for i, (label, values) in enumerate(data_dict.items()):
-            # Replace None with 0 for plotting, but track which are missing
-            plot_values = [v if v is not None else 0 for v in values]
-            color = colors[i % len(colors)]
-            bar = ax.bar(x + offsets[i], plot_values, width, label=label, color=color, alpha=0.85)
-            bars.append(bar)
-            # Add value labels on bars (skip None/missing data)
-            max_val = max((v for v in values if v is not None), default=1)
-            for rect, val in zip(bar, values):
-                if val is not None and val > 0:
-                    ax.text(
-                        rect.get_x() + rect.get_width() / 2,
-                        rect.get_height() + max_val * 0.02,
-                        f"{val:{fmt}}",
-                        ha="center",
-                        va="bottom",
-                        fontsize=9,
-                        fontweight="bold",
-                    )
-        ax.set_xlabel("Concurrency", fontsize=12)
-        ax.set_ylabel(ylabel, fontsize=12)
-        ax.set_title(title, fontsize=13, fontweight="bold")
-        ax.set_xticks(x)
-        ax.set_xticklabels([str(c) for c in concurrencies])
-        ax.legend(fontsize=10)
-        ax.grid(axis="y", alpha=0.3)
-        ax.set_axisbelow(True)
-
-    plot_metric(axes[0, 0], ttfp_data, "TTFP (ms)", "Time to First Audio Packet (TTFP)")
-    plot_metric(axes[0, 1], e2e_data, "E2E Latency (ms)", "End-to-End Latency (E2E)")
-    plot_metric(axes[1, 0], rtf_data, "RTF", "Real-Time Factor (RTF)", fmt=".3f")
-    plot_metric(axes[1, 1], throughput_data, "Audio-sec / Wall-sec", "Audio Throughput", fmt=".2f")
-
-    plt.tight_layout()
-    plt.savefig(output_path, dpi=150, bbox_inches="tight")
-    print(f"Plot saved to {output_path}")
-    plt.close()
-
-
-def plot_single_summary(results: list[dict], label: str, output_path: str):
-    """Generate a single-config summary with percentile breakdown."""
-    concurrencies = [r["concurrency"] for r in results]
-
-    fig, axes = plt.subplots(1, 3, figsize=(16, 5))
-    fig.suptitle(f"Qwen3-TTS Benchmark - {label}", fontsize=15, fontweight="bold")
-
-    # TTFP breakdown
-    ax = axes[0]
-    means = [r["mean_ttfp_ms"] for r in results]
-    medians = [r["median_ttfp_ms"] for r in results]
-    p90s = [r["p90_ttfp_ms"] for r in results]
-    p99s = [r["p99_ttfp_ms"] for r in results]
-    x = np.arange(len(concurrencies))
-    w = 0.2
-    ax.bar(x - 1.5 * w, means, w, label="mean", color="#2196F3")
-    ax.bar(x - 0.5 * w, medians, w, label="median", color="#4CAF50")
-    ax.bar(x + 0.5 * w, p90s, w, label="p90", color="#FF9800")
-    ax.bar(x + 1.5 * w, p99s, w, label="p99", color="#F44336")
-    ax.set_xticks(x)
-    ax.set_xticklabels([str(c) for c in concurrencies])
-    ax.set_xlabel("Concurrency")
-    ax.set_ylabel("TTFP (ms)")
-    ax.set_title("Time to First Audio Packet")
-    ax.legend(fontsize=9)
-    ax.grid(axis="y", alpha=0.3)
-
-    # E2E breakdown
-    ax = axes[1]
-    means = [r["mean_e2e_ms"] for r in results]
-    medians = [r["median_e2e_ms"] for r in results]
-    p90s = [r["p90_e2e_ms"] for r in results]
-    p99s = [r["p99_e2e_ms"] for r in results]
-    ax.bar(x - 1.5 * w, means, w, label="mean", color="#2196F3")
-    ax.bar(x - 0.5 * w, medians, w, label="median", color="#4CAF50")
-    ax.bar(x + 0.5 * w, p90s, w, label="p90", color="#FF9800")
-    ax.bar(x + 1.5 * w, p99s, w, label="p99", color="#F44336")
-    ax.set_xticks(x)
-    ax.set_xticklabels([str(c) for c in concurrencies])
-    ax.set_xlabel("Concurrency")
-    ax.set_ylabel("E2E Latency (ms)")
-    ax.set_title("End-to-End Latency")
-    ax.legend(fontsize=9)
-    ax.grid(axis="y", alpha=0.3)
-
-    # RTF
-    ax = axes[2]
-    means = [r["mean_rtf"] for r in results]
-    medians = [r["median_rtf"] for r in results]
-    ax.bar(x - 0.15, means, 0.3, label="mean", color="#2196F3")
-    ax.bar(x + 0.15, medians, 0.3, label="median", color="#4CAF50")
-    ax.set_xticks(x)
-    ax.set_xticklabels([str(c) for c in concurrencies])
-    ax.set_xlabel("Concurrency")
-    ax.set_ylabel("RTF")
-    ax.set_title("Real-Time Factor")
-    ax.legend(fontsize=9)
-    ax.grid(axis="y", alpha=0.3)
-
-    plt.tight_layout()
-    plt.savefig(output_path, dpi=150, bbox_inches="tight")
-    print(f"Plot saved to {output_path}")
-    plt.close()
-
-
-def print_comparison_table(all_results: list[list[dict]], labels: list[str]):
-    """Print a markdown-formatted comparison table."""
-    concurrencies = sorted(set(r["concurrency"] for r in all_results[0]))
-
-    print("\n## Benchmark Results\n")
-    header = "| Metric | Concurrency |"
-    sep = "| --- | --- |"
-    for label in labels:
-        header += f" {label} |"
-        sep += " --- |"
-    print(header)
-    print(sep)
-
-    for metric, key, fmt in [
-        ("TTFP (ms)", "mean_ttfp_ms", ".1f"),
-        ("E2E (ms)", "mean_e2e_ms", ".1f"),
-        ("RTF", "mean_rtf", ".3f"),
-        ("Throughput (audio-s/s)", "audio_throughput", ".2f"),
-    ]:
-        for c in concurrencies:
-            row = f"| {metric} | {c} |"
-            for results in all_results:
-                conc_map = {r["concurrency"]: r for r in results}
-                val = conc_map.get(c, {}).get(key, 0)
-                row += f" {val:{fmt}} |"
-            print(row)
-
-    # Improvement calculation (only if 2 configs)
-    if len(all_results) == 2:
-        print(f"\n## Improvement ({labels[0]} vs {labels[1]})\n")
-        print("| Metric | Concurrency | Improvement |")
-        print("| --- | --- | --- |")
-        for metric, key in [("TTFP", "mean_ttfp_ms"), ("E2E", "mean_e2e_ms"), ("RTF", "mean_rtf")]:
-            for c in concurrencies:
-                m0 = {r["concurrency"]: r for r in all_results[0]}
-                m1 = {r["concurrency"]: r for r in all_results[1]}
-                v0 = m0.get(c, {}).get(key, 0)
-                v1 = m1.get(c, {}).get(key, 0)
-                if v1 > 0:
-                    pct = (v1 - v0) / v1 * 100
-                    print(f"| {metric} | {c} | {pct:+.1f}% |")
-
-
-def parse_args():
-    parser = argparse.ArgumentParser(description="Plot Qwen3-TTS benchmark results")
-    parser.add_argument(
-        "--results", type=str, nargs="+", required=True, help="Path(s) to result JSON files (one per config)"
-    )
-    parser.add_argument(
-        "--labels", type=str, nargs="+", required=True, help="Labels for each config (must match --results count)"
-    )
-    parser.add_argument("--output", type=str, default="results/qwen3_tts_benchmark.png", help="Output image path")
-    parser.add_argument("--title", type=str, default="Qwen3-TTS", help="Title prefix for the plot")
-    return parser.parse_args()
-
-
-if __name__ == "__main__":
-    args = parse_args()
-    assert len(args.results) == len(args.labels), "--results and --labels must have the same count"
-
-    all_results = load_results(args.results)
-    print_comparison_table(all_results, args.labels)
-
-    Path(args.output).parent.mkdir(parents=True, exist_ok=True)
-
-    if len(all_results) == 1:
-        plot_single_summary(all_results[0], args.labels[0], args.output)
-    else:
-        plot_comparison(all_results, args.labels, args.output, title_prefix=args.title)

benchmarks/qwen3-tts/results/.gitignore

@@ -1,3 +0,0 @@
-# Benchmark results are machine-specific - do not commit
-*
-!.gitignore

benchmarks/qwen3-tts/run_benchmark.sh

@@ -1,272 +0,0 @@
-#!/bin/bash
-# Qwen3-TTS Benchmark Runner
-#
-# Compares vllm-omni streaming serving vs HuggingFace transformers offline inference.
-# Produces JSON results and comparison plots.
-#
-# Usage:
-#   # Full comparison (vllm-omni + HF):
-#   bash run_benchmark.sh
-#
-#   # Only vllm-omni async_chunk config:
-#   bash run_benchmark.sh --async-only
-#
-#   # Only HuggingFace baseline:
-#   bash run_benchmark.sh --hf-only
-#
-#   # vllm-omni only (skip HF):
-#   bash run_benchmark.sh --skip-hf
-#
-#   # Custom settings:
-#   GPU_DEVICE=1 NUM_PROMPTS=20 CONCURRENCY="1 4" bash run_benchmark.sh
-#
-#   # Use 1.7B model:
-#   MODEL=Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice bash run_benchmark.sh --async-only
-#
-#   # Use batch_size=4 config:
-#   STAGE_CONFIG=vllm_omni/configs/qwen3_tts_bs4.yaml bash run_benchmark.sh --async-only
-#
-# Environment variables:
-#   GPU_DEVICE       - GPU index to use (default: 0)
-#   NUM_PROMPTS      - Number of prompts per concurrency level (default: 50)
-#   CONCURRENCY      - Space-separated concurrency levels (default: "1 4 10")
-#   MODEL            - Model name (default: Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice)
-#   PORT             - Server port (default: 8000)
-#   GPU_MEM_TALKER   - gpu_memory_utilization for talker stage (default: 0.3)
-#   GPU_MEM_CODE2WAV - gpu_memory_utilization for code2wav stage (default: 0.2)
-#   STAGE_CONFIG     - Path to stage config YAML (default: configs/qwen3_tts_bs1.yaml)
-
-set -euo pipefail
-
-SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
-PROJECT_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)"
-
-# Defaults
-GPU_DEVICE="${GPU_DEVICE:-0}"
-NUM_PROMPTS="${NUM_PROMPTS:-50}"
-CONCURRENCY="${CONCURRENCY:-1 4 10}"
-MODEL="${MODEL:-Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice}"
-PORT="${PORT:-8000}"
-GPU_MEM_TALKER="${GPU_MEM_TALKER:-0.3}"
-GPU_MEM_CODE2WAV="${GPU_MEM_CODE2WAV:-0.2}"
-NUM_WARMUPS="${NUM_WARMUPS:-3}"
-STAGE_CONFIG="${STAGE_CONFIG:-vllm_omni/configs/qwen3_tts_bs1.yaml}"
-RESULT_DIR="${SCRIPT_DIR}/results"
-TIMESTAMP="$(date +%Y%m%d_%H%M%S)"
-
-# Parse args
-RUN_ASYNC=true
-RUN_HF=true
-for arg in "$@"; do
-    case "$arg" in
-        --async-only) RUN_HF=false ;;
-        --hf-only) RUN_ASYNC=false ;;
-        --skip-hf) RUN_HF=false ;;
-    esac
-done
-
-mkdir -p "${RESULT_DIR}"
-
-echo "============================================================"
-echo " Qwen3-TTS Benchmark"
-echo "============================================================"
-echo " GPU:          ${GPU_DEVICE}"
-echo " Model:        ${MODEL}"
-echo " Prompts:      ${NUM_PROMPTS}"
-echo " Concurrency:  ${CONCURRENCY}"
-echo " Port:         ${PORT}"
-echo " Stage config: ${STAGE_CONFIG}"
-echo " Results:      ${RESULT_DIR}"
-echo "============================================================"
-
-# Prepare stage config with correct GPU device and memory settings
-prepare_config() {
-    local config_template="$1"
-    local config_name="$2"
-    local output_path="${RESULT_DIR}/${config_name}_stage_config.yaml"
-
-    # Use sed to patch GPU device and memory utilization
-    sed \
-        -e "s/devices: \"0\"/devices: \"${GPU_DEVICE}\"/g" \
-        -e "s/gpu_memory_utilization: 0.3/gpu_memory_utilization: ${GPU_MEM_TALKER}/g" \
-        -e "s/gpu_memory_utilization: 0.2/gpu_memory_utilization: ${GPU_MEM_CODE2WAV}/g" \
-        "${config_template}" > "${output_path}"
-
-    echo "${output_path}"
-}
-
-# Start server and wait for it to be ready
-start_server() {
-    local stage_config="$1"
-    local config_name="$2"
-    local log_file="${RESULT_DIR}/server_${config_name}_${TIMESTAMP}.log"
-
-    echo ""
-    echo "Starting server with config: ${config_name}"
-    echo "  Stage config: ${stage_config}"
-    echo "  Log file: ${log_file}"
-
-    VLLM_WORKER_MULTIPROC_METHOD=spawn \
-    CUDA_VISIBLE_DEVICES="${GPU_DEVICE}" \
-    python -m vllm_omni.entrypoints.cli.main serve "${MODEL}" \
-        --omni \
-        --host 127.0.0.1 \
-        --port "${PORT}" \
-        --stage-configs-path "${stage_config}" \
-        --stage-init-timeout 120 \
-        --trust-remote-code \
-        --disable-log-stats \
-        > "${log_file}" 2>&1 &
-
-    SERVER_PID=$!
-    echo "  Server PID: ${SERVER_PID}"
-
-    # Wait for server to be ready
-    echo "  Waiting for server to be ready..."
-    local max_wait=300
-    local waited=0
-    while [ ${waited} -lt ${max_wait} ]; do
-        if curl -sf "http://127.0.0.1:${PORT}/v1/models" > /dev/null 2>&1; then
-            echo "  Server is ready! (waited ${waited}s)"
-            return 0
-        fi
-        # Check if process is still alive
-        if ! kill -0 ${SERVER_PID} 2>/dev/null; then
-            echo "  ERROR: Server process died. Check log: ${log_file}"
-            tail -20 "${log_file}"
-            return 1
-        fi
-        sleep 2
-        waited=$((waited + 2))
-    done
-
-    echo "  ERROR: Server did not start within ${max_wait}s. Check log: ${log_file}"
-    kill ${SERVER_PID} 2>/dev/null || true
-    return 1
-}
-
-# Stop the server
-stop_server() {
-    if [ -n "${SERVER_PID:-}" ]; then
-        echo "  Stopping server (PID: ${SERVER_PID})..."
-        kill ${SERVER_PID} 2>/dev/null || true
-        wait ${SERVER_PID} 2>/dev/null || true
-        # Kill any remaining child processes on the port
-        local pids
-        pids=$(lsof -ti:${PORT} 2>/dev/null || true)
-        if [ -n "${pids}" ]; then
-            echo "  Cleaning up remaining processes on port ${PORT}..."
-            echo "${pids}" | xargs kill -9 2>/dev/null || true
-        fi
-        echo "  Server stopped."
-        SERVER_PID=""
-    fi
-}
-
-# Cleanup on exit
-trap 'stop_server' EXIT
-
-# Run benchmark for a given config
-run_bench() {
-    local config_name="$1"
-    local config_template="$2"
-
-    echo ""
-    echo "============================================================"
-    echo " Benchmarking: ${config_name}"
-    echo "============================================================"
-
-    local stage_config
-    stage_config=$(prepare_config "${config_template}" "${config_name}")
-
-    start_server "${stage_config}" "${config_name}"
-
-    # Convert concurrency string to args
-    local conc_args=""
-    for c in ${CONCURRENCY}; do
-        conc_args="${conc_args} ${c}"
-    done
-
-    cd "${PROJECT_ROOT}"
-    python "${SCRIPT_DIR}/vllm_omni/bench_tts_serve.py" \
-        --host 127.0.0.1 \
-        --port "${PORT}" \
-        --num-prompts "${NUM_PROMPTS}" \
-        --max-concurrency ${conc_args} \
-        --num-warmups "${NUM_WARMUPS}" \
-        --config-name "${config_name}" \
-        --result-dir "${RESULT_DIR}"
-
-    stop_server
-
-    # Allow GPU memory to settle
-    sleep 5
-}
-
-# Run vllm-omni benchmark
-if [ "${RUN_ASYNC}" = true ]; then
-    run_bench "async_chunk" "${SCRIPT_DIR}/${STAGE_CONFIG}"
-fi
-
-# Run HuggingFace baseline benchmark
-if [ "${RUN_HF}" = true ]; then
-    echo ""
-    echo "============================================================"
-    echo " Benchmarking: HuggingFace transformers (offline)"
-    echo "============================================================"
-
-    cd "${PROJECT_ROOT}"
-    python "${SCRIPT_DIR}/transformers/bench_tts_hf.py" \
-        --model "${MODEL}" \
-        --num-prompts "${NUM_PROMPTS}" \
-        --num-warmups "${NUM_WARMUPS}" \
-        --gpu-device "${GPU_DEVICE}" \
-        --config-name "hf_transformers" \
-        --result-dir "${RESULT_DIR}"
-
-    # Allow GPU memory to settle
-    sleep 5
-fi
-
-# Plot results
-echo ""
-echo "============================================================"
-echo " Generating plots..."
-echo "============================================================"
-
-RESULT_FILES=""
-LABELS=""
-
-if [ "${RUN_ASYNC}" = true ]; then
-    ASYNC_FILE=$(ls -t "${RESULT_DIR}"/bench_async_chunk_*.json 2>/dev/null | head -1)
-    if [ -n "${ASYNC_FILE}" ]; then
-        RESULT_FILES="${ASYNC_FILE}"
-        LABELS="async_chunk"
-    fi
-fi
-
-if [ "${RUN_HF}" = true ]; then
-    HF_FILE=$(ls -t "${RESULT_DIR}"/bench_hf_transformers_*.json 2>/dev/null | head -1)
-    if [ -n "${HF_FILE}" ]; then
-        if [ -n "${RESULT_FILES}" ]; then
-            RESULT_FILES="${RESULT_FILES} ${HF_FILE}"
-            LABELS="${LABELS} hf_transformers"
-        else
-            RESULT_FILES="${HF_FILE}"
-            LABELS="hf_transformers"
-        fi
-    fi
-fi
-
-if [ -n "${RESULT_FILES}" ]; then
-    python "${SCRIPT_DIR}/plot_results.py" \
-        --results ${RESULT_FILES} \
-        --labels ${LABELS} \
-        --output "${RESULT_DIR}/qwen3_tts_benchmark_${TIMESTAMP}.png"
-fi
-
-echo ""
-echo "============================================================"
-echo " Benchmark complete!"
-echo " Results: ${RESULT_DIR}"
-echo "============================================================"

benchmarks/qwen3-tts/transformers/bench_tts_hf.py

@@ -1,281 +0,0 @@
-"""Benchmark Qwen3-TTS using HuggingFace transformers (qwen_tts library).
-
-Measures E2E latency, RTF, and audio duration for offline (non-serving) inference.
-Results are saved in the same JSON format as bench_tts_serve.py for unified plotting.
-
-Usage:
-    python bench_tts_hf.py \
-        --model Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice \
-        --num-prompts 50 \
-        --num-warmups 3 \
-        --gpu-device 0 \
-        --result-dir results/
-"""
-
-import argparse
-import json
-import time
-from dataclasses import asdict, dataclass, field
-from datetime import datetime
-from pathlib import Path
-
-import numpy as np
-import soundfile as sf
-import torch
-
-PROMPTS = [
-    "Hello, welcome to the voice synthesis benchmark test.",
-    "She said she would be here by noon, but nobody showed up.",
-    "The quick brown fox jumps over the lazy dog near the riverbank.",
-    "I can't believe how beautiful the sunset looks from up here on the mountain.",
-    "Please remember to bring your identification documents to the appointment tomorrow morning.",
-    "Have you ever wondered what it would be like to travel through time and visit ancient civilizations?",
-    "The restaurant on the corner serves the best pasta I have ever tasted in my entire life.",
-    "After the meeting, we should discuss the quarterly results and plan for the next phase.",
-    "Learning a new language takes patience, practice, and a genuine curiosity about other cultures.",
-    "The train leaves at half past seven, so we need to arrive at the station before then.",
-    "Could you please turn down the music a little bit, I'm trying to concentrate on my work.",
-    "It was a dark and stormy night when the old lighthouse keeper heard a knock at the door.",
-]
-
-
-@dataclass
-class BenchmarkResult:
-    config_name: str = ""
-    concurrency: int = 1  # always 1 for offline
-    num_prompts: int = 0
-    completed: int = 0
-    failed: int = 0
-    duration_s: float = 0.0
-    # TTFP stats - not applicable for HF offline, set to E2E for compatibility
-    mean_ttfp_ms: float = 0.0
-    median_ttfp_ms: float = 0.0
-    std_ttfp_ms: float = 0.0
-    p90_ttfp_ms: float = 0.0
-    p95_ttfp_ms: float = 0.0
-    p99_ttfp_ms: float = 0.0
-    # E2E stats (ms)
-    mean_e2e_ms: float = 0.0
-    median_e2e_ms: float = 0.0
-    std_e2e_ms: float = 0.0
-    p90_e2e_ms: float = 0.0
-    p95_e2e_ms: float = 0.0
-    p99_e2e_ms: float = 0.0
-    # RTF stats
-    mean_rtf: float = 0.0
-    median_rtf: float = 0.0
-    std_rtf: float = 0.0
-    p99_rtf: float = 0.0
-    # Audio stats
-    mean_audio_duration_s: float = 0.0
-    total_audio_duration_s: float = 0.0
-    audio_throughput: float = 0.0
-    request_throughput: float = 0.0
-    # Per-request details
-    per_request: list = field(default_factory=list)
-
-
-def run_benchmark(args):
-    from qwen_tts import Qwen3TTSModel
-
-    device = f"cuda:{args.gpu_device}"
-    print(f"Loading model: {args.model} on {device}")
-    model = Qwen3TTSModel.from_pretrained(
-        args.model,
-        device_map=device,
-        dtype=torch.bfloat16,
-    )
-    print("Model loaded.")
-
-    # Build prompt list
-    prompts = [PROMPTS[i % len(PROMPTS)] for i in range(args.num_prompts)]
-
-    # Warmup
-    if args.num_warmups > 0:
-        print(f"Warming up with {args.num_warmups} requests...")
-        for i in range(args.num_warmups):
-            p = PROMPTS[i % len(PROMPTS)]
-            wavs, sr = model.generate_custom_voice(
-                text=p,
-                language=args.language,
-                speaker=args.voice,
-            )
-        # Sync GPU
-        torch.cuda.synchronize(device)
-        print("Warmup done.")
-
-    # Benchmark
-    print(f"Running {args.num_prompts} requests sequentially...")
-    e2e_times = []
-    rtfs = []
-    audio_durations = []
-    per_request = []
-    failed = 0
-
-    audio_dir = None
-    if args.save_audio:
-        audio_dir = Path(args.result_dir) / "audio_hf"
-        audio_dir.mkdir(parents=True, exist_ok=True)
-
-    total_start = time.perf_counter()
-
-    for i, prompt in enumerate(prompts):
-        try:
-            torch.cuda.synchronize(device)
-            st = time.perf_counter()
-
-            wavs, sr = model.generate_custom_voice(
-                text=prompt,
-                language=args.language,
-                speaker=args.voice,
-            )
-
-            torch.cuda.synchronize(device)
-            elapsed = time.perf_counter() - st
-
-            # Compute audio duration
-            audio_samples = wavs[0]
-            if isinstance(audio_samples, torch.Tensor):
-                audio_samples = audio_samples.cpu().numpy()
-            audio_dur = len(audio_samples) / sr
-
-            rtf = elapsed / audio_dur if audio_dur > 0 else 0.0
-
-            e2e_times.append(elapsed)
-            rtfs.append(rtf)
-            audio_durations.append(audio_dur)
-            per_request.append(
-                {
-                    "e2e_ms": elapsed * 1000,
-                    "ttfp_ms": elapsed * 1000,  # no streaming, TTFP = E2E
-                    "rtf": rtf,
-                    "audio_duration_s": audio_dur,
-                    "prompt": prompt,
-                }
-            )
-
-            if audio_dir:
-                sf.write(str(audio_dir / f"output_{i:04d}.wav"), audio_samples, sr)
-
-            if (i + 1) % 10 == 0 or i == 0:
-                print(
-                    f"  [{i + 1}/{args.num_prompts}] e2e={elapsed * 1000:.0f}ms  rtf={rtf:.3f}  audio={audio_dur:.2f}s"
-                )
-
-        except Exception as e:
-            print(f"  [{i + 1}/{args.num_prompts}] FAILED: {e}")
-            failed += 1
-
-    total_duration = time.perf_counter() - total_start
-    completed = len(e2e_times)
-
-    # Compute stats
-    result = BenchmarkResult(
-        config_name=args.config_name,
-        concurrency=1,
-        num_prompts=args.num_prompts,
-        completed=completed,
-        failed=failed,
-        duration_s=total_duration,
-    )
-
-    if e2e_times:
-        e2e_ms = [t * 1000 for t in e2e_times]
-
-        result.mean_e2e_ms = float(np.mean(e2e_ms))
-        result.median_e2e_ms = float(np.median(e2e_ms))
-        result.std_e2e_ms = float(np.std(e2e_ms))
-        result.p90_e2e_ms = float(np.percentile(e2e_ms, 90))
-        result.p95_e2e_ms = float(np.percentile(e2e_ms, 95))
-        result.p99_e2e_ms = float(np.percentile(e2e_ms, 99))
-
-        # For HF offline, TTFP = E2E (no streaming)
-        result.mean_ttfp_ms = result.mean_e2e_ms
-        result.median_ttfp_ms = result.median_e2e_ms
-        result.std_ttfp_ms = result.std_e2e_ms
-        result.p90_ttfp_ms = result.p90_e2e_ms
-        result.p95_ttfp_ms = result.p95_e2e_ms
-        result.p99_ttfp_ms = result.p99_e2e_ms
-
-        result.mean_rtf = float(np.mean(rtfs))
-        result.median_rtf = float(np.median(rtfs))
-        result.std_rtf = float(np.std(rtfs))
-        result.p99_rtf = float(np.percentile(rtfs, 99))
-
-        result.mean_audio_duration_s = float(np.mean(audio_durations))
-        result.total_audio_duration_s = float(np.sum(audio_durations))
-        result.audio_throughput = result.total_audio_duration_s / total_duration
-        result.request_throughput = completed / total_duration
-        result.per_request = per_request
-
-    # Print summary in standardized performance template
-    W = 50
-    print("")
-    print(f"{'=' * W}")
-    print(f"{'Serving Benchmark Result':^{W}}")
-    print(f"{'=' * W}")
-    print(f"{'Successful requests:':<40}{completed:<10}")
-    print(f"{'Failed requests:':<40}{failed:<10}")
-    print(f"{'Maximum request concurrency:':<40}{1:<10}")
-    print(f"{'Benchmark duration (s):':<40}{total_duration:<10.2f}")
-    print(f"{'Request throughput (req/s):':<40}{result.request_throughput:<10.2f}")
-    print(f"{'-' * W}")
-    print(f"{'End-to-end Latency':^{W}}")
-    print(f"{'-' * W}")
-    print(f"{'Mean E2EL (ms):':<40}{result.mean_e2e_ms:<10.2f}")
-    print(f"{'Median E2EL (ms):':<40}{result.median_e2e_ms:<10.2f}")
-    print(f"{'P99 E2EL (ms):':<40}{result.p99_e2e_ms:<10.2f}")
-    print(f"{'=' * W}")
-    print(f"{'Audio Result':^{W}}")
-    print(f"{'=' * W}")
-    print(f"{'Total audio duration generated (s):':<40}{result.total_audio_duration_s:<10.2f}")
-    print(f"{'Audio throughput (audio duration/s):':<40}{result.audio_throughput:<10.2f}")
-    print(f"{'-' * W}")
-    print(f"{'Time to First Packet':^{W}}")
-    print(f"{'-' * W}")
-    print(f"{'Mean AUDIO_TTFP (ms):':<40}{result.mean_ttfp_ms:<10.2f}")
-    print(f"{'Median AUDIO_TTFP (ms):':<40}{result.median_ttfp_ms:<10.2f}")
-    print(f"{'P99 AUDIO_TTFP (ms):':<40}{result.p99_ttfp_ms:<10.2f}")
-    print(f"{'-' * W}")
-    print(f"{'Real Time Factor':^{W}}")
-    print(f"{'-' * W}")
-    print(f"{'Mean AUDIO_RTF:':<40}{result.mean_rtf:<10.3f}")
-    print(f"{'Median AUDIO_RTF:':<40}{result.median_rtf:<10.3f}")
-    print(f"{'P99 AUDIO_RTF:':<40}{result.p99_rtf:<10.3f}")
-    print(f"{'=' * W}")
-    print("")
-
-    # Save results (as a list with single concurrency=1 entry, matching serve format)
-    result_dir = Path(args.result_dir)
-    result_dir.mkdir(parents=True, exist_ok=True)
-    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-    result_file = result_dir / f"bench_{args.config_name}_{timestamp}.json"
-
-    with open(result_file, "w") as f:
-        json.dump([asdict(result)], f, indent=2)
-    print(f"Results saved to {result_file}")
-
-    return result
-
-
-def parse_args():
-    parser = argparse.ArgumentParser(description="Qwen3-TTS HuggingFace Benchmark")
-    parser.add_argument(
-        "--model", type=str, default="Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice", help="HuggingFace model name or path"
-    )
-    parser.add_argument("--num-prompts", type=int, default=50)
-    parser.add_argument("--num-warmups", type=int, default=3)
-    parser.add_argument("--gpu-device", type=int, default=0)
-    parser.add_argument("--voice", type=str, default="Vivian")
-    parser.add_argument("--language", type=str, default="English")
-    parser.add_argument(
-        "--config-name", type=str, default="hf_transformers", help="Label for this config (used in filenames)"
-    )
-    parser.add_argument("--result-dir", type=str, default="results")
-    parser.add_argument("--save-audio", action="store_true", help="Save generated audio files")
-    return parser.parse_args()
-
-
-if __name__ == "__main__":
-    args = parse_args()
-    run_benchmark(args)

benchmarks/qwen3-tts/vllm_omni/bench_async_chunk.py

@@ -1,301 +0,0 @@
-"""Benchmark comparing async_chunk on vs off for Qwen3-TTS.
-
-Measures TTFP (Time-to-First-Packet), E2E latency, and RTF across
-concurrency levels for both async_chunk modes. Saves results as JSON.
-
-Usage:
-    # Run against a server already serving with a given config:
-    python bench_async_chunk.py \
-        --host 127.0.0.1 --port 8000 \
-        --config-name async_chunk_on \
-        --num-prompts 50 \
-        --max-concurrency 1 10 \
-        --result-dir results/
-"""
-
-import argparse
-import asyncio
-import json
-import time
-from dataclasses import asdict, dataclass, field
-from datetime import datetime
-from pathlib import Path
-
-import aiohttp
-import numpy as np
-from tqdm.asyncio import tqdm
-
-PROMPTS = [
-    "Hello, welcome to the voice synthesis benchmark test.",
-    "She said she would be here by noon, but nobody showed up.",
-    "The quick brown fox jumps over the lazy dog near the riverbank.",
-    "I can't believe how beautiful the sunset looks from up here on the mountain.",
-    "Please remember to bring your identification documents to the appointment tomorrow morning.",
-    "Have you ever wondered what it would be like to travel through time and visit ancient civilizations?",
-    "The restaurant on the corner serves the best pasta I have ever tasted in my entire life.",
-    "After the meeting, we should discuss the quarterly results and plan for the next phase.",
-    "Learning a new language takes patience, practice, and a genuine curiosity about other cultures.",
-    "The train leaves at half past seven, so we need to arrive at the station before then.",
-    "Could you please turn down the music a little bit, I'm trying to concentrate on my work.",
-    "It was a dark and stormy night when the old lighthouse keeper heard a knock at the door.",
-]
-
-
-@dataclass
-class RequestResult:
-    success: bool = False
-    ttfp: float = 0.0
-    e2e: float = 0.0
-    audio_bytes: int = 0
-    audio_duration: float = 0.0
-    rtf: float = 0.0
-    prompt: str = ""
-    error: str = ""
-
-
-@dataclass
-class BenchmarkResult:
-    config_name: str = ""
-    concurrency: int = 0
-    num_prompts: int = 0
-    completed: int = 0
-    failed: int = 0
-    duration_s: float = 0.0
-    mean_ttfp_ms: float = 0.0
-    median_ttfp_ms: float = 0.0
-    std_ttfp_ms: float = 0.0
-    p90_ttfp_ms: float = 0.0
-    p95_ttfp_ms: float = 0.0
-    p99_ttfp_ms: float = 0.0
-    mean_e2e_ms: float = 0.0
-    median_e2e_ms: float = 0.0
-    std_e2e_ms: float = 0.0
-    p90_e2e_ms: float = 0.0
-    p95_e2e_ms: float = 0.0
-    p99_e2e_ms: float = 0.0
-    mean_rtf: float = 0.0
-    median_rtf: float = 0.0
-    std_rtf: float = 0.0
-    mean_audio_duration_s: float = 0.0
-    total_audio_duration_s: float = 0.0
-    audio_throughput: float = 0.0
-    request_throughput: float = 0.0
-    per_request: list = field(default_factory=list)
-
-
-def pcm_bytes_to_duration(num_bytes: int, sample_rate: int = 24000, sample_width: int = 2) -> float:
-    return num_bytes / sample_width / sample_rate
-
-
-async def send_tts_request(
-    session: aiohttp.ClientSession,
-    api_url: str,
-    prompt: str,
-    voice: str = "vivian",
-    language: str = "English",
-    stream: bool = True,
-    pbar: tqdm | None = None,
-) -> RequestResult:
-    payload = {
-        "input": prompt,
-        "voice": voice,
-        "language": language,
-        "stream": stream,
-        "response_format": "pcm",
-    }
-
-    result = RequestResult(prompt=prompt)
-    st = time.perf_counter()
-
-    try:
-        async with session.post(api_url, json=payload) as response:
-            if response.status != 200:
-                result.error = f"HTTP {response.status}: {await response.text()}"
-                return result
-
-            first_chunk = True
-            total_bytes = 0
-
-            async for chunk in response.content.iter_any():
-                if first_chunk and len(chunk) > 0:
-                    result.ttfp = time.perf_counter() - st
-                    first_chunk = False
-                total_bytes += len(chunk)
-
-            result.e2e = time.perf_counter() - st
-            result.audio_bytes = total_bytes
-            result.audio_duration = pcm_bytes_to_duration(total_bytes)
-            if result.audio_duration > 0:
-                result.rtf = result.e2e / result.audio_duration
-            result.success = True
-
-    except Exception as e:
-        result.error = str(e)
-        result.e2e = time.perf_counter() - st
-
-    if pbar:
-        pbar.update(1)
-    return result
-
-
-async def run_benchmark(
-    host: str,
-    port: int,
-    num_prompts: int,
-    max_concurrency: int,
-    num_warmups: int = 3,
-    voice: str = "vivian",
-    language: str = "English",
-    stream: bool = True,
-) -> BenchmarkResult:
-    api_url = f"http://{host}:{port}/v1/audio/speech"
-
-    connector = aiohttp.TCPConnector(limit=max_concurrency, limit_per_host=max_concurrency, keepalive_timeout=60)
-    session = aiohttp.ClientSession(connector=connector, timeout=aiohttp.ClientTimeout(total=600))
-
-    if num_warmups > 0:
-        print(f"  Warming up with {num_warmups} requests...")
-        warmup_tasks = [
-            send_tts_request(session, api_url, PROMPTS[i % len(PROMPTS)], voice, language, stream)
-            for i in range(num_warmups)
-        ]
-        await asyncio.gather(*warmup_tasks)
-        print("  Warmup done.")
-
-    request_prompts = [PROMPTS[i % len(PROMPTS)] for i in range(num_prompts)]
-
-    print(f"  Running {num_prompts} requests with concurrency={max_concurrency}...")
-    semaphore = asyncio.Semaphore(max_concurrency)
-    pbar = tqdm(total=num_prompts, desc=f"  concurrency={max_concurrency}")
-
-    async def limited_request(prompt):
-        async with semaphore:
-            return await send_tts_request(session, api_url, prompt, voice, language, stream, pbar)
-
-    start_time = time.perf_counter()
-    tasks = [asyncio.create_task(limited_request(p)) for p in request_prompts]
-    results: list[RequestResult] = await asyncio.gather(*tasks)
-    duration = time.perf_counter() - start_time
-    pbar.close()
-
-    await session.close()
-
-    successful = [r for r in results if r.success]
-    failed = [r for r in results if not r.success]
-
-    bench = BenchmarkResult(
-        concurrency=max_concurrency,
-        num_prompts=num_prompts,
-        completed=len(successful),
-        failed=len(failed),
-        duration_s=duration,
-    )
-
-    if successful:
-        ttfps = [r.ttfp * 1000 for r in successful]
-        e2es = [r.e2e * 1000 for r in successful]
-        rtfs = [r.rtf for r in successful]
-        audio_durs = [r.audio_duration for r in successful]
-
-        bench.mean_ttfp_ms = float(np.mean(ttfps))
-        bench.median_ttfp_ms = float(np.median(ttfps))
-        bench.std_ttfp_ms = float(np.std(ttfps))
-        bench.p90_ttfp_ms = float(np.percentile(ttfps, 90))
-        bench.p95_ttfp_ms = float(np.percentile(ttfps, 95))
-        bench.p99_ttfp_ms = float(np.percentile(ttfps, 99))
-
-        bench.mean_e2e_ms = float(np.mean(e2es))
-        bench.median_e2e_ms = float(np.median(e2es))
-        bench.std_e2e_ms = float(np.std(e2es))
-        bench.p90_e2e_ms = float(np.percentile(e2es, 90))
-        bench.p95_e2e_ms = float(np.percentile(e2es, 95))
-        bench.p99_e2e_ms = float(np.percentile(e2es, 99))
-
-        bench.mean_rtf = float(np.mean(rtfs))
-        bench.median_rtf = float(np.median(rtfs))
-        bench.std_rtf = float(np.std(rtfs))
-
-        bench.mean_audio_duration_s = float(np.mean(audio_durs))
-        bench.total_audio_duration_s = float(np.sum(audio_durs))
-        bench.audio_throughput = bench.total_audio_duration_s / duration
-        bench.request_throughput = len(successful) / duration
-
-        bench.per_request = [
-            {
-                "ttfp_ms": r.ttfp * 1000,
-                "e2e_ms": r.e2e * 1000,
-                "rtf": r.rtf,
-                "audio_duration_s": r.audio_duration,
-                "prompt": r.prompt,
-            }
-            for r in successful
-        ]
-
-    print(f"\n{'=' * 60}")
-    print(f"  Concurrency: {max_concurrency}  |  Completed: {bench.completed}  |  Failed: {bench.failed}")
-    print(f"  Duration: {duration:.2f}s  |  Throughput: {bench.request_throughput:.2f} req/s")
-    print(
-        f"  TTFP (ms):  mean={bench.mean_ttfp_ms:.1f}  median={bench.median_ttfp_ms:.1f}"
-        f"  p90={bench.p90_ttfp_ms:.1f}  p99={bench.p99_ttfp_ms:.1f}"
-    )
-    print(
-        f"  E2E (ms):   mean={bench.mean_e2e_ms:.1f}  median={bench.median_e2e_ms:.1f}"
-        f"  p90={bench.p90_e2e_ms:.1f}  p99={bench.p99_e2e_ms:.1f}"
-    )
-    print(f"  RTF:        mean={bench.mean_rtf:.3f}  median={bench.median_rtf:.3f}")
-    print(f"  Throughput: {bench.audio_throughput:.2f} audio-sec/wall-sec")
-    print(f"{'=' * 60}\n")
-
-    if failed:
-        for r in failed[:3]:
-            print(f"  [ERROR] {r.error[:200]}")
-
-    return bench
-
-
-async def main(args):
-    all_results = []
-
-    for concurrency in args.max_concurrency:
-        result = await run_benchmark(
-            host=args.host,
-            port=args.port,
-            num_prompts=args.num_prompts,
-            max_concurrency=concurrency,
-            num_warmups=args.num_warmups,
-            voice=args.voice,
-            language=args.language,
-            stream=args.stream,
-        )
-        result.config_name = args.config_name
-        all_results.append(asdict(result))
-
-    result_dir = Path(args.result_dir)
-    result_dir.mkdir(parents=True, exist_ok=True)
-    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-    result_file = result_dir / f"bench_{args.config_name}_{timestamp}.json"
-
-    with open(result_file, "w") as f:
-        json.dump(all_results, f, indent=2)
-    print(f"Results saved to {result_file}")
-
-
-def parse_args():
-    parser = argparse.ArgumentParser(description="Qwen3-TTS async_chunk benchmark client")
-    parser.add_argument("--host", type=str, default="127.0.0.1")
-    parser.add_argument("--port", type=int, default=8000)
-    parser.add_argument("--num-prompts", type=int, default=50)
-    parser.add_argument("--max-concurrency", type=int, nargs="+", default=[1, 10])
-    parser.add_argument("--num-warmups", type=int, default=3)
-    parser.add_argument("--voice", type=str, default="vivian")
-    parser.add_argument("--language", type=str, default="English")
-    parser.add_argument("--stream", action="store_true", default=True)
-    parser.add_argument("--no-stream", dest="stream", action="store_false")
-    parser.add_argument("--config-name", type=str, default="async_chunk_on")
-    parser.add_argument("--result-dir", type=str, default="results")
-    return parser.parse_args()
-
-
-if __name__ == "__main__":
-    args = parse_args()
-    asyncio.run(main(args))

benchmarks/qwen3-tts/vllm_omni/bench_tts_serve.py

@@ -1,348 +0,0 @@
-"""Benchmark client for Qwen3-TTS via /v1/audio/speech endpoint.
-
-Measures TTFP (Time-to-First-Packet), E2E latency, and RTF (Real-Time Factor)
-across configurable concurrency levels. Saves results as JSON for plotting.
-
-Usage:
-    python bench_tts_serve.py \
-        --host 127.0.0.1 --port 8000 \
-        --num-prompts 50 \
-        --max-concurrency 1 4 10 \
-        --result-dir results/
-"""
-
-import argparse
-import asyncio
-import json
-import time
-from dataclasses import asdict, dataclass, field
-from datetime import datetime
-from pathlib import Path
-
-import aiohttp
-import numpy as np
-from tqdm.asyncio import tqdm
-
-PROMPTS = [
-    "Hello, welcome to the voice synthesis benchmark test.",
-    "She said she would be here by noon, but nobody showed up.",
-    "The quick brown fox jumps over the lazy dog near the riverbank.",
-    "I can't believe how beautiful the sunset looks from up here on the mountain.",
-    "Please remember to bring your identification documents to the appointment tomorrow morning.",
-    "Have you ever wondered what it would be like to travel through time and visit ancient civilizations?",
-    "The restaurant on the corner serves the best pasta I have ever tasted in my entire life.",
-    "After the meeting, we should discuss the quarterly results and plan for the next phase.",
-    "Learning a new language takes patience, practice, and a genuine curiosity about other cultures.",
-    "The train leaves at half past seven, so we need to arrive at the station before then.",
-    "Could you please turn down the music a little bit, I'm trying to concentrate on my work.",
-    "It was a dark and stormy night when the old lighthouse keeper heard a knock at the door.",
-]
-
-
-@dataclass
-class RequestResult:
-    success: bool = False
-    ttfp: float = 0.0  # Time to first audio packet (seconds)
-    e2e: float = 0.0  # End-to-end latency (seconds)
-    audio_bytes: int = 0  # Total audio bytes received
-    audio_duration: float = 0.0  # Audio duration in seconds (estimated from PCM)
-    rtf: float = 0.0  # Real-time factor = e2e / audio_duration
-    prompt: str = ""
-    error: str = ""
-
-
-@dataclass
-class BenchmarkResult:
-    config_name: str = ""
-    concurrency: int = 0
-    num_prompts: int = 0
-    completed: int = 0
-    failed: int = 0
-    duration_s: float = 0.0
-    # TTFP stats (ms)
-    mean_ttfp_ms: float = 0.0
-    median_ttfp_ms: float = 0.0
-    std_ttfp_ms: float = 0.0
-    p90_ttfp_ms: float = 0.0
-    p95_ttfp_ms: float = 0.0
-    p99_ttfp_ms: float = 0.0
-    # E2E stats (ms)
-    mean_e2e_ms: float = 0.0
-    median_e2e_ms: float = 0.0
-    std_e2e_ms: float = 0.0
-    p90_e2e_ms: float = 0.0
-    p95_e2e_ms: float = 0.0
-    p99_e2e_ms: float = 0.0
-    # RTF stats
-    mean_rtf: float = 0.0
-    median_rtf: float = 0.0
-    std_rtf: float = 0.0
-    p99_rtf: float = 0.0
-    # Audio stats
-    mean_audio_duration_s: float = 0.0
-    total_audio_duration_s: float = 0.0
-    audio_throughput: float = 0.0  # audio_duration / wall_time
-    request_throughput: float = 0.0  # requests / second
-    # Per-request details
-    per_request: list = field(default_factory=list)
-
-
-def pcm_bytes_to_duration(num_bytes: int, sample_rate: int = 24000, sample_width: int = 2) -> float:
-    """Convert raw PCM byte count to duration in seconds."""
-    num_samples = num_bytes / sample_width
-    return num_samples / sample_rate
-
-
-async def send_tts_request(
-    session: aiohttp.ClientSession,
-    api_url: str,
-    prompt: str,
-    voice: str = "vivian",
-    language: str = "English",
-    pbar: tqdm | None = None,
-) -> RequestResult:
-    """Send a streaming TTS request and measure latency metrics."""
-    payload = {
-        "input": prompt,
-        "voice": voice,
-        "language": language,
-        "stream": True,
-        "response_format": "pcm",
-    }
-
-    result = RequestResult(prompt=prompt)
-    st = time.perf_counter()
-
-    try:
-        async with session.post(api_url, json=payload) as response:
-            if response.status != 200:
-                result.error = f"HTTP {response.status}: {await response.text()}"
-                result.success = False
-                return result
-
-            first_chunk = True
-            total_bytes = 0
-
-            async for chunk in response.content.iter_any():
-                if first_chunk and len(chunk) > 0:
-                    result.ttfp = time.perf_counter() - st
-                    first_chunk = False
-                total_bytes += len(chunk)
-
-            result.e2e = time.perf_counter() - st
-            result.audio_bytes = total_bytes
-            result.audio_duration = pcm_bytes_to_duration(total_bytes)
-
-            if result.audio_duration > 0:
-                result.rtf = result.e2e / result.audio_duration
-            result.success = True
-
-    except Exception as e:
-        result.error = str(e)
-        result.success = False
-        result.e2e = time.perf_counter() - st
-
-    if pbar:
-        pbar.update(1)
-    return result
-
-
-async def run_benchmark(
-    host: str,
-    port: int,
-    num_prompts: int,
-    max_concurrency: int,
-    num_warmups: int = 3,
-    voice: str = "vivian",
-    language: str = "English",
-) -> BenchmarkResult:
-    """Run benchmark at a given concurrency level."""
-    api_url = f"http://{host}:{port}/v1/audio/speech"
-
-    connector = aiohttp.TCPConnector(
-        limit=max_concurrency,
-        limit_per_host=max_concurrency,
-        keepalive_timeout=60,
-    )
-    session = aiohttp.ClientSession(
-        connector=connector,
-        timeout=aiohttp.ClientTimeout(total=600),
-    )
-
-    # Warmup
-    if num_warmups > 0:
-        print(f"  Warming up with {num_warmups} requests...")
-        warmup_tasks = []
-        for i in range(num_warmups):
-            prompt = PROMPTS[i % len(PROMPTS)]
-            warmup_tasks.append(send_tts_request(session, api_url, prompt, voice, language))
-        await asyncio.gather(*warmup_tasks)
-        print("  Warmup done.")
-
-    # Build request list
-    request_prompts = [PROMPTS[i % len(PROMPTS)] for i in range(num_prompts)]
-
-    # Run benchmark
-    print(f"  Running {num_prompts} requests with concurrency={max_concurrency}...")
-    semaphore = asyncio.Semaphore(max_concurrency)
-    pbar = tqdm(total=num_prompts, desc=f"  concurrency={max_concurrency}")
-
-    async def limited_request(prompt):
-        async with semaphore:
-            return await send_tts_request(session, api_url, prompt, voice, language, pbar)
-
-    start_time = time.perf_counter()
-    tasks = [asyncio.create_task(limited_request(p)) for p in request_prompts]
-    results: list[RequestResult] = await asyncio.gather(*tasks)
-    duration = time.perf_counter() - start_time
-    pbar.close()
-
-    await session.close()
-
-    # Compute stats
-    successful = [r for r in results if r.success]
-    failed = [r for r in results if not r.success]
-
-    bench = BenchmarkResult(
-        concurrency=max_concurrency,
-        num_prompts=num_prompts,
-        completed=len(successful),
-        failed=len(failed),
-        duration_s=duration,
-    )
-
-    if successful:
-        ttfps = [r.ttfp * 1000 for r in successful]  # convert to ms
-        e2es = [r.e2e * 1000 for r in successful]
-        rtfs = [r.rtf for r in successful]
-        audio_durs = [r.audio_duration for r in successful]
-
-        bench.mean_ttfp_ms = float(np.mean(ttfps))
-        bench.median_ttfp_ms = float(np.median(ttfps))
-        bench.std_ttfp_ms = float(np.std(ttfps))
-        bench.p90_ttfp_ms = float(np.percentile(ttfps, 90))
-        bench.p95_ttfp_ms = float(np.percentile(ttfps, 95))
-        bench.p99_ttfp_ms = float(np.percentile(ttfps, 99))
-
-        bench.mean_e2e_ms = float(np.mean(e2es))
-        bench.median_e2e_ms = float(np.median(e2es))
-        bench.std_e2e_ms = float(np.std(e2es))
-        bench.p90_e2e_ms = float(np.percentile(e2es, 90))
-        bench.p95_e2e_ms = float(np.percentile(e2es, 95))
-        bench.p99_e2e_ms = float(np.percentile(e2es, 99))
-
-        bench.mean_rtf = float(np.mean(rtfs))
-        bench.median_rtf = float(np.median(rtfs))
-        bench.std_rtf = float(np.std(rtfs))
-        bench.p99_rtf = float(np.percentile(rtfs, 99))
-
-        bench.mean_audio_duration_s = float(np.mean(audio_durs))
-        bench.total_audio_duration_s = float(np.sum(audio_durs))
-        bench.audio_throughput = bench.total_audio_duration_s / duration
-        bench.request_throughput = len(successful) / duration
-
-        bench.per_request = [
-            {
-                "ttfp_ms": r.ttfp * 1000,
-                "e2e_ms": r.e2e * 1000,
-                "rtf": r.rtf,
-                "audio_duration_s": r.audio_duration,
-                "prompt": r.prompt,
-            }
-            for r in successful
-        ]
-
-    # Print summary in standardized performance template
-    W = 50
-    print("")
-    print(f"{'=' * W}")
-    print(f"{'Serving Benchmark Result':^{W}}")
-    print(f"{'=' * W}")
-    print(f"{'Successful requests:':<40}{bench.completed:<10}")
-    print(f"{'Failed requests:':<40}{bench.failed:<10}")
-    print(f"{'Maximum request concurrency:':<40}{max_concurrency:<10}")
-    print(f"{'Benchmark duration (s):':<40}{duration:<10.2f}")
-    print(f"{'Request throughput (req/s):':<40}{bench.request_throughput:<10.2f}")
-    print(f"{'-' * W}")
-    print(f"{'End-to-end Latency':^{W}}")
-    print(f"{'-' * W}")
-    print(f"{'Mean E2EL (ms):':<40}{bench.mean_e2e_ms:<10.2f}")
-    print(f"{'Median E2EL (ms):':<40}{bench.median_e2e_ms:<10.2f}")
-    print(f"{'P99 E2EL (ms):':<40}{bench.p99_e2e_ms:<10.2f}")
-    print(f"{'=' * W}")
-    print(f"{'Audio Result':^{W}}")
-    print(f"{'=' * W}")
-    print(f"{'Total audio duration generated (s):':<40}{bench.total_audio_duration_s:<10.2f}")
-    print(f"{'Audio throughput (audio duration/s):':<40}{bench.audio_throughput:<10.2f}")
-    print(f"{'-' * W}")
-    print(f"{'Time to First Packet':^{W}}")
-    print(f"{'-' * W}")
-    print(f"{'Mean AUDIO_TTFP (ms):':<40}{bench.mean_ttfp_ms:<10.2f}")
-    print(f"{'Median AUDIO_TTFP (ms):':<40}{bench.median_ttfp_ms:<10.2f}")
-    print(f"{'P99 AUDIO_TTFP (ms):':<40}{bench.p99_ttfp_ms:<10.2f}")
-    print(f"{'-' * W}")
-    print(f"{'Real Time Factor':^{W}}")
-    print(f"{'-' * W}")
-    print(f"{'Mean AUDIO_RTF:':<40}{bench.mean_rtf:<10.3f}")
-    print(f"{'Median AUDIO_RTF:':<40}{bench.median_rtf:<10.3f}")
-    print(f"{'P99 AUDIO_RTF:':<40}{bench.p99_rtf:<10.3f}")
-    print(f"{'=' * W}")
-    print("")
-
-    if failed:
-        for r in failed[:3]:
-            print(f"  [ERROR] {r.error[:200]}")
-
-    return bench
-
-
-async def main(args):
-    all_results = []
-
-    for concurrency in args.max_concurrency:
-        result = await run_benchmark(
-            host=args.host,
-            port=args.port,
-            num_prompts=args.num_prompts,
-            max_concurrency=concurrency,
-            num_warmups=args.num_warmups,
-            voice=args.voice,
-            language=args.language,
-        )
-        result.config_name = args.config_name
-        all_results.append(asdict(result))
-
-    # Save results
-    result_dir = Path(args.result_dir)
-    result_dir.mkdir(parents=True, exist_ok=True)
-    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-    result_file = result_dir / f"bench_{args.config_name}_{timestamp}.json"
-
-    with open(result_file, "w") as f:
-        json.dump(all_results, f, indent=2)
-    print(f"Results saved to {result_file}")
-
-    return all_results
-
-
-def parse_args():
-    parser = argparse.ArgumentParser(description="Qwen3-TTS Benchmark Client")
-    parser.add_argument("--host", type=str, default="127.0.0.1")
-    parser.add_argument("--port", type=int, default=8000)
-    parser.add_argument("--num-prompts", type=int, default=50, help="Number of prompts per concurrency level")
-    parser.add_argument(  # noqa: E501
-        "--max-concurrency", type=int, nargs="+", default=[1, 4, 10], help="Concurrency levels to test"
-    )
-    parser.add_argument("--num-warmups", type=int, default=3)
-    parser.add_argument("--voice", type=str, default="vivian")
-    parser.add_argument("--language", type=str, default="English")
-    parser.add_argument(
-        "--config-name", type=str, default="async_chunk", help="Label for this config (used in filenames)"
-    )
-    parser.add_argument("--result-dir", type=str, default="results")
-    return parser.parse_args()
-
-
-if __name__ == "__main__":
-    args = parse_args()
-    asyncio.run(main(args))

benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs16.yaml

@@ -1,94 +0,0 @@
-# Qwen3-TTS max_num_seqs=16 config (streaming with async_chunk)
-# High-throughput concurrent request processing
-# 2-stage pipeline: Talker -> Code2Wav
-async_chunk: true
-stage_args:
-  - stage_id: 0
-    stage_type: llm
-    is_comprehension: true
-    runtime:
-      devices: "0"
-    engine_args:
-      max_num_seqs: 16
-      model_stage: qwen3_tts
-      model_arch: Qwen3TTSTalkerForConditionalGeneration
-      worker_type: ar
-      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      enforce_eager: false
-      trust_remote_code: true
-      async_scheduling: true
-      enable_prefix_caching: false
-      engine_output_type: latent
-      gpu_memory_utilization: 0.3
-      distributed_executor_backend: "mp"
-      max_num_batched_tokens: 4096
-      max_model_len: 4096
-      custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_tts.talker2code2wav_async_chunk
-    output_connectors:
-      to_stage_1: connector_of_shared_memory
-    default_sampling_params:
-      temperature: 0.9
-      top_k: 50
-      max_tokens: 4096
-      seed: 42
-      detokenize: false
-      repetition_penalty: 1.05
-      stop_token_ids: [2150]
-
-  - stage_id: 1
-    stage_type: llm
-    runtime:
-      devices: "0"
-    engine_args:
-      max_num_seqs: 16
-      model_stage: code2wav
-      model_arch: Qwen3TTSCode2Wav
-      worker_type: generation
-      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
-      enforce_eager: true
-      trust_remote_code: true
-      async_scheduling: true
-      enable_prefix_caching: false
-      engine_output_type: audio
-      gpu_memory_utilization: 0.2
-      distributed_executor_backend: "mp"
-      max_num_batched_tokens: 16384
-      max_model_len: 32768
-    engine_input_source: [0]
-    final_output: true
-    final_output_type: audio
-    input_connectors:
-      from_stage_0: connector_of_shared_memory
-    tts_args:
-      max_instructions_length: 500
-    default_sampling_params:
-      temperature: 0.0
-      top_p: 1.0
-      top_k: -1
-      max_tokens: 65536
-      seed: 42
-      detokenize: true
-      repetition_penalty: 1.0
-
-runtime:
-  enabled: true
-  defaults:
-    window_size: -1
-    max_inflight: 16
-
-  connectors:
-    connector_of_shared_memory:
-      name: SharedMemoryConnector
-      extra:
-        shm_threshold_bytes: 65536
-        codec_streaming: true
-        connector_get_sleep_s: 0.01
-        connector_get_max_wait_first_chunk: 3000
-        connector_get_max_wait: 300
-        codec_chunk_frames: 25
-        codec_left_context_frames: 25
-
-  edges:
-    - from: 0
-      to: 1
-      window_size: -1

benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs4.yaml

@@ -1,94 +0,0 @@
-# Qwen3-TTS batch_size=4 config (streaming with async_chunk)
-# Enables concurrent request processing
-# 2-stage pipeline: Talker -> Code2Wav
-async_chunk: true
-stage_args:
-  - stage_id: 0
-    stage_type: llm
-    is_comprehension: true
-    runtime:
-      devices: "0"
-    engine_args:
-      max_num_seqs: 4
-      model_stage: qwen3_tts
-      model_arch: Qwen3TTSTalkerForConditionalGeneration
-      worker_type: ar
-      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      enforce_eager: false
-      trust_remote_code: true
-      async_scheduling: true
-      enable_prefix_caching: false
-      engine_output_type: latent
-      gpu_memory_utilization: 0.3
-      distributed_executor_backend: "mp"
-      max_num_batched_tokens: 512
-      max_model_len: 4096
-      custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_tts.talker2code2wav_async_chunk
-    output_connectors:
-      to_stage_1: connector_of_shared_memory
-    default_sampling_params:
-      temperature: 0.9
-      top_k: 50
-      max_tokens: 4096
-      seed: 42
-      detokenize: false
-      repetition_penalty: 1.05
-      stop_token_ids: [2150]
-
-  - stage_id: 1
-    stage_type: llm
-    runtime:
-      devices: "0"
-    engine_args:
-      max_num_seqs: 4
-      model_stage: code2wav
-      model_arch: Qwen3TTSCode2Wav
-      worker_type: generation
-      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
-      enforce_eager: true
-      trust_remote_code: true
-      async_scheduling: true
-      enable_prefix_caching: false
-      engine_output_type: audio
-      gpu_memory_utilization: 0.2
-      distributed_executor_backend: "mp"
-      max_num_batched_tokens: 8192
-      max_model_len: 32768
-    engine_input_source: [0]
-    final_output: true
-    final_output_type: audio
-    input_connectors:
-      from_stage_0: connector_of_shared_memory
-    tts_args:
-      max_instructions_length: 500
-    default_sampling_params:
-      temperature: 0.0
-      top_p: 1.0
-      top_k: -1
-      max_tokens: 65536
-      seed: 42
-      detokenize: true
-      repetition_penalty: 1.0
-
-runtime:
-  enabled: true
-  defaults:
-    window_size: -1
-    max_inflight: 4
-
-  connectors:
-    connector_of_shared_memory:
-      name: SharedMemoryConnector
-      extra:
-        shm_threshold_bytes: 65536
-        codec_streaming: true
-        connector_get_sleep_s: 0.01
-        connector_get_max_wait_first_chunk: 3000
-        connector_get_max_wait: 300
-        codec_chunk_frames: 25
-        codec_left_context_frames: 25
-
-  edges:
-    - from: 0
-      to: 1
-      window_size: -1

benchmarks/qwen3-tts/vllm_omni/plot_async_chunk.py

@@ -1,249 +0,0 @@
-"""Plot TTFP comparison: async_chunk off vs on.
-
-Generates a bar chart with improvement arrows, matching the Qwen3-Omni
-async_chunk benchmark figure style.
-
-Usage:
-    python plot_async_chunk.py \
-        --off results/bench_async_chunk_off_*.json \
-        --on results/bench_async_chunk_on_*.json \
-        --output results/qwen3_tts_async_chunk_ttfp.png
-
-    # Also supports E2E and RTF metrics:
-    python plot_async_chunk.py \
-        --off results/bench_async_chunk_off_*.json \
-        --on results/bench_async_chunk_on_*.json \
-        --metric e2e \
-        --output results/qwen3_tts_async_chunk_e2e.png
-"""
-
-import argparse
-import json
-from pathlib import Path
-
-import matplotlib.pyplot as plt
-import numpy as np
-
-METRIC_CONFIG = {
-    "ttfp": {
-        "key": "mean_ttfp_ms",
-        "ylabel": "TTFP (s)",
-        "title": "TTFP (Time to First Audio Packet) - Qwen3-TTS, by concurrency",
-        "to_seconds": True,
-    },
-    "e2e": {
-        "key": "mean_e2e_ms",
-        "ylabel": "E2E (s)",
-        "title": "E2E Latency - Qwen3-TTS, by concurrency",
-        "to_seconds": True,
-    },
-    "rtf": {
-        "key": "mean_rtf",
-        "ylabel": "RTF",
-        "title": "Real-Time Factor - Qwen3-TTS, by concurrency",
-        "to_seconds": False,
-    },
-}
-
-
-def load_results(path: str) -> list[dict]:
-    with open(path) as f:
-        return json.load(f)
-
-
-def plot_ttfp_comparison(
-    off_results: list[dict],
-    on_results: list[dict],
-    metric: str,
-    output_path: str,
-    title_override: str | None = None,
-):
-    cfg = METRIC_CONFIG[metric]
-    key = cfg["key"]
-    to_seconds = cfg["to_seconds"]
-
-    off_map = {r["concurrency"]: r for r in off_results}
-    on_map = {r["concurrency"]: r for r in on_results}
-    concurrencies = sorted(set(off_map.keys()) & set(on_map.keys()))
-
-    off_vals = []
-    on_vals = []
-    for c in concurrencies:
-        v_off = off_map[c][key]
-        v_on = on_map[c][key]
-        if to_seconds:
-            v_off /= 1000.0
-            v_on /= 1000.0
-        off_vals.append(v_off)
-        on_vals.append(v_on)
-
-    fig, ax = plt.subplots(figsize=(8, 6))
-
-    x = np.arange(len(concurrencies))
-    width = 0.3
-
-    ax.bar(x - width / 2, off_vals, width, label="async_chunk off", color="#87CEEB", edgecolor="none")
-    ax.bar(x + width / 2, on_vals, width, label="async_chunk on", color="#FFF8DC", edgecolor="#DDD8B8")
-
-    # Draw improvement arrows and labels
-    for i in range(len(concurrencies)):
-        v_off = off_vals[i]
-        v_on = on_vals[i]
-        if v_on > 0:
-            improvement = v_off / v_on
-        else:
-            improvement = float("inf")
-
-        # Arrow from top of off-bar to top of on-bar
-        arrow_start_x = x[i] - width / 2
-        arrow_start_y = v_off * 0.95
-        arrow_end_x = x[i] + width / 2
-        arrow_end_y = v_on * 1.05
-
-        ax.annotate(
-            "",
-            xy=(arrow_end_x, arrow_end_y),
-            xytext=(arrow_start_x, arrow_start_y),
-            arrowprops=dict(arrowstyle="->", color="red", lw=1.5),
-        )
-
-        # Improvement label
-        label_x = (arrow_start_x + arrow_end_x) / 2
-        label_y = arrow_start_y + (v_off - v_on) * 0.15
-        ax.text(
-            label_x,
-            label_y,
-            f"{improvement:.1f}x improvement",
-            ha="center",
-            va="bottom",
-            fontsize=10,
-            color="red",
-            fontweight="bold",
-        )
-
-    title = title_override or cfg["title"]
-    ax.set_title(title, fontsize=13, fontweight="bold")
-    ax.set_ylabel(cfg["ylabel"], fontsize=12)
-    ax.set_xlabel("Max concurrency", fontsize=12)
-    ax.set_xticks(x)
-    ax.set_xticklabels([str(c) for c in concurrencies])
-    ax.set_yscale("log")
-    ax.legend(loc="upper left", fontsize=11)
-    ax.grid(axis="y", alpha=0.3, linestyle="--")
-    ax.set_axisbelow(True)
-
-    plt.tight_layout()
-    Path(output_path).parent.mkdir(parents=True, exist_ok=True)
-    plt.savefig(output_path, dpi=150, bbox_inches="tight")
-    print(f"Plot saved to {output_path}")
-    plt.close()
-
-
-def plot_all_metrics(off_results: list[dict], on_results: list[dict], output_path: str):
-    """Generate a 1x3 subplot with TTFP, E2E, and RTF comparisons."""
-    off_map = {r["concurrency"]: r for r in off_results}
-    on_map = {r["concurrency"]: r for r in on_results}
-    concurrencies = sorted(set(off_map.keys()) & set(on_map.keys()))
-
-    fig, axes = plt.subplots(1, 3, figsize=(18, 6))
-    fig.suptitle("Qwen3-TTS: async_chunk on vs off", fontsize=15, fontweight="bold")
-
-    for ax, metric in zip(axes, ["ttfp", "e2e", "rtf"]):
-        cfg = METRIC_CONFIG[metric]
-        key = cfg["key"]
-        to_seconds = cfg["to_seconds"]
-
-        off_vals = []
-        on_vals = []
-        for c in concurrencies:
-            v_off = off_map[c][key]
-            v_on = on_map[c][key]
-            if to_seconds:
-                v_off /= 1000.0
-                v_on /= 1000.0
-            off_vals.append(v_off)
-            on_vals.append(v_on)
-
-        x = np.arange(len(concurrencies))
-        width = 0.3
-        ax.bar(x - width / 2, off_vals, width, label="async_chunk off", color="#87CEEB")
-        ax.bar(x + width / 2, on_vals, width, label="async_chunk on", color="#FFF8DC", edgecolor="#DDD8B8")
-
-        for i in range(len(concurrencies)):
-            if on_vals[i] > 0:
-                improvement = off_vals[i] / on_vals[i]
-                ax.annotate(
-                    "",
-                    xy=(x[i] + width / 2, on_vals[i] * 1.05),
-                    xytext=(x[i] - width / 2, off_vals[i] * 0.95),
-                    arrowprops=dict(arrowstyle="->", color="red", lw=1.5),
-                )
-                label_y = off_vals[i] * 0.85
-                ax.text(x[i], label_y, f"{improvement:.1f}x", ha="center", fontsize=10, color="red", fontweight="bold")
-
-        ax.set_title(cfg["title"].split(" - ")[0], fontsize=12, fontweight="bold")
-        ax.set_ylabel(cfg["ylabel"], fontsize=11)
-        ax.set_xlabel("Max concurrency", fontsize=11)
-        ax.set_xticks(x)
-        ax.set_xticklabels([str(c) for c in concurrencies])
-        if metric != "rtf":
-            ax.set_yscale("log")
-        ax.legend(fontsize=9)
-        ax.grid(axis="y", alpha=0.3, linestyle="--")
-        ax.set_axisbelow(True)
-
-    plt.tight_layout()
-    Path(output_path).parent.mkdir(parents=True, exist_ok=True)
-    plt.savefig(output_path, dpi=150, bbox_inches="tight")
-    print(f"Plot saved to {output_path}")
-    plt.close()
-
-
-def print_table(off_results: list[dict], on_results: list[dict]):
-    off_map = {r["concurrency"]: r for r in off_results}
-    on_map = {r["concurrency"]: r for r in on_results}
-    concurrencies = sorted(set(off_map.keys()) & set(on_map.keys()))
-
-    print("\n## Benchmark Results: async_chunk off vs on\n")
-    print("| Metric | Concurrency | async_chunk off | async_chunk on | Improvement |")
-    print("| --- | --- | --- | --- | --- |")
-
-    for name, key, fmt in [
-        ("TTFP (ms)", "mean_ttfp_ms", ".1f"),
-        ("E2E (ms)", "mean_e2e_ms", ".1f"),
-        ("RTF", "mean_rtf", ".3f"),
-        ("Throughput", "audio_throughput", ".2f"),
-    ]:
-        for c in concurrencies:
-            v_off = off_map[c].get(key, 0)
-            v_on = on_map[c].get(key, 0)
-            if v_on > 0 and key != "audio_throughput":
-                ratio = f"{v_off / v_on:.1f}x"
-            elif v_off > 0 and key == "audio_throughput":
-                ratio = f"{v_on / v_off:.1f}x"
-            else:
-                ratio = "N/A"
-            print(f"| {name} | {c} | {v_off:{fmt}} | {v_on:{fmt}} | {ratio} |")
-
-
-def parse_args():
-    parser = argparse.ArgumentParser(description="Plot async_chunk comparison for Qwen3-TTS")
-    parser.add_argument("--off", type=str, required=True, help="JSON results for async_chunk off")
-    parser.add_argument("--on", type=str, required=True, help="JSON results for async_chunk on")
-    parser.add_argument("--metric", type=str, default="ttfp", choices=["ttfp", "e2e", "rtf", "all"])
-    parser.add_argument("--output", type=str, default="results/qwen3_tts_async_chunk.png")
-    parser.add_argument("--title", type=str, default=None, help="Custom title override")
-    return parser.parse_args()
-
-
-if __name__ == "__main__":
-    args = parse_args()
-    off_results = load_results(args.off)
-    on_results = load_results(args.on)
-
-    print_table(off_results, on_results)
-
-    if args.metric == "all":
-        plot_all_metrics(off_results, on_results, args.output)
-    else:
-        plot_ttfp_comparison(off_results, on_results, args.metric, args.output, args.title)

benchmarks/qwen3-tts/vllm_omni/run_async_chunk_benchmark.sh

@@ -1,167 +0,0 @@
-#!/bin/bash
-# Qwen3-TTS async_chunk on vs off Benchmark
-#
-# Starts two servers (async_chunk on and off), benchmarks both,
-# and generates comparison plots.
-#
-# Usage:
-#   bash run_async_chunk_benchmark.sh
-#
-# Environment variables:
-#   GPU_DEVICE       - GPU index (default: 0)
-#   MODEL            - Model path (default: Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice)
-#   NUM_PROMPTS      - Prompts per concurrency level (default: 50)
-#   CONCURRENCY      - Space-separated concurrency levels (default: "1 10")
-#   PORT_ON          - Port for async_chunk on server (default: 8000)
-#   PORT_OFF         - Port for async_chunk off server (default: 8001)
-
-set -euo pipefail
-
-SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
-PROJECT_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)"
-cd "$PROJECT_ROOT"
-
-GPU_DEVICE="${GPU_DEVICE:-0}"
-MODEL="${MODEL:-Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice}"
-NUM_PROMPTS="${NUM_PROMPTS:-50}"
-CONCURRENCY="${CONCURRENCY:-1 10}"
-NUM_WARMUPS="${NUM_WARMUPS:-3}"
-PORT_ON="${PORT_ON:-8000}"
-PORT_OFF="${PORT_OFF:-8001}"
-RESULT_DIR="${SCRIPT_DIR}/results"
-TIMESTAMP="$(date +%Y%m%d_%H%M%S)"
-
-STAGE_CONFIG_ON="vllm_omni/model_executor/stage_configs/qwen3_tts.yaml"
-STAGE_CONFIG_OFF="vllm_omni/model_executor/stage_configs/qwen3_tts_no_async_chunk.yaml"
-
-mkdir -p "${RESULT_DIR}"
-
-echo "============================================================"
-echo " Qwen3-TTS async_chunk Benchmark"
-echo "============================================================"
-echo " GPU:            ${GPU_DEVICE}"
-echo " Model:          ${MODEL}"
-echo " Prompts:        ${NUM_PROMPTS}"
-echo " Concurrency:    ${CONCURRENCY}"
-echo " Port (on/off):  ${PORT_ON} / ${PORT_OFF}"
-echo " Results:        ${RESULT_DIR}"
-echo "============================================================"
-
-cleanup() {
-    echo "Cleaning up servers..."
-    kill "$PID_ON" 2>/dev/null || true
-    kill "$PID_OFF" 2>/dev/null || true
-    wait "$PID_ON" 2>/dev/null || true
-    wait "$PID_OFF" 2>/dev/null || true
-}
-trap cleanup EXIT
-
-wait_for_server() {
-    local port=$1
-    local name=$2
-    local max_wait=300
-    local elapsed=0
-    echo "Waiting for ${name} server on port ${port}..."
-    while ! curl -s "http://localhost:${port}/health" >/dev/null 2>&1; do
-        sleep 5
-        elapsed=$((elapsed + 5))
-        if [ $elapsed -ge $max_wait ]; then
-            echo "ERROR: ${name} server failed to start within ${max_wait}s"
-            exit 1
-        fi
-    done
-    echo "${name} server ready (${elapsed}s)"
-}
-
-# ---- Phase 1: Start async_chunk ON server ----
-echo ""
-echo "[Phase 1] Starting async_chunk ON server on port ${PORT_ON}..."
-CUDA_VISIBLE_DEVICES=${GPU_DEVICE} vllm-omni serve "${MODEL}" \
-    --stage-configs-path "${STAGE_CONFIG_ON}" \
-    --host 0.0.0.0 --port "${PORT_ON}" \
-    --trust-remote-code --enforce-eager --omni \
-    > "${RESULT_DIR}/server_on_${TIMESTAMP}.log" 2>&1 &
-PID_ON=$!
-
-wait_for_server "${PORT_ON}" "async_chunk_on"
-
-echo "[Phase 1] Benchmarking async_chunk ON..."
-# shellcheck disable=SC2086
-python "${SCRIPT_DIR}/bench_async_chunk.py" \
-    --host 127.0.0.1 --port "${PORT_ON}" \
-    --config-name "async_chunk_on" \
-    --num-prompts "${NUM_PROMPTS}" \
-    --max-concurrency ${CONCURRENCY} \
-    --num-warmups "${NUM_WARMUPS}" \
-    --result-dir "${RESULT_DIR}"
-
-echo "[Phase 1] Stopping async_chunk ON server..."
-kill "$PID_ON" 2>/dev/null || true
-wait "$PID_ON" 2>/dev/null || true
-sleep 5
-
-# ---- Phase 2: Start async_chunk OFF server ----
-echo ""
-echo "[Phase 2] Starting async_chunk OFF server on port ${PORT_OFF}..."
-CUDA_VISIBLE_DEVICES=${GPU_DEVICE} vllm-omni serve "${MODEL}" \
-    --stage-configs-path "${STAGE_CONFIG_OFF}" \
-    --host 0.0.0.0 --port "${PORT_OFF}" \
-    --trust-remote-code --enforce-eager --omni \
-    > "${RESULT_DIR}/server_off_${TIMESTAMP}.log" 2>&1 &
-PID_OFF=$!
-
-wait_for_server "${PORT_OFF}" "async_chunk_off"
-
-echo "[Phase 2] Benchmarking async_chunk OFF (non-streaming)..."
-# shellcheck disable=SC2086
-python "${SCRIPT_DIR}/bench_async_chunk.py" \
-    --host 127.0.0.1 --port "${PORT_OFF}" \
-    --config-name "async_chunk_off" \
-    --num-prompts "${NUM_PROMPTS}" \
-    --max-concurrency ${CONCURRENCY} \
-    --num-warmups "${NUM_WARMUPS}" \
-    --no-stream \
-    --result-dir "${RESULT_DIR}"
-
-echo "[Phase 2] Stopping async_chunk OFF server..."
-kill "$PID_OFF" 2>/dev/null || true
-wait "$PID_OFF" 2>/dev/null || true
-
-# ---- Phase 3: Plot results ----
-echo ""
-echo "[Phase 3] Generating plots..."
-
-# Find the latest result files
-RESULT_ON=$(ls -t "${RESULT_DIR}"/bench_async_chunk_on_*.json 2>/dev/null | head -1)
-RESULT_OFF=$(ls -t "${RESULT_DIR}"/bench_async_chunk_off_*.json 2>/dev/null | head -1)
-
-if [ -z "$RESULT_ON" ] || [ -z "$RESULT_OFF" ]; then
-    echo "ERROR: Could not find result files. Check logs in ${RESULT_DIR}/"
-    exit 1
-fi
-
-echo "  ON results:  ${RESULT_ON}"
-echo "  OFF results: ${RESULT_OFF}"
-
-# TTFP comparison (main figure)
-python "${SCRIPT_DIR}/plot_async_chunk.py" \
-    --off "${RESULT_OFF}" \
-    --on "${RESULT_ON}" \
-    --metric ttfp \
-    --output "${RESULT_DIR}/qwen3_tts_async_chunk_ttfp.png"
-
-# All metrics comparison
-python "${SCRIPT_DIR}/plot_async_chunk.py" \
-    --off "${RESULT_OFF}" \
-    --on "${RESULT_ON}" \
-    --metric all \
-    --output "${RESULT_DIR}/qwen3_tts_async_chunk_all.png"
-
-echo ""
-echo "============================================================"
-echo " Benchmark complete!"
-echo " Results: ${RESULT_DIR}/"
-echo " Plots:"
-echo "   - ${RESULT_DIR}/qwen3_tts_async_chunk_ttfp.png"
-echo "   - ${RESULT_DIR}/qwen3_tts_async_chunk_all.png"
-echo "============================================================"

benchmarks/tts/README.md

@@ -0,0 +1,227 @@
+# TTS Universal Benchmark
+
+A model-agnostic serving benchmark for TTS models in vllm-omni. One CLI
+(`bench_tts.py`) + one YAML registry (`model_configs.yaml`) drive perf and
+quality runs for every registered checkpoint: **Qwen3-TTS** (Base / CustomVoice)
+and **VoxCPM2** today, more to come.
+
+The same three task types — `voice_clone`, `default_voice`, `voice_design` —
+are wired into both the manual CLI and the DFX nightly CI matrix
+(`tests/dfx/perf/tests/test_tts.json`).
+
+## Quick start
+
+### 1. Start the server
+
+```bash
+vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base --omni --port 8000
+```
+
+The server auto-loads its Deploy YAML from `vllm_omni/deploy/qwen3_tts.yaml`
+(Pipeline + Deploy schema introduced in #2383). No `--stage-configs-path` or
+`--deploy-config` flag is needed for any registered model.
+
+### 2. Run the benchmark (`vllm bench serve --omni`)
+
+The primary, directly-controllable path. Copy-paste one of these and tweak
+any bench flag (sampling params, endpoint, extra body, warmups, etc.):
+
+#### voice_clone (Qwen3-TTS-Base, seed-tts dataset)
+
+```bash
+vllm bench serve --omni \
+    --host 127.0.0.1 --port 8000 \
+    --model Qwen/Qwen3-TTS-12Hz-1.7B-Base \
+    --backend openai-audio-speech \
+    --endpoint /v1/audio/speech \
+    --dataset-name seed-tts \
+    --dataset-path /path/to/seed-tts-eval \
+    --seed-tts-locale en \
+    --num-prompts 20 --num-warmups 2 \
+    --extra-body '{"task_type":"Base"}' \
+    --max-concurrency 1 --request-rate inf \
+    --percentile-metrics ttft,e2el,audio_rtf,audio_ttfp,audio_duration \
+    --save-result --result-dir ./results
+```
+
+#### default_voice (Qwen3-TTS-CustomVoice, bundled seed_tts_smoke)
+
+```bash
+vllm bench serve --omni \
+    --host 127.0.0.1 --port 8000 \
+    --model Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
+    --backend openai-audio-speech \
+    --endpoint /v1/audio/speech \
+    --dataset-name seed-tts-text \
+    --dataset-path benchmarks/build_dataset/seed_tts_smoke \
+    --seed-tts-locale en \
+    --num-prompts 20 --num-warmups 2 \
+    --extra-body '{"voice":"Vivian","language":"English","task_type":"CustomVoice"}' \
+    --max-concurrency 1 --request-rate inf \
+    --percentile-metrics ttft,e2el,audio_rtf,audio_ttfp,audio_duration \
+    --save-result --result-dir ./results
+```
+
+#### voice_design (Qwen3-TTS-CustomVoice, bundled seed_tts_design)
+
+```bash
+vllm bench serve --omni \
+    --host 127.0.0.1 --port 8000 \
+    --model Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
+    --backend openai-audio-speech \
+    --endpoint /v1/audio/speech \
+    --dataset-name seed-tts-design \
+    --dataset-path benchmarks/build_dataset/seed_tts_design \
+    --seed-tts-locale en \
+    --num-prompts 20 --num-warmups 2 \
+    --extra-body '{"task_type":"VoiceDesign","language":"English"}' \
+    --max-concurrency 1 --request-rate inf \
+    --percentile-metrics ttft,e2el,audio_rtf,audio_ttfp,audio_duration \
+    --save-result --result-dir ./results
+```
+
+#### Add WER / SIM / UTMOS to any of the above
+
+Append `--seed-tts-wer-eval` (and optionally `SEED_TTS_EVAL_DEVICE=cuda:0`
+in the env, per PR #2558). This triggers the seed-tts-eval protocol:
+Whisper-large-v3 ASR → WER, WavLM embeddings → SIM, balacoon/utmos → UTMOS.
+
+### 3. Convenience wrapper (`bench_tts.py`)
+
+If you're running the **canonical** configuration for a registered model,
+`bench_tts.py` loads the right defaults from `model_configs.yaml` and
+emits the exact `vllm bench serve --omni` command above — useful for
+concurrency sweeps and multi-task runs:
+
+```bash
+# Smallest smoke — 5 prompts, concurrency=1
+python benchmarks/tts/bench_tts.py \
+    --model Qwen/Qwen3-TTS-12Hz-1.7B-Base \
+    --task voice_clone \
+    --dataset-path /path/to/seed-tts-eval \
+    --concurrency 1 --num-prompts 5 \
+    --output-dir ./results
+
+# Full concurrency sweep
+python benchmarks/tts/bench_tts.py \
+    --model Qwen/Qwen3-TTS-12Hz-1.7B-Base \
+    --task voice_clone \
+    --dataset-path /path/to/seed-tts-eval \
+    --concurrency 1 2 4 8 16 32 \
+    --num-prompts 20 \
+    --output-dir ./results
+
+# With WER / SIM / UTMOS quality eval (adds ASR + embedding compute)
+python benchmarks/tts/bench_tts.py \
+    --model Qwen/Qwen3-TTS-12Hz-1.7B-Base \
+    --task voice_clone \
+    --dataset-path /path/to/seed-tts-eval \
+    --wer-eval \
+    --concurrency 4 --num-prompts 200 \
+    --output-dir ./results
+```
+
+### 4. Plot a sweep
+
+```bash
+python benchmarks/tts/plot_results.py \
+    --results ./results/*.json \
+    --output ./results/curve.png
+```
+
+Outputs TTFP / RTF / throughput curves (and a markdown table) for every
+`(task, concurrency)` combination in the result set.
+
+## Task types
+
+| Task            | Dataset           | Request body                                        | Checkpoints that support it              |
+|-----------------|-------------------|-----------------------------------------------------|------------------------------------------|
+| `voice_clone`   | `seed-tts`        | `ref_audio` + `ref_text` + `task_type=Base`         | `Qwen3-TTS-*-Base`, `VoxCPM2`            |
+| `default_voice` | `seed-tts-text`   | `voice=Vivian` + `task_type=CustomVoice`            | `Qwen3-TTS-*-CustomVoice`                |
+| `voice_design`  | `seed-tts-design` | `instructions=<natural-language description>` + `task_type=VoiceDesign` | `Qwen3-TTS-*-CustomVoice` |
+
+**`-CustomVoice` checkpoints do NOT ship `speaker_encoder` weights**, so
+voice_clone requests raise `ValueError` at model runtime. Use `-Base` for
+voice_clone.
+
+## Adding a new TTS model
+
+Drop an entry into `model_configs.yaml` — no Python changes required:
+
+```yaml
+models:
+  <org>/<model-id>:
+    supported_tasks: [voice_clone]          # or default_voice / voice_design
+    backend: openai-audio-speech            # vllm bench serve backend
+    endpoint: /v1/audio/speech              # OpenAI-compatible endpoint
+    task_extra_body:                        # merged into every request's body
+      voice_clone:
+        task_type: Base
+```
+
+Then add the model's Deploy YAML under `vllm_omni/deploy/<model>.yaml`
+(Pipeline + Deploy schema) and it's immediately benchable.
+
+## Datasets
+
+| Dataset            | Bundled? | Format            | Source                                                         |
+|--------------------|----------|-------------------|----------------------------------------------------------------|
+| `seed-tts-design`  | ✅       | 5-field meta.lst  | `benchmarks/build_dataset/seed_tts_design/en/meta.lst` (20 prompts) |
+| `seed_tts_smoke`   | ✅       | 4-field meta.lst  | `benchmarks/build_dataset/seed_tts_smoke/en/meta.lst` (20 text-only) |
+| `seed-tts`         | ❌       | 4-field meta.lst + WAVs | Google-Drive: [BytedanceSpeech/seed-tts-eval][seedtts] (~1.2 GB) |
+| `seed-tts-text`    | ❌       | 4-field meta.lst  | Same archive as `seed-tts` (wav column unused)                 |
+
+[seedtts]: https://github.com/BytedanceSpeech/seed-tts-eval
+
+For manual voice_clone / default_voice runs against the full corpus, follow
+`benchmarks/build_dataset/download_process_data_seedtts.md` and point
+`--dataset-path` at the extracted `seedtts_testset` directory.
+
+## DFX nightly CI
+
+`tests/dfx/perf/tests/test_tts.json` wires three perf regimes plus quality:
+
+| eval_phase    | concurrency | purpose                                                 | Baseline metrics                        |
+|---------------|-------------|---------------------------------------------------------|-----------------------------------------|
+| `latency`     | 1           | Single-request TTFP / RTF SLO                           | `median_audio_ttfp_ms`, `median_audio_rtf` |
+| `throughput`  | 8           | Codec-batching cliff sentinel (PDF #272 concurrency≥8)  | `median_audio_ttfp_ms`, `median_audio_rtf` |
+| `quality`     | 4           | WER / SIM / UTMOS regression (disabled in CI by default)| `mean_audio_rtf`                        |
+
+Why `median_*` for latency/throughput and `mean_*` for quality: latency
+distributions have cold-start tails that drag the mean; quality aggregates
+over 200 prompts so single-request outliers don't matter.
+
+Quality entries are `enabled: false` in CI because seed-tts-eval is not
+staged in the Buildkite container (matches the precedent in
+PR #2558 — quality runs are manual / release-validation, not nightly).
+
+## Concurrency cliff regression sentinel
+
+Observed on H20-3e, Qwen3-TTS-1.7B (measured pre-merge on this branch):
+
+| Task          | Model         | c=1    | c=4    | **c=8**    | c=16   | c=32   |
+|---------------|---------------|--------|--------|------------|--------|--------|
+| voice_clone   | 1.7B-Base     | RTF 0.15 / TTFP 165ms | 0.28 / 412ms | **0.49 / 1701ms** | 0.72 / 3355ms | 0.77 / 3772ms |
+| voice_design  | 1.7B-CustomVoice | RTF 0.08 / TTFP 53ms  | 0.11 / 154ms | **0.21 / 872ms**  | 0.33 / 1801ms | 0.38 / 1989ms |
+
+Both models show a **4–6× TTFP jump from c=4 to c=8** while audio throughput
+saturates around c=4–8 — the codec-bs=1 bottleneck documented in
+vllm-project/vllm-omni#272. The `throughput` CI regime at c=8 is the
+sentinel for regressions in this area.
+
+## File layout
+
+```
+benchmarks/tts/
+├── README.md                  (this file)
+├── bench_tts.py               CLI — serve-mode benchmark driver
+├── bench_voxcpm_offline.py    CLI — offline VoxCPM benchmark (sync + streaming)
+├── plot_results.py            Generate per-task / per-concurrency curves
+└── model_configs.yaml         Model registry (supported tasks + extra body)
+```
+
+## Related
+
+- Upstream seed-tts-eval integration: vllm-project/vllm-omni#2558
+- Pipeline + Deploy schema: vllm-project/vllm-omni#2383
+- Concurrency cliff RFC: vllm-project/vllm-omni#272

benchmarks/tts/bench_tts.py

@@ -0,0 +1,308 @@
+#!/usr/bin/env python3
+"""Universal TTS benchmark CLI for vllm-omni.
+
+Runs ``vllm bench serve --omni`` with model-aware defaults loaded from
+``model_configs.yaml``.  Supports Qwen3-TTS, VoxCPM2, and any future TTS
+model registered in the config file -- no code changes needed to add models.
+
+Usage::
+
+    python benchmarks/tts/bench_tts.py \\
+        --model Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \\
+        --task voice_clone \\
+        --locale en \\
+        --concurrency 1 4 \\
+        --num-prompts 20 \\
+        --dataset-path /path/to/seed-tts-eval \\
+        --host localhost --port 8000
+
+See ``--help`` for full option list.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import os
+import subprocess
+import sys
+from datetime import datetime
+from pathlib import Path
+from typing import Any
+
+import yaml
+
+
+def _vllm_omni_bin() -> str:
+    """Return the vllm-omni (or vllm) binary co-located with the current Python."""
+    bin_dir = Path(sys.executable).parent
+    for candidate in ("vllm-omni", "vllm"):
+        p = bin_dir / candidate
+        if p.is_file():
+            return str(p)
+    return "vllm-omni"  # fall back and let the shell resolve it
+
+
+_REPO_ROOT = Path(__file__).resolve().parent.parent.parent
+_SCRIPT_DIR = Path(__file__).resolve().parent
+_DEFAULT_MODEL_CONFIGS = _SCRIPT_DIR / "model_configs.yaml"
+
+# Maps task name to the dataset_name used with vllm bench serve
+_TASK_TO_DATASET: dict[str, str] = {
+    "voice_clone": "seed-tts",
+    "default_voice": "seed-tts-text",
+    "voice_design": "seed-tts-design",
+}
+
+# Default design dataset path (bundled with the repo)
+_DEFAULT_DESIGN_DATASET_PATH = str(_REPO_ROOT / "benchmarks" / "build_dataset" / "seed_tts_design")
+
+
+def load_model_configs(path: Path) -> dict[str, Any]:
+    """Load model registry from YAML file."""
+    with open(path, encoding="utf-8") as f:
+        data = yaml.safe_load(f)
+    return data.get("models", {})
+
+
+def build_bench_args(
+    *,
+    host: str,
+    port: int,
+    model: str,
+    task: str,
+    model_cfg: dict[str, Any],
+    locale: str,
+    num_prompts: int,
+    concurrency: int | None,
+    dataset_path: str | None,
+    wer_eval: bool,
+    output_dir: str | None,
+    result_filename: str | None,
+    extra_cli_args: list[str],
+) -> list[str]:
+    """Build the ``vllm bench serve --omni`` command for one (task, concurrency) run."""
+    dataset_name = _TASK_TO_DATASET[task]
+    backend: str = model_cfg["backend"]
+    endpoint: str = model_cfg["endpoint"]
+    task_extra_body: dict[str, Any] = (model_cfg.get("task_extra_body") or {}).get(task) or {}
+
+    # Resolve dataset path
+    if dataset_path:
+        resolved_dataset_path = dataset_path
+    elif task == "voice_design":
+        resolved_dataset_path = _DEFAULT_DESIGN_DATASET_PATH
+    else:
+        resolved_dataset_path = None
+
+    cmd = [
+        _vllm_omni_bin(),
+        "bench",
+        "serve",
+        "--omni",
+        "--host",
+        host,
+        "--port",
+        str(port),
+        "--model",
+        model,
+        "--backend",
+        backend,
+        "--endpoint",
+        endpoint,
+        "--dataset-name",
+        dataset_name,
+        "--num-prompts",
+        str(num_prompts),
+        "--num-warmups",
+        "2",
+        "--percentile-metrics",
+        "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
+    ]
+
+    if resolved_dataset_path:
+        cmd += ["--dataset-path", resolved_dataset_path]
+
+    if locale:
+        cmd += ["--seed-tts-locale", locale]
+
+    if task_extra_body:
+        cmd += ["--extra-body", json.dumps(task_extra_body, separators=(",", ":"))]
+
+    if concurrency is not None:
+        cmd += ["--max-concurrency", str(concurrency), "--request-rate", "inf"]
+
+    if wer_eval:
+        cmd.append("--seed-tts-wer-eval")
+
+    if output_dir or result_filename:
+        out_dir = output_dir or "."
+        os.makedirs(out_dir, exist_ok=True)
+        cmd += ["--save-result", "--result-dir", out_dir]
+        if result_filename:
+            cmd += ["--result-filename", result_filename]
+
+    cmd += extra_cli_args
+    return cmd
+
+
+def run_one_benchmark(cmd: list[str]) -> dict[str, Any] | None:
+    """Run a single benchmark subprocess and return parsed JSON result if available."""
+    print(f"\n{'=' * 60}")
+    print("Running:", " ".join(cmd))
+    print("=" * 60)
+    result = subprocess.run(cmd, check=False)
+    if result.returncode != 0:
+        print(f"[bench_tts] WARNING: benchmark exited with code {result.returncode}")
+        return None
+    # If --save-result was used, find the result file
+    try:
+        result_dir_idx = cmd.index("--result-dir")
+        result_dir = Path(cmd[result_dir_idx + 1])
+        if "--result-filename" in cmd:
+            fname_idx = cmd.index("--result-filename")
+            result_file = result_dir / cmd[fname_idx + 1]
+        else:
+            # find most recently modified json
+            jsons = sorted(result_dir.glob("result_*.json"), key=lambda p: p.stat().st_mtime)
+            result_file = jsons[-1] if jsons else None
+        if result_file and result_file.is_file():
+            return json.loads(result_file.read_text(encoding="utf-8"))
+    except (ValueError, IndexError, OSError):
+        pass
+    return None
+
+
+def print_summary_table(results: list[dict[str, Any]]) -> None:
+    """Print a unified metrics table across all (task, concurrency) runs."""
+    if not results:
+        return
+    header = (
+        f"{'Task':<16} {'Concurrency':>11} {'RTF mean':>10} "
+        f"{'TTFP (ms)':>10} {'Throughput':>12} {'WER':>7} {'SIM':>7} {'UTMOS':>7}"
+    )
+    print(f"\n{'=' * len(header)}")
+    print("BENCHMARK SUMMARY")
+    print("=" * len(header))
+    print(header)
+    print("-" * len(header))
+    for r in results:
+        task = r.get("_task", "?")
+        conc = r.get("_concurrency", "?")
+        rtf = r.get("mean_audio_rtf", float("nan"))
+        ttfp = r.get("mean_audio_ttfp_ms", float("nan"))
+        throughput = r.get("audio_throughput", float("nan"))
+        wer = r.get("seed_tts_mean_wer", float("nan"))
+        sim = r.get("seed_tts_mean_sim", float("nan"))
+        utmos = r.get("seed_tts_mean_utmos", float("nan"))
+
+        def fmt(v: float, digits: int = 3) -> str:
+            return f"{v:.{digits}f}" if not math.isnan(v) else "  n/a"
+
+        print(
+            f"{task:<16} {str(conc):>11} {fmt(rtf):>10} {fmt(ttfp, 0):>10} "
+            f"{fmt(throughput):>12} {fmt(wer):>7} {fmt(sim):>7} {fmt(utmos):>7}"
+        )
+    print("=" * len(header))
+
+
+def main() -> None:
+    """Entry point for the universal TTS benchmark CLI."""
+    parser = argparse.ArgumentParser(
+        description=__doc__,
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    parser.add_argument(
+        "--model", required=True, help="HuggingFace model ID (e.g. Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice)"
+    )
+    parser.add_argument("--task", default="all", help="Task type: voice_clone | default_voice | voice_design | all")
+    parser.add_argument("--locale", default="en", choices=["en", "zh"])
+    parser.add_argument("--concurrency", type=int, nargs="+", default=[1, 4], metavar="N")
+    parser.add_argument(
+        "--num-prompts",
+        type=int,
+        nargs="+",
+        default=[20],
+        metavar="N",
+        help="Number of prompts per run. If one value, applied to all concurrency levels.",
+    )
+    parser.add_argument(
+        "--dataset-path", default=None, help="Root of seed-tts-eval dataset (required for voice_clone/default_voice)"
+    )
+    parser.add_argument("--wer-eval", action="store_true", help="Enable WER/SIM/UTMOS quality eval")
+    parser.add_argument("--output-dir", default=None, help="Directory to save result JSON files")
+    parser.add_argument("--host", default="localhost")
+    parser.add_argument("--port", type=int, default=8000)
+    parser.add_argument("--model-configs", default=str(_DEFAULT_MODEL_CONFIGS), help="Path to model_configs.yaml")
+    parser.add_argument("extra", nargs=argparse.REMAINDER, help="Extra args passed directly to vllm bench serve")
+    args = parser.parse_args()
+
+    model_configs = load_model_configs(Path(args.model_configs))
+    if args.model not in model_configs:
+        known = "\n  ".join(model_configs.keys())
+        print(f"[bench_tts] ERROR: model '{args.model}' not in model_configs.yaml.\nKnown models:\n  {known}")
+        sys.exit(1)
+
+    model_cfg = model_configs[args.model]
+    supported_tasks: list[str] = model_cfg.get("supported_tasks", [])
+
+    tasks_to_run: list[str]
+    if args.task == "all":
+        tasks_to_run = supported_tasks
+    elif args.task in supported_tasks:
+        tasks_to_run = [args.task]
+    else:
+        print(
+            f"[bench_tts] ERROR: task '{args.task}' not supported by {args.model}.\nSupported tasks: {supported_tasks}"
+        )
+        sys.exit(1)
+
+    # Align num_prompts list with concurrency list
+    num_prompts_list: list[int] = args.num_prompts
+    if len(num_prompts_list) == 1:
+        num_prompts_list = num_prompts_list * len(args.concurrency)
+    elif len(num_prompts_list) != len(args.concurrency):
+        print(
+            f"[bench_tts] ERROR: --num-prompts ({len(num_prompts_list)} values) must be "
+            f"length 1 or match --concurrency ({len(args.concurrency)} values)."
+        )
+        sys.exit(1)
+
+    all_results: list[dict[str, Any]] = []
+
+    for task in tasks_to_run:
+        for concurrency, num_prompts in zip(args.concurrency, num_prompts_list):
+            ts = datetime.now().strftime("%Y%m%d-%H%M%S")
+            result_filename = f"bench_tts_{args.model.replace('/', '_')}_{task}_c{concurrency}_{ts}.json"
+            cmd = build_bench_args(
+                host=args.host,
+                port=args.port,
+                model=args.model,
+                task=task,
+                model_cfg=model_cfg,
+                locale=args.locale,
+                num_prompts=num_prompts,
+                concurrency=concurrency,
+                dataset_path=args.dataset_path,
+                wer_eval=args.wer_eval,
+                output_dir=args.output_dir,
+                result_filename=result_filename,
+                extra_cli_args=args.extra or [],
+            )
+            result = run_one_benchmark(cmd)
+            if result is not None:
+                result["_task"] = task
+                result["_concurrency"] = concurrency
+                all_results.append(result)
+                # Persist the metadata so plot_results.py can pick it up.
+                if args.output_dir and result_filename:
+                    result_path = Path(args.output_dir) / result_filename
+                    if result_path.is_file():
+                        result_path.write_text(json.dumps(result, indent=2), encoding="utf-8")
+
+    print_summary_table(all_results)
+
+
+if __name__ == "__main__":
+    main()

benchmarks/tts/bench_voxcpm_offline.py

@@ -0,0 +1,922 @@
+"""Offline VoxCPM benchmark for vLLM Omni.
+
+Supports both:
+- sync one-shot (Omni.generate)
+- streaming (AsyncOmni.generate with async_chunk config)
+- text-only synthesis
+- voice cloning
+- text/clone batch inputs from txt or jsonl
+
+Usage::
+
+    # Sync (default voice)
+    python benchmarks/tts/bench_voxcpm_offline.py \\
+        --model /path/to/VoxCPM \\
+        --text "Hello world" \\
+        --output-dir results/audio/
+
+    # Streaming (async_chunk)
+    python benchmarks/tts/bench_voxcpm_offline.py \\
+        --model /path/to/VoxCPM \\
+        --stage-configs-path vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml \\
+        --txt-prompts prompts.txt \\
+        --output-dir results/audio/
+
+    # Voice cloning batch via JSONL
+    python benchmarks/tts/bench_voxcpm_offline.py \\
+        --model /path/to/VoxCPM \\
+        --jsonl-prompts prompts.jsonl \\
+        --output-dir results/audio/
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+import os
+import tempfile
+import time
+import uuid
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import torch
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+from vllm_omni import AsyncOmni, Omni
+
+
+def _find_repo_root(start: Path) -> Path:
+    """Walk up from ``start`` until a repo marker is found.
+
+    Falls back to ``parents[2]`` for backwards compatibility if no marker hits
+    (which can only happen in unusual checkouts — the tree should always have
+    pyproject.toml + vllm_omni/ at the top level).
+    """
+    for candidate in [start, *start.parents]:
+        if (candidate / "pyproject.toml").is_file() and (candidate / "vllm_omni").is_dir():
+            return candidate
+    return start.parents[2]
+
+
+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"
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True, slots=True)
+class PromptSpec:
+    text: str
+    label: str
+    ref_audio: str | None = None
+    ref_text: str | None = None
+
+
+def _require_soundfile():
+    try:
+        import soundfile as sf  # type: ignore
+    except ModuleNotFoundError as exc:
+        raise RuntimeError(
+            "soundfile is required to write VoxCPM benchmark WAV outputs. Install it with: pip install soundfile"
+        ) from exc
+    return sf
+
+
+def _build_prompt(
+    args,
+    *,
+    text: str,
+    ref_audio: str | None = None,
+    ref_text: str | None = None,
+    global_request_id: str | None = None,
+) -> dict[str, Any]:
+    additional_information: dict[str, list[Any]] = {
+        "text": [text],
+        "cfg_value": [args.cfg_value],
+        "inference_timesteps": [args.inference_timesteps],
+        "min_len": [args.min_len],
+        "max_new_tokens": [args.max_new_tokens],
+    }
+    if args.streaming_prefix_len is not None:
+        additional_information["streaming_prefix_len"] = [args.streaming_prefix_len]
+
+    if ref_audio:
+        additional_information["ref_audio"] = [ref_audio]
+    if ref_text:
+        additional_information["ref_text"] = [ref_text]
+    if global_request_id is not None:
+        additional_information["global_request_id"] = [global_request_id]
+
+    return {
+        "prompt_token_ids": [1],
+        "additional_information": additional_information,
+    }
+
+
+def _extract_audio_tensor(mm: dict[str, Any]) -> torch.Tensor:
+    audio = mm.get("audio", mm.get("model_outputs"))
+    if audio is None:
+        raise ValueError("No audio output found in multimodal output.")
+    if isinstance(audio, list):
+        parts = [torch.as_tensor(a).float().cpu().reshape(-1) for a in audio]
+        audio = torch.cat(parts, dim=-1) if parts else torch.zeros(0)
+    if not isinstance(audio, torch.Tensor):
+        audio = torch.as_tensor(audio)
+    return audio.float().cpu().reshape(-1)
+
+
+def _extract_sample_rate(mm: dict[str, Any]) -> int:
+    sr_raw = mm.get("sr", 24000)
+    if isinstance(sr_raw, list) and sr_raw:
+        sr_raw = sr_raw[-1]
+    if hasattr(sr_raw, "item"):
+        return int(sr_raw.item())
+    return int(sr_raw)
+
+
+def _emit_offline_metrics(
+    *,
+    request_id: str,
+    elapsed_s: float,
+    first_audio_elapsed: float | None,
+    audio_duration_s: float,
+) -> None:
+    metrics = {
+        "request_id": request_id,
+        "ttfp_ms": round(first_audio_elapsed * 1000.0, 3) if first_audio_elapsed is not None else None,
+        "audio_duration_s": round(audio_duration_s, 6),
+        "rtf": round(elapsed_s / audio_duration_s, 6) if audio_duration_s > 0 else None,
+    }
+    print(f"[OfflineMetrics] {metrics}")
+
+
+def _write_audio_tensor(output_path: Path, audio_tensor: Any, sample_rate: int) -> None:
+    sf = _require_soundfile()
+    if isinstance(audio_tensor, torch.Tensor):
+        audio_np = audio_tensor.float().cpu().clamp(-1.0, 1.0).numpy()
+    else:
+        audio_np = torch.as_tensor(audio_tensor).float().cpu().clamp(-1.0, 1.0).numpy()
+    sf.write(
+        output_path,
+        audio_np,
+        sample_rate,
+        format="WAV",
+        subtype="PCM_16",
+    )
+
+
+def _save_wav(mm: dict[str, Any], output_dir: Path, request_id: str) -> Path:
+    output_dir.mkdir(parents=True, exist_ok=True)
+    output_path = output_dir / f"output_{request_id}.wav"
+    _write_audio_tensor(output_path, _extract_audio_tensor(mm), _extract_sample_rate(mm))
+    return output_path
+
+
+def _iter_request_multimodal_outputs(request_output: Any):
+    outputs = getattr(request_output, "outputs", None)
+    if outputs:
+        for output in outputs:
+            mm = getattr(output, "multimodal_output", None)
+            if isinstance(mm, dict):
+                yield mm
+
+    mm = getattr(request_output, "multimodal_output", None)
+    if isinstance(mm, dict):
+        yield mm
+
+
+def _read_non_empty_lines(path: str) -> list[str]:
+    with open(path, encoding="utf-8") as f:
+        return [line.strip() for line in f if line.strip()]
+
+
+def _load_prompt_specs(args) -> list[PromptSpec]:
+    specs: list[PromptSpec] = []
+
+    if args.txt_prompts is not None:
+        texts = _read_non_empty_lines(args.txt_prompts)
+        if not texts:
+            raise ValueError(f"No prompts found in {args.txt_prompts}")
+        for idx, text in enumerate(texts, start=1):
+            specs.append(
+                PromptSpec(
+                    text=text,
+                    label=f"item{idx:03d}",
+                    ref_audio=args.ref_audio,
+                    ref_text=args.ref_text,
+                )
+            )
+        return specs
+
+    if args.jsonl_prompts is not None:
+        with open(args.jsonl_prompts, encoding="utf-8") as f:
+            for line_no, raw_line in enumerate(f, start=1):
+                line = raw_line.strip()
+                if not line:
+                    continue
+                try:
+                    item = json.loads(line)
+                except json.JSONDecodeError as exc:
+                    raise ValueError(f"{args.jsonl_prompts}:{line_no} is not valid JSON: {exc}") from exc
+                if not isinstance(item, dict):
+                    raise ValueError(f"{args.jsonl_prompts}:{line_no} must be a JSON object")
+
+                text = item.get("text")
+                if not isinstance(text, str) or not text.strip():
+                    raise ValueError(f"{args.jsonl_prompts}:{line_no} requires non-empty string field 'text'")
+
+                ref_audio = item.get("ref_audio", args.ref_audio)
+                ref_text = item.get("ref_text", args.ref_text)
+                if (ref_audio is None) != (ref_text is None):
+                    raise ValueError(
+                        f"{args.jsonl_prompts}:{line_no} must provide both 'ref_audio' and 'ref_text' together"
+                    )
+
+                specs.append(
+                    PromptSpec(
+                        text=text.strip(),
+                        label=f"item{len(specs) + 1:03d}",
+                        ref_audio=ref_audio,
+                        ref_text=ref_text,
+                    )
+                )
+
+        if not specs:
+            raise ValueError(f"No prompts found in {args.jsonl_prompts}")
+        return specs
+
+    specs.append(
+        PromptSpec(
+            text=args.text,
+            label="item001",
+            ref_audio=args.ref_audio,
+            ref_text=args.ref_text,
+        )
+    )
+    return specs
+
+
+def _build_prompt_for_spec(args, spec: PromptSpec, *, global_request_id: str | None = None) -> dict[str, Any]:
+    return _build_prompt(
+        args,
+        text=spec.text,
+        ref_audio=spec.ref_audio,
+        ref_text=spec.ref_text,
+        global_request_id=global_request_id,
+    )
+
+
+def _count_voice_clone_prompts(prompt_specs: list[PromptSpec]) -> int:
+    return sum(1 for spec in prompt_specs if spec.ref_audio is not None)
+
+
+def _get_warmup_specs(prompt_specs: list[PromptSpec]) -> list[PromptSpec]:
+    return prompt_specs[:1]
+
+
+def _extract_stream_finished(stage_output: Any) -> bool:
+    request_output = getattr(stage_output, "request_output", None)
+    request_finished = getattr(request_output, "finished", None)
+    if request_finished is not None:
+        return bool(request_finished)
+    return bool(getattr(stage_output, "finished", False))
+
+
+def _build_profiled_stage_config(
+    stage_configs_path: str,
+    profiler_dir: str,
+) -> str:
+    stage_config_path = Path(stage_configs_path)
+    yaml_text = stage_config_path.read_text(encoding="utf-8")
+    injected_lines: list[str] = []
+    injected_count = 0
+
+    for line in yaml_text.splitlines():
+        injected_lines.append(line)
+        if line.strip() != "engine_args:":
+            continue
+        indent = line[: len(line) - len(line.lstrip())]
+        child_indent = indent + "  "
+        grandchild_indent = child_indent + "  "
+        injected_lines.extend(
+            [
+                f"{child_indent}profiler_config:",
+                f'{grandchild_indent}profiler: "torch"',
+                f'{grandchild_indent}torch_profiler_dir: "{profiler_dir}"',
+                f"{grandchild_indent}torch_profiler_with_stack: true",
+            ]
+        )
+        injected_count += 1
+
+    if injected_count == 0:
+        raise ValueError(f"No engine_args block found in stage config: {stage_configs_path}")
+
+    tmp = tempfile.NamedTemporaryFile(
+        mode="w",
+        encoding="utf-8",
+        delete=False,
+        suffix=".yaml",
+        prefix=f"{stage_config_path.stem}_profile_",
+    )
+    tmp.write("\n".join(injected_lines) + "\n")
+    tmp.close()
+    return tmp.name
+
+
+def parse_args():
+    parser = FlexibleArgumentParser(
+        description="Offline split-stage VoxCPM inference with vLLM Omni (auto sync/streaming by stage config)"
+    )
+    parser.add_argument(
+        "--model",
+        type=str,
+        default=os.environ.get("VOXCPM_MODEL"),
+        help="Local VoxCPM model directory. Defaults to $VOXCPM_MODEL.",
+    )
+    parser.add_argument(
+        "--text",
+        type=str,
+        default="This is a split-stage VoxCPM synthesis example running on vLLM Omni.",
+        help="Text to synthesize. Ignored when --txt-prompts or --jsonl-prompts is used.",
+    )
+    parser.add_argument(
+        "--txt-prompts",
+        type=str,
+        default=None,
+        help="Path to a .txt file with one synthesis text per line.",
+    )
+    parser.add_argument(
+        "--jsonl-prompts",
+        type=str,
+        default=None,
+        help=(
+            "Path to a .jsonl file. Each line must contain at least {'text': ...}; "
+            "clone rows can also set ref_audio/ref_text, and ref_text must be the "
+            "real transcript of ref_audio."
+        ),
+    )
+    parser.add_argument(
+        "--ref-audio",
+        type=str,
+        default=None,
+        help=(
+            "Optional reference audio path for voice cloning. With --txt-prompts, "
+            "the same reference is applied to every line."
+        ),
+    )
+    parser.add_argument(
+        "--ref-text",
+        type=str,
+        default=None,
+        help=(
+            "Real transcript of the reference audio. Placeholder text or mismatched "
+            "text will usually produce noisy/electronic clone audio."
+        ),
+    )
+    parser.add_argument(
+        "--stage-configs-path",
+        type=str,
+        default=str(DEFAULT_STAGE_SYNC),
+        help="Stage config YAML path. Routing is selected only from this path.",
+    )
+    parser.add_argument(
+        "--cfg-value",
+        type=float,
+        default=2.0,
+        help="Classifier-free guidance value for VoxCPM.",
+    )
+    parser.add_argument(
+        "--inference-timesteps",
+        type=int,
+        default=10,
+        help="Number of inference timesteps.",
+    )
+    parser.add_argument(
+        "--min-len",
+        type=int,
+        default=2,
+        help="Minimum generated token length.",
+    )
+    parser.add_argument(
+        "--max-new-tokens",
+        type=int,
+        default=4096,
+        help="Maximum generated token length.",
+    )
+    parser.add_argument(
+        "--streaming-prefix-len",
+        type=int,
+        default=None,
+        help="VoxCPM streaming window (optional, streaming mode only).",
+    )
+    parser.add_argument(
+        "--output-dir",
+        type=str,
+        default=None,
+        help="Directory for output WAV files.",
+    )
+    parser.add_argument(
+        "--stage-init-timeout",
+        type=int,
+        default=600,
+        help="Stage initialization timeout in seconds.",
+    )
+    parser.add_argument(
+        "--log-stats",
+        dest="log_stats",
+        action="store_true",
+        help="Enable vLLM Omni stats logging.",
+    )
+    parser.add_argument(
+        "--no-log-stats",
+        dest="log_stats",
+        action="store_false",
+        help="Disable vLLM Omni stats logging.",
+    )
+    parser.set_defaults(log_stats=True)
+    parser.add_argument(
+        "--num-runs",
+        type=int,
+        default=1,
+        help="Number of full inference runs (same prompt each time). Default 1.",
+    )
+    parser.add_argument(
+        "--warmup-runs",
+        type=int,
+        default=0,
+        help=(
+            "Optional number of warmup passes before measured runs. Warmup uses only "
+            "the first prompt and does not save outputs."
+        ),
+    )
+    parser.add_argument(
+        "--enable-profiler",
+        action="store_true",
+        help=(
+            "Enable torch profiler for the configured stages. A temporary profiled "
+            "stage config is generated automatically."
+        ),
+    )
+    parser.add_argument(
+        "--profiler-dir",
+        type=str,
+        default=None,
+        help="Directory for profiler traces. Defaults to <output-dir>/profiler when profiling is enabled.",
+    )
+    parser.add_argument(
+        "--profiler-stages",
+        type=int,
+        nargs="*",
+        default=None,
+        help="Optional stage ids to profile. Defaults to all stages that have profiler_config.",
+    )
+    parser.add_argument(
+        "--profiler-wait-seconds",
+        type=float,
+        default=30.0,
+        help="Seconds to wait after stop_profile for trace files to flush.",
+    )
+    args = parser.parse_args()
+
+    if not args.model:
+        parser.error("--model is required unless $VOXCPM_MODEL is set")
+    if args.txt_prompts is not None and args.jsonl_prompts is not None:
+        parser.error("--txt-prompts and --jsonl-prompts are mutually exclusive")
+    if (args.ref_audio is None) != (args.ref_text is None):
+        parser.error("--ref-audio and --ref-text must be provided together")
+    if args.num_runs < 1:
+        parser.error("--num-runs must be >= 1")
+    if args.warmup_runs < 0:
+        parser.error("--warmup-runs must be >= 0")
+    if args.output_dir is None:
+        args.output_dir = (
+            "output_audio_streaming" if _is_streaming_stage_config(args.stage_configs_path) else "output_audio"
+        )
+    if args.enable_profiler and args.profiler_dir is None:
+        args.profiler_dir = str(Path(args.output_dir) / "profiler")
+    try:
+        args.prompt_specs = _load_prompt_specs(args)
+    except ValueError as exc:
+        parser.error(str(exc))
+
+    return args
+
+
+def _is_streaming_stage_config(stage_configs_path: str) -> bool:
+    cfg_name = Path(stage_configs_path).name.lower()
+    return "async_chunk" in cfg_name
+
+
+async def _collect_streaming_audio(
+    omni: AsyncOmni,
+    args: Any,
+    spec: PromptSpec,
+    request_id: str,
+    *,
+    phase_label: str,
+    prompt_index: int,
+    prompt_count: int,
+    print_prompt: bool = False,
+) -> tuple[torch.Tensor, int, float, float | None]:
+    prompt = _build_prompt_for_spec(args, spec, global_request_id=request_id)
+    delta_chunks: list[torch.Tensor] = []
+    sample_rate = 24000
+    chunk_i = 0
+    prev_total_samples = 0
+    t_start = time.perf_counter()
+    first_audio_elapsed: float | None = None
+
+    if print_prompt:
+        print(f"---prompt---:{prompt}")
+
+    async for stage_output in omni.generate(prompt, request_id=request_id):
+        mm = getattr(stage_output, "multimodal_output", None)
+        if not isinstance(mm, dict):
+            ro = getattr(stage_output, "request_output", None)
+            if ro is None:
+                continue
+            mm = getattr(ro, "multimodal_output", None)
+            if not isinstance(mm, dict) and getattr(ro, "outputs", None):
+                seq = ro.outputs[0]
+                mm = getattr(seq, "multimodal_output", None)
+        if not isinstance(mm, dict):
+            continue
+        sample_rate = _extract_sample_rate(mm)
+        try:
+            w = _extract_audio_tensor(mm)
+            n = int(w.numel())
+            if n == 0:
+                continue
+            finished = _extract_stream_finished(stage_output)
+            if n > prev_total_samples:
+                delta = w.reshape(-1)[prev_total_samples:]
+                prev_total_samples = n
+            elif finished and n == prev_total_samples:
+                delta = w.reshape(-1)[:0]
+            else:
+                delta = w.reshape(-1)
+                prev_total_samples += int(delta.numel())
+            if int(delta.numel()) > 0:
+                delta_chunks.append(delta)
+            if first_audio_elapsed is None and int(delta.numel()) > 0:
+                first_audio_elapsed = time.perf_counter() - t_start
+            logger.info(
+                "%s prompt=%d/%d chunk=%d delta_samples=%d buf_len=%d finished=%s",
+                phase_label,
+                prompt_index + 1,
+                prompt_count,
+                chunk_i,
+                int(delta.numel()),
+                n,
+                finished,
+            )
+            chunk_i += 1
+        except ValueError:
+            if not _extract_stream_finished(stage_output):
+                logger.debug("skip non-audio partial output chunk=%d", chunk_i)
+
+    if not delta_chunks:
+        raise RuntimeError("No audio chunks received; check stage config and logs.")
+
+    audio_cat = torch.cat([c.reshape(-1) for c in delta_chunks], dim=0)
+    elapsed = time.perf_counter() - t_start
+    return audio_cat, sample_rate, elapsed, first_audio_elapsed
+
+
+async def _abort_streaming_residual_work(
+    omni: AsyncOmni,
+    request_id: str,
+    *,
+    settle_seconds: float = 0.1,
+) -> None:
+    """Stop any late stage-0 work once the final audio has been collected."""
+    await omni.engine.abort_async([request_id])
+    if settle_seconds > 0:
+        await asyncio.sleep(settle_seconds)
+
+
+async def _run_streaming_single(
+    omni: AsyncOmni,
+    args: Any,
+    spec: PromptSpec,
+    output_dir: Path,
+    request_id: str,
+    *,
+    run_index: int,
+    num_runs: int,
+    prompt_index: int,
+    prompt_count: int,
+) -> Path:
+    audio_cat, sample_rate, elapsed, first_audio_elapsed = await _collect_streaming_audio(
+        omni,
+        args,
+        spec,
+        request_id,
+        phase_label=f"run={run_index + 1}/{num_runs}",
+        prompt_index=prompt_index,
+        prompt_count=prompt_count,
+        print_prompt=(run_index == 0 and prompt_index == 0),
+    )
+    await _abort_streaming_residual_work(omni, request_id)
+    output_path = output_dir / f"output_run{run_index + 1}_{spec.label}.wav"
+    _write_audio_tensor(output_path, audio_cat, sample_rate)
+    audio_duration_s = float(audio_cat.numel()) / float(sample_rate) if sample_rate > 0 else 0.0
+    ttfp_text = f", ttfp={first_audio_elapsed:.2f}s" if first_audio_elapsed is not None else ""
+    rtf_text = f", rtf={elapsed / audio_duration_s:.3f}" if audio_duration_s > 0 else ""
+    print(
+        f"Saved (streaming) run {run_index + 1}/{num_runs}, "
+        f"prompt {prompt_index + 1}/{prompt_count}: {output_path} ({elapsed:.2f}s{ttfp_text}{rtf_text})"
+    )
+    _emit_offline_metrics(
+        request_id=request_id,
+        elapsed_s=elapsed,
+        first_audio_elapsed=first_audio_elapsed,
+        audio_duration_s=audio_duration_s,
+    )
+    return output_path
+
+
+async def _run_streaming_warmup(args, omni: AsyncOmni) -> None:
+    if args.warmup_runs == 0:
+        return
+
+    warmup_specs = _get_warmup_specs(args.prompt_specs)
+    print(
+        f"Warmup: {args.warmup_runs} run(s) using the first prompt "
+        f"({len(warmup_specs)} prompt(s)); outputs will be discarded."
+    )
+    for warmup_index in range(args.warmup_runs):
+        t_warmup = time.perf_counter()
+        tasks = []
+        request_ids: list[str] = []
+        for prompt_index, spec in enumerate(warmup_specs):
+            request_id = f"warmup_stream_{warmup_index + 1}_{spec.label}_{uuid.uuid4().hex[:8]}"
+            request_ids.append(request_id)
+            tasks.append(
+                _collect_streaming_audio(
+                    omni,
+                    args,
+                    spec,
+                    request_id,
+                    phase_label=f"warmup={warmup_index + 1}/{args.warmup_runs}",
+                    prompt_index=prompt_index,
+                    prompt_count=len(warmup_specs),
+                )
+            )
+        results = await asyncio.gather(*tasks)
+        for request_id in request_ids:
+            await _abort_streaming_residual_work(omni, request_id)
+        total_samples = sum(int(audio.numel()) for audio, _, _, _ in results)
+        warmup_ttfps = [ttfp for _, _, _, ttfp in results if ttfp is not None]
+        ttfp_text = f", ttfp={min(warmup_ttfps):.2f}s" if warmup_ttfps else ""
+        print(
+            f"Warmup (streaming) {warmup_index + 1}/{args.warmup_runs} finished: "
+            f"{len(results)} prompt(s), {total_samples} sample(s) "
+            f"({time.perf_counter() - t_warmup:.2f}s{ttfp_text})"
+        )
+
+
+async def _run_streaming(args) -> list[Path]:
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    omni = AsyncOmni(
+        model=args.model,
+        stage_configs_path=args.stage_configs_path,
+        log_stats=args.log_stats,
+        stage_init_timeout=args.stage_init_timeout,
+    )
+
+    await _run_streaming_warmup(args, omni)
+    profiler_started = False
+    if args.enable_profiler:
+        profile_prefix = f"voxcpm_streaming_{int(time.time())}"
+        stages_text = args.profiler_stages if args.profiler_stages is not None else "all-configured"
+        print(f"Starting profiler (streaming): stages={stages_text}, dir={args.profiler_dir}")
+        await omni.start_profile(profile_prefix=profile_prefix, stages=args.profiler_stages)
+        profiler_started = True
+    t_total = time.perf_counter()
+    total_elapsed = 0.0
+    paths: list[Path] = []
+    prompt_specs: list[PromptSpec] = args.prompt_specs
+    try:
+        for run in range(args.num_runs):
+            for prompt_index, spec in enumerate(prompt_specs):
+                request_id = f"stream_{run + 1}_{spec.label}_{uuid.uuid4().hex[:8]}"
+                paths.append(
+                    await _run_streaming_single(
+                        omni,
+                        args,
+                        spec,
+                        output_dir,
+                        request_id,
+                        run_index=run,
+                        num_runs=args.num_runs,
+                        prompt_index=prompt_index,
+                        prompt_count=len(prompt_specs),
+                    )
+                )
+        total_elapsed = time.perf_counter() - t_total
+    finally:
+        if profiler_started:
+            print("Stopping profiler (streaming)...")
+            await omni.stop_profile(stages=args.profiler_stages)
+            if args.profiler_wait_seconds > 0:
+                print(f"Waiting {args.profiler_wait_seconds:.1f}s for profiler traces to flush...")
+                await asyncio.sleep(args.profiler_wait_seconds)
+
+    print(
+        f"All streaming runs finished: {args.num_runs} run(s), "
+        f"{len(prompt_specs)} prompt(s), {len(paths)} file(s) in {total_elapsed:.2f}s total"
+    )
+    return paths
+
+
+def _run_sync(args) -> list[Path]:
+    output_dir = Path(args.output_dir)
+
+    omni = Omni(
+        model=args.model,
+        stage_configs_path=args.stage_configs_path,
+        log_stats=args.log_stats,
+        stage_init_timeout=args.stage_init_timeout,
+    )
+
+    def _run_sync_single(
+        spec: PromptSpec,
+        *,
+        request_prefix: str,
+        save_outputs: bool,
+        run_index: int | None = None,
+    ) -> tuple[list[Path], int, float | None, float, float, str]:
+        global_request_id = f"{request_prefix}_{spec.label}"
+        prompt = _build_prompt_for_spec(args, spec, global_request_id=global_request_id)
+        if save_outputs and run_index == 0 and spec.label == "item001":
+            print(f"---prompt---:{prompt}")
+
+        saved_paths: list[Path] = []
+        output_count = 0
+        first_audio_elapsed: float | None = None
+        total_audio_duration_s = 0.0
+        metrics_request_id = global_request_id
+        t_start = time.perf_counter()
+        for stage_outputs in omni.generate(prompt):
+            request_output = stage_outputs.request_output
+            if request_output is None:
+                continue
+            request_output_id = getattr(request_output, "request_id", None)
+            if isinstance(request_output_id, str) and request_output_id:
+                metrics_request_id = request_output_id
+            for j, mm in enumerate(_iter_request_multimodal_outputs(request_output)):
+                output_count += 1
+                if first_audio_elapsed is None:
+                    try:
+                        audio_tensor = _extract_audio_tensor(mm)
+                        if int(audio_tensor.numel()) > 0:
+                            first_audio_elapsed = time.perf_counter() - t_start
+                        total_audio_duration_s += float(audio_tensor.numel()) / float(_extract_sample_rate(mm))
+                    except ValueError:
+                        pass
+                else:
+                    try:
+                        audio_tensor = _extract_audio_tensor(mm)
+                        total_audio_duration_s += float(audio_tensor.numel()) / float(_extract_sample_rate(mm))
+                    except ValueError:
+                        pass
+                if not save_outputs:
+                    continue
+                save_stem = f"run{run_index + 1}_{spec.label}" if j == 0 else f"run{run_index + 1}_{spec.label}_{j}"
+                saved_paths.append(_save_wav(mm, output_dir, save_stem))
+
+        if output_count == 0:
+            raise RuntimeError("No output from Omni.generate")
+        elapsed_s = time.perf_counter() - t_start
+        return saved_paths, output_count, first_audio_elapsed, elapsed_s, total_audio_duration_s, metrics_request_id
+
+    if args.warmup_runs:
+        warmup_specs = _get_warmup_specs(args.prompt_specs)
+        print(
+            f"Warmup: {args.warmup_runs} run(s) using the first prompt "
+            f"({len(warmup_specs)} prompt(s)); outputs will be discarded."
+        )
+        for warmup_index in range(args.warmup_runs):
+            t_warmup = time.perf_counter()
+            _, output_count, first_audio_elapsed, elapsed_s, audio_duration_s, _ = _run_sync_single(
+                warmup_specs[0],
+                request_prefix=f"warmup_sync{warmup_index + 1}",
+                save_outputs=False,
+            )
+            ttfp_text = f", ttfp={first_audio_elapsed:.2f}s" if first_audio_elapsed is not None else ""
+            rtf_text = f", rtf={elapsed_s / audio_duration_s:.3f}" if audio_duration_s > 0 else ""
+            print(
+                f"Warmup (sync) {warmup_index + 1}/{args.warmup_runs} finished: "
+                f"{output_count} output(s) ({time.perf_counter() - t_warmup:.2f}s{ttfp_text}{rtf_text})"
+            )
+
+    profiler_started = False
+    if args.enable_profiler:
+        profile_prefix = f"voxcpm_sync_{int(time.time())}"
+        stages_text = args.profiler_stages if args.profiler_stages is not None else "all-configured"
+        print(f"Starting profiler (sync): stages={stages_text}, dir={args.profiler_dir}")
+        omni.start_profile(profile_prefix=profile_prefix, stages=args.profiler_stages)
+        profiler_started = True
+
+    t_total = time.perf_counter()
+    total_elapsed = 0.0
+    saved_paths: list[Path] = []
+    prompt_specs: list[PromptSpec] = args.prompt_specs
+    try:
+        for run in range(args.num_runs):
+            t_run = time.perf_counter()
+            run_paths: list[Path] = []
+            for prompt_index, spec in enumerate(prompt_specs):
+                prompt_paths, _, first_audio_elapsed, elapsed_s, audio_duration_s, metrics_request_id = (
+                    _run_sync_single(
+                        spec,
+                        request_prefix=f"sync_run{run + 1}_{prompt_index + 1:03d}",
+                        save_outputs=True,
+                        run_index=run,
+                    )
+                )
+                run_paths.extend(prompt_paths)
+                ttfp_text = f", ttfp={first_audio_elapsed:.2f}s" if first_audio_elapsed is not None else ""
+                rtf_text = f", rtf={elapsed_s / audio_duration_s:.3f}" if audio_duration_s > 0 else ""
+                print(
+                    f"Saved (sync) run {run + 1}/{args.num_runs}, "
+                    f"prompt {prompt_index + 1}/{len(prompt_specs)}: {len(prompt_paths)} file(s){ttfp_text}{rtf_text}"
+                )
+                _emit_offline_metrics(
+                    request_id=metrics_request_id,
+                    elapsed_s=elapsed_s,
+                    first_audio_elapsed=first_audio_elapsed,
+                    audio_duration_s=audio_duration_s,
+                )
+
+            saved_paths.extend(run_paths)
+            print(
+                f"Run {run + 1}/{args.num_runs} finished: {len(run_paths)} file(s) ({time.perf_counter() - t_run:.2f}s)"
+            )
+            for path in run_paths:
+                print(f"  {path}")
+
+        total_elapsed = time.perf_counter() - t_total
+    finally:
+        if profiler_started:
+            print("Stopping profiler (sync)...")
+            omni.stop_profile(stages=args.profiler_stages)
+            if args.profiler_wait_seconds > 0:
+                print(f"Waiting {args.profiler_wait_seconds:.1f}s for profiler traces to flush...")
+                time.sleep(args.profiler_wait_seconds)
+
+    print(
+        f"All sync runs finished: {args.num_runs} run(s), "
+        f"{len(prompt_specs)} prompt(s), {len(saved_paths)} file(s) in {total_elapsed:.2f}s total"
+    )
+    return saved_paths
+
+
+def main(args) -> int:
+    logging.basicConfig(level=logging.INFO)
+    profiled_stage_config_path: str | None = None
+    original_stage_config_path = args.stage_configs_path
+    if args.enable_profiler:
+        Path(args.profiler_dir).mkdir(parents=True, exist_ok=True)
+        profiled_stage_config_path = _build_profiled_stage_config(
+            args.stage_configs_path,
+            str(Path(args.profiler_dir).resolve()),
+        )
+        args.stage_configs_path = profiled_stage_config_path
+
+    is_streaming = _is_streaming_stage_config(args.stage_configs_path)
+    voice_clone_count = _count_voice_clone_prompts(args.prompt_specs)
+    print(f"Model: {args.model}")
+    print(f"Stage config: {original_stage_config_path}")
+    print(f"Route: {'streaming' if is_streaming else 'sync'} (from stage-configs-path)")
+    print(f"Prompt count: {len(args.prompt_specs)}")
+    print("Batch mode: sequential (aligned with native VoxCPM)")
+    print(f"Warmup runs: {args.warmup_runs}")
+    print(f"Voice cloning prompts: {voice_clone_count}/{len(args.prompt_specs)}")
+    if args.enable_profiler:
+        print(f"Profiler: enabled (dir={args.profiler_dir}, stages={args.profiler_stages or 'all-configured'})")
+        print(f"Profiled stage config: {args.stage_configs_path}")
+    if voice_clone_count:
+        print("Voice cloning note: --ref-text/ref_text must match the spoken content of the reference audio.")
+    print(f"Num runs: {args.num_runs}")
+    try:
+        if is_streaming:
+            asyncio.run(_run_streaming(args))
+        else:
+            _run_sync(args)
+    finally:
+        if profiled_stage_config_path is not None and os.path.exists(profiled_stage_config_path):
+            os.unlink(profiled_stage_config_path)
+    return 0
+
+
+if __name__ == "__main__":
+    os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+    raise SystemExit(main(parse_args()))

benchmarks/tts/model_configs.yaml

@@ -0,0 +1,39 @@
+# Universal TTS benchmark model registry.
+# Maps HuggingFace model ID → supported tasks + per-task extra body fields.
+# To add a new TTS model: add an entry here. No code changes required.
+#
+# The server auto-loads its Deploy YAML from vllm_omni/deploy/<model>.yaml via
+# the Pipeline + Deploy schema introduced in #2383, so no stage_config path
+# is tracked here.
+
+models:
+  # -CustomVoice checkpoints lack speaker_encoder weights, so voice_clone is
+  # NOT supported (an attempt raises ValueError from _extract_speaker_embedding
+  # at model runtime).  Use -Base for voice_clone.
+  Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice:
+    supported_tasks: [default_voice, voice_design]
+    backend: openai-audio-speech
+    endpoint: /v1/audio/speech
+    task_extra_body:
+      default_voice:
+        voice: Vivian
+        language: English
+        task_type: CustomVoice
+      voice_design:
+        task_type: VoiceDesign
+        language: English
+
+  Qwen/Qwen3-TTS-12Hz-1.7B-Base:
+    supported_tasks: [voice_clone]
+    backend: openai-audio-speech
+    endpoint: /v1/audio/speech
+    task_extra_body:
+      voice_clone:
+        task_type: Base
+
+  openbmb/VoxCPM2:
+    supported_tasks: [voice_clone]
+    backend: openai-audio-speech
+    endpoint: /v1/audio/speech
+    task_extra_body:
+      voice_clone: {}

benchmarks/tts/plot_results.py

@@ -0,0 +1,324 @@
+"""Plot universal TTS benchmark results.
+
+Reads JSON files saved by ``bench_tts.py`` (via ``vllm bench serve --omni``)
+and generates comparison bar charts grouped by task type.
+
+Metrics plotted:
+- AUDIO_TTFP  (mean audio time-to-first-packet, ms)
+- E2EL        (mean end-to-end latency, ms)
+- Audio RTF   (mean real-time factor)
+- Audio throughput (audio-seconds / wall-second)
+
+Quality metrics (WER / SIM / UTMOS) are printed in a table when present.
+
+Usage::
+
+    # Single run — one JSON per task, all in results/
+    python benchmarks/tts/plot_results.py \\
+        --results results/bench_tts_*.json \\
+        --output results/tts_benchmark.png
+
+    # Compare two runs (e.g. async_chunk on vs off)
+    python benchmarks/tts/plot_results.py \\
+        --results run_a/bench_tts_*.json \\
+        --results run_b/bench_tts_*.json \\
+        --labels "async_chunk_on" "async_chunk_off" \\
+        --output results/comparison.png
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import math
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+# ---------------------------------------------------------------------------
+# JSON loading
+# ---------------------------------------------------------------------------
+
+
+def load_run(paths: list[str]) -> list[dict]:
+    """Load and merge all JSON files for one run into a flat list of records.
+
+    Each record is expected to have at least ``_concurrency`` (int) and
+    ``_task`` (str) keys injected by ``bench_tts.py``.  Records that come
+    from a file that contains a list are flattened.
+    """
+    records: list[dict] = []
+    for p in paths:
+        raw = json.loads(Path(p).read_text(encoding="utf-8"))
+        if isinstance(raw, list):
+            records.extend(raw)
+        elif isinstance(raw, dict):
+            records.append(raw)
+    return records
+
+
+def _get(record: dict, key: str) -> float:
+    v = record.get(key, float("nan"))
+    if v is None or (isinstance(v, float) and math.isnan(v)):
+        return float("nan")
+    try:
+        return float(v)
+    except (TypeError, ValueError):
+        return float("nan")
+
+
+# ---------------------------------------------------------------------------
+# Plotting helpers
+# ---------------------------------------------------------------------------
+
+
+def _bar_group(
+    ax: plt.Axes,
+    x: np.ndarray,
+    data_per_label: dict[str, list[float]],
+    width: float,
+    colors: list[str],
+    ylabel: str,
+    title: str,
+    concurrency_labels: list[str],
+    fmt: str = ".1f",
+) -> None:
+    n = len(data_per_label)
+    offsets = np.linspace(-(n - 1) * width / 2, (n - 1) * width / 2, n) if n > 1 else [0.0]
+
+    for i, (label, values) in enumerate(data_per_label.items()):
+        plot_vals = [0.0 if math.isnan(v) else v for v in values]
+        bar = ax.bar(x + offsets[i], plot_vals, width, label=label, color=colors[i % len(colors)], alpha=0.85)
+        max_val = max((v for v in values if not math.isnan(v)), default=1.0)
+        for rect, val in zip(bar, values):
+            if not math.isnan(val) and val > 0:
+                ax.text(
+                    rect.get_x() + rect.get_width() / 2,
+                    rect.get_height() + max_val * 0.02,
+                    f"{val:{fmt}}",
+                    ha="center",
+                    va="bottom",
+                    fontsize=8,
+                    fontweight="bold",
+                )
+
+    ax.set_xlabel("Concurrency", fontsize=11)
+    ax.set_ylabel(ylabel, fontsize=11)
+    ax.set_title(title, fontsize=12, fontweight="bold")
+    ax.set_xticks(x)
+    ax.set_xticklabels(concurrency_labels)
+    ax.legend(fontsize=9)
+    ax.grid(axis="y", alpha=0.3)
+    ax.set_axisbelow(True)
+
+
+COLORS = ["#2196F3", "#FF5722", "#4CAF50", "#FFC107", "#9C27B0"]
+
+
+# ---------------------------------------------------------------------------
+# Comparison plot (multiple labels / runs)
+# ---------------------------------------------------------------------------
+
+
+def plot_comparison(
+    all_runs: list[list[dict]],
+    labels: list[str],
+    output_path: str,
+    task_filter: str | None = None,
+    title_prefix: str = "TTS",
+) -> None:
+    """One 2×2 subplot per task found in the data."""
+    # Determine tasks to plot
+    tasks: list[str] = []
+    for run in all_runs:
+        for r in run:
+            t = r.get("_task", "unknown")
+            if t not in tasks:
+                tasks.append(t)
+    if task_filter:
+        tasks = [t for t in tasks if t == task_filter]
+
+    n_tasks = len(tasks)
+    if n_tasks == 0:
+        print("[plot_results] No tasks found in data.")
+        return
+
+    fig, axes_grid = plt.subplots(n_tasks, 4, figsize=(18, 4.5 * n_tasks))
+    fig.suptitle(f"{title_prefix} Benchmark", fontsize=15, fontweight="bold")
+
+    # Ensure axes_grid is always 2D
+    if n_tasks == 1:
+        axes_grid = [axes_grid]
+
+    for row_idx, task in enumerate(tasks):
+        # Collect concurrencies across all runs for this task
+        all_concs: set[int] = set()
+        for run in all_runs:
+            for r in run:
+                if r.get("_task") == task:
+                    c = r.get("_concurrency")
+                    if c is not None:
+                        all_concs.add(int(c))
+        concurrencies = sorted(all_concs)
+        x = np.arange(len(concurrencies))
+        conc_labels = [str(c) for c in concurrencies]
+
+        def _series(run: list[dict], metric_key: str) -> list[float]:
+            conc_map = {int(r["_concurrency"]): r for r in run if r.get("_task") == task and "_concurrency" in r}
+            return [_get(conc_map.get(c, {}), metric_key) for c in concurrencies]
+
+        metrics = [
+            ("mean_audio_ttfp_ms", "TTFP (ms)", "Time-to-First-Packet", ".0f"),
+            ("mean_e2el_ms", "E2E Latency (ms)", "End-to-End Latency", ".0f"),
+            ("mean_audio_rtf", "RTF", "Real-Time Factor (RTF)", ".3f"),
+            ("audio_throughput", "audio-s / wall-s", "Audio Throughput", ".2f"),
+        ]
+
+        axes_row = axes_grid[row_idx]
+        for col_idx, (key, ylabel, subtitle, fmt) in enumerate(metrics):
+            data_per_label = {lbl: _series(run, key) for lbl, run in zip(labels, all_runs)}
+            _bar_group(
+                axes_row[col_idx],
+                x,
+                data_per_label,
+                width=0.3 if len(labels) > 1 else 0.5,
+                colors=COLORS,
+                ylabel=ylabel,
+                title=f"{task} — {subtitle}",
+                concurrency_labels=conc_labels,
+                fmt=fmt,
+            )
+
+    plt.tight_layout()
+    Path(output_path).parent.mkdir(parents=True, exist_ok=True)
+    plt.savefig(output_path, dpi=150, bbox_inches="tight")
+    print(f"Plot saved to {output_path}")
+    plt.close()
+
+
+# ---------------------------------------------------------------------------
+# Markdown comparison table
+# ---------------------------------------------------------------------------
+
+
+def print_comparison_table(all_runs: list[list[dict]], labels: list[str]) -> None:
+    tasks: list[str] = []
+    for run in all_runs:
+        for r in run:
+            t = r.get("_task", "unknown")
+            if t not in tasks:
+                tasks.append(t)
+
+    perf_metrics = [
+        ("TTFP (ms)", "mean_audio_ttfp_ms", ".1f"),
+        ("E2E (ms)", "mean_e2el_ms", ".1f"),
+        ("RTF", "mean_audio_rtf", ".3f"),
+        ("Throughput (a-s/s)", "audio_throughput", ".2f"),
+    ]
+    quality_metrics = [
+        ("WER (%)", "seed_tts_mean_wer", ".1f"),
+        ("SIM", "seed_tts_mean_sim", ".3f"),
+        ("UTMOS", "seed_tts_mean_utmos", ".2f"),
+    ]
+
+    for task in tasks:
+        all_concs: set[int] = set()
+        for run in all_runs:
+            for r in run:
+                if r.get("_task") == task:
+                    c = r.get("_concurrency")
+                    if c is not None:
+                        all_concs.add(int(c))
+        concurrencies = sorted(all_concs)
+
+        print(f"\n## {task}\n")
+        col_header = "| Metric | Concurrency |" + "".join(f" {lbl} |" for lbl in labels)
+        sep = "| --- | --- |" + " --- |" * len(labels)
+        print(col_header)
+        print(sep)
+
+        for metric, key, fmt in perf_metrics + quality_metrics:
+            for c in concurrencies:
+                row = f"| {metric} | {c} |"
+                for run in all_runs:
+                    conc_map = {
+                        int(r["_concurrency"]): r for r in run if r.get("_task") == task and "_concurrency" in r
+                    }
+                    val = _get(conc_map.get(c, {}), key)
+                    row += f" {val:{fmt}} |" if not math.isnan(val) else " n/a |"
+                print(row)
+
+        # Improvement column (2-run comparison only)
+        if len(all_runs) == 2:
+            print(f"\n### Improvement ({labels[0]} vs {labels[1]})\n")
+            print("| Metric | Concurrency | Change |")
+            print("| --- | --- | --- |")
+            for metric, key, _ in perf_metrics:
+                for c in concurrencies:
+                    conc_map0 = {
+                        int(r["_concurrency"]): r for r in all_runs[0] if r.get("_task") == task and "_concurrency" in r
+                    }
+                    conc_map1 = {
+                        int(r["_concurrency"]): r for r in all_runs[1] if r.get("_task") == task and "_concurrency" in r
+                    }
+                    v0 = _get(conc_map0.get(c, {}), key)
+                    v1 = _get(conc_map1.get(c, {}), key)
+                    if not math.isnan(v0) and not math.isnan(v1) and v1 > 0:
+                        pct = (v1 - v0) / v1 * 100
+                        print(f"| {metric} | {c} | {pct:+.1f}% |")
+
+
+# ---------------------------------------------------------------------------
+# CLI
+# ---------------------------------------------------------------------------
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description=__doc__,
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    parser.add_argument(
+        "--results",
+        type=str,
+        nargs="+",
+        action="append",
+        required=True,
+        metavar="FILE",
+        help="JSON result file(s) for one run. Repeat --results for multiple runs to compare.",
+    )
+    parser.add_argument(
+        "--labels",
+        type=str,
+        nargs="+",
+        default=None,
+        help="Label for each --results group (must match the number of --results groups).",
+    )
+    parser.add_argument("--output", type=str, default="results/tts_benchmark.png", help="Output image path.")
+    parser.add_argument("--title", type=str, default="TTS", help="Title prefix for the plot.")
+    parser.add_argument("--task", type=str, default=None, help="Filter to a single task (e.g. voice_clone).")
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+
+    # args.results is a list-of-lists due to action="append"
+    all_runs: list[list[dict]] = [load_run(group) for group in args.results]
+    n_runs = len(all_runs)
+
+    labels: list[str]
+    if args.labels:
+        if len(args.labels) != n_runs:
+            raise SystemExit(f"--labels count ({len(args.labels)}) must match --results groups ({n_runs})")
+        labels = args.labels
+    else:
+        labels = [f"run{i + 1}" for i in range(n_runs)]
+
+    print_comparison_table(all_runs, labels)
+    plot_comparison(all_runs, labels, args.output, task_filter=args.task, title_prefix=args.title)
+
+
+if __name__ == "__main__":
+    main()

docker/Dockerfile.ci

@@ -1,5 +1,5 @@
 ARG VLLM_BASE_IMAGE=vllm/vllm-openai
-ARG VLLM_BASE_TAG=v0.18.0
+ARG VLLM_BASE_TAG=v0.20.0
 FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG}
 ARG APP_DIR=/workspace/vllm-omni
 WORKDIR ${APP_DIR}
@@ -7,13 +7,11 @@ COPY . .
 
 # Install system dependencies
 RUN apt-get update && \
-    apt-get install -y espeak-ng ffmpeg git sox libsox-fmt-all jq && \
+    apt-get install -y espeak-ng git jq && \
     apt-get clean && \
     rm -rf /var/lib/apt/lists/*
 
-# Install vllm-omni into the same uv-managed Python environment used by the base image.
-# Use bash -c so that $(python3 -c ...) is expanded inside the container.
-RUN uv pip install --system --no-cache-dir ".[dev]"
+RUN uv pip install --system ".[dev]"
 
 RUN ln -sf /usr/bin/python3 /usr/bin/python
 

docker/Dockerfile.cuda

@@ -0,0 +1,22 @@
+ARG BASE_IMAGE=vllm/vllm-openai:v0.20.0
+FROM ${BASE_IMAGE}
+
+ARG COMMON_WORKDIR=/app
+
+WORKDIR ${COMMON_WORKDIR}
+
+# Step 1: Setup - Install system dependencies
+RUN apt-get update && \
+    apt-get install -y git jq && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/*
+
+RUN mkdir -p ${COMMON_WORKDIR}/vllm-omni
+
+# Step 2: Copy vllm-omni code and install
+COPY . ${COMMON_WORKDIR}/vllm-omni
+RUN cd ${COMMON_WORKDIR}/vllm-omni && uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" --no-cache-dir "."
+
+RUN ln -sf /usr/bin/python3 /usr/bin/python
+
+ENTRYPOINT []

docker/Dockerfile.npu

@@ -1,20 +1,20 @@
 ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
-ARG VLLM_ASCEND_TAG=v0.17.0rc1
+ARG VLLM_ASCEND_TAG=v0.18.0rc1
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 
-WORKDIR /vllm-workspace/vllm
-RUN git fetch origin --tags && git checkout v0.18.0
+# WORKDIR /vllm-workspace/vllm
+# RUN git fetch origin --tags && git checkout v0.18.0
 
-WORKDIR /vllm-workspace/vllm-ascend
-RUN git fetch origin releases/v0.18.0 && git checkout d781902ce9dbda8ab1e11bb0f2f0c1bc508fee7a
-# Install vllm-ascend
-# Append `libascend_hal.so` path (devlib) to LD_LIBRARY_PATH
-RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
-    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 -v -e /vllm-workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/ && \
-    python3 -m pip cache purge
+# WORKDIR /vllm-workspace/vllm-ascend
+# RUN git fetch origin releases/v0.18.0 && git checkout d781902ce9dbda8ab1e11bb0f2f0c1bc508fee7a
+# # Install vllm-ascend
+# # Append `libascend_hal.so` path (devlib) to LD_LIBRARY_PATH
+# RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
+#     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 -v -e /vllm-workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/ && \
+#     python3 -m pip cache purge
 
 ARG APP_DIR=/vllm-workspace/vllm-omni
 WORKDIR ${APP_DIR}

docker/Dockerfile.npu.a3

@@ -1,20 +1,20 @@
 ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
-ARG VLLM_ASCEND_TAG=v0.17.0rc1-a3
+ARG VLLM_ASCEND_TAG=v0.18.0rc1-a3
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 
-WORKDIR /vllm-workspace/vllm
-RUN git fetch origin --tags && git checkout v0.18.0
+# WORKDIR /vllm-workspace/vllm
+# RUN git fetch origin --tags && git checkout v0.18.0
 
-WORKDIR /vllm-workspace/vllm-ascend
-RUN git fetch origin releases/v0.18.0 && git checkout d781902ce9dbda8ab1e11bb0f2f0c1bc508fee7a
-# Install vllm-ascend
-# Append `libascend_hal.so` path (devlib) to LD_LIBRARY_PATH
-RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
-    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 -v -e /vllm-workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/ && \
-    python3 -m pip cache purge
+# WORKDIR /vllm-workspace/vllm-ascend
+# RUN git fetch origin releases/v0.18.0 && git checkout d781902ce9dbda8ab1e11bb0f2f0c1bc508fee7a
+# # Install vllm-ascend
+# # Append `libascend_hal.so` path (devlib) to LD_LIBRARY_PATH
+# RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
+#     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 -v -e /vllm-workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/ && \
+#     python3 -m pip cache purge
 
 ARG APP_DIR=/vllm-workspace/vllm-omni
 WORKDIR ${APP_DIR}

docker/Dockerfile.npu.ci

@@ -0,0 +1,42 @@
+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

@@ -0,0 +1,42 @@
+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 []

docker/Dockerfile.rocm

@@ -1,4 +1,4 @@
-ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.18.0
+ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.20.0
 FROM ${BASE_IMAGE} AS base
 
 # Declare a variable to know if we want to use the nightly build or the stable build.
@@ -18,8 +18,10 @@ ARG COMMON_WORKDIR=/app
 WORKDIR ${COMMON_WORKDIR}
 
 # Step 1: Setup - Install system dependencies
+# Need to include ffmpeg because vllm rocm upstream docker image
+# does not include it.
 RUN apt-get update && \
-    apt-get install -y espeak-ng ffmpeg git sox libsox-fmt-all jq && \
+    apt-get install -y espeak-ng ffmpeg git jq && \
     apt-get clean && \
     rm -rf /var/lib/apt/lists/*
 
@@ -39,6 +41,24 @@ RUN if [ "${USE_NIGHTLY_BUILD}" = "1" ]; then \
 # Step 3: Copy vllm-omni code and install without uv
 RUN mkdir -p ${COMMON_WORKDIR}/vllm-omni
 COPY . ${COMMON_WORKDIR}/vllm-omni
+
+# This is a workaround to ensure pytest exits with the correct status code in CI tests.
+RUN printf '%s\n' \
+    'import os' \
+    '' \
+    '_exit_code = 1' \
+    '' \
+    'def pytest_sessionfinish(session, exitstatus):' \
+    '    global _exit_code' \
+    '    _exit_code = int(exitstatus)' \
+    '' \
+    'def pytest_unconfigure(config):' \
+    '    import sys' \
+    '    sys.stdout.flush()' \
+    '    sys.stderr.flush()' \
+    '    os._exit(_exit_code)' \
+    > ${COMMON_WORKDIR}/vllm-omni/conftest.py
+
 RUN cd ${COMMON_WORKDIR}/vllm-omni && uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" --no-cache-dir ".[dev]" --no-build-isolation
 
 RUN ln -sf /usr/bin/python3 /usr/bin/python

docker/Dockerfile.xpu

@@ -14,7 +14,6 @@ RUN wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRO
 RUN apt clean && apt-get update -y && \
     apt-get install -y --no-install-recommends --fix-missing \
     curl \
-    espeak-ng \
     ffmpeg \
     git \
     libsndfile1 \
@@ -30,8 +29,20 @@ RUN apt clean && apt-get update -y && \
     python3.12-dev \
     python3-pip
 
-RUN apt update && apt upgrade -y && \
-    apt install -y intel-oneapi-compiler-dpcpp-cpp-2025.3
+# Add oneAPI repo, pin oneAPI to 2025.3, then install pinned packages in one layer.
+RUN wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB | gpg --dearmor | tee /usr/share/keyrings/oneapi-archive-keyring.gpg > /dev/null && \
+    echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" | tee /etc/apt/sources.list.d/oneAPI.list && \
+    printf '%s\n' \
+    'Package: intel-oneapi-* intel-deep-learning-essentials* intel-pti*' \
+    'Pin: version 2025.3*' \
+    'Pin-Priority: 1001' \
+    > /etc/apt/preferences.d/oneapi-2025.3.pref && \
+    apt-get update -y && \
+    apt-get install -y --no-install-recommends \
+    intel-oneapi-compiler-dpcpp-cpp-2025.3 \
+    intel-oneapi-mkl-devel-2025.3 \
+    intel-oneapi-dnnl-devel-2025.3 && \
+    rm -rf /var/lib/apt/lists/*
 
 # Install UMD with fixed version
 RUN mkdir neo && \
@@ -54,9 +65,9 @@ RUN curl -LsSf https://astral.sh/uv/install.sh | sh
 RUN uv venv --python ${PYTHON_VERSION} --seed ${VIRTUAL_ENV}
 ENV PATH="$VIRTUAL_ENV/bin:$PATH"
 
-# This oneccl contains the BMG support which is not the case for default version of oneapi 2025.2.
-ARG ONECCL_INSTALLER="intel-oneccl-2021.15.7.8_offline.sh"
-RUN wget "https://github.com/uxlfoundation/oneCCL/releases/download/2021.15.7/${ONECCL_INSTALLER}" && \
+# This oneccl contains the BMG support which is not the case for default version of oneapi 2025.3.
+ARG ONECCL_INSTALLER="intel-oneccl-2021.15.9.14_offline.sh"
+RUN wget "https://github.com/uxlfoundation/oneCCL/releases/download/2021.15.9/${ONECCL_INSTALLER}" && \
     bash "${ONECCL_INSTALLER}" -a --silent --eula accept && \
     rm "${ONECCL_INSTALLER}" && \
     echo "source /opt/intel/oneapi/setvars.sh --force" >> /root/.bashrc && \
@@ -68,6 +79,7 @@ SHELL ["bash", "-c"]
 CMD ["bash", "-c", "source /root/.bashrc && exec bash"]
 
 WORKDIR /workspace/
+
 ENV UV_HTTP_TIMEOUT=500
 
 # Configure package index for XPU
@@ -76,24 +88,18 @@ ENV UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
 ENV UV_INDEX_STRATEGY="unsafe-best-match"
 ENV UV_LINK_MODE="copy"
 
-ARG VLLM_VERSION=v0.18.0
+ARG VLLM_VERSION=v0.20.0
 RUN git clone -b ${VLLM_VERSION} https://github.com/vllm-project/vllm
 WORKDIR /workspace/vllm
 
 RUN --mount=type=cache,target=/root/.cache/uv \
     uv pip install --upgrade pip && \
-    uv pip install -r requirements/xpu.txt
-
- # used for suffix method speculative decoding
- # build deps for proto + nanobind-based extensions to set up the build environment
-RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install grpcio-tools protobuf nanobind
- # arctic-inference is built from source which needs torch-xpu properly installed first
-RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip install -r requirements/xpu.txt && \
+    uv pip install setuptools_scm &&\
+    uv pip install grpcio-tools protobuf nanobind && \
     source /opt/intel/oneapi/setvars.sh --force && \
     source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force && \
-    export CMAKE_PREFIX_PATH="$(python -c 'import site; print(site.getsitepackages()[0])'):${CMAKE_PREFIX_PATH}" && \
-    uv pip install --no-build-isolation arctic-inference==0.1.1
+    export CMAKE_PREFIX_PATH="$(python3 -c 'import site; print(site.getsitepackages()[0])'):${CMAKE_PREFIX_PATH}"
 
 ENV LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/usr/local/lib/"
 
@@ -107,6 +113,9 @@ CMD ["/bin/bash"]
 
 FROM vllm-base AS vllm-openai
 
+# install nixl from source code
+RUN uv pip install nixl
+
 # install additional dependencies for openai api server
 RUN --mount=type=cache,target=/root/.cache/uv \
     uv pip install accelerate hf_transfer pytest pytest_asyncio lm_eval[api] modelscope
@@ -114,10 +123,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \
 # install development dependencies (for testing)
 RUN uv pip install -e tests/vllm_test_utils
 
-# install nixl from source code
-ENV NIXL_VERSION=0.7.0
-RUN python /workspace/vllm/tools/install_nixl_from_source_ubuntu.py
-
 # ensure vllm is properly installed
 RUN python -c "import vllm, inspect; print(vllm.__file__)"
 RUN uv pip show vllm
@@ -137,7 +142,7 @@ RUN uv pip install --no-cache-dir ".[dev]" --no-build-isolation
 # FIX triton
 RUN --mount=type=cache,target=/root/.cache/uv \
     uv pip uninstall triton triton-xpu && \
-    uv pip install triton-xpu==3.6.0
+    uv pip install triton-xpu==3.7.0
 
 # remove torch bundled oneccl to avoid conflicts
 RUN --mount=type=cache,target=/root/.cache/uv \

docs/.nav.yml

@@ -10,38 +10,38 @@ nav:
       - Image Generation: serving/image_generation_api.md
       - Image Edit: serving/image_edit_api.md
       - Text to Speech: serving/speech_api.md
+      - Audio Generation: serving/audio_generate_api.md
+      - Streaming Video Input: serving/video_stream_api.md
   - Examples:
     - examples/README.md
     - Offline Inference:
       - BAGEL-7B-MoT: user_guide/examples/offline_inference/bagel.md
-      - CosyVoice3: user_guide/examples/offline_inference/cosyvoice3.md
-      - Fish Speech S2 Pro: user_guide/examples/offline_inference/fish_speech.md
       - GLM-Image Multistage End-to-End Inference: user_guide/examples/offline_inference/glm_image.md
       - Helios Video Generation: user_guide/examples/offline_inference/helios.md
       - HunyuanImage-3.0 Image-to-Text Inference: user_guide/examples/offline_inference/hunyuan_image3.md
       - Image-To-Image: user_guide/examples/offline_inference/image_to_image.md
       - Image-To-Video: user_guide/examples/offline_inference/image_to_video.md
+      - InternVLA-A1: user_guide/examples/offline_inference/internvla_a1.md
       - MammothModa2-Preview: user_guide/examples/offline_inference/mammothmodal2_preview.md
       - MiMo-Audio Offline Inference: user_guide/examples/offline_inference/mimo_audio.md
       - Qwen2.5-Omni: user_guide/examples/offline_inference/qwen2_5_omni.md
       - Qwen3-Omni: user_guide/examples/offline_inference/qwen3_omni.md
-      - Qwen3-TTS: user_guide/examples/offline_inference/qwen3_tts.md
       - Text-To-Audio: user_guide/examples/offline_inference/text_to_audio.md
       - Text-To-Image: user_guide/examples/offline_inference/text_to_image.md
+      - Text-To-Speech: user_guide/examples/offline_inference/text_to_speech.md
       - Text-To-Video: user_guide/examples/offline_inference/text_to_video.md
-      - Voxtral TTS Offline Inference: user_guide/examples/offline_inference/voxtral_tts.md
       - X-To-Video-Audio: user_guide/examples/offline_inference/x_to_video_audio.md
     - Online Serving:
       - BAGEL-7B-MoT: user_guide/examples/online_serving/bagel.md
       - vLLM-Omni Helm Chart: user_guide/examples/online_serving/chart-helm.md
-      - Fish Speech S2 Pro: user_guide/examples/online_serving/fish_speech.md
+      - Diffusers Backend Adapter: user_guide/examples/online_serving/diffusers_pipeline_adapter.md
       - GLM-Image Online Serving: user_guide/examples/online_serving/glm_image.md
       - Image-To-Image: user_guide/examples/online_serving/image_to_image.md
       - Image-To-Video: user_guide/examples/online_serving/image_to_video.md
       - Online serving Example of vLLM-Omni for MiMo-Audio: user_guide/examples/online_serving/mimo_audio.md
       - Qwen2.5-Omni: user_guide/examples/online_serving/qwen2_5_omni.md
       - Qwen3-Omni: user_guide/examples/online_serving/qwen3_omni.md
-      - Qwen3-TTS: user_guide/examples/online_serving/qwen3_tts.md
+      - Text-To-Speech: user_guide/examples/online_serving/text_to_speech.md
       - Text-To-Image: user_guide/examples/online_serving/text_to_image.md
       - Text-To-Video: user_guide/examples/online_serving/text_to_video.md
   - General:
@@ -59,11 +59,8 @@ nav:
       - Cache Acceleration:
         - TeaCache: user_guide/diffusion/cache_acceleration/teacache.md
         - Cache-DiT: user_guide/diffusion/cache_acceleration/cache_dit.md
-      - Quantization:
-        - Overview: user_guide/diffusion/quantization/overview.md
-        - FP8: user_guide/diffusion/quantization/fp8.md
-        - Int8: user_guide/diffusion/quantization/int8.md
-        - GGUF: user_guide/diffusion/quantization/gguf.md
+      - Attention Backends: user_guide/diffusion/attention_backends.md
+      - Frame Interpolation: user_guide/diffusion/frame_interpolation.md
       - Parallelism:
         - Overview: user_guide/diffusion/parallelism/overview.md
         - CFG Parallel: user_guide/diffusion/parallelism/cfg_parallel.md
@@ -76,11 +73,20 @@ nav:
       - LoRA: user_guide/diffusion/lora.md
       - Custom Pipeline: features/custom_pipeline.md
       - Step Execution: user_guide/diffusion/step_execution.md
+    - 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
     - ComfyUI: features/comfyui.md
 - Developer Guide:
   - General:
     - contributing/README.md
-    - pr_reviewer.md
     - glob: contributing/*
       flatten_single_child_sections: true
   - Model Implementation:
@@ -88,7 +94,6 @@ nav:
     - contributing/model/adding_omni_model.md
     - contributing/model/adding_tts_model.md
     - contributing/model/adding_diffusion_model.md
-    - contributing/model/adding_quantization_model.md
   - CI: contributing/ci
   - Design Documents:
     - design/index.md
@@ -97,16 +102,19 @@ nav:
       - design/feature/disaggregated_inference.md
       - design/feature/ray_based_execution.md
       - design/feature/omni_connectors/
+      - design/feature/prefix_caching.md
       - design/feature/cfg_parallel.md
+      - design/feature/expert_parallel.md
       - design/feature/sequence_parallel.md
       - design/feature/tensor_parallel.md
       - design/feature/vae_parallel.md
       - design/feature/hsdp.md
       - design/feature/cache_dit.md
       - design/feature/teacache.md
-      - design/feature/async_chunk_design.md
+      - design/feature/async_chunk.md
       - design/feature/vae_parallel.md
       - design/feature/diffusion_step_execution.md
+      - design/feature/diffusion_continuous_batching.md
     - Module Design:
       - design/module/ar_module.md
       - design/module/dit_module.md

docs/api/README.md

@@ -5,8 +5,7 @@
 Main entry points for vLLM-Omni inference and serving.
 
 - [vllm_omni.entrypoints.async_omni.AsyncOmni][]
-- [vllm_omni.entrypoints.async_omni_diffusion.AsyncOmniDiffusion][]
-- [vllm_omni.entrypoints.cfg_companion_tracker.CfgCompanionTracker][]
+- [vllm_omni.engine.cfg_companion_tracker.CfgCompanionTracker][]
 - [vllm_omni.entrypoints.cli.benchmark.base.OmniBenchmarkSubcommandBase][]
 - [vllm_omni.entrypoints.cli.benchmark.main.OmniBenchmarkSubcommand][]
 - [vllm_omni.entrypoints.cli.benchmark.serve.OmniBenchmarkServingSubcommand][]

docs/cli/serve.md

@@ -1,5 +1,61 @@
 # vllm-omni serve
 
+## Stage-based CLI quickstart
+
+The stage-based CLI is designed for deployments that require launching each pipeline stage in an isolated process
+(e.g., across separate operating system processes, distinct GPUs, or distributed hosts).
+
+- For **migrated models** that utilize the bundled deployment YAML configurations located in
+  `vllm_omni/deploy/`, the `--deploy-config` flag is only required to override the default configuration. By default, executing `vllm serve MODEL --omni ...`
+  automatically loads the bundled deployment configuration.
+- For **legacy models** utilizing configuration files located in
+  `vllm_omni/model_executor/stage_configs/`, the `--stage-configs-path` parameter remains mandatory.
+
+Example: Initializing Stage 0 (Orchestrator and API Server):
+The commands below show a common device mapping where Stage 0 uses GPU 0 and
+worker stages use GPU 1 via `CUDA_VISIBLE_DEVICES`.
+
+```bash
+CUDA_VISIBLE_DEVICES=0 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
+    --port 8091 \
+    --stage-id 0 \
+    --omni-master-address 127.0.0.1 \
+    --omni-master-port 26000
+```
+
+Example: Initializing a Headless Worker Stage (Stage 1):
+
+```bash
+CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
+    --stage-id 1 \
+    --headless \
+    --omni-master-address 127.0.0.1 \
+    --omni-master-port 26000
+```
+
+When utilizing a custom deployment YAML based on the new schema, append `--deploy-config /path/to/override.yaml` to each command execution. Conversely, for legacy models, substitute this parameter with `--stage-configs-path /path/to/stage_configs.yaml`.
+
+In the standard execution paradigm, the `--stage-overrides` argument is utilized to apply stage-specific configurations from a single CLI command.
+However, under the **stage-based CLI** paradigm, where each process strictly encapsulates a single stage, it is recommended to specify tuning parameters directly via discrete command-line flags for the respective stage, rather than constructing a composite `--stage-overrides` JSON string.
+
+For example, as an alternative to the following composite configuration:
+
+```bash
+vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
+    --stage-overrides '{"1": {"gpu_memory_utilization": 0.5}}'
+```
+
+the stage-based CLI permits the direct initialization of Stage 1 with explicit parameters:
+
+```bash
+CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
+    --stage-id 1 \
+    --headless \
+    --gpu-memory-utilization 0.5 \
+    --omni-master-address 127.0.0.1 \
+    --omni-master-port 26000
+```
+
 ## JSON CLI Arguments
 
 --8<-- "docs/cli/json_tip.inc.md"

docs/community/governance.md

@@ -37,10 +37,12 @@ Committers have write access and merge rights. They typically have deep expertis
 - [@gcanlin](https://github.com/gcanlin): Hardware plugin and NPU integration
 - [@Isotr0py](https://github.com/Isotr0py): Diffusion and Quantization
 - [@linyueqian](https://github.com/linyueqian): TTS and Omni Support
+- [@lishunyang12](https://github.com/lishunyang12): Quantization and Configuration
 - [@princepride](https://github.com/princepride): Diffusion and Omni Support
 - [@SamitHuang](https://github.com/SamitHuang): RL and Diffusion
 - [@tzhouam](https://github.com/tzhouam): Engine and New Model Support
-- [@wtomin](https://github.com/wtomin):
+- [@wtomin](https://github.com/wtomin): Diffusion and Parallelism
+- [@ZeldaHuang](https://github.com/ZeldaHuang): Omni Support
 - [@ZJY0516](https://github.com/ZJY0516): Diffusion and CustomOp
 
 ## Meetings

docs/configuration/README.md

@@ -6,7 +6,7 @@ For options within a vLLM Engine. Please refer to [vLLM Configuration](https://d
 
 Currently, the main options are maintained by stage configs for each model.
 
-For specific example, please refer to [Qwen2.5-omni stage config](stage_configs/qwen2_5_omni.yaml)
+For a specific example, see the [Qwen2.5-Omni deploy config](gh-file:vllm_omni/deploy/qwen2_5_omni.yaml). The matching frozen pipeline topology lives at [vllm_omni/model_executor/models/qwen2_5_omni/pipeline.py](gh-file:vllm_omni/model_executor/models/qwen2_5_omni/pipeline.py).
 
 For introduction, please check [Introduction for stage config](./stage_configs.md)
 

docs/configuration/pd_disaggregation.md

@@ -11,7 +11,7 @@ deployment-specific values usually change per environment:
 - connector backend and connector ports
 - connector IPs or bootstrap addresses
 
-Start from the [default Qwen3-Omni stage config](gh-file:vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml)
+Start from the [default Qwen3-Omni stage config](gh-file:vllm_omni/deploy/qwen3_omni_moe.yaml)
 and copy it to your own file, for example `qwen3_omni_pd.yaml`. Then apply the
 changes below.
 
@@ -145,19 +145,13 @@ Compared with the default Qwen3-Omni config:
 ```yaml
 runtime:
   enabled: true
-  defaults:
-    window_size: -1
-    max_inflight: 1
   edges:
     - from: 0
       to: 1
-      window_size: -1
     - from: 1
       to: 2
-      window_size: -1
     - from: 2
       to: 3
-      window_size: -1
 ```
 
 ## 4. Launch with your custom config