Revert "Add online serving to Stable Audio Diffusion and introduce `v1/audio/…"

This reverts commit 7a560492a2.

.buildkite/bootstrap-amd-omni.sh

@@ -56,88 +56,6 @@ check_run_all_label() {
     fi
 }
 
-is_docs_only_change() {
-    local file_path
-    local has_any=0
-
-    while IFS= read -r file_path; do
-        [[ -z "${file_path}" ]] && continue
-        has_any=1
-
-        if [[ "${file_path}" == docs/* ]]; then
-            continue
-        fi
-        if [[ "${file_path}" == *.md ]]; then
-            continue
-        fi
-        if [[ "${file_path}" == "mkdocs.yml" ]]; then
-            continue
-        fi
-        return 1
-    done
-
-    [[ "${has_any}" -eq 1 ]]
-}
-
-resolve_skip_ci() {
-    local is_pr_build=0
-    local files
-    local base_branch base_ref
-
-    if [[ "${BUILDKITE_PULL_REQUEST:-false}" != "false" && -n "${BUILDKITE_PULL_REQUEST:-}" ]]; then
-        is_pr_build=1
-    fi
-
-    if [[ "${is_pr_build}" -eq 1 ]]; then
-        base_branch="${BUILDKITE_PULL_REQUEST_BASE_BRANCH:-main}"
-        if ! git rev-parse --verify "origin/${base_branch}" >/dev/null 2>&1; then
-            echo "resolve_skip_ci: origin/${base_branch} not found locally; trying fetch" >&2
-            git fetch --depth=200 origin "${base_branch}" >/dev/null 2>&1 || true
-        fi
-
-        base_ref=""
-        if git rev-parse --verify "origin/${base_branch}" >/dev/null 2>&1; then
-            base_ref="origin/${base_branch}"
-        elif git rev-parse --verify "${base_branch}" >/dev/null 2>&1; then
-            base_ref="${base_branch}"
-        else
-            echo "resolve_skip_ci: cannot resolve PR base ${base_branch}; skip-ci=0" >&2
-            echo -n 0
-            return 0
-        fi
-
-        if ! files="$(git diff --name-only "${base_ref}...${BUILDKITE_COMMIT}" 2>/dev/null)"; then
-            echo "resolve_skip_ci: failed to compute PR changed files; skip-ci=0" >&2
-            echo -n 0
-            return 0
-        fi
-    elif [[ "${BUILDKITE_BRANCH:-}" == "main" ]]; then
-        if ! git rev-parse --verify "${BUILDKITE_COMMIT}^" >/dev/null 2>&1; then
-            echo "resolve_skip_ci: commit has no parent on main; skip-ci=0" >&2
-            echo -n 0
-            return 0
-        fi
-        if ! files="$(git diff --name-only "${BUILDKITE_COMMIT}^..${BUILDKITE_COMMIT}" 2>/dev/null)"; then
-            echo "resolve_skip_ci: failed to compute main changed files; skip-ci=0" >&2
-            echo -n 0
-            return 0
-        fi
-    else
-        echo "resolve_skip_ci: not PR/main build; skip-ci=0" >&2
-        echo -n 0
-        return 0
-    fi
-
-    if is_docs_only_change <<< "${files}"; then
-        echo "resolve_skip_ci: docs-only change detected; skip-ci=1" >&2
-        echo -n 1
-        return 0
-    fi
-
-    echo "resolve_skip_ci: non-doc changes detected; skip-ci=0" >&2
-    echo -n 0
-}
-
 if [[ -z "${COV_ENABLED:-}" ]]; then
     COV_ENABLED=0
 fi
@@ -172,34 +90,10 @@ upload_pipeline() {
     FAIL_FAST=$(fail_fast)
 
     cd .buildkite
-
-    # Select test definition file: merge suite for main, ready suite for PRs.
-    # For debugging, DEBUG_TEST_YAML can override the selection — accepts
-    # "merge" or "ready" (case-insensitive).
-    if [[ -n "${DEBUG_TEST_YAML:-}" ]]; then
-        case "${DEBUG_TEST_YAML,,}" in
-            merge)
-                TEST_YAML="test-amd-merge.yml"
-                ;;
-            ready)
-                TEST_YAML="test-amd-ready.yaml"
-                ;;
-            *)
-                echo "ERROR: DEBUG_TEST_YAML must be 'merge' or 'ready', got '$DEBUG_TEST_YAML'" >&2
-                exit 1
-                ;;
-        esac
-        echo "DEBUG_TEST_YAML override: using $TEST_YAML"
-    elif [[ $BUILDKITE_BRANCH == "main" ]]; then
-        TEST_YAML="test-amd-merge.yml"
-    else
-        TEST_YAML="test-amd-ready.yaml"
-    fi
-
     (
         set -x
         # Output pipeline.yaml with all blank lines removed
-        minijinja-cli test-template.j2 "$TEST_YAML" \
+        minijinja-cli test-template.j2 test-amd.yaml \
             -D branch="$BUILDKITE_BRANCH" \
             -D list_file_diff="$LIST_FILE_DIFF" \
             -D run_all="$RUN_ALL" \
@@ -238,13 +132,30 @@ fi
 # Early exit start: skip pipeline if conditions are met
 # ----------------------------------------------------------------------
 
-# Match CUDA Buildkite skip-ci behavior for docs/markdown-only changes.
+# skip pipeline if all changed files are under docs/
 if [[ "${DOCS_ONLY_DISABLE}" != "1" ]]; then
-    SKIP_CI="$(resolve_skip_ci)"
-    if [[ "${SKIP_CI}" == "1" ]]; then
-        echo "[docs-only] Docs/markdown-only changes detected. Exiting before pipeline upload."
-        exit 0
+  if [[ -n "${file_diff:-}" ]]; then
+    docs_only=1
+    # Robust iteration over newline-separated file_diff
+    while IFS= read -r f; do
+      [[ -z "$f" ]] && continue
+      # **Policy:** only skip if *every* path starts with docs/
+      if [[ "$f" != docs/* ]]; then
+        docs_only=0
+        break
+      fi
+    done < <(printf '%s\n' "$file_diff" | tr ' ' '\n' | tr -d '\r')
+
+    if [[ "$docs_only" -eq 1 ]]; then
+      buildkite-agent annotate ":memo: CI skipped — docs/** only changes detected
+
+\`\`\`
+${file_diff}
+\`\`\`" --style "info" || true
+      echo "[docs-only] All changes are under docs/. Exiting before pipeline upload."
+      exit 0
     fi
+  fi
 fi
 
 # ----------------------------------------------------------------------

.buildkite/bootstrap-intel-omni.sh

@@ -1,51 +0,0 @@
-#!/bin/bash
-# vllm-omni Intel bootstrap
-# Uses static pipeline-intel.yaml for Intel XPU tests
-# Based on: bootstrap-amd-omni.sh
-
-set -euo pipefail
-
-if [[ -z "${DOCS_ONLY_DISABLE:-}" ]]; then
-    DOCS_ONLY_DISABLE=0
-fi
-
-upload_pipeline() {
-    echo "--- 🛠 Preparing Intel pipeline"
-    buildkite-agent pipeline upload .buildkite/pipeline-intel.yaml
-}
-
-get_diff() {
-    $(git add .)
-    echo $(git diff --name-only --diff-filter=ACMDR $(git merge-base origin/main HEAD))
-}
-
-get_diff_main() {
-    $(git add .)
-    echo $(git diff --name-only --diff-filter=ACMDR HEAD~1)
-}
-
-file_diff=$(get_diff)
-if [[ $BUILDKITE_BRANCH == "main" ]]; then
-    file_diff=$(get_diff_main)
-fi
-
-# ----------------------------------------------------------------------
-# Early exit: skip pipeline if all changed files are under docs/
-# ----------------------------------------------------------------------
-if [[ "${DOCS_ONLY_DISABLE}" != "1" ]] && [[ -n "${file_diff:-}" ]]; then
-    docs_only=1
-    while IFS= read -r f; do
-        [[ -z "$f" ]] && continue
-        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" --style "info"
-        exit 0
-    fi
-fi
-
-upload_pipeline

.buildkite/pipeline-intel.yaml

@@ -1,16 +0,0 @@
-# .buildkite/pipeline-intel.yaml
-steps:
-  - group: "INTEL Tests"
-    steps:
-      - label: "Intel: XPU Test"
-        command: "bash .buildkite/scripts/hardware_ci/run-xpu-test.sh"
-        agents:
-          queue: "intel-gpu-omni"
-        env:
-          DOCKER_BUILDKIT: "1"
-          # Buildkite will automatically replace this with the actual commit hash
-          VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}"
-          VLLM_VERSION: "v0.22.0"
-        priority: 100
-        timeout_in_minutes: 60
-        soft_fail: false

.buildkite/pipeline-npu-a3.yaml

@@ -1,110 +0,0 @@
-# steps:
-#   - label: ":buildkit: Build and Push NPU Test Image - A3"
-#     key: image-build-a3
-#     agents:
-#       queue: "ascend-a3"
-#       resource_class: "npu-2"
-#     plugins:
-#       - kubernetes:
-#           metadata:
-#             annotations:
-#               vault.hashicorp.com/agent-init-first: "true"
-#               vault.hashicorp.com/agent-inject: "true"
-#               vault.hashicorp.com/agent-inject-perms-ca.pem: "0400"
-#               vault.hashicorp.com/agent-inject-perms-cert.pem: "0400"
-#               vault.hashicorp.com/agent-inject-perms-config.json: "0400"
-#               vault.hashicorp.com/agent-inject-perms-key.pem: "0400"
-#               vault.hashicorp.com/agent-inject-secret-ca.pem: internal/data/ascend/buildkitd
-#               vault.hashicorp.com/agent-inject-secret-cert.pem: internal/data/ascend/buildkitd
-#               vault.hashicorp.com/agent-inject-secret-config.json: internal/data/ascend/buildkitd
-#               vault.hashicorp.com/agent-inject-secret-key.pem: internal/data/ascend/buildkitd
-#               vault.hashicorp.com/agent-inject-template-ca.pem: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.RootCA }}\n{{- end }}"
-#               vault.hashicorp.com/agent-inject-template-cert.pem: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.ClientCaCert }}\n{{- end }}"
-#               vault.hashicorp.com/agent-inject-template-config.json: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.dockerConfig }}\n{{- end }}"
-#               vault.hashicorp.com/agent-inject-template-key.pem: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.ClientCaKey }}\n{{- end }}"
-#               vault.hashicorp.com/agent-pre-populate-only: "true"
-#               vault.hashicorp.com/agent-run-as-group: "1000"
-#               vault.hashicorp.com/agent-run-as-user: "1000"
-#               vault.hashicorp.com/agent-service-account-token-volume-name: token-vol
-#               vault.hashicorp.com/role: ascend-gha-runners
-#               vault.hashicorp.com/secret-volume-path: /home/user/.docker/
-#               vault.hashicorp.com/tls-skip-verify: "true"
-#           podSpecPatch:
-#             volumes:
-#             - name: token-vol
-#               projected:
-#                 defaultMode: 420
-#                 sources:
-#                 - serviceAccountToken:
-#                     audience: api
-#                     expirationSeconds: 600
-#                     path: token
-#     env:
-#       VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}"
-#       IMAGE_NAME: "vllm-omni-ci-npu"
-#       IMAGE_REGISTRY: "swr.cn-southwest-2.myhuaweicloud.com/modelfoundry"
-#       BUILDKITD_ADDR: "tcp://buildkitd-service.buildkitd:1234"
-#     command: |
-#       set -ex
-
-#       echo "--- Building and pushing NPU Test Image"
-#       echo "Image: $${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG}"
-#       echo "buildkitd address: $${BUILDKITD_ADDR}"
-
-#       git config --global url."https://gh-proxy.test.osinfra.cn/https://github.com/".insteadOf "https://github.com/"
-
-#       if ! command -v buildctl &> /dev/null; then
-#         echo "Installing buildctl..."
-#         mkdir -p /tmp/buildkit
-#         BUILDKIT_VERSION="v0.29.0"
-#         wget -q "https://gh-proxy.test.osinfra.cn/https://github.com/moby/buildkit/releases/download/$${BUILDKIT_VERSION}/buildkit-$${BUILDKIT_VERSION}.linux-arm64.tar.gz" -O /tmp/buildkit.tar.gz
-#         tar -xzf /tmp/buildkit.tar.gz -C /tmp/buildkit
-#         cp /tmp/buildkit/bin/buildctl /usr/local/bin/
-#       fi
-
-#       export DOCKER_CONFIG=/home/user/.docker
-#       buildctl \
-#         --addr="$${BUILDKITD_ADDR}" \
-#         --tlscacert=/home/user/.docker/ca.pem \
-#         --tlscert=/home/user/.docker/cert.pem \
-#         --tlskey=/home/user/.docker/key.pem \
-#         build \
-#         --frontend dockerfile.v0 \
-#         --local context=. \
-#         --local dockerfile=./docker \
-#         --opt filename=Dockerfile.npu.ci.a3 \
-#         --secret id=dockerconfig,src=/home/user/.docker/config.json \
-#         --output type=image,name=$${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG},push=true \
-#         --progress=plain
-
-#       echo "--- Image pushed successfully"
-#       echo "$${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG}"
-
-#   # - label: "Upload Ready Pipeline - A3"
-#   #   depends_on: image-build-a3
-#   #   key: upload-ready-pipeline-a3
-#   #   if: build.branch != "main" && build.pull_request.labels includes "npu-test"
-#   #   agents:
-#   #     queue: "ascend-a3"
-#   #     resource_class: "npu-2"
-#   #   command: |
-#   #     buildkite-agent pipeline upload .buildkite/test-npu-ready.yaml
-
-  # # L4 Test — main+NIGHTLY=1 (scheduled), or PR with specific label (e.g. add label then Rebuild)
-  # - label: "Upload Nightly Pipeline"
-  #   depends_on: image-build
-  #   key: upload-nightly-pipeline
-  #   if: >-
-  #     (build.branch == "main" && build.env("NIGHTLY") == "1") ||
-  #     (build.branch != "main" && (
-  #       build.pull_request.labels includes "nightly-test" ||
-  #       build.pull_request.labels includes "omni-test" ||
-  #       build.pull_request.labels includes "tts-test" ||
-  #       build.pull_request.labels includes "diffusion-x2iat-test" ||
-  #       build.pull_request.labels includes "diffusion-x2v-test"
-  #     ))
-  #   commands:
-  #     - buildkite-agent pipeline upload .buildkite/test-npu-nightly.yaml
-  #   agents:
-  #     queue: "ascend-a3"
-  #     resource_class: "npu-2"

.buildkite/pipeline-npu.yaml

@@ -1,112 +0,0 @@
-steps:
-  - label: ":buildkit: Build and Push NPU Test Image - A2"
-    key: image-build
-    agents:
-      queue: "ascend-a2b3"
-      resource_class: "npu-1"
-    plugins:
-      - kubernetes:
-          metadata:
-            annotations:
-              vault.hashicorp.com/agent-init-first: "true"
-              vault.hashicorp.com/agent-inject: "true"
-              vault.hashicorp.com/agent-inject-perms-ca.pem: "0400"
-              vault.hashicorp.com/agent-inject-perms-cert.pem: "0400"
-              vault.hashicorp.com/agent-inject-perms-config.json: "0400"
-              vault.hashicorp.com/agent-inject-perms-key.pem: "0400"
-              vault.hashicorp.com/agent-inject-secret-ca.pem: internal/data/ascend/buildkitd
-              vault.hashicorp.com/agent-inject-secret-cert.pem: internal/data/ascend/buildkitd
-              vault.hashicorp.com/agent-inject-secret-config.json: internal/data/ascend/buildkitd
-              vault.hashicorp.com/agent-inject-secret-key.pem: internal/data/ascend/buildkitd
-              vault.hashicorp.com/agent-inject-template-ca.pem: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.RootCA }}\n{{- end }}"
-              vault.hashicorp.com/agent-inject-template-cert.pem: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.ClientCaCert }}\n{{- end }}"
-              vault.hashicorp.com/agent-inject-template-config.json: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.dockerConfig }}\n{{- end }}"
-              vault.hashicorp.com/agent-inject-template-key.pem: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.ClientCaKey }}\n{{- end }}"
-              vault.hashicorp.com/agent-pre-populate-only: "true"
-              vault.hashicorp.com/agent-run-as-group: "1000"
-              vault.hashicorp.com/agent-run-as-user: "1000"
-              vault.hashicorp.com/agent-service-account-token-volume-name: token-vol
-              vault.hashicorp.com/role: ascend-gha-runners
-              vault.hashicorp.com/secret-volume-path: /home/user/.docker/
-              vault.hashicorp.com/tls-skip-verify: "true"
-          podSpecPatch:
-            volumes:
-            - name: token-vol
-              projected:
-                defaultMode: 420
-                sources:
-                - serviceAccountToken:
-                    audience: api
-                    expirationSeconds: 600
-                    path: token
-    env:
-      VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}"
-      IMAGE_NAME: "vllm-omni-ci-npu"
-      IMAGE_REGISTRY: "swr.cn-southwest-2.myhuaweicloud.com/modelfoundry"
-      BUILDKITD_ADDR: "tcp://buildkitd-service.buildkitd:1234"
-    command: |
-      set -ex
-
-      echo "--- Building and pushing NPU Test Image"
-      echo "Image: $${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG}"
-      echo "buildkitd address: $${BUILDKITD_ADDR}"
-
-      if ! command -v buildctl &> /dev/null; then
-        echo "Installing buildctl..."
-        mkdir -p /tmp/buildkit
-        BUILDKIT_VERSION="v0.29.0"
-        wget -q "https://github.com/moby/buildkit/releases/download/$${BUILDKIT_VERSION}/buildkit-$${BUILDKIT_VERSION}.linux-arm64.tar.gz" -O /tmp/buildkit.tar.gz
-        tar -xzf /tmp/buildkit.tar.gz -C /tmp/buildkit
-        cp /tmp/buildkit/bin/buildctl /usr/local/bin/
-      fi
-
-      export DOCKER_CONFIG=/home/user/.docker
-      buildctl \
-        --addr="$${BUILDKITD_ADDR}" \
-        --tlscacert=/home/user/.docker/ca.pem \
-        --tlscert=/home/user/.docker/cert.pem \
-        --tlskey=/home/user/.docker/key.pem \
-        build \
-        --frontend dockerfile.v0 \
-        --local context=. \
-        --local dockerfile=./docker \
-        --opt filename=Dockerfile.npu.ci \
-        --secret id=dockerconfig,src=/home/user/.docker/config.json \
-        --output type=image,name=$${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG},push=true \
-        --progress=plain
-
-      echo "--- Image pushed successfully"
-      echo "$${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG}"
-
-  - label: "Upload Ready Pipeline"
-    depends_on: image-build
-    key: upload-ready-pipeline
-    if: >-
-      build.branch != "main" && (
-        build.pull_request.labels includes "npu-test" ||
-        build.pull_request.labels includes "ready"
-      )
-    agents:
-      queue: "ascend-a2b3"
-      resource_class: "npu-1"
-    command: |
-      buildkite-agent pipeline upload .buildkite/test-npu-ready.yaml
-
-  # L4 Test — main+NIGHTLY=1 (scheduled), or PR with specific label (e.g. add label then Rebuild)
-  - label: "Upload Nightly Pipeline"
-    depends_on: image-build
-    key: upload-nightly-pipeline
-    if: >-
-      (build.branch == "main" && build.env("NIGHTLY") == "1") ||
-      (build.branch != "main" && (
-        build.pull_request.labels includes "nightly-test" ||
-        build.pull_request.labels includes "omni-test" ||
-        build.pull_request.labels includes "tts-test" ||
-        build.pull_request.labels includes "diffusion-x2iat-test" ||
-        build.pull_request.labels includes "diffusion-x2v-test"
-      ))
-    commands:
-      - buildkite-agent pipeline upload .buildkite/test-npu-nightly.yaml
-    agents:
-      queue: "ascend-a2b3"
-      resource_class: "npu-1"

.buildkite/pipeline.yml

@@ -1,22 +1,6 @@
-# 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 ."
@@ -29,7 +13,7 @@ steps:
   - label: "Upload Ready Pipeline"
     depends_on: image-build
     key: upload-ready-pipeline
-    if: __UPLOAD_READY_IF__
+    if: build.branch != "main"
     commands:
       - buildkite-agent pipeline upload .buildkite/test-ready.yml
     agents:
@@ -39,38 +23,18 @@ steps:
   - label: "Upload Merge Pipeline"
     depends_on: image-build
     key: upload-merge-pipeline
-    if: __UPLOAD_MERGE_IF__
+    if: build.branch == "main"
     commands:
       - buildkite-agent pipeline upload .buildkite/test-merge.yml
     agents:
       queue: "cpu_queue_premerge"
 
-  # L4 Test — main+NIGHTLY=1 (scheduled), or PR with specific label (e.g. add label then Rebuild)
+  # L4 Test
   - 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"
-      ))
+    if: build.env("NIGHTLY") == "1"
     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/rebase-pipeline.yaml

@@ -1,38 +0,0 @@
-steps:
-  - label: ":docker: Build image"
-    key: image-build
-    commands:
-      - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-      - "docker build --progress=plain --file docker/Dockerfile.ci -t vllm-omni-ci ."
-      - "docker tag vllm-omni-ci public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
-      - "docker push public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
-    agents:
-      queue: "cpu_queue_premerge"
-
-  # L2 Test
-  - label: "Upload Ready Pipeline"
-    depends_on: image-build
-    key: upload-ready-pipeline
-    commands:
-      - buildkite-agent pipeline upload .buildkite/test-ready.yml
-    agents:
-      queue: "cpu_queue_premerge"
-
-  # L3 Test
-  - label: "Upload Merge Pipeline"
-    depends_on: image-build
-    key: upload-merge-pipeline
-    commands:
-      - buildkite-agent pipeline upload .buildkite/test-merge.yml
-    agents:
-      queue: "cpu_queue_premerge"
-
-  # L4 Test
-  - label: "Upload Nightly Pipeline"
-    depends_on: image-build
-    key: upload-nightly-pipeline
-    if: build.env("NIGHTLY") == "1"
-    commands:
-      - buildkite-agent pipeline upload .buildkite/test-nightly.yml
-    agents:
-      queue: "cpu_queue_premerge"

.buildkite/release-pipeline.yaml

@@ -1,110 +1,38 @@
 steps:
-  - input: "Provide Release version here"
-    id: input-release-version
-    fields:
-      - text: "What is the release version?"
-        key: release-version
-
-  - label: "Build and upload wheel"
-    depends_on: ~
-    key: "build-wheel"
-    agents:
-      queue: cpu_queue_release
+  - label: ":docker: Build image"
+    key: image-build
     commands:
-      - "curl -LsSf https://astral.sh/uv/install.sh | sh"
-      - 'export PATH="$HOME/.local/bin:$PATH"'
-      - "uv venv --python=3.12 && source .venv/bin/activate"
-      - "uv pip install --upgrade build"
-      - "python3 -m build"
-      - "bash .buildkite/scripts/upload-nightly-wheels.sh"
-
-  - label: "Generate and upload wheel indices"
-    depends_on: "build-wheel"
-    allow_dependency_failure: true
+      - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
+      - "docker build --progress=plain --file docker/Dockerfile.ci -t vllm-omni-ci ."
+      - "docker tag vllm-omni-ci public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
+      - "docker push public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
     agents:
-      queue: small_cpu_queue_release
+      queue: "cpu_queue_premerge"
+
+  # L2 Test
+  - label: "Upload Ready Pipeline"
+    depends_on: image-build
+    key: upload-ready-pipeline
     commands:
-      - "bash .buildkite/scripts/generate-and-upload-nightly-index.sh"
-
-  - block: "Upload release wheels to PyPI"
-    depends_on: "build-wheel"
-    key: block-upload-release-wheels
-
-  - label: "Upload release wheels to PyPI"
-    depends_on:
-      - block-upload-release-wheels
-      - input-release-version
-    key: upload-release-wheels
+      - buildkite-agent pipeline upload .buildkite/test-ready.yml
     agents:
-      queue: small_cpu_queue_release
+      queue: "cpu_queue_premerge"
+
+  # L3 Test
+  - label: "Upload Merge Pipeline"
+    depends_on: image-build
+    key: upload-merge-pipeline
     commands:
-      - "bash .buildkite/scripts/publish-release-wheels-pypi.sh"
-
-  - block: "Unblock to build release Docker images"
-    depends_on: ~
-    key: block-build-release-images
-
-  - group: "Build release Docker images"
-    key: "build-release-images"
-    depends_on: block-build-release-images
-    allow_dependency_failure: true
-    steps:
-      - label: "Build release image - x86_64"
-        key: build-release-image-x86
-        depends_on: ~
-        agents:
-          queue: cpu_queue_release
-        commands:
-          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-          - |
-            DOCKER_BUILDKIT=1 docker build \
-              --build-arg max_jobs=16 \
-              --tag public.ecr.aws/q9t5s3a7/vllm-omni-release-repo:$BUILDKITE_COMMIT-$(uname -m) \
-              --progress plain \
-              -f docker/Dockerfile.cuda .
-          - "docker push public.ecr.aws/q9t5s3a7/vllm-omni-release-repo:$BUILDKITE_COMMIT-$(uname -m)"
-          # re-tag to default image tag and push, just in case arm64 build fails
-          - "docker tag public.ecr.aws/q9t5s3a7/vllm-omni-release-repo:$BUILDKITE_COMMIT-$(uname -m) public.ecr.aws/q9t5s3a7/vllm-omni-release-repo:$BUILDKITE_COMMIT"
-          - "docker push public.ecr.aws/q9t5s3a7/vllm-omni-release-repo:$BUILDKITE_COMMIT"
-        env:
-          DOCKER_BUILDKIT: "1"
-
-      - label: "Build release image - aarch64"
-        key: build-release-image-arm64
-        depends_on: ~
-        agents:
-          queue: arm64_cpu_queue_release
-        commands:
-          - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-          - |
-            DOCKER_BUILDKIT=1 docker build \
-              --build-arg max_jobs=16 \
-              --tag public.ecr.aws/q9t5s3a7/vllm-omni-release-repo:$BUILDKITE_COMMIT-$(uname -m) \
-              --progress plain \
-              -f docker/Dockerfile.cuda .
-          - "docker push public.ecr.aws/q9t5s3a7/vllm-omni-release-repo:$BUILDKITE_COMMIT-$(uname -m)"
-        env:
-          DOCKER_BUILDKIT: "1"
-
-  - block: "Publish release images to DockerHub"
-    key: block-publish-release-images
-    depends_on:
-      - build-release-image-x86
-      - build-release-image-arm64
-      - input-release-version
-
-  - label: "Publish release images to DockerHub"
-    depends_on:
-      - block-publish-release-images
-    key: publish-release-images-dockerhub
+      - buildkite-agent pipeline upload .buildkite/test-merge.yml
     agents:
-      queue: small_cpu_queue_release
+      queue: "cpu_queue_premerge"
+
+  # L4 Test
+  - label: "Upload Nightly Pipeline"
+    depends_on: image-build
+    key: upload-nightly-pipeline
+    if: build.env("NIGHTLY") == "1"
     commands:
-      - "bash .buildkite/scripts/publish-release-images.sh"
-    plugins:
-      - docker-login#v3.0.0:
-          username: vllmbot
-          password-env: DOCKERHUB_TOKEN
-    env:
-      DOCKER_BUILDKIT: "1"
-      DOCKERHUB_USERNAME: "vllmbot"
+      - buildkite-agent pipeline upload .buildkite/test-nightly.yml
+    agents:
+      queue: "cpu_queue_premerge"

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

@@ -1,88 +0,0 @@
-#!/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

@@ -1,192 +0,0 @@
-#!/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,6 +11,15 @@ 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
 
@@ -42,24 +51,28 @@ 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-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-omni:${BUILDKITE_COMMIT}"
+image_name="public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:${BUILDKITE_COMMIT}-rocm-omni"
 container_name="rocm_${BUILDKITE_COMMIT}_$(tr -dc A-Za-z0-9 < /dev/urandom | head -c 10; echo)"
 
-# TODO: @tjtanaa uncomment this once the amd agents are configured to use ECR Public Gallery.
-# # Install AWS CLI to authenticate to ECR Public Gallery to get higher rate limit for pulling images
-# sudo apt-get update && sudo apt-get install -y awscli
-##  Use safe docker login helper to prevent race conditions
-# source "$(dirname "${BASH_SOURCE[0]}")/../docker_login_ecr_public.sh"
-# safe_docker_login_ecr_public
-
-## Pull the container from AMD Docker Hub
+# Install AWS CLI to authenticate to ECR Public Gallery to get higher rate limit for pulling images
+sudo apt-get update && sudo apt-get install -y awscli
+# Use safe docker login helper to prevent race conditions
+source "$(dirname "${BASH_SOURCE[0]}")/../docker_login_ecr_public.sh"
+safe_docker_login_ecr_public
+# Pull the container from ECR Public Gallery
 
 docker pull "${image_name}"
 
@@ -74,11 +87,7 @@ HF_CACHE="$(realpath ~)/huggingface"
 mkdir -p "${HF_CACHE}"
 HF_MOUNT="/root/.cache/huggingface"
 
-if [[ -n "${TEST_COMMAND:-}" ]]; then
-    commands="$TEST_COMMAND"
-else
-    commands="$@"
-fi
+commands=$@
 echo "Commands:$commands"
 
 PARALLEL_JOB_COUNT=8
@@ -93,7 +102,6 @@ if [[ -z "$render_gid" ]]; then
 fi
 
 # check if the command contains shard flag, we will run all shards in parallel because the host have 8 GPUs.
-# TODO: @tjtanaa reenable to run VLLM_ROCM_USE_AITER=1 when AITER is shipped with prebuilt kernels.
 if [[ $commands == *"--shard-id="* ]]; then
   # assign job count as the number of shards used
   commands=$(echo "$commands" | sed -E "s/--num-shards[[:blank:]]*=[[:blank:]]*[0-9]*/--num-shards=${PARALLEL_JOB_COUNT} /g" | sed 's/ \\ / /g')
@@ -110,9 +118,11 @@ if [[ $commands == *"--shard-id="* ]]; then
         --rm \
         -e MIOPEN_DEBUG_CONV_DIRECT=0 \
         -e MIOPEN_DEBUG_CONV_GEMM=0 \
-        -e VLLM_ROCM_USE_AITER=0 \
+        -e VLLM_ROCM_USE_AITER=1 \
         -e HIP_VISIBLE_DEVICES="${GPU}" \
         -e HF_TOKEN \
+        -e AWS_ACCESS_KEY_ID \
+        -e AWS_SECRET_ACCESS_KEY \
         -v "${HF_CACHE}:${HF_MOUNT}" \
         -e "HF_HOME=${HF_MOUNT}" \
         -e "PYTHONPATH=${MYPYTHONPATH}" \
@@ -143,8 +153,10 @@ else
           --rm \
           -e MIOPEN_DEBUG_CONV_DIRECT=0 \
           -e MIOPEN_DEBUG_CONV_GEMM=0 \
-          -e VLLM_ROCM_USE_AITER=0 \
+          -e VLLM_ROCM_USE_AITER=1 \
           -e HF_TOKEN \
+          -e AWS_ACCESS_KEY_ID \
+          -e AWS_SECRET_ACCESS_KEY \
           -v "${HF_CACHE}:${HF_MOUNT}" \
           -e "HF_HOME=${HF_MOUNT}" \
           -e "PYTHONPATH=${MYPYTHONPATH}" \

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

@@ -1,58 +0,0 @@
-#!/bin/bash
-
-# This script build the XPU docker image and run the offline inference inside the container.
-set -ex
-
-omni_source_dir=$(git rev-parse --show-toplevel)
-
-base_image_name="xpu/vllm-omni-ci-base:${VLLM_VERSION:?VLLM_VERSION must be set}"
-image_name="xpu/vllm-omni-ci:${BUILDKITE_COMMIT:?BUILDKITE_COMMIT must be set}"
-container_name="xpu_${BUILDKITE_COMMIT}_$(
-    tr -dc A-Za-z0-9 </dev/urandom | head -c 10
-    echo
-)"
-
-cd "${omni_source_dir}"
-if [ -z "$(docker images -q "${base_image_name}")" ]; then
-    docker build --target vllm-base -t "${base_image_name}" --build-arg "VLLM_VERSION=${VLLM_VERSION}" -f docker/Dockerfile.xpu .
-fi
-
-# Try building the docker image
-docker build --build-arg "VLLM_BASE=${base_image_name}" --build-arg "VLLM_VERSION=${VLLM_VERSION}" -t "${image_name}" -f docker/Dockerfile.xpu .
-
-# Setup cleanup
-remove_docker_container() {
-    docker rm -f "${container_name}" || true
-    docker image rm -f "${image_name}" || true
-    docker system prune -f || true
-}
-trap remove_docker_container EXIT
-
-HF_CACHE="${HF_CACHE:-$(realpath ~)/.cache/huggingface}"
-mkdir -p "${HF_CACHE}"
-HF_MOUNT="/root/.cache/huggingface"
-
-time timeout -k 30 30m docker run \
-    --device /dev/dri:/dev/dri \
-    --net=host \
-    --ipc=host \
-    -v /dev/dri/by-path:/dev/dri/by-path \
-    -v "${HF_CACHE}:${HF_MOUNT}" \
-    --security-opt seccomp=unconfined \
-    --entrypoint="" \
-    -e VLLM_LOGGING_LEVEL \
-    -e VLLM_OMNI_LOGGING_LEVEL \
-    -e HF_TOKEN \
-    -e ZE_AFFINITY_MASK \
-    --name "${container_name}" \
-    "${image_name}" \
-    bash -c '
-    set -e
-    echo $ZE_AFFINITY_MASK
-    pip install tblib==3.1.0
-    cd /workspace/vllm-omni
-    pytest -v -s -m "core_model and xpu and B60"
-    pytest -v -s tests/diffusion/quantization/test_mxfp8_config.py
-    export VLLM_XPU_USE_SAMPLER_KERNEL=0    # NOTE: Remove this after vLLM v0.21.1 is merged. Fixes Qwen2-5 omni-expansion tests.
-    pytest -v -s -m "advanced_model and xpu and B60"
-'

.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_expansion.py
+    VLLM_USE_MODELSCOPE=True pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
 '

.buildkite/scripts/publish-release-images.sh

@@ -1,50 +0,0 @@
-#!/bin/bash
-# SPDX-License-Identifier: Apache-2.0
-#
-# Publish release Docker images from ECR to DockerHub.
-# Pulls per-arch images, tags with `latest` and versioned tags, pushes them,
-# then creates and pushes a multi-arch manifest.
-
-set -euo pipefail
-
-RELEASE_VERSION=$(buildkite-agent meta-data get release-version --default "" | sed 's/^v//')
-if [ -z "${RELEASE_VERSION}" ]; then
-  echo "ERROR: release-version metadata not set"
-  exit 1
-fi
-
-COMMIT="$BUILDKITE_COMMIT"
-DOCKERHUB_REPO="vllm/vllm-omni"
-ECR_REPO="public.ecr.aws/q9t5s3a7/vllm-omni-release-repo"
-
-echo "========================================"
-echo "Publishing release images v${RELEASE_VERSION}"
-echo "  Commit: ${COMMIT}"
-echo "========================================"
-
-# Login to ECR to pull staging images
-aws ecr-public get-login-password --region us-east-1 | \
-  docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
-
-docker pull "${ECR_REPO}:${COMMIT}-x86_64"
-docker pull "${ECR_REPO}:${COMMIT}-aarch64"
-
-docker tag "${ECR_REPO}:${COMMIT}-x86_64" "${DOCKERHUB_REPO}:latest-x86_64"
-docker tag "${ECR_REPO}:${COMMIT}-x86_64" "${DOCKERHUB_REPO}:v${RELEASE_VERSION}-x86_64"
-docker push "${DOCKERHUB_REPO}:latest-x86_64"
-docker push "${DOCKERHUB_REPO}:v${RELEASE_VERSION}-x86_64"
-
-docker tag "${ECR_REPO}:${COMMIT}-aarch64" "${DOCKERHUB_REPO}:latest-aarch64"
-docker tag "${ECR_REPO}:${COMMIT}-aarch64" "${DOCKERHUB_REPO}:v${RELEASE_VERSION}-aarch64"
-docker push "${DOCKERHUB_REPO}:latest-aarch64"
-docker push "${DOCKERHUB_REPO}:v${RELEASE_VERSION}-aarch64"
-
-docker manifest rm "${DOCKERHUB_REPO}:latest" || true
-docker manifest rm "${DOCKERHUB_REPO}:v${RELEASE_VERSION}" || true
-docker manifest create "${DOCKERHUB_REPO}:latest" "${DOCKERHUB_REPO}:latest-x86_64" "${DOCKERHUB_REPO}:latest-aarch64"
-docker manifest create "${DOCKERHUB_REPO}:v${RELEASE_VERSION}" "${DOCKERHUB_REPO}:v${RELEASE_VERSION}-x86_64" "${DOCKERHUB_REPO}:v${RELEASE_VERSION}-aarch64"
-docker manifest push "${DOCKERHUB_REPO}:latest"
-docker manifest push "${DOCKERHUB_REPO}:v${RELEASE_VERSION}"
-
-echo ""
-echo "Successfully published release images for v${RELEASE_VERSION}"

.buildkite/scripts/publish-release-wheels-pypi.sh

@@ -1,52 +0,0 @@
-#!/usr/bin/env bash
-# SPDX-License-Identifier: Apache-2.0
-
-set -e
-
-BUCKET="vllm-wheels"
-SUBPATH="omni/$BUILDKITE_COMMIT"
-S3_COMMIT_PREFIX="s3://$BUCKET/$SUBPATH/"
-
-RELEASE_VERSION=$(buildkite-agent meta-data get release-version | sed 's/^v//')
-if [[ -z "$RELEASE_VERSION" ]]; then
-  echo "[FATAL] release-version metadata not set."
-  exit 1
-fi
-echo "Release version: $RELEASE_VERSION"
-
-if [[ -z "$PYPI_TOKEN" ]]; then
-  echo "[FATAL] PYPI_TOKEN is not set."
-  exit 1
-else
-  export TWINE_USERNAME="__token__"
-  export TWINE_PASSWORD="$PYPI_TOKEN"
-fi
-
-set -x
-
-if ! command -v uv &> /dev/null; then
-  curl -LsSf https://astral.sh/uv/install.sh | sh
-  export PATH="$HOME/.local/bin:$PATH"
-fi
-
-uv venv --python 3.12 /tmp/vllm-omni-release-env
-source /tmp/vllm-omni-release-env/bin/activate
-uv pip install twine
-
-DIST_DIR=/tmp/vllm-omni-release-dist
-mkdir -p "$DIST_DIR"
-
-echo "Downloading wheels from S3:"
-aws s3 ls "$S3_COMMIT_PREFIX"
-aws s3 cp --recursive --exclude "*" --include "vllm_omni-${RELEASE_VERSION}*.whl" --exclude "*dev*" "$S3_COMMIT_PREFIX" "$DIST_DIR"
-ls -la "$DIST_DIR"
-
-PYPI_WHEEL_FILES=$(find "$DIST_DIR" -name "vllm_omni-${RELEASE_VERSION}*.whl" -not -name "*+*")
-if [[ -z "$PYPI_WHEEL_FILES" ]]; then
-  echo "[FATAL] No wheels found for version ${RELEASE_VERSION}"
-  exit 1
-fi
-
-python3 -m twine check $PYPI_WHEEL_FILES
-python3 -m twine upload --non-interactive --verbose $PYPI_WHEEL_FILES
-echo "Wheels uploaded to PyPI"

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

@@ -1,33 +0,0 @@
-#!/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

@@ -1,146 +0,0 @@
-#!/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

@@ -1,222 +0,0 @@
-steps:
-
-- label: "Simple · Diffusion & Model Executor Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  timeout_in_minutes: 50
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - export VLLM_ROCM_USE_AITER=0
-    # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
-    # that is not supported on CPU
-    - "timeout 40m pytest -sv tests/diffusion tests/model_executor -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: "Simple · Other Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  timeout_in_minutes: 50
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - export VLLM_ROCM_USE_AITER=0
-    # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
-    # that is not supported on CPU
-    - "timeout 40m pytest -sv tests/diffusion tests/model_executor -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
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - timeout 20m pytest -s -v tests/diffusion/lora/test_images_generations_lora.py
-
-- label: "Diffusion Model CPU offloading Test"
-  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
-    - |
-      timeout 20m bash -c '
-        set +e
-        pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py
-        EXIT1=\$?
-        pytest -s -v tests/diffusion/offloader/test_diffusion_layerwise_offload.py
-        EXIT2=\$?
-        exit \$((EXIT1 | EXIT2))
-      '
-
-## ISSUE depends on `diffusers` package: https://github.com/huggingface/diffusers/issues/13274
-# - label: "Audio Generation 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
-#     - 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 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/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"
-  agent_pool: mi325_2
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - pytest -s -v tests/diffusion/distributed/test_zimage_parallelism.py
-
-- label: "Diffusion Model Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - timeout 30m pytest -s -v tests/e2e/offline_inference/test_z_image.py
-
-- label: "Diffusion GPU Worker Test"
-  agent_pool: mi325_2
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py
-
-- 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 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_4
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  commands:
-    - 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
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py -m 'advanced_model and cuda' --run-level 'advanced_model'
-
-- label: "Qwen3-TTS CustomVoice E2E Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - |
-      timeout 20m bash -c '
-        export VLLM_LOGGING_LEVEL=DEBUG
-        export VLLM_WORKER_MULTIPROC_METHOD=spawn
-        export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
-        pytest -s -v tests/e2e/online_serving/test_qwen3_tts_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_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
-  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 tests/e2e/offline_inference/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
-      '

.buildkite/test-amd-ready.yaml

@@ -1,238 +0,0 @@
-steps:
-
-- label: "Simple · Diffusion & Model Executor Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - export VLLM_ROCM_USE_AITER=0
-    # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
-    # that is not supported on CPU
-    - "timeout 40m pytest -sv tests/diffusion tests/model_executor -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: "Simple · Other Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - export VLLM_ROCM_USE_AITER=0
-    # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
-    # that is not supported on CPU
-    - "timeout 40m pytest -sv tests/diffusion tests/model_executor -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: "Voxtral TTS CUDA Unit Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - timeout 10m pytest -s -v tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py
-
-- label: "Diffusion · Batch Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - timeout 20m pytest -sv tests/diffusion/batching -m 'core_model and cuda' --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
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - export VLLM_LOGGING_LEVEL=DEBUG
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - |
-      timeout 20m bash -c '
-        set +e
-        pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py
-        EXIT1=\$?
-        pytest -s -v tests/diffusion/offloader/test_diffusion_layerwise_offload.py
-        EXIT2=\$?
-        exit \$((EXIT1 | EXIT2))
-      '
-
-## ISSUE depends on `diffusers` package: https://github.com/huggingface/diffusers/issues/13274
-# - label: "Audio Generation 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
-#     - 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
-  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 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
-  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/diffusion/distributed/test_sequence_parallel.py -m core_model
-
-- label: "Diffusion GPU Worker Test"
-  agent_pool: mi325_2
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py
-
-- label: "Diffusion · Model Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - timeout 15m pytest -sv tests/diffusion/models/ -m 'core_model and cuda' --run-level "core_model"
-
-- label: "Diffusion · Wan22 Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - |
-      timeout 40m bash -c "
-        export VLLM_TEST_CLEAN_GPU_MEMORY=1
-        export VLLM_IMAGE_FETCH_TIMEOUT=60
-        pytest -s -v tests/e2e/online_serving/test_wan22_t2v.py -m 'core_model' --run-level 'core_model'
-      "
-
-- label: "Engine Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - |
-      timeout 15m bash -c '
-              pytest -s -v tests/engine/test_async_omni_engine_abort.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: "MiMo-Audio E2E Test with H100"
-  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
-        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_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - export VLLM_LOGGING_LEVEL=DEBUG
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
-    - timeout 30m pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m "core_model" --run-level "core_model"
-
-- label: "Voxtral-TTS E2E Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - |
-      timeout 20m bash -c '
-        export VLLM_LOGGING_LEVEL=DEBUG
-        export VLLM_WORKER_MULTIPROC_METHOD=spawn
-        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"
-      '
-
-# 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:
-    - |
-      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

@@ -1,17 +1,17 @@
 steps:
 
 - label: "Diffusion Model Test"
-  timeout_in_minutes: 30
+  timeout_in_minutes: 20
   agent_pool: mi325_2
   depends_on: amd-build
   mirror_hardwares: [amdproduction]
   grade: Blocking
   commands:
     - export GPU_ARCHS=gfx942
-    - pytest -s -v tests/e2e/offline_inference/test_z_image.py
+    - pytest -s -v tests/e2e/offline_inference/test_t2i_model.py
 
 - label: "Diffusion Images API LoRA E2E"
-  timeout_in_minutes: 30
+  timeout_in_minutes: 20
   agent_pool: mi325_1
   depends_on: amd-build
   mirror_hardwares: [amdproduction]
@@ -20,7 +20,7 @@ steps:
     - export GPU_ARCHS=gfx942
     - export VLLM_LOGGING_LEVEL=DEBUG
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - pytest -s -v tests/diffusion/lora/test_images_generations_lora.py
+    - pytest -s -v tests/e2e/online_serving/test_images_generations_lora.py
 
 - label: "Diffusion Model CPU offloading Test"
   timeout_in_minutes: 20
@@ -32,8 +32,8 @@ steps:
     - export GPU_ARCHS=gfx942
     - export VLLM_LOGGING_LEVEL=DEBUG
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py
-    - pytest -s -v tests/diffusion/offloader/test_diffusion_layerwise_offload.py
+    - pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
+    - pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
 
 - label: "Diffusion Cache Backend Test"
   timeout_in_minutes: 15
@@ -57,7 +57,7 @@ steps:
     - export GPU_ARCHS=gfx942
     - export VLLM_LOGGING_LEVEL=DEBUG
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - pytest -s -v tests/diffusion/distributed/test_sequence_parallel.py -m core_model
+    - pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py -m core_model
 
 - label: "Diffusion Tensor Parallelism Test"
   timeout_in_minutes: 20
@@ -67,7 +67,7 @@ steps:
     - export GPU_ARCHS=gfx942
     - export VLLM_LOGGING_LEVEL=DEBUG
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - pytest -s -v tests/diffusion/distributed/test_zimage_parallelism.py
+    - pytest -s -v tests/e2e/offline_inference/test_zimage_parallelism.py
 
 - label: "Diffusion GPU Worker Test"
   timeout_in_minutes: 20
@@ -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_expansion.py
+    - pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
     - pytest -s -v tests/engine/test_async_omni_engine_abort.py
 
 - label: "Omni Model Test Qwen3-Omni"
@@ -100,14 +100,14 @@ steps:
   commands:
     - export VLLM_LOGGING_LEVEL=DEBUG
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - export VLLM_ROCM_USE_AITER=0
+    - export VLLM_TEST_CLEAN_GPU_MEMORY="1"
     - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
     - pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py
 
 
-- label: "Omni Sleep Mode Test"
-  timeout_in_minutes: 40
-  agent_pool: mi325_2
+- label: "Diffusion Image Edit Test"
+  timeout_in_minutes: 15
+  agent_pool: mi325_1
   depends_on: amd-build
   mirror_hardwares: [amdproduction]
   grade: Blocking
@@ -115,4 +115,4 @@ steps:
     - export GPU_ARCHS=gfx942
     - export VLLM_LOGGING_LEVEL=DEBUG
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - pytest -s -v tests/entrypoints/test_omni_sleep_mode.py -m "advanced_model and omni and MI325" --run-level "advanced_model"
+    - pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py

.buildkite/test-npu-nightly.yaml

@@ -1,24 +0,0 @@
-steps:
-  - group: ":card_index_dividers: Diffusion X2V Model Test"
-    key: nightly-diffusion-x2v-group
-    depends_on: upload-nightly-pipeline
-    if: >-
-      build.env("NIGHTLY") == "1" ||
-      build.pull_request.labels includes "nightly-test" ||
-      build.pull_request.labels includes "diffusion-x2v-test"
-    steps:
-      - label: ":full_moon: Diffusion X2V · Function Test"
-        agents:
-          queue: "ascend-a2b3"
-          resource_class: "npu-4"
-        image: "swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/vllm-omni-ci-npu:${BUILDKITE_COMMIT}"
-        plugins:
-          - kubernetes:
-              podSpecPatch:
-                imagePullSecrets:
-                  - name: swr-secret
-        env:
-          VLLM_WORKER_MULTIPROC_METHOD: spawn
-          HF_TOKEN: "${HF_TOKEN}"
-        commands:
-          - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k npu_i2v --run-level "full_model"

.buildkite/test-npu-ready.yaml

@@ -1,21 +0,0 @@
-steps:
-  - group: ":card_index_dividers: TTS Model Test"
-    key: ready-tts-group
-    depends_on: upload-ready-pipeline
-    steps:
-      - label: ":sound: TTS · Qwen3-TTS CustomVoice Test"
-        agents:
-          queue: "ascend-a2b3"
-          resource_class: "npu-1"
-        image: "swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/vllm-omni-ci-npu:${BUILDKITE_COMMIT}"
-        plugins:
-          - kubernetes:
-              podSpecPatch:
-                imagePullSecrets:
-                  - name: swr-secret
-        env:
-          VLLM_WORKER_MULTIPROC_METHOD: spawn
-          VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
-          HF_TOKEN: "${HF_TOKEN}"
-        commands:
-          - pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m 'core_model' --run-level 'core_model'

.buildkite/test-ready.yml

@@ -1,49 +1,8 @@
-env:
-  VLLM_WORKER_MULTIPROC_METHOD: spawn
-  VLLM_USE_FLASHINFER_MOE_FP16: 0
-  HF_HUB_DOWNLOAD_TIMEOUT: 300
-  HF_HUB_ETAG_TIMEOUT: 60
-
 steps:
-  - group: ":card_index_dividers: Simple Test"
-    depends_on: upload-ready-pipeline
-    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: "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"
+  - label: "Simple Unit Test"
     depends_on: upload-ready-pipeline
     commands:
-      - timeout 20m pytest -s -v tests/e2e/offline_inference/custom_pipeline/ -m "core_model"
+      - "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:
@@ -56,116 +15,101 @@ steps:
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
-  - group: ":card_index_dividers: Diffusion Test"
+  - label: "Diffusion Model Test"
     depends_on: upload-ready-pipeline
-    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"
+    commands:
+      - timeout 20m 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"
+          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 · 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 · 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 · 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"
+  - label: "Diffusion Model CPU offloading Test"
     depends_on: upload-ready-pipeline
     commands:
       - |
-        timeout 15m bash -c "
-                  pytest -sv tests/engine/ tests/model_executor/ -m 'core_model and cuda' --run-level 'core_model'
-        "
+        timeout 20m bash -c '
+          set +e
+          pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
+          EXIT1=$$?
+          pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
+          EXIT2=$$?
+          exit $$((EXIT1 | EXIT2))
+        '
     agents:
-      queue: "gpu_1_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
+      queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Audio Generation Model Test"
+    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"
+          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"
+          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"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Diffusion GPU Worker Test"
+    depends_on: upload-ready-pipeline
+    commands:
+      - timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py
+    agents:
+      queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
     plugins:
       - docker#v5.2.0:
           image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
@@ -174,450 +118,190 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
-            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
 
-  - label: "Distributed Test"
+  - label: "Benchmark & Engine Test"
     depends_on: upload-ready-pipeline
     commands:
       - |
-        timeout 20m bash -c "
-                  pytest -sv tests/distributed/ -m 'core_model and cuda' --run-level 'core_model'
-        "
+        timeout 15m bash -c '
+                export VLLM_WORKER_MULTIPROC_METHOD=spawn
+                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))
+        '
     agents:
-      queue: "mithril-h100-pool"
+      queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
     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
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          shm-size: "8gb"
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
 
-  - group: ":card_index_dividers: E2E Test"
+
+  - label: "Omni Model Test"
     depends_on: upload-ready-pipeline
-    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"
+    commands:
+      - |
+        timeout 17m bash -c '
+          export VLLM_LOGGING_LEVEL=DEBUG
+          export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
+        '
+    agents:
+      queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
 
-      - label: "Omni · 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
-                    emptyDir:
-                      medium: Memory
-                  - name: hf-cache
-                    hostPath:
-                      path: /mnt/hf-cache
-                      type: DirectoryOrCreate
+  # - label: "Omni Model Test with H100"
+  #   depends_on: upload-ready-pipeline
+  #   commands:
+  #     - |
+  #       timeout 20m bash -c '
+  #         export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  #         export VLLM_TEST_CLEAN_GPU_MEMORY="1"
+  #         # - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
+  #         pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "core_model" --run-level "core_model"
+  #       '
+  #   agents:
+  #     queue: "mithril-h100-pool"
+  #   plugins:
+  #     - kubernetes:
+  #         podSpec:
+  #           containers:
+  #             - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+  #               resources:
+  #                 limits:
+  #                   nvidia.com/gpu: 2
+  #               volumeMounts:
+  #                 - name: devshm
+  #                   mountPath: /dev/shm
+  #                 - name: hf-cache
+  #                   mountPath: /root/.cache/huggingface
+  #               env:
+  #                 - name: HF_HOME
+  #                   value: /root/.cache/huggingface
+  #           nodeSelector:
+  #             node.kubernetes.io/instance-type: gpu-h100-sxm
+  #           volumes:
+  #             - name: devshm
+  #               emptyDir:
+  #                 medium: Memory
+  #             - name: hf-cache
+  #               hostPath:
+  #                 path: /mnt/hf-cache
+  #                 type: DirectoryOrCreate
 
-      - label: "Omni · Qwen3-Omni Realtime WebSocket"
-        commands:
-          - |
-            timeout 30m bash -c "
-              export VLLM_WORKER_MULTIPROC_METHOD=spawn
-              pytest -s -v tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py::TestQwen3OmniRealtimeWebSocket::test_streaming_audio_input_pcm_output_async_chunk_smoke -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: "Qwen3-TTS E2E Test"
+    depends_on: upload-ready-pipeline
+    commands:
+      - |
+        timeout 20m bash -c '
+          huggingface-cli download Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice
+          export VLLM_LOGGING_LEVEL=DEBUG
+          export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py -k "not NoAsyncChunk"
+        '
+    agents:
+      queue: "gpu_4_queue"
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          shm-size: "8gb"
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
 
-      - label: "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: "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: "TTS · VoxCPM2 Online Test"
-        timeout_in_minutes: 30
-        commands:
-          - |
-            timeout 30m bash -c "
-              export VLLM_LOGGING_LEVEL=DEBUG
-              export VLLM_WORKER_MULTIPROC_METHOD=spawn
-              pytest -s -v tests/e2e/online_serving/test_voxcpm2_tts.py -m 'core_model and tts' --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_4_queue"
-        plugins:
-          - docker#v5.2.0:
-              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-              always-pull: true
-              propagate-environment: true
-              shm-size: "8gb"
-              environment:
-                - "HF_HOME=/fsx/hf_cache"
-                - "HF_TOKEN"
-              volumes:
-                - "/fsx/hf_cache:/fsx/hf_cache"
-
-      - label: "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_4_queue"
-        plugins:
-          - docker#v5.2.0:
-              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-              always-pull: true
-              propagate-environment: true
-              shm-size: "8gb"
-              environment:
-                - "HF_HOME=/fsx/hf_cache"
-                - "HF_TOKEN"
-              volumes:
-                - "/fsx/hf_cache:/fsx/hf_cache"
-
-      - label: "TTS · Higgs-Audio-V3 Test"
-        commands:
-          - |
-            timeout 20m bash -c "
-              export VLLM_WORKER_MULTIPROC_METHOD=spawn
-              export VLLM_USE_DEEP_GEMM=0
-              export VLLM_MOE_USE_DEEP_GEMM=0
-              pytest -s -v tests/e2e/online_serving/test_higgs_audio_v3.py -m 'core_model' --run-level 'core_model'
-            "
-        agents:
-          queue: "gpu_4_queue"
-        plugins:
-          - docker#v5.2.0:
-              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-              always-pull: true
-              propagate-environment: true
-              shm-size: "8gb"
-              environment:
-                - "HF_HOME=/fsx/hf_cache"
-                - "HF_TOKEN"
-              volumes:
-                - "/fsx/hf_cache:/fsx/hf_cache"
-
-      - label: "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
-                    emptyDir:
-                      medium: Memory
-                  - name: hf-cache
-                    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
+  # - label: "Bagel Text2Img Model Test with H100"
+  #   depends_on: upload-ready-pipeline
+  #   commands:
+  #     - |
+  #       timeout 30m bash -c '
+  #         export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  #         pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py
+  #       '
+  #   agents:
+  #     queue: "mithril-h100-pool"
+  #   plugins:
+  #     - kubernetes:
+  #         podSpec:
+  #           containers:
+  #             - image: 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

.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-omni:$BUILDKITE_COMMIT" %}
+{% set docker_image_amd = "public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT-rocm-omni" %}
 {% set default_working_dir = "/app/vllm-omni" %}
 
   - group: "AMD Tests"
@@ -13,6 +13,7 @@
         depends_on: ~
         soft_fail: false
         commands:
+          - "source .buildkite/scripts/docker_login_ecr_public.sh && safe_docker_login_ecr_public"
           - "docker build -f docker/Dockerfile.rocm -t {{ docker_image_amd }} --target test --build-arg ARG_PYTORCH_ROCM_ARCH=gfx942 --progress plain ."
           - "docker push {{ docker_image_amd }}"
         key: "amd-build"
@@ -21,15 +22,13 @@
         retry:
           automatic:
             - exit_status: -1  # Agent was lost
-              limit: 2
+              limit: 1
             - exit_status: -10  # Agent was lost
-              limit: 2
-            - exit_status: 128  # Git connectivity issues
-              limit: 2
+              limit: 1
             - exit_status: 1  # Machine occasionally fail
               limit: 1
         agents:
-          queue: amd-cpu
+          queue: cpu_queue_premerge
 
     {% for step in steps %}
     {% if step.mirror_hardwares and mirror_hw in step.mirror_hardwares %}
@@ -41,21 +40,10 @@
           {% else %}
           queue: amd_mi325_1
           {% endif %}
-{% set cmd_body = (step.command or (step.commands | join("\n"))) | trim %}
-{% set indented_cmd = cmd_body | replace("\n", "\n            ") %}
-        command: bash .buildkite/scripts/hardware_ci/run-amd-test.sh
+        command: bash .buildkite/scripts/hardware_ci/run-amd-test.sh "(command rocm-smi || true) && cd {{ (step.working_dir or default_working_dir) | safe  }} ; {{ step.command  or (step.commands | join(" && ")) | safe }}"
         env:
           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.timeout_in_minutes %}
-        timeout_in_minutes: {{ step.timeout_in_minutes }}
-        {% endif %}
         {% if step.grade and step.grade == "Blocking" %}
         soft_fail: false
         {% else %}

.buildkite/test-weekly.yml

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

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

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

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

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

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

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

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

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

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

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

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

@@ -1,108 +0,0 @@
-# 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/diffusion-perf-opt/SKILL.md

@@ -1,510 +0,0 @@
----
-name: diffusion-perf-opt
-description: Diagnose and optimize vLLM Omni diffusion workloads, especially Wan/Qwen/Flux-style image and video generation. Use when Codex is asked to analyze profiling traces, choose parallel strategies, inspect torch profiler trace.json or trace.json.gz timelines, estimate optimization ROI, investigate GPU idle/free bubbles, compare USP/CFG/HSDP/VAE parallelism, or design operator/host/quantization optimizations for vLLM Omni.
----
-
-# vLLM Omni Performance Optimization
-
-Use this skill to run a disciplined optimization loop for vLLM Omni diffusion workloads. Keep two ideas separate: real performance baselines are collected with low overhead, while torch profiler traces are diagnostic artifacts and may distort latency.
-
-## First Questions
-
-Before proposing changes, ask for the optimization scene if it is not already known:
-
-- GPU model, card count, topology, and whether NVLink is present.
-- Model and pipeline, for example Wan2.2 I2V A14B.
-- User workload: resolution, frames, steps, batch/concurrency, CFG scales, prompt/image inputs.
-- Three runnable commands for each model/workload, or scripts that generate them:
-  - **Server startup command**: exact `vllm serve` command, environment variables, model path, port, parallelism flags, profiler flags, and precision/compile settings.
-  - **User request command**: exact single-request client command (for example curl or Python client) with fixed prompt/media/seed/size/steps. Use this to validate correctness and collect per-request stage timings.
-  - **User benchmark command**: exact repeatable benchmark command or script with warmup count, measured iteration count, concurrency/batch policy, output directory, and summary format.
-- Current enabled strategies: USP/SP, CFG parallel, HSDP/FSDP, VAE patch parallel, torch.compile, profiler options.
-- Optimization target: latency, throughput, memory, cost, or quality-preserving speed.
-- Precision/quality tolerance: bf16/fp8/quantization/sparsity/approximate attention allowed or not.
-- Quality validation method: exact-output comparison when feasible, or metrics
-  such as SSIM, PSNR, LPIPS, cosine similarity, MAE, MSE, temporal flicker
-  checks, baseline self-run variance, and reviewable output artifacts.
-
-## Workflow
-
-1. **Freeze the measurement protocol and commands.**
-   - Before analyzing parallel strategies or traces, establish the three commands for the target model/workload: server startup, single user request, and user benchmark.
-   - Prefer checked-in scripts or one-off shell scripts over free-form commands in chat. The scripts should make workload variables explicit, including model path, port/base URL, prompt/media inputs, resolution, frames, steps, seed, warmup, measured iterations, concurrency, and output path.
-   - Measurement reports must preserve the concrete process, not just the final numbers. For each tested configuration, write down the server startup command, the user request command, the user benchmark command or polling loop, the final server response metrics, and the client-side timing/HTTP result.
-   - Present all duration and latency values in **milliseconds (ms)** in tables and summaries. If an API returns seconds, convert to ms before reporting. Keep raw units only inside quoted raw logs or source snippets.
-   - Keep profiler modes separate in the command set:
-     - Baseline/benchmark commands should avoid torch profiler and stack collection.
-     - Diagnostic commands may enable torch profiler after a baseline identifies a bottleneck.
-   - If any of the three commands is missing, ask for it or create a proposed script before proceeding.
-   - Make clear which command is authoritative for correctness validation and which command is authoritative for performance numbers.
-
-2. **Collect a real baseline.**
-   - Disable PyTorch profiler and stack collection.
-   - Disable or fix torch.compile state so A/B is fair.
-   - Avoid `--enforce-eager` for production-speed baselines unless eager is the target.
-   - Prefer `--log-stats` and `--enable-diffusion-pipeline-profiler` for low-overhead stage timing. PR 3069 has the relevant metrics/log-stats changes; if local code does not include them, fetch or cherry-pick the minimal metrics changes rather than merging unrelated PR drift.
-   - Run warmup requests and exclude first-request lazy init.
-
-3. **Model the parallel strategy before testing.**
-   - Estimate compute and communication for self-attention, cross-attention, FFN, CFG, VAE encode/decode, and HSDP.
-   - Select a small candidate matrix rather than exhaustively testing everything.
-   - Typical 2-card candidates: `USP=2`, `CFG=2`, `USP=1/HSDP on-off`, VAE parallel on-off if memory allows.
-   - Typical 4-card candidates: `USP=4`, `CFG=2 x USP=2`, `USP=2 x HSDP`, VAE parallel on-off.
-   - Typical 8-card candidates for official CFG-enabled video diffusion models:
-     - Primary candidate: `CFG=2 x USP=4` with VAE patch parallel across all 8 ranks.
-     - Compare against: `CFG=1 x USP=8` to test whether larger sequence parallel groups beat CFG branch parallelism.
-     - Isolate VAE parallelism: keep the DiT strategy fixed and compare VAE patch parallel on/off or different VAE patch sizes.
-     - If long-video `diffuse` remains dominant and Ulysses all-to-all is suspected, test a hybrid sequence strategy such as `CFG=2 x USP=2 x Ring=2`.
-     - Test HSDP off only if the model fits without it; treat HSDP primarily as a memory strategy until A/B proves otherwise.
-   - Prefer USP/SP for long video token sequences; prefer CFG parallel when CFG doubles transformer forwards and sequence length is modest.
-   - Convert workload shape to latent/patch token counts before choosing candidates. For video models, record frames, latent frames, latent HxW, patch size, approximate token count, and which stages should scale with the token count.
-
-4. **Search for the best parallel configuration.**
-   - Start with a small matrix that answers one question per comparison:
-     - `CFG parallel vs larger USP`: compare `CFG=2 x USP=world/2` against `CFG=1 x USP=world` for CFG-enabled workloads.
-     - `VAE patch parallel value`: compare the best DiT strategy with VAE patch parallel enabled and disabled.
-     - `HSDP cost`: compare HSDP on/off only when both configurations fit memory.
-     - `Ulysses vs Ring`: test a Ring/Ulysses hybrid only after long-sequence `diffuse` is confirmed dominant.
-   - For each configuration, create or record a stable config id, for example `A_cfg2_usp4_vaepp8_hsdp_tiling`.
-   - For each config id, capture:
-     - exact server command and environment variables,
-     - observed distributed setup from logs, such as SP groups, CFG groups, HSDP shard/replicate sizes, VAE patch size,
-     - exact request command for every scenario,
-     - final server response metrics and client-side elapsed time,
-     - output artifact paths and any failed/empty responses.
-   - Use one warmup request per scenario, then at least three measured repeats for the shortlisted config. For exploratory matrix pruning, one measured request is acceptable only if the margin is large; label it as one-shot.
-   - Compare configurations by stage, not only end-to-end latency. A configuration can improve `diffuse` while hurting `vae.decode`; record both effects.
-   - Select the best config only after checking the target metric, dominant stages, memory headroom, and output correctness.
-
-5. **Run targeted A/B tests.**
-   - Change one variable per test.
-   - Keep model, input, seed, request parameters, GPU placement, and warmup policy fixed.
-   - Record latency, stage timings, memory, output quality, and logs.
-   - Report comparison tables in ms. Include at least end-to-end client time, server `inference_time`, server stage generation time if available, `vae.encode`, `diffuse`, `vae.decode`, and peak memory.
-
-6. **Enforce the quality and precision gate for every optimization.**
-   - Do not mark an optimization ready until quality is checked on the same
-     workload: model, prompt/media, seed, shape, steps, scheduler, dtype,
-     backend, parallelism, and output encoding.
-   - For math-preserving changes, require exact/near-exact agreement or show
-     that differences are within baseline self-run variance.
-   - For precision, quantization, approximate attention, sparse/custom kernels,
-     or backend changes, include reviewable artifacts and metrics such as SSIM,
-     PSNR, LPIPS, cosine similarity, MAE, MSE, and temporal flicker checks.
-   - If quality tolerance is not stated, default to quality-preserving behavior.
-     Failed or inconclusive quality validation blocks a ready-to-merge claim.
-
-7. **Collect diagnostic trace only after narrowing hypotheses.**
-   - Use torch profiler for a small number of requests.
-   - Run two separate diagnostic traces instead of mixing concerns:
-     - **Operator/shape trace**: enable `torch_profiler_record_shapes=True` and keep stack collection disabled. Use this to rank CUDA kernels, NCCL collectives, attention/MLP/norm/RoPE work, and shape-specific hot operators.
-     - **Host-stack trace**: enable `torch_profiler_with_stack=True` and normally keep shape collection disabled. Use this to map CPU/Python host gaps, synchronization points, scheduler paths, and request handling overhead.
-   - Keep both trace commands and reports separate from baseline/benchmark commands. Torch profiler latency is diagnostic only and must not be used as the final latency claim.
-   - Prefer profiling only the narrowed dominant scenario, for example the highest-resolution/video-length case where `diffuse` dominates.
-   - If profiler endpoints are available, run one warmup request first, then call `/start_profile`, run one profiled request, and call `/stop_profile`. This keeps model initialization and warmup out of the diagnostic trace.
-   - Start by analyzing rank 0 only. Expand to more ranks only if rank 0 suggests imbalance, unclear GPU idle/free bubbles, high NCCL wait, server timing mismatch, CFG branch imbalance, or USP group stragglers.
-   - Diagnostic reports must be written to disk and preserve: server command, profiler config, warmup/request/polling commands, trace artifact paths, rank analyzed, analyzer output or summary, and the decision about whether additional ranks are necessary.
-   - Analyze both rank-level balance and device-level free bubbles when additional ranks are opened.
-
-8. **Analyze host, communication, and operators.**
-   - Find GPU idle/free intervals and map each large gap to the enclosing CPU/Python code.
-   - Separate real GPU idle from profiler overhead such as CUPTI `Command Buffer Full`.
-   - Compare NCCL kernel time to user annotations; annotations can overcount nested intervals.
-   - Rank operator work by total CUDA time and by repeated small-kernel launch count.
-
-9. **Produce an optimization plan.**
-   - Classify candidates as P0/P1/P2.
-   - For each candidate, state necessity, expected benefit, implementation path, validation plan, and quality risk.
-   - Include a concrete quality/precision validation gate for each candidate.
-   - Do not implement high-risk operator rewrites before proving the operator is a bottleneck for the target shapes.
-   - End the plan with a user-facing candidate selection table. The assistant
-     should not automatically choose a risky optimization just because it is
-     technically possible. Present the options clearly and let the user decide
-     which item is worth implementing based on latency target, engineering
-     budget, memory headroom, and quality tolerance.
-   - Organize the plan by optimization layer, not by a flat list of ideas:
-     - host/runtime optimization,
-     - measurement/benchmark reliability,
-     - parallelism and communication,
-     - VAE encode/decode and media pre/post processing,
-     - operator fusion and layout cleanup,
-     - attention main-path optimization,
-     - algorithmic, precision, or approximation changes.
-   - For each layer, tie every candidate back to evidence from baseline
-     metrics, diagnostic trace, source code, or output quality requirements.
-     Do not list generic optimizations without a trace or workload reason.
-
-## Priority Rules
-
-- **P0:** low risk, likely useful, or required for trustworthy measurement. Examples: real baseline, warmup, targeted parallel A/B, disabling avoidable `empty_cache`, scheduler coefficient caching.
-- **P1:** meaningful code changes with contained risk. Examples: cross-attention KV caching, VAE gather/broadcast reduction, AdaLayerNorm/RMSNorm/RoPE fusion after trace evidence.
-- **P2:** high implementation or quality risk. Examples: FA to LA replacement, custom Triton/CUDA fused kernels, FP8/quantization, sparsity/Rainfusion-style acceleration.
-
-Every implemented optimization needs A/B validation and a passing
-quality/precision gate, including math-preserving changes such as padding trim,
-layout cleanup, cache reuse, or host/runtime cleanup. Use objective metrics
-when available, keep reviewable artifacts, and compare against baseline
-self-run variance for precision, quantization, approximate kernels, or backend
-changes. If quality validation fails or is inconclusive, do not present the
-optimization as ready to merge.
-
-## Optimization Layers
-
-After baseline, parallel-search, and diagnostic traces, summarize optimization
-opportunities by layer. This is the core of the performance analysis: the goal
-is to connect evidence to a scoped implementation and a validation plan.
-
-### Host and Runtime Optimization
-
-Purpose: remove CPU/Python stalls, synchronization points, allocator overhead,
-and request-path overhead that leave GPU lanes empty.
-
-Evidence to look for:
-
-- High `idle_pct` or large `GAP` blocks in `trace_analyzer.py`.
-- Host-stack trace lines such as `torch.cuda.empty_cache`,
-  `cudaStreamSynchronize`, `cudaDeviceSynchronize`, Python locks, scheduler
-  waits, image/video preprocessing, or repeated small allocation paths.
-- Difference between client wall-clock time and server `inference_time_s`.
-
-Typical candidates:
-
-- Make avoidable `torch.cuda.empty_cache()` optional or guard it by memory
-  headroom.
-- Cache scheduler coefficients, timesteps, masks, or other tiny repeated CPU
-  computations when the request shape/steps are fixed.
-- Remove avoidable host-device synchronizations and blocking logging/stat calls.
-- Move expensive preprocessing out of the critical path or cache fixed prompt,
-  image, and transform work for benchmark scenarios.
-- Ensure benchmark scripts record client-side elapsed time, HTTP status, output
-  path, and server response metrics.
-
-Priority guidance:
-
-- Usually P0 when the change is measurement reliability or an obvious removable
-  synchronization.
-- Usually P1 when it changes scheduling, memory lifetime, or request execution
-  order.
-
-Validation:
-
-- Re-run non-profiler baseline with same workload and seed.
-- Confirm peak memory headroom if disabling cache cleanup.
-- Confirm generated output exists and quality/seed behavior is unchanged.
-
-### Parallelism and Communication Optimization
-
-Purpose: choose the right decomposition for CFG branches, sequence tokens, model
-weights, VAE tiles, and rank topology.
-
-Evidence to look for:
-
-- Baseline A/B across `CFG`, `USP/SP`, `Ring`, `HSDP/FSDP`, and VAE patch
-  parallelism.
-- Stage timing shifts: `diffuse`, `vae.encode`, `vae.decode`, and server
-  end-to-end.
-- NCCL kernel time from trace, not only `user_annotation` time.
-- Rank imbalance across SP group ranks or CFG branch ranks.
-- Memory headroom and OOM risk.
-
-Typical candidates:
-
-- `CFG=2 x USP=world/2` versus `CFG=1 x USP=world` for CFG-enabled models.
-- VAE patch parallel on/off or patch size tuning.
-- HSDP/FSDP on/off only if both configurations fit memory.
-- Ulysses versus Ulysses+Ring only after long-sequence `diffuse` is confirmed
-  dominant and all-to-all is suspected.
-- Rank mapping/topology changes if all-rank traces show stragglers or NCCL wait.
-- Buffer reuse or preallocation for FSDP/HSDP all-gather paths.
-
-Priority guidance:
-
-- P0/P1 for configuration-only changes with strong measured wins.
-- P1 for buffer reuse or rank mapping changes.
-- P2 for invasive distributed algorithm changes.
-
-Validation:
-
-- Measure by stage and memory, not only end-to-end.
-- Use one-variable A/B with identical prompt/media/seed/shape/steps.
-- When communication is suspected, compare rank0-3 in one USP group and rank0
-  versus rank4 across CFG branches for `CFG=2 x USP=4`.
-
-### VAE and Media Pipeline Optimization
-
-Purpose: reduce encode/decode, tiling, split/gather, and media conversion time.
-
-Evidence to look for:
-
-- Large `vae.encode` or `vae.decode` in low-overhead stage timings.
-- Host-stack gaps in VAE tile split, gather, merge, broadcast, or image/video
-  transforms.
-- VAE kernels or cuDNN convolution in operator trace.
-- Whether every rank needs the final decoded tensor.
-
-Typical candidates:
-
-- Keep VAE patch parallel enabled when it has clear measured benefit.
-- Reduce VAE gather/broadcast to only ranks that need the final media output.
-- Reuse tile metadata, split buffers, or gather buffers.
-- Evaluate bf16/autocast behavior for VAE only with visual quality checks.
-- Avoid redundant image conversion, resize, or tensor construction in repeated
-  benchmark runs.
-
-Priority guidance:
-
-- P0/P1 if VAE is a large share of the target workload or if a host gap is
-  obvious and low risk.
-- Lower priority when `diffuse` dominates and VAE is already patch-parallelized.
-
-Validation:
-
-- Compare `vae.encode`, `vae.decode`, server end-to-end, and peak memory.
-- Check output video integrity, artifacts, flicker, and seed stability.
-
-### Operator Fusion and Layout Cleanup
-
-Purpose: reduce high-frequency small kernels, memory bandwidth pressure, layout
-conversions, and launch overhead in transformer and VAE blocks.
-
-Evidence to look for:
-
-- Top operator tables showing many `aten::copy_`, `aten::cat`,
-  `split_with_sizes_copy`, `aten::add`, `aten::mul`, `aten::div`, norm,
-  activation, RoPE, or reshape/layout kernels.
-- `ops_rankN.xlsx` `by_shape` sheet showing repeated small shapes inside the
-  same block path.
-- Trace lanes showing many short kernels between larger GEMM/attention kernels.
-- Source code patterns with repeated elementwise chains or layout conversions.
-
-Typical fusion targets:
-
-- AdaLayerNorm / RMSNorm / LayerNorm plus scale/shift fusion.
-- RoPE fusion with Q/K layout preparation when shapes are stable.
-- Residual add, scale, gate, and elementwise chains.
-- MLP gate/up/down cleanup, such as fusing activation and multiply around
-  GEMM outputs when feasible.
-- QKV projection and reshape/split/cat path cleanup.
-- Attention pre/post layout cleanup to avoid unnecessary copies, cats, and
-  splits around sequence parallel all-to-all.
-
-Priority guidance:
-
-- P1 when implemented with existing PyTorch/Triton/local helper patterns and
-  validated against exact outputs or tolerances.
-- P2 when it requires custom CUDA/Triton kernels, changes numerics, or touches
-  attention math directly.
-
-Validation:
-
-- First prove the operator family is material for the target shape.
-- Use non-profiler A/B for latency and stage timing.
-- Use quality regression checks for generated video stability.
-- Check compile behavior and graph breaks if using `torch.compile`.
-
-### Attention Main-Path Optimization
-
-Purpose: address the dominant self-attention cost when FlashAttention or other
-attention kernels dominate CUDA time.
-
-Evidence to look for:
-
-- Operator trace where FlashAttention/SDPA kernels dominate total CUDA time.
-- Attention shape from `ops_rankN.xlsx` `by_shape`, model code, or trace
-  metadata.
-- Whether attention cost scales with latent frames, latent H/W, patch size, or
-  CFG duplication.
-- Layout/copy/all-to-all work around attention.
-
-Typical candidates:
-
-- Verify the attention backend and shape are on the intended fast kernel path.
-- Compare supported attention backends only with identical workload and quality
-  settings.
-- Reduce attention input size by safe model/config choices when allowed:
-  latent resolution, frame count, patching, boundary ratio, or windowing.
-- Remove avoidable layout conversions before/after attention.
-- Reuse condition-side KV or other static inputs if the model structure allows.
-- Consider custom kernels, sparse/window/linear attention, or approximation only
-  after quality risk is accepted.
-
-Priority guidance:
-
-- P1 for backend/config/layout changes with preserved math.
-- P2 for approximate attention, sparsity, custom kernels, or any change that can
-  alter quality/temporal consistency.
-
-Validation:
-
-- Always include output quality and seed behavior checks.
-- Compare `diffuse`, server end-to-end, peak memory, and attention kernel time
-  in diagnostic traces if needed.
-
-### Algorithmic, Precision, and Approximation Optimization
-
-Purpose: reduce mathematical work or precision cost beyond local code cleanup.
-
-Evidence to look for:
-
-- A single operator family dominates even after low-risk runtime, parallel, and
-  fusion work.
-- Memory bandwidth or compute utilization suggests precision or quantization
-  could matter.
-- The user explicitly allows quality-preserving or approximate methods.
-
-Typical candidates:
-
-- FP8/quantization for transformer or selected projections.
-- Sparsity or Rainfusion-style acceleration.
-- Reduced steps, scheduler changes, distillation, or caching across frames.
-- Approximate attention or linear attention.
-
-Priority guidance:
-
-- Usually P2 because quality, numerics, and implementation risk are high.
-
-Validation:
-
-- Requires strict A/B, visual quality review, temporal flicker checks, seed
-  stability, and possibly human evaluation.
-
-### Interpolation, Super-Resolution, and E2E Pipeline Optimization
-
-Purpose: optimize the whole user-visible video product, not only the base
-diffusion invocation. Some deployments trade base-model latency against
-post-processing, interpolation, or super-resolution stages.
-
-Evidence to look for:
-
-- E2E latency breakdown across base generation, interpolation, super-resolution,
-  encoding, storage, and response streaming.
-- Fast/slow GPU or fast/slow stage analysis across multiple cards and pipeline
-  stages.
-- User quality target: resolution, FPS, temporal smoothness, and acceptable
-  post-processing artifacts.
-
-Typical candidates:
-
-- Add or optimize a frame interpolation stage when it reduces required base
-  model frames for the same perceived FPS.
-- Add or optimize a super-resolution model when generating lower base
-  resolution plus SR is faster for the target quality.
-- Analyze E2E2 pipeline behavior: client request, service scheduling, diffusion,
-  VAE/media, post-process, file write, and response.
-- Identify fast/slow cards or stages and rebalance pipeline placement.
-
-Priority guidance:
-
-- P1 when using proven interpolation/SR components without changing diffusion
-  math.
-- P2 when quality risk is high or the pipeline adds significant operational
-  complexity.
-
-Validation:
-
-- Measure E2E wall-clock, per-stage server timings, output FPS/resolution,
-  artifacts, flicker, and user-visible quality.
-
-### Optimization Candidate Library
-
-Use this table as a compact menu, not as automatic recommendations. Pick items
-only when baseline metrics, trace evidence, source inspection, or quality
-tolerance supports them.
-
-| Layer | Candidate | Evidence | Priority | Validation focus |
-|---|---|---|---|---|
-| Measurement | Freeze server/request/benchmark commands | Missing or drifting commands | P0 | Repeatable non-profiler A/B |
-| Measurement | Separate baseline and diagnostic profiler runs | Profiler used for latency claims | P0 | Low-overhead stage timings |
-| Host/runtime | Guard or remove avoidable `empty_cache` | Host-stack gaps or sync stalls | P0/P1 | Latency, peak memory, OOM safety |
-| Host/runtime | Cache scheduler coefficients/timesteps | Repeated tiny CPU/GPU work | P0/P1 | Same seed/output, stage timing |
-| Host/runtime | Reduce framework scheduling overhead | Client time exceeds server time | P1 | E2E latency, throughput |
-| Parallel | `CFG=2 x USP=world/2` vs `USP=world` | CFG doubles forward work | P0/P1 | `diffuse`, NCCL, memory |
-| Parallel | Tune VAE patch parallelism | VAE encode/decode is material | P0/P1 | VAE time, output correctness |
-| Parallel | HSDP on/off or buffer reuse | HSDP affects memory/all-gather | P1 | Memory, latency, OOM risk |
-| Parallel | Ulysses vs Ulysses+Ring | Long sequence all-to-all suspected | P1/P2 | Rank balance, NCCL kernels |
-| Cross-attn | Disable SP for short condition tokens | Cross-attn comm exceeds compute | P1 | `diffuse`, correctness |
-| VAE/media | Reduce VAE gather/broadcast | Rank traces show VAE wait | P1 | Rank balance, output file |
-| VAE/media | Reuse tile metadata/buffers | Tile split/merge host gaps | P1 | `vae.encode/decode`, memory |
-| VAE/media | VAE bf16/autocast | VAE float kernels are slow | P1/P2 | Artifacts, flicker, seed stability |
-| Operator fusion | AdaLayerNorm/LayerNorm fusion | Norm plus scale/shift kernels | P1 | Numeric tolerance, latency |
-| Operator fusion | RMSNorm fusion | Many small RMSNorm kernels | P1 | Numeric tolerance, latency |
-| Operator fusion | RoPE cache/fuse/layout cleanup | RoPE copy/reshape kernels | P1/P2 | Kernel count, correctness |
-| Operator layout | QKV or attention layout cleanup | Copy/cat/split around attention | P1 | Copy kernels, compile behavior |
-| Attention | Verify backend fast path | FA/SDPA dominates trace | P1 | `diffuse`, attention kernels |
-| Attention | FA to LA or selected-head LA | Attention remains dominant | P2 | Quality, temporal stability |
-| Precision | Transformer FP8/quantization | Compute/bandwidth bound and allowed | P2 | Quality, speed, stability |
-| Sparsity | Rainfusion-style acceleration | DiT compute remains dominant | P2 | Prompt diversity, quality |
-| Pipeline | Frame interpolation | Fewer base frames can meet FPS | P1/P2 | E2E latency, motion artifacts |
-| Pipeline | Super-resolution | Lower base res plus SR may win | P1/P2 | Detail quality, artifacts |
-| E2E | Fast/slow-card analysis | Multi-card stragglers | P0/P1 | Per-rank/stage wall-clock |
-
-### Optimization Plan Template
-
-Use this table shape when reporting the next work items:
-
-| Priority | Layer | Candidate | Evidence | Expected benefit | Implementation path | Validation | Quality risk |
-|---|---|---|---|---|---|---|---|
-| P0 | Host/runtime | Guard `empty_cache` | Host-stack gap points to `torch.cuda.empty_cache` | Small latency reduction, less idle | Add config/env guard | Non-profiler A/B, memory check | Low |
-| P1 | Operator fusion | RMSNorm/AdaLayerNorm fusion | High-frequency norm/elementwise kernels | Lower launch/bandwidth overhead | Use existing fusion helper or targeted Triton | A/B + output check | Medium |
-| P1/P2 | Attention | Attention layout/backend investigation | FA kernel dominates CUDA time | Potentially large | Inspect shapes/backend and remove layout copies | A/B + trace + quality | Medium/high |
-
-Then present a short selection prompt using the same rows:
-
-```text
-Which candidate should we implement next?
-
-1. P0 Host/runtime: guard empty_cache
-   - Expected benefit: small but low-risk latency reduction.
-   - Risk: possible memory increase/OOM if memory headroom is insufficient.
-
-2. P1 Operator fusion: inspect by_shape and implement first norm/RoPE/layout fusion
-   - Expected benefit: medium if high-frequency small kernels are confirmed.
-   - Risk: numerical/compile/quality validation needed.
-
-3. P1/P2 Attention: FA/LA/backend/layout investigation
-   - Expected benefit: potentially large.
-   - Risk: high quality and implementation risk.
-```
-
-If the user has not chosen an item, default to explaining tradeoffs and asking
-which candidate to execute. Only proceed autonomously on low-risk P0 measurement
-or instrumentation fixes.
-
-## Analysis Helpers
-
-PyTorch profiler traces are Chrome/Perfetto-compatible JSON files, usually
-`trace_rankN.json` or `trace_rankN.json.gz`. They normally contain a top-level
-`traceEvents` list, though some exporters emit the raw event list directly.
-
-Use the checked-in analyzer from the repository root:
-
-```bash
-python3 .claude/skills/diffusion-perf-opt/scripts/trace_analyzer.py \
-  vllm_profile/.../trace_rank0.json.gz \
-  --min-gap-ms 5 \
-  --topn 20
-```
-
-For rank imbalance or communication questions, pass all relevant rank traces in
-one command. For host gaps, lower `--min-gap-ms` to `1` and use a host-stack
-trace. Read `gpu_span_s`, `busy_union_s`, `idle_union_s`, `idle_pct`, `GAP`
-blocks, `Top GPU/operator events by total duration`, and `Top NCCL-like events
-by category`. Treat `cat=user_annotation` NCCL ranges as enclosing annotations;
-prefer `cat=kernel` or `cat=gpu_user_annotation` for real device work.
-
-The analyzer summarizes timing only. It does not parse tensor shapes, attribute
-overlap to individual streams, prove quality, or provide final latency claims.
-Use `ops_rankN.xlsx` or PyTorch key averages for shape analysis, and re-test any
-optimization with non-profiler baseline commands.
-
-Read `references/optimization-playbook.md` when drafting the optimization table or comparing candidate techniques.
-
-## vLLM Omni Heuristics
-
-- Cross-attention usually should not use USP/SP when text/image condition token count is much smaller than latent video tokens. Confirm via trace; in Wan2.2 I2V, self-attention dominates cross-attention.
-- VAE bf16/autocast is often worthwhile but requires visual quality checks.
-- VAE patch parallel can help decode/encode but may add gather/merge/broadcast overhead. Check whether all ranks need the final decoded tensor.
-- HSDP/FSDP is primarily a memory strategy. If the model fits without it, run an on/off latency comparison.
-- Scheduler work can create small host/device gaps; cache tiny solve coefficients when timesteps/order are known.
-- `torch.cuda.empty_cache()` can prevent OOM but creates synchronization/idle. Make it optional if memory headroom is sufficient.
-- `Command Buffer Full` in profiler output is profiler overhead, not a model optimization target.

.claude/skills/diffusion-perf-opt/agents/openai.yaml

@@ -1,4 +0,0 @@
-interface:
-  display_name: "vLLM Omni Perf Opt"
-  short_description: "Profile and optimize vLLM Omni diffusion performance."
-  default_prompt: "Analyze and optimize a vLLM Omni diffusion workload using baselines, traces, parallel strategy modeling, and prioritized optimization plans."

.claude/skills/diffusion-perf-opt/references/optimization-playbook.md

@@ -1,82 +0,0 @@
-# Optimization Playbook
-
-Use this reference when turning trace evidence into prioritized vLLM Omni optimizations.
-
-## Parallel Strategy
-
-Start from workload shape:
-
-- Long video tokens, high resolution, many frames: self-attention and FFN dominate. USP/SP is usually valuable.
-- CFG enabled with positive and negative prompts: transformer forward count doubles. CFG parallel may beat USP on small/moderate sequence lengths.
-- Cross-attention: skip SP/USP when condition tokens are much smaller than latent tokens; extra all-to-all can exceed compute saved.
-- HSDP/FSDP: use to fit memory first. If memory allows, compare without HSDP because all-gathers and lazy init add overhead.
-- VAE parallel: useful when encode/decode is material. Check gather/merge/broadcast and whether every rank needs final decoded output.
-
-Candidate matrices:
-
-- 2 cards: `USP=2`, `CFG=2`, HSDP on/off, VAE PP on/off.
-- 4 cards: `USP=4`, `CFG=2 x USP=2`, `USP=2 x HSDP`, VAE PP on/off.
-- 8 cards: add pipeline/stage parallel only after measuring orchestration and inter-stage transfer overhead.
-
-## Host And Free Bubble Analysis
-
-Look for device intervals where no kernel, memcpy, or memset is active.
-
-Common causes:
-
-- First request FSDP/HSDP lazy init.
-- Python preprocessing: image load, resize, `to_tensor`, tokenizer, request marshaling.
-- `torch.cuda.empty_cache()` or explicit synchronizations.
-- Scheduler small linear algebra such as tiny `torch.linalg.solve`.
-- VAE tile split/merge Python loops.
-- Profiler overhead such as CUPTI command buffer flush/full.
-
-Also inspect trace-external time:
-
-- API request parsing and multipart upload.
-- Engine/orchestrator dispatch.
-- Worker RPC and queue wait.
-- Output serialization, video write, base64, HTTP response.
-
-## Operator Optimization Candidates
-
-Evaluate necessity before implementation:
-
-- **FA to LA / LA for head:** high potential only if attention is the main CUDA hotspot and target shapes fit LA well.
-- **Cross-attention BSND/no transpose:** do only if transpose/copy around cross-attn is measurable.
-- **AdaLayerNorm fusion:** fuse norm, scale/shift, gate, residual, and dtype conversion if repeated elementwise kernels are significant.
-- **RMSNorm fusion:** useful for Q/K norms and VAE norms when small-kernel launch count is high.
-- **FP32 LayerNorm to bf16:** quality-sensitive; run visual regression.
-- **LA preprocess:** integrate only if preprocess overhead is cached or amortized.
-- **RoPE move/cache/fuse:** useful when RoPE small kernels are visible around every self-attn block.
-- **VAE RMSNorm replacement:** only after VAE internal trace shows norm as a real component of decode/encode.
-- **Scheduler fusion/cache:** cache coefficient solves and fuse CFG combine with scheduler arithmetic when possible.
-
-Before writing a custom kernel, check:
-
-- Local vLLM Omni implementations.
-- Upstream vLLM kernels and attention backends.
-- SGLang, FlashInfer, FlashAttention, and Triton examples.
-- Shape constraints and layout expectations.
-
-## Quantization And Sparsity
-
-Treat these as high-risk optimizations:
-
-- Rainfusion/sparse diffusion can be high reward but needs model-specific quality validation.
-- VAE decode quantization can reduce decode time but usually has smaller end-to-end upside than transformer attention/FFN work.
-- FP8 or approximate kernels require seed-stable visual regression, temporal flicker checks, and prompt diversity.
-
-## Output Format For Recommendations
-
-Return a table with:
-
-- Priority: P0/P1/P2.
-- Optimization.
-- Evidence from trace or code.
-- Expected benefit.
-- Risk/quality impact.
-- Implementation path.
-- Validation command or A/B test.
-
-State clearly when a recommendation is a hypothesis requiring new trace with shapes.

.claude/skills/diffusion-perf-opt/scripts/trace_analyzer.py

@@ -1,144 +0,0 @@
-#!/usr/bin/env python3
-"""Summarize torch profiler Chrome traces for vLLM Omni diffusion optimization."""
-
-from __future__ import annotations
-
-import argparse
-import collections
-import gzip
-import json
-from pathlib import Path
-from typing import Any
-
-GPU_CATS = {"kernel", "gpu_memcpy", "gpu_memset"}
-CPU_CATS = {"python_function", "user_annotation", "cpu_op", "cuda_runtime", "cuda_driver"}
-
-
-def open_trace(path: Path):
-    if path.suffix == ".gz":
-        return gzip.open(path, "rt")
-    return path.open("rt")
-
-
-def event_name(event: dict[str, Any]) -> str:
-    return str(event.get("name", ""))
-
-
-def summarize_trace(path: Path, min_gap_us: float, topn: int) -> None:
-    with open_trace(path) as f:
-        data = json.load(f)
-
-    events = data if isinstance(data, list) else data.get("traceEvents", [])
-    gpu: list[tuple[float, float, float, str, str, Any, Any]] = []
-    cpu: list[tuple[float, float, float, str, str, Any, Any]] = []
-    by_gpu_name: dict[str, list[float]] = collections.defaultdict(lambda: [0, 0.0, 0.0])
-    by_name: dict[str, list[float]] = collections.defaultdict(lambda: [0, 0.0, 0.0])
-    nccl: dict[tuple[str, str], list[float]] = collections.defaultdict(lambda: [0, 0.0, 0.0])
-
-    for event in events:
-        dur = event.get("dur")
-        ts = event.get("ts")
-        if dur is None or ts is None or dur <= 0:
-            continue
-        cat = str(event.get("cat", ""))
-        name = event_name(event)
-        row = (float(ts), float(ts + dur), float(dur), name, cat, event.get("pid"), event.get("tid"))
-        stat = by_name[name]
-        stat[0] += 1
-        stat[1] += dur
-        stat[2] = max(stat[2], dur)
-        if cat in GPU_CATS:
-            gpu.append(row)
-            gstat = by_gpu_name[name]
-            gstat[0] += 1
-            gstat[1] += dur
-            gstat[2] = max(gstat[2], dur)
-        elif cat in CPU_CATS:
-            cpu.append(row)
-        if "nccl" in name.lower():
-            nstat = nccl[(cat, name)]
-            nstat[0] += 1
-            nstat[1] += dur
-            nstat[2] = max(nstat[2], dur)
-
-    print(f"\n== {path}")
-    print(f"events={len(events)} gpu_events={len(gpu)} cpu_events={len(cpu)}")
-    if not gpu:
-        print("No GPU events found.")
-        return
-
-    gpu.sort()
-    merged: list[list[Any]] = []
-    for start, end, dur, name, cat, pid, tid in gpu:
-        if not merged or start > merged[-1][1]:
-            merged.append([start, end, [(start, end, dur, name, cat, pid, tid)]])
-        else:
-            merged[-1][1] = max(merged[-1][1], end)
-            merged[-1][2].append((start, end, dur, name, cat, pid, tid))
-
-    span = max(end for _, end, *_ in gpu) - min(start for start, *_ in gpu)
-    busy = sum(end - start for start, end, _ in merged)
-    idle = span - busy
-    print(
-        f"gpu_span_s={span / 1e6:.3f} "
-        f"busy_union_s={busy / 1e6:.3f} "
-        f"idle_union_s={idle / 1e6:.3f} "
-        f"idle_pct={idle / span * 100:.2f}"
-    )
-
-    interesting_cpu = [
-        row
-        for row in cpu
-        if row[2] >= 1000
-        and (
-            row[4] in {"python_function", "user_annotation"}
-            or "cudaStreamSynchronize" in row[3]
-            or "cudaDeviceSynchronize" in row[3]
-            or "cudaLaunch" in row[3]
-            or "cudaMemcpy" in row[3]
-        )
-    ]
-
-    gaps = []
-    for idx in range(1, len(merged)):
-        gap_start = merged[idx - 1][1]
-        gap_end = merged[idx][0]
-        gap_dur = gap_end - gap_start
-        if gap_dur >= min_gap_us:
-            prev_event = max(merged[idx - 1][2], key=lambda x: x[1])
-            next_event = min(merged[idx][2], key=lambda x: x[0])
-            mid = (gap_start + gap_end) / 2
-            containers = [row for row in interesting_cpu if row[0] <= mid <= row[1]]
-            containers = sorted(containers, key=lambda x: x[2])[:8]
-            gaps.append((gap_dur, gap_start, gap_end, prev_event, next_event, containers))
-
-    print(f"gaps_ge_{min_gap_us / 1000:.3f}ms count={len(gaps)} sum_s={sum(g[0] for g in gaps) / 1e6:.3f}")
-    for gap_dur, gap_start, gap_end, prev_event, next_event, containers in sorted(gaps, reverse=True)[:topn]:
-        print(f"\nGAP {gap_dur / 1000:.3f} ms ts={gap_start:.0f}->{gap_end:.0f}")
-        print(f"  prev {prev_event[4]} {prev_event[2] / 1000:.3f} ms {prev_event[3][:160]}")
-        print(f"  next {next_event[4]} {next_event[2] / 1000:.3f} ms {next_event[3][:160]}")
-        for row in containers:
-            print(f"  in   {row[4]} {row[2] / 1000:.3f} ms {row[3][:180]}")
-
-    print("\nTop GPU/operator events by total duration:")
-    for name, (count, total, max_dur) in sorted(by_gpu_name.items(), key=lambda kv: kv[1][1], reverse=True)[:topn]:
-        print(f"  {int(count):8d} total={total / 1e6:9.3f}s max={max_dur / 1000:9.3f}ms {name[:180]}")
-
-    print("\nTop NCCL-like events by category:")
-    for (cat, name), (count, total, max_dur) in sorted(nccl.items(), key=lambda kv: kv[1][1], reverse=True)[:topn]:
-        print(f"  {int(count):8d} total={total / 1e6:9.3f}s max={max_dur / 1000:9.3f}ms cat={cat} {name[:160]}")
-
-
-def main() -> None:
-    parser = argparse.ArgumentParser(description=__doc__)
-    parser.add_argument("traces", nargs="+", type=Path, help="trace.json or trace.json.gz files")
-    parser.add_argument("--min-gap-ms", type=float, default=5.0, help="minimum GPU idle gap to print")
-    parser.add_argument("--topn", type=int, default=20, help="number of gaps/hotspots to print")
-    args = parser.parse_args()
-
-    for trace in args.traces:
-        summarize_trace(trace, args.min_gap_ms * 1000.0, args.topn)
-
-
-if __name__ == "__main__":
-    main()

.claude/skills/precheck-pr/SKILL.md

@@ -1,97 +0,0 @@
----
-name: precheck-pr
-description: Self-check your branch before creating a PR — catch dead code, verify accuracy/perf claims, validate PR title format, and confirm merge readiness. Use when the user says "precheck", "self review", "pre-submit check", or "check my PR before I open it." Never posts to GitHub.
----
-
-# PR Pre-Check
-
-Self-review your branch before creating a PR against `vllm-project/vllm-omni`. Two modes: **quick** catches showstoppers, **full** does a thorough maintainer-grade review. Never posts to GitHub; the report is for the contributor's terminal only.
-
-## Mode Selection
-
-| Mode | When | Time |
-|------|------|------|
-| **Quick** | About to push, final sanity check | ~3 min |
-| **Full** | Ready for review, want maintainer-level scan | ~10 min |
-
-Default to quick if unsure. Run full before marking a PR "ready for review."
-
-## Workflow
-
-### Step 1: Detect Base Branch
-
-```bash
-BASE_SHA=$(git merge-base HEAD origin/main 2>/dev/null \
-         || git merge-base HEAD main 2>/dev/null \
-         || echo origin/main)
-echo "diffing against ${BASE_SHA}"
-git diff --name-only ${BASE_SHA}...HEAD
-```
-
-### Step 2: Validate PR Title
-
-Check the most recent commit message (or branch name if no commit yet) against the project convention. Valid prefixes:
-
-| Prefix | Applies to |
-|--------|-----------|
-| `[Bugfix]` | Bug fixes |
-| `[CI/Build]` | Build or CI improvements |
-| `[Doc]` | Documentation changes |
-| `[Model]` | New/improved models (include model name) |
-| `[Frontend]` | Frontend changes (API server, OmniLLM class, etc.) |
-| `[Kernel]` | CUDA/kernel changes |
-| `[Core]` | Core logic changes (OmniProcessor, OmniARScheduler, etc.) |
-| `[Hardware][Vendor]` | Hardware-specific (e.g., `[Hardware][Ascend]`) |
-| `[Misc]` | Other changes (use sparingly) |
-
-✗ if: missing prefix, wrong case (`[bugfix]`), or WIP/Draft in title.
-⚠ if: `[Model]` prefix without the model identifier (e.g., `[Model] Add new model` — should be `[Model] Add <ModelName> ...`).
-
-### Step 3: Categorize the PR
-
-| Diff contains | PR type |
-|---------------|---------|
-| New files under `vllm_omni/model_executor/models/<name>/` | **New Model** |
-| Changes to `vllm_omni/diffusion/` | **Diffusion Model** |
-| `[Bugfix]` prefix or single-file fix | **Bug Fix** |
-| Perf/benchmark/throughput claims in commit msg or diff | **Performance** |
-| Everything else | **General** |
-
-If multiple rows apply (e.g., a diffusion model is also a new model), union the checklists.
-
-### Step 4: Run Checklist
-
-Ask: "Quick mode or full mode?" Then walk the checklist for the detected PR type from [references/checklists.md](references/checklists.md). Each item produces ✓, ✗, or ⚠.
-
-### Step 5: Print Report
-
-```
-Pre-check report for <branch>
-
-  Mode: quick | full
-  Type: <new-model | diffusion-model | bug-fix | perf | general>
-
-  Dimension          Result
-  ─────────────────  ──────
-  PR title format    ✓
-  PR desc integrity  ✓
-  Registry/config    ✓
-  Dead code          ⚠ 2 warnings
-  Accuracy           ✓
-  Benchmark          ✗ missing software versions
-
-  Verdict: 1 blocking | 2 warnings | recommend fixing ✗ before PR
-```
-
-**Severity:**
-
-| Mark | Meaning |
-|------|---------|
-| ✗ | Blocking — fix before opening PR |
-| ⚠ | Warning — consider fixing |
-| ✓ | Pass |
-| — | Skipped (not applicable) |
-
-### Stop Here
-
-Do not post comments, open PRs, or modify files. The report is for the contributor's terminal only.

.claude/skills/precheck-pr/references/checklists.md

@@ -1,135 +0,0 @@
-# Pre-Check Checklists
-
-Each item is a concrete check to run against the diff. ✗ = fix before PR. ⚠ = consider fixing.
-
----
-
-## All PRs: Title Format (check first)
-
-PR title must follow the project convention documented in `docs/contributing/README.md`:
-
-| Prefix | Applies to |
-|--------|-----------|
-| `[Bugfix]` | Bug fixes |
-| `[CI/Build]` | Build or CI improvements |
-| `[Doc]` | Documentation changes |
-| `[Model]` | New/improved models (include model name) |
-| `[Frontend]` | Frontend changes (API server, OmniLLM class, etc.) |
-| `[Kernel]` | CUDA/kernel changes |
-| `[Core]` | Core logic changes (OmniProcessor, OmniARScheduler, etc.) |
-| `[Hardware][Vendor]` | Hardware-specific (e.g., `[Hardware][Ascend]`) |
-| `[Misc]` | Other changes (use sparingly) |
-
-**Examples:**
-```
-✓ [Model] Add MiniCPM-o 4.5 support
-✓ [Bugfix] Fix Ovis image text encoder dtype
-✓ [Core] Qwen3-Omni performance optimization
-✓ [Hardware][Ascend] Fix VoxCPM2 on Ascend
-✗ Fix Ovis image text encoder dtype              ← missing prefix
-✗ [bugfix] Fix Ovis image text encoder           ← wrong case
-✗ [WIP] Add model                                ← WIP in title
-```
-
-- [ ] **Title has valid prefix** — ✗ if missing prefix, wrong case, or WIP/Draft in title. For `[Model]` PRs, the title MUST also include the model identifier (e.g., `[Model] Add <ModelName> ...`). ⚠ if missing the model name.
-
----
-
-## Bug Fix PRs
-
-### Quick
-
-- [ ] **PR body has 4 sections:** what broke (error message), repro steps, root cause, fix description
-- [ ] **Fix is minimal:** diff touches only files directly related to the bug. No unrelated refactors or style changes.
-- [ ] **Repro is runnable:** the PR body includes a copy-paste command that reproduces the bug
-- [ ] **Branch is rebased:** `git merge-base HEAD origin/main` is recent, no merge conflicts
-
-### Full (adds)
-
-- [ ] **Regression test exists:** `git diff --name-only` includes at least one `tests/` file
-- [ ] **Regression test reproduces the original error:** the test would fail on main, pass on this branch
-- [ ] **Fix matches root cause exactly:** no "fix the symptom + something else" — if the root cause is dtype, the fix is dtype, not dtype + reformatting
-- [ ] **No silent failure risk:** no bare `except: pass`, no `try: ... except: return None`, no empty fallback added
-- [ ] **Upstream pattern match:** the fix follows the same pattern as existing code for similar cases (grep for analogous `from_pretrained` calls, etc.)
-- [ ] **Environment documented:** torch/transformers/vllm versions listed if the bug is version-dependent
-
----
-
-## Performance PRs
-
-### Quick
-
-- [ ] **Before/after numbers in PR body:** at minimum, one metric (RTF, RPS, latency, VRAM) with main vs PR values
-- [ ] **Hardware specified:** GPU model, count, VRAM. "L4" or "H20" is sufficient. "tested on GPU" is not.
-- [ ] **Concurrency scaling:** if the optimization claims better throughput, numbers at 2+ concurrency levels
-- [ ] **No unexplained regressions:** if any metric goes down at any concurrency level, explain why
-- [ ] **Branch is rebased:** `git merge-base HEAD origin/main` is recent, no merge conflicts
-
-### Full (adds)
-
-- [ ] **Software versions:** torch, CUDA, vllm, vllm-omni versions in PR body
-- [ ] **Warmup stated:** "1 warmup + 3 measured" or equivalent
-- [ ] **Stddev or range:** not just a single best run. "average of 2 rounds" is minimum.
-- [ ] **VRAM measured:** peak GPU memory before and after
-- [ ] **Benchmark script checked in:** under `tests/` or `examples/`, runnable by reviewer
-- [ ] **Exact command line in PR:** not pseudocode, not "run the benchmark"
-- [ ] **Pytest summary included:** `36 passed, 3 skipped, 17 warnings in 2897.66s` — proves the whole suite passes
-- [ ] **Config changes documented:** if `max_num_seqs`, `gpu_memory_utilization`, or deploy YAML changed, old and new values stated with reasoning
-
----
-
-## New Model PRs
-
-### Quick
-
-- [ ] **PR body file list matches `git diff --name-only`:** every file path in the description exists in the diff
-- [ ] **Registry entries match claimed count:** count `_OMNI_MODELS` or `_DIFFUSION_MODELS` entries added vs claimed in body
-- [ ] **Pipeline config `custom_process_input_func` paths resolve:** each referenced module.function exists
-- [ ] **`__init__.py` exports match registry:** each registered class appears in `__init__.py` and `__all__`
-- [ ] **Shell script modes match README:** `run_curl.sh` supports every mode the README documents
-- [ ] **Branch is rebased:** `git merge-base HEAD origin/main` is recent, no merge conflicts
-
-### Full (adds)
-
-- [ ] **Dead code scan (5 patterns):**
-  1. Dead `forward()` — defined but never called through the pipeline path
-  2. Dead factory/builder functions — instantiated nowhere
-  3. Dead wrappers — thin wrappers with no callers
-  4. Dead branch guards — `if version > X` where X is always satisfied
-  5. Unused parameters — `__init__` args never accessed
-- [ ] **Copy-paste detection:** no string constant defined in 3+ files, no cross-module validation duplication, no near-identical shape coercion fns
-- [ ] **Import hygiene:** no re-export-only imports, no module-level side effects without docs, no `import os` in function body when already at top. Function-body imports are acceptable when the import is only reachable from a closure or factory return value (i.e., `import` would otherwise execute at module load even though the user may never trigger that path). Either lift to top-of-file or annotate with `# noqa: PLC0415`.
-- [ ] **Accuracy:** at least one test compares output against reference; audio = valid WAV at correct sample rate; determinism verified if seed claimed. ⚠ if missing and not explicitly deferred to a follow-up PR or RFC section; — if explicitly deferred (the deferral target SHOULD be linked in the PR body).
-- [ ] **Performance:** hardware, software versions, warmup, and metrics (RTF, VRAM, RPS, latency) stated. ⚠ if missing and not explicitly deferred; — if explicitly deferred.
-- [ ] **Benchmark settings:** model config (TP, PP, max_model_len, enforce_eager, quant), runtime config (batch, input/output spec), environment versions
-- [ ] **Benchmark script checked in + exact command line + pytest summary**
-
----
-
-## Diffusion Model PRs
-
-Diffusion models live under `vllm_omni/diffusion/`. In addition to the New Model checklist above, verify:
-
-- [ ] **Transformer adapter follows contract:** inherits from `nn.Module`, implements `load_weights()`, uses `vllm_omni.diffusion.attention.layer.Attention`
-- [ ] **Pipeline has both offline and online paths:** `forward()` for batch generation, streaming path for online serving
-- [ ] **Cache-DiT integration:** at least one acceleration (TeaCache, FBCache, etc.) is wired up. ✗ for new **top-level** pipelines. For **subclass** pipelines that inherit `__call__` from a parent where Cache-DiT is already absent, downgrade to ⚠ and recommend filing a follow-up issue against the parent pipeline family.
-- [ ] **Parallelism config:** TP/SP/USP/CFG-Parallel options exposed in pipeline config
-- [ ] **Docs table updated:** model added to `docs/models/supported_models.md` with correct metadata
-- [ ] **E2E test exists:** at least one test exercises the full generate path
-
----
-
-## General PRs
-
-### Quick
-
-- [ ] **PR body matches diff:** description doesn't claim changes not in the diff
-- [ ] **Branch is rebased:** `git merge-base HEAD origin/main` is recent, no merge conflicts
-- [ ] **CI gates passing:** DCO, pre-commit, build — check `gh pr view --json statusCheckRollup` if a PR already exists
-
-### Full (adds)
-
-- [ ] **No new dead code in changed files:** scan for unreachable paths, unused imports
-- [ ] **Test coverage:** if core code changed (`engine/`, `stages/`, `connectors/`), tests are added or existing tests cover the path
-- [ ] **Import hygiene:** no redundant imports added
-- [ ] **No unrelated changes:** diff scope matches PR description

.claude/skills/readme.md

@@ -1,42 +0,0 @@
-# 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`
-- `diffusion-perf-opt`: guides diffusion model performance optimization,
-  including profiling traces, parallel strategies, stage timing analysis, and
-  benchmark-driven tuning
-- `add-omni-model`: covers addition of new omni-modality model support
-- `add-tts-model`: covers integration of new TTS models and related serving
-  workflows
-- `generate-release-note`: helps prepare release notes for repository changes
-- `precheck-pr`: self-check a branch before creating a PR — validates PR title
-  format, catches dead code, verifies accuracy/perf claims, and confirms merge
-  readiness
-- `review-pr`: provides a structured workflow for reviewing pull requests
-- `vllm-omni-npu-model-runner-upgrade`: upgrades NPU model runners to align with the
-  latest vllm-ascend NPUModelRunner
-
-## 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

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

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

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

@@ -1,222 +0,0 @@
-# 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,7 +27,6 @@
 *.mov binary
 *.mp4 binary
 *.mp3 binary
-*.wav binary
 *.flv binary
 *.fla binary
 *.swf binary

.github/CODEOWNERS

@@ -3,64 +3,15 @@
 # More info: https://docs.github.com/en/repositories/managing-your-repositorys-settings-and-features/customizing-your-repository/about-code-owners
 
 # Default owners (fallback reviewers for all files)
-* @Gaohan123 @hsliuustc0106 @ywang96
+# later, we can add more reviewers here
+* @hsliuustc0106
 
-# Repository operations
-/.github/ @hsliuustc0106
+# Assign reviewers by directory
+/.github/ @owner1
+/tests/ @owner1
+/docs/ @owner1
 
-# Documentation and Recipes
-/docs/ @hsliuustc0106 @Gaohan123 @david6666666
-/recipes @hsliuustc0106 @Gaohan123 @david6666666
-
-# Skills
-/.claude/ @hsliuustc0106
-
-# Tests
-/tests/ @yenuo26
-
-# CI
-/.buildkite/ @yenuo26 @congw729
-
-# Engine and model_executor
-/vllm_omni/core/ @tzhouam @yinpeiqi @fake0fan
-/vllm_omni/engine/ @tzhouam @yinpeiqi @fake0fan
-/vllm_omni/entrypoints/ @tzhouam @yinpeiqi @fake0fan
-/vllm_omni/model_executor/ @tzhouam @gcanlin
-/vllm_omni/worker/ @tzhouam @gcanlin
-
-# Diffusion
-/vllm_omni/diffusion/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @david6666666
-/vllm_omni/lora/ @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @yuanheng-zhao
-
-# Diffusion CPU offloading
-/vllm_omni/diffusion/offloader/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @yuanheng-zhao
-
-# Omni and TTS support
-/vllm_omni/transformers_utils/ @tzhouam @yuanheng-zhao
-/vllm_omni/model_executor/models/ @linyueqian @ZeldaHuang @princepride @yuanheng-zhao @amy-why-3459 @Bounty-hunter
-
-# Benchmarks
-/benchmarks/accuracy/ @david6666666
-/benchmarks/build_dataset/ @linyueqian @yuanheng-zhao
-/benchmarks/diffusion/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa
-/benchmarks/distributed/ @princepride
-/benchmarks/glm_image/ @princepride @JaredforReal
-/benchmarks/tts/ @linyueqian @ZeldaHuang @princepride
-
-# Usability and Observability
-/vllm_omni/deploy/ @tzhouam @lishunyang12
-/vllm_omni/metrics/ @lishunyang12
-/vllm_omni/profiler/ @gcanlin
-
-# ComfyUI
-/apps/ @fhfuih
-
-# Large Scale Serving
-/vllm_omni/distributed/ @princepride @yuanheng-zhao @xuechendi @natureofnature @fake0fan
-
-# Quantization and configuration
-/vllm_omni/config/ @lishunyang12
-/vllm_omni/quantization/ @david6666666 @Isotr0py @lishunyang12
-
-# Hardware plugins and accelerators
-/vllm_omni/platforms/ @gcanlin @tjtanaa @xuechendi
+# Example: reviewers for specific directories
+# /examples/offline_inference/ @offline-inference-team
+# /vllm_omni/core/ @core-team
+# /vllm_omni/model_executor/ @model-team

.github/ISSUE_TEMPLATE/400-bug-report.yml

@@ -74,21 +74,28 @@ body:
       If relevant, add a minimal example so that we can reproduce the error by running the code. It is very important for the snippet to be as succinct (minimal) as possible, so please take time to trim down any irrelevant code to help us debug efficiently. We are going to copy-paste your code and we expect to get the same result as you did: avoid any external data, and include the relevant imports, etc. For example:
 
       ```python
-      from vllm_omni.entrypoints.omni import Omni
-      from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-      from vllm import SamplingParams
+      from vllm_omni import OmniLLM, create_ar_stage_config, create_dit_stage_config
 
-      omni = Omni(
-          model="Qwen/Qwen-Image",
-          stage_configs_path="/path/to/stage_configs.yaml",
+      # Create stage configurations
+      ar_config = create_ar_stage_config(
+          stage_id=0,
+          model_path="Qwen/Qwen3-0.6B",
+          input_modalities=["text"],
+          output_modalities=["text"]
       )
 
-      prompts = [{"prompt": "A scenic watercolor painting of a lighthouse at sunset"}]
-      sampling_params_list = [
-          SamplingParams(max_tokens=1),
-          OmniDiffusionSamplingParams(num_outputs_per_prompt=1),
-      ]
-      outputs = omni.generate(prompts=prompts, sampling_params_list=sampling_params_list)
+      dit_config = create_dit_stage_config(
+          stage_id=1,
+          model_path="stabilityai/stable-diffusion-2-1",
+          input_modalities=["text"],
+          output_modalities=["image"]
+      )
+
+      # Initialize OmniLLM
+      omni = OmniLLM([ar_config, dit_config])
+
+      # Generate
+      outputs = omni.generate(prompt="A scenic watercolor painting of a lighthouse at sunset")
       ```
 
       If the code is too long (hopefully, it isn't), feel free to put it in a public gist and link it in the issue: https://gist.github.com.

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

.github/PULL_REQUEST_TEMPLATE.md

@@ -1,6 +1,4 @@
 <!-- markdownlint-disable -->
-**Before submitting:** run the precheck-pr skill with code agent for a self-check against project conventions — catches dead code, missing benchmarks, and title format issues.
-
 PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTTOM) HAVE BEEN CONSIDERED.
 
 ## Purpose
@@ -20,5 +18,4 @@ PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTT
 - [ ] (Optional) Release notes update. If your change is user-facing, please update the release notes draft.
 </details>
 
-**BEFORE SUBMITTING,** please read <https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md>
-(anything written below this line will be removed by GitHub Actions)
+**BEFORE SUBMITTING, PLEASE READ <https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md>** (anything written below this line will be removed by GitHub Actions)

.github/mergify.yml

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

@@ -0,0 +1,629 @@
+#!/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/build_wheel.yml

@@ -39,10 +39,10 @@ jobs:
         python-version: ["3.11", "3.12"]
 
     steps:
-      - uses: actions/checkout@v6
+      - uses: actions/checkout@v4
 
       - name: Set up Python ${{ matrix.python-version }}
-        uses: actions/setup-python@v6
+        uses: actions/setup-python@v5
         with:
           python-version: ${{ matrix.python-version }}
 
@@ -57,7 +57,7 @@ jobs:
           bash scripts/build_wheel.sh --python python
 
       - name: Upload Wheel Artifact
-        uses: actions/upload-artifact@v7
+        uses: actions/upload-artifact@v4
         with:
           name: vllm-omni-wheel-py${{ matrix.python-version }}
           path: dist/*.whl

.github/workflows/pr-reviewer.yml

@@ -0,0 +1,62 @@
+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@v4.2.2
+        with:
+          fetch-depth: 0
+
+      - name: Set up Python
+        uses: actions/setup-python@v5.3.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@v4.4.3
+        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

.github/workflows/pre-commit.yml

@@ -18,8 +18,8 @@ jobs:
   pre-commit:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v6
-      - uses: actions/setup-python@v6
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
         with:
           python-version: "3.12"
           cache: 'pip'

.gitignore

@@ -83,9 +83,6 @@ 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:
@@ -158,23 +155,10 @@ cython_debug/
 
 # Claude
 CLAUDE.md
-/.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
+.claude/
 
 # Codex
 AGENTS.md
-.codex
 .codex/
 
 # cursor
@@ -204,7 +188,6 @@ checkpoints/
 # Cache directories
 cache/
 !vllm_omni/diffusion/cache/
-!tests/diffusion/cache/
 .cache/
 diffusion_cache/
 kv_cache/
@@ -239,7 +222,6 @@ data/
 datasets/
 *.csv
 *.json
-!tests/dfx/perf/tests/*.json
 !apps/ComfyUI-vLLM-Omni/example_workflows/*.json
 *.jsonl
 *.parquet
@@ -265,11 +247,3 @@ 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/
-
-# Local Lance precision-matching dump scripts (not part of the repo)
-lance_compare/

.readthedocs.yml

@@ -10,11 +10,6 @@ build:
   jobs:
     post_checkout:
       - git fetch --unshallow || true
-    pre_install:
-      # Pre-install CPU torch so the subsequent `pip install .[docs]` does not
-      # pull ~2 GB of CUDA 13 libs (cuBLAS/cuDNN/NCCL/...) that exceed RTD's
-      # 1 GB memory cap and time out the mkdocs build.
-      - python -m pip install --index-url https://download.pytorch.org/whl/cpu torch
 
 mkdocs:
   configuration: mkdocs.yml

README.md

@@ -9,20 +9,18 @@ Easy, fast, and cheap omni-modality model serving for everyone
 </h3>
 
 <p align="center">
-| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://deepwiki.com/vllm-project/vllm-omni"><b>DeepWiki</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> | <a href="docs/assets/WeChat.jpg"><b>WeChat</b></a> | <a href="https://arxiv.org/abs/2602.02204"><b>Paper</b></a> | <a href="https://docs.google.com/presentation/d/111-L8zF7A1j_YI_cR8JsblofdScdRr2f/edit?usp=sharing&ouid=110473603432222024453&rtpof=true&sd=true"><b>Slides</b></a> |
+| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> | <a href="docs/assets/WeChat.jpg"><b>WeChat</b></a> | <a href="https://arxiv.org/abs/2602.02204"><b>Paper</b></a> | <a href="https://drive.google.com/file/d/12k_DOo2xoEROvd4xQxYNuEi6C2GECoRx/view?usp=drive_link"><b>Slides</b></a> |
 </p>
 
 
 ---
 
 *Latest News* 🔥
-- [2026/06] We released [0.22.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.22.0) - an **omnimodal world-model** release aligned with vLLM 0.22, featuring [Nvidia Cosmos3](recipes/cosmos3/Cosmos3-Nano.md)/DreamZero world model support, expanded quantization coverage across Blackwell/NPU/XPU, TTS production improvements, new models including MiniCPM-o 4.5, MOSS-TTS, and Lance, plus RL integration with [VeRL-Omni](https://github.com/verl-project/verl-omni).
-- [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.
 
 ---

apps/ComfyUI-vLLM-Omni/README.md

@@ -47,7 +47,6 @@ If no, check your shell running the ComfyUI process. There may be some error mes
 This extension offers the following nodes based on the output modalities (at **ComfyUI sidebar -> Node Library**):
 
 - **Generate Image** for text-to-image and image-to-image tasks
-- **Generate Video** for text-to-video and image-to-video tasks
 - **Multimodality Understanding** for multimodality-to-text and multimodality-to-audio tasks
 - **TTS** and **TTS Voice Clone** for TTS tasks
 
@@ -60,7 +59,7 @@ To build a simple workflow yourself,
 
 - Drag a generation node onto the canvas.
 - Depending on your need, grab built-in multimedia file loader nodes, such as **image->Load Image**, **image->video->Load Video**, **audio->Load Audio**
-- Depending on your need, grab built-in multimedia file preview nodes, such as **image->Preview Image**, **image->video->Save Video**, **audio->Preview Audio**, **utils->Preview as Text**.
+- Depending on your need, grab built-in multimedia file preview nodes, such as **image->Preview Image**, **image->video->Save Video**, **audio->Preview Audio**. For text output, you can install [ComfyUI-Custom-Scripts plugin](https://github.com/pythongosssss/ComfyUI-Custom-Scripts/) and grab its **utils->Show Text 🐍** node.
 - If you want to tune sampling parameters, grab corresponding nodes from **vLLM-Omni-> Sampling Params**.
     - For multi-stage models, you can connect multiple **AR Sampling Params** and **Diffusion Sampling Params** nodes to a **Multi-Stage Sampling Params List** node, and connect this node to the generation node.
     - For some multi-stage models like BAGEL, [only one stage's sampling parameters are exposed and tunable via vLLM-Omni's online serving API](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/examples/online_serving/bagel/). Thus, these models are treated as single-stage ones. Please check the vLLM-Omni documentation on how to correctly set each model's sampling parameters.
@@ -105,20 +104,6 @@ You can configure per-stage sampling parameters for multi-stage models.
 > [!TIP]
 > The node automatically choose text-to-image or image-to-image API endpoints depending on whether you connect an image input or not.
 
-### Text-to-video and image-to-video generation (e.g., Wan)
-
-(Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni Video Generation**)
-
-<p align="center">
-  <picture>
-    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-video-generation.jpg">
-    <img alt="vLLM-Omni video generation" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-video-generation.jpg" width=55%>
-  </picture>
-</p>
-
-> [!TIP]
-> The node automatically choose text-to-video or image-to-video API endpoints depending on whether you connect an image input or not.
-
 ### TTS (e.g., Qwen TTS series)
 
 (Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni TTS**)

apps/ComfyUI-vLLM-Omni/__init__.py

@@ -14,21 +14,17 @@ from .comfyui_vllm_omni.nodes import (
     VLLMOmniARSampling,
     VLLMOmniDiffusionSampling,
     VLLMOmniGenerateImage,
-    VLLMOmniGenerateVideo,
     VLLMOmniQwenTTSParams,
-    VLLMOmniRemoteLoRA,
     VLLMOmniSamplingParamsList,
     VLLMOmniTTS,
     VLLMOmniUnderstanding,
     VLLMOmniVoiceClone,
-    VLLMOmniWanParams,
 )
 
 # A dictionary that contains all nodes you want to export with their names
 NODE_CLASS_MAPPINGS = {
     # === Generation ===
     "VLLMOmniGenerateImage": VLLMOmniGenerateImage,
-    "VLLMOmniGenerateVideo": VLLMOmniGenerateVideo,
     "VLLMOmniUnderstanding": VLLMOmniUnderstanding,
     "VLLMOmniTTS": VLLMOmniTTS,
     "VLLMOmniVoiceClone": VLLMOmniVoiceClone,
@@ -36,16 +32,13 @@ NODE_CLASS_MAPPINGS = {
     "VLLMOmniARSampling": VLLMOmniARSampling,
     "VLLMOmniDiffusionSampling": VLLMOmniDiffusionSampling,
     "VLLMOmniSamplingParamsList": VLLMOmniSamplingParamsList,
-    "VLLMOmniRemoteLoRA": VLLMOmniRemoteLoRA,
     "VLLMOmniQwenTTSParams": VLLMOmniQwenTTSParams,
-    "VLLMOmniWanParams": VLLMOmniWanParams,
 }
 
 # A dictionary that contains the friendly/humanly readable titles for the nodes
 NODE_DISPLAY_NAME_MAPPINGS = {
     # === Generation ===
     "VLLMOmniGenerateImage": "Generate Image",
-    "VLLMOmniGenerateVideo": "Generate Video",
     "VLLMOmniUnderstanding": "Multimodality Understanding",
     "VLLMOmniTTS": "TTS (Text to Speech)",
     "VLLMOmniVoiceClone": "TTS Voice Cloning",
@@ -53,9 +46,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "VLLMOmniARSampling": "AR Sampling Params",
     "VLLMOmniDiffusionSampling": "Diffusion Sampling Params",
     "VLLMOmniSamplingParamsList": "Multi-Stage Sampling Params List",
-    "VLLMOmniRemoteLoRA": "LoRA",
     "VLLMOmniQwenTTSParams": "Qwen TTS Params",
-    "VLLMOmniWanParams": "Wan Video Params",
 }
 
 WEB_DIRECTORY = "./web"

apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/nodes.py

@@ -11,7 +11,6 @@ from .utils.types import (
     AutoregressionSamplingParams,
     DiffusionSamplingParams,
     QwenTTSModelSpecificParams,
-    WanModelSpecificParams,
 )
 from .utils.validators import (
     add_sampling_parameters_to_stage,
@@ -57,7 +56,6 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
                 # "video": ("VIDEO",),
                 # "audio": ("AUDIO",),
                 "sampling_params": ("SAMPLING_PARAMS",),
-                "lora": ("REMOTE_LORA",),
             },
         }
 
@@ -78,11 +76,9 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
         audio: AudioInput | None = None,  # Hidden & unused
         video: VideoInput | None = None,  # Hidden & unused
         sampling_params: dict | list[dict] | None = None,
-        lora: dict | None = None,
         **kwargs,
     ):
-        if kwargs:
-            logger.info("Uncaught kwargs: %s", kwargs)
+        logger.info("Uncaught kwargs: %s", kwargs)
         logger.debug("Got sampling params: %s", sampling_params)
         validate_model_and_sampling_params_types(model, sampling_params)
         if image is None and mask is not None:
@@ -109,7 +105,6 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
                     height=height,
                     negative_prompt=negative_prompt,
                     sampling_params=sampling_params,
-                    lora=lora,
                 )
                 return (output,)
             elif image is not None and audio is None and video is None:
@@ -124,7 +119,6 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
                     negative_prompt=negative_prompt,
                     mask=mask,
                     sampling_params=sampling_params,
-                    lora=lora,
                 )
                 return (output,)
 
@@ -142,91 +136,11 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
             audio=audio,
             video=video,
             sampling_params=sampling_params,
-            lora=lora,
         )
 
         return (output,)
 
 
-class VLLMOmniGenerateVideo(_VLLMOmniGenerateBase):
-    @classmethod
-    def INPUT_TYPES(cls):
-        return {
-            "required": {
-                "url": ("STRING", {"default": "http://localhost:8000/v1"}),
-                "model": ("STRING", {"default": "Wan-AI/Wan2.2-T2V-A14B-Diffusers"}),
-                "prompt": ("STRING", {"multiline": True}),
-                "negative_prompt": ("STRING", {"multiline": True, "default": ""}),
-                "width": ("INT", {"default": 832, "min": 1}),
-                "height": ("INT", {"default": 480, "min": 1}),
-                "fps": ("INT", {"default": 16, "min": 1}),
-                "num_frames": ("INT", {"default": 41, "min": 1}),
-            },
-            "optional": {
-                "image": ("IMAGE",),
-                "sampling_params": ("SAMPLING_PARAMS",),
-                "lora": ("REMOTE_LORA",),
-                "model_params": ("VIDEO_PARAMS",),
-            },
-        }
-
-    RETURN_TYPES = ("VIDEO",)
-    RETURN_NAMES = ("video",)
-    FUNCTION = "generate"
-
-    async def generate(
-        self,
-        url: str,
-        model: str,
-        prompt: str,
-        width: int,
-        height: int,
-        fps: int,
-        num_frames: int,
-        negative_prompt: str | None = None,
-        image: torch.Tensor | None = None,
-        sampling_params: dict | list[dict] | None = None,
-        model_params: dict | None = None,
-        lora: dict | None = None,
-        **kwargs,
-    ):
-        if kwargs:
-            logger.info("Uncaught kwargs: %s", kwargs)
-        logger.debug("Got sampling params: %s", sampling_params)
-        logger.debug("Got model params: %s", model_params)
-        validate_model_and_sampling_params_types(model, sampling_params)
-
-        # Currently, all video generation models are single-stage diffusion models
-        if isinstance(sampling_params, list):
-            if len(sampling_params) != 1:
-                raise ValueError(
-                    "Video generation expects a single sampling params group. "
-                    "Please provide one Diffusion sampling node."
-                )
-            sampling_params = sampling_params[0]
-
-        if sampling_params is not None:
-            sampling_params.pop("type", None)  # internal fields
-        if model_params is not None:
-            model_params.pop("type", None)  # internal fields
-
-        client = VLLMOmniClient(url)
-        output = await client.generate_video(
-            model=model,
-            prompt=prompt,
-            image=image,  # image present => i2v, absent => t2v
-            width=width,
-            height=height,
-            num_frames=num_frames,
-            fps=fps,
-            negative_prompt=negative_prompt,
-            sampling_params=sampling_params,
-            lora=lora,
-            model_params=model_params,
-        )
-        return (output,)
-
-
 class VLLMOmniUnderstanding(_VLLMOmniGenerateBase):
     @classmethod
     def INPUT_TYPES(cls):
@@ -274,8 +188,7 @@ class VLLMOmniUnderstanding(_VLLMOmniGenerateBase):
         use_audio_in_video: bool = True,
         **kwargs,
     ) -> tuple[str, AudioInput]:
-        if kwargs:
-            logger.info("Uncaught kwargs: %s", kwargs)
+        logger.info("Uncaught kwargs: %s", kwargs)
         logger.debug("Got sampling params: %s", sampling_params)
         validate_model_and_sampling_params_types(model, sampling_params)
 
@@ -634,52 +547,6 @@ class VLLMOmniSamplingParamsList:
         return (params,)
 
 
-class VLLMOmniRemoteLoRA:
-    @classmethod
-    def INPUT_TYPES(cls):
-        return {
-            "required": {
-                "local_path": ("STRING", {"default": ""}),
-                "name": ("STRING", {"default": ""}),
-                "scale": (
-                    "FLOAT",
-                    {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.1},
-                ),
-                "int_id": (
-                    "INT",
-                    {
-                        "default": 0,
-                        "min": 0,
-                        "step": 1,
-                        "tooltip": "0 means it is not set and the server can derive it.",
-                    },
-                ),
-            }
-        }
-
-    RETURN_TYPES = ("REMOTE_LORA",)
-    RETURN_NAMES = ("lora",)
-    FUNCTION = "get_lora"
-    CATEGORY = "vLLM-Omni"
-
-    @classmethod
-    def VALIDATE_INPUTS(cls, local_path, name) -> str | Literal[True]:
-        if not local_path.strip() or not name.strip():
-            return "Both local_path and name must be provided."
-        return True
-
-    def get_lora(self, local_path: str, name: str, scale: float, int_id: int):
-        local_path = local_path.strip()
-        name = name.strip()
-        lora = {
-            "local_path": local_path or None,
-            "name": name or None,
-            "scale": float(scale),
-            "int_id": int(int_id) if int_id > 0 else None,
-        }
-        return (lora,)
-
-
 class VLLMOmniQwenTTSParams:
     @classmethod
     def INPUT_TYPES(cls):
@@ -705,32 +572,3 @@ class VLLMOmniQwenTTSParams:
 
     def get_params(self, **kwargs):
         return (QwenTTSModelSpecificParams(kwargs),)
-
-
-class VLLMOmniWanParams:
-    @classmethod
-    def INPUT_TYPES(cls):
-        return {
-            "required": {
-                "guidance_scale_2": (
-                    "FLOAT",
-                    {"default": 4.0, "min": 0.0, "max": 20.0, "step": 0.1},
-                ),
-                "boundary_ratio": (
-                    "FLOAT",
-                    {"default": 0.875, "min": 0.0, "max": 1.0, "step": 0.001},
-                ),
-                "flow_shift": (
-                    "FLOAT",
-                    {"default": 5.0, "min": 0.0, "max": 100.0, "step": 0.1},
-                ),
-            }
-        }
-
-    RETURN_TYPES = ("VIDEO_PARAMS",)
-    RETURN_NAMES = ("Wan video params",)
-    FUNCTION = "get_params"
-    CATEGORY = "vLLM-Omni/Video Params"
-
-    def get_params(self, **kwargs):
-        return (WanModelSpecificParams(kwargs),)

apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/api_client.py

@@ -6,8 +6,6 @@ The image generation part is derived from dougbtv/comfyui-vllm-omni by Doug (@do
 Original source at https://github.com/dougbtv/comfyui-vllm-omni, distributed under the MIT License.
 """
 
-import asyncio
-import json
 from typing import Any
 
 import aiohttp
@@ -20,7 +18,6 @@ from .format import (
     base64_to_audio,
     base64_to_image_tensor,
     bytes_to_audio,
-    bytes_to_video,
     image_tensor_to_base64,
     image_tensor_to_png_bytes,
     video_to_base64,
@@ -32,43 +29,10 @@ from .types import AudioFormat
 logger = get_logger(__name__)
 
 
-async def url_json(session: aiohttp.ClientSession, url: str, verb: str = "get", **kwargs) -> dict[str, Any]:
-    try:
-        async with getattr(session, verb)(url, **kwargs) as response:
-            if not response.ok:
-                error_text = await response.text()
-                raise (ValueError if response.status < 500 else RuntimeError)(
-                    f"vLLM-Omni API returned status {response.status}: {error_text}"
-                )
-            try:
-                return await response.json()
-            except aiohttp.ContentTypeError as e:
-                raise RuntimeError(f"Invalid JSON response from vLLM-Omni: {e}")
-    except aiohttp.ClientError as e:
-        raise RuntimeError(f"Network error connecting to vLLM-Omni at {url}: {e}")
-
-
-async def url_bytes(session: aiohttp.ClientSession, url: str, verb: str = "get", **kwargs) -> bytes:
-    try:
-        async with getattr(session, verb)(url, **kwargs) as response:
-            if not response.ok:
-                error_text = await response.text()
-                raise (ValueError if response.status < 500 else RuntimeError)(
-                    f"vLLM-Omni API returned status {response.status}: {error_text}"
-                )
-            return await response.read()
-    except aiohttp.ClientError as e:
-        raise RuntimeError(f"Network error connecting to vLLM-Omni at {url}: {e}")
-
-
 class VLLMOmniClient:
-    def __init__(
-        self, base_url: str, timeout: float | None = None, poll_interval: float = 5.0, max_poll_duration: float = 60 * 5
-    ):
+    def __init__(self, base_url: str, timeout: float = 300.0):
         self.base_url = base_url
         self.timeout = aiohttp.ClientTimeout(total=timeout)
-        self.poll_interval = poll_interval
-        self.max_poll_duration = max_poll_duration
 
     async def generate_image(
         self,
@@ -79,13 +43,12 @@ class VLLMOmniClient:
         height: int,
         negative_prompt: str | None = None,
         sampling_params: dict | None = None,
-        lora: dict | None = None,
     ) -> torch.Tensor:
         """Run text-to-image generation via DALLE API"""
         await self._check_model_exist(model)
 
         size = f"{width}x{height}"
-        payload: dict[str, Any] = {
+        payload = {
             "model": model,
             "prompt": prompt,
             "size": size,
@@ -95,8 +58,6 @@ class VLLMOmniClient:
             payload["negative_prompt"] = negative_prompt
         if sampling_params is not None:
             payload.update(sampling_params)
-        if lora is not None:
-            payload["lora"] = lora
         logger.debug("img gen payload: %s", payload)
 
         url = self.base_url + "/images/generations"
@@ -149,7 +110,6 @@ class VLLMOmniClient:
         negative_prompt: str | None = None,
         mask: torch.Tensor | None = None,
         sampling_params: dict | None = None,
-        lora: dict | None = None,
     ) -> torch.Tensor:
         """Run image editing via DALLE API"""
         await self._check_model_exist(model)
@@ -171,8 +131,6 @@ class VLLMOmniClient:
         if sampling_params is not None:
             for k, v in sampling_params.items():
                 form.add_field(k, str(v))
-        if lora is not None:
-            form.add_field("lora", json.dumps(lora, ensure_ascii=False))
         if mask is not None:
             mask_filename = "mask.png"
             form.add_field(
@@ -225,7 +183,6 @@ class VLLMOmniClient:
         audio: AudioInput | None = None,
         video: VideoInput | None = None,
         sampling_params: dict | list[dict] | None = None,
-        lora: dict | None = None,
     ) -> torch.Tensor:
         payload = VLLMOmniClient._prepare_chat_completion_messages(
             model=model,
@@ -236,8 +193,6 @@ class VLLMOmniClient:
             video=video,
             sampling_params=sampling_params,
             modalities=["image"],
-            # === below are additional `extra_body` fields, handled by **kwargs ===
-            lora=lora,
         )
         choices = await self._generate_base_chat_completion(model, payload)
 
@@ -251,87 +206,6 @@ class VLLMOmniClient:
 
         return torch.stack(image_tensors, dim=0)
 
-    async def generate_video(
-        self,
-        *,
-        model: str,
-        prompt: str,
-        width: int,
-        height: int,
-        num_frames: int,
-        fps: int,
-        negative_prompt: str | None = None,
-        image: torch.Tensor | None = None,
-        sampling_params: dict | None = None,
-        model_params: dict | None = None,
-        lora: dict | None = None,
-        **extra_body,
-    ) -> VideoInput:
-        form = aiohttp.FormData()
-        form.add_field("model", model)
-        form.add_field("prompt", prompt)
-        form.add_field("width", str(width))
-        form.add_field("height", str(height))
-        form.add_field("num_frames", str(num_frames))
-        form.add_field("fps", str(fps))
-        if negative_prompt:
-            form.add_field("negative_prompt", negative_prompt)
-        if sampling_params is not None:
-            for k, v in sampling_params.items():
-                form.add_field(k, str(v))
-        if model_params is not None:
-            for k, v in model_params.items():
-                form.add_field(k, str(v))
-        if lora is not None:
-            form.add_field("lora", json.dumps(lora, ensure_ascii=False))
-        if extra_body:
-            form.add_field("extra_body", json.dumps(extra_body, ensure_ascii=False))
-
-        if image is not None:
-            image_filename = "image.png"  # Required for multipart form
-            form.add_field(
-                "input_reference",
-                image_tensor_to_png_bytes(image, image_filename),
-                filename=image_filename,
-                content_type="image/png",
-            )
-
-        async with aiohttp.ClientSession(timeout=self.timeout) as session:
-            # Start the video generation job
-            url = f"{self.base_url}/videos"
-            data = await url_json(session, url, "post", data=form)
-            if (job_id := data.get("id", None)) is None:
-                raise RuntimeError("API response missing job 'id' field - expected OpenAI compliant format")
-            if (job_status := data.get("status", None)) is None:
-                raise RuntimeError("API response missing job 'status' field - expected OpenAI compliant format")
-
-            # Poll for video generation job completion
-            deadline = asyncio.get_running_loop().time() + self.max_poll_duration
-            url = f"{self.base_url}/videos/{job_id}"
-            while job_status not in {"completed", "failed"}:
-                await asyncio.sleep(self.poll_interval)
-
-                data = await url_json(session, url)
-                if (job_status := data.get("status", None)) is None:
-                    raise RuntimeError("API response missing job 'status' field - expected OpenAI compliant format")
-                if asyncio.get_running_loop().time() >= deadline:
-                    raise RuntimeError(f"Timed out waiting for video job {job_id} to complete")
-
-            if job_status == "failed":
-                raise RuntimeError(f"Video job failed: {data}")
-
-            # Retrieve completed content
-            video_bytes = await url_bytes(session, f"{url}/content")
-
-            # Decode video and make a best effort at cleaning up server resources
-            try:
-                return bytes_to_video(video_bytes)
-            finally:
-                try:
-                    await url_json(session, url, "delete")
-                except Exception as exc:
-                    logger.warning("Failed to clean up video job %s: %s", job_id, exc)
-
     async def generate_understanding_chat_completion(
         self,
         *,

apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/format.py

@@ -6,17 +6,12 @@ Original source at https://github.com/dougbtv/comfyui-vllm-omni, distributed und
 
 import base64
 import mimetypes
-from fractions import Fraction
 from io import BytesIO
 
 import av
 import numpy as np
 import torch
-from av.audio.frame import AudioFrame
-from av.audio.resampler import AudioResampler
-from av.video.frame import VideoFrame
 from comfy_api.input import AudioInput, VideoInput
-from comfy_api.latest import InputImpl, Types
 from comfy_extras import nodes_audio
 from PIL import Image
 
@@ -143,75 +138,6 @@ def video_to_base64(video: VideoInput, filename: str = "video.mp4") -> str:
     return f"data:{mime_type};base64,{base64_str}"
 
 
-def bytes_to_video(video_bytes: bytes) -> VideoInput:
-    video_buffer = BytesIO(video_bytes)
-
-    try:
-        with av.open(video_buffer, mode="r") as container:
-            video_stream = next((s for s in container.streams if s.type == "video"), None)
-            if video_stream is None:
-                raise ValueError("No video stream found in decoded payload.")
-
-            frames = []
-            for frame in container.decode(video_stream):
-                if not isinstance(frame, VideoFrame):
-                    continue
-                image = frame.to_ndarray(format="rgb24")
-                frames.append(torch.from_numpy(image).float() / 255.0)
-
-            if len(frames) == 0:
-                raise ValueError("No video frames found in decoded payload.")
-
-            images = torch.stack(frames, dim=0)
-            frame_rate = Fraction(video_stream.average_rate) if video_stream.average_rate else Fraction(1)
-
-            audio: AudioInput | None = None
-            if len(container.streams.audio):
-                audio_stream = container.streams.audio[-1]
-                audio_frames = []
-                resampler = AudioResampler(format="fltp")
-                for frame in container.decode(audio_stream):
-                    if not isinstance(frame, AudioFrame):
-                        continue
-                    resampled = resampler.resample(frame)
-                    if not isinstance(resampled, list):
-                        resampled = [resampled]
-                    for audio_frame in resampled:
-                        if audio_frame is not None:
-                            audio_frames.append(audio_frame.to_ndarray())
-
-                if len(audio_frames) > 0:
-                    audio_data = np.concatenate(audio_frames, axis=1)
-                    sample_rate = int(audio_stream.sample_rate) if audio_stream.sample_rate else 1
-                    audio = {
-                        "waveform": torch.from_numpy(audio_data).unsqueeze(0),
-                        "sample_rate": sample_rate,
-                    }
-
-            components = Types.VideoComponents(
-                images=images,
-                frame_rate=frame_rate,
-                audio=audio,
-                metadata=container.metadata if container.metadata else None,
-            )
-    except Exception as e:
-        raise RuntimeError(f"Failed to decode video: {e}")
-
-    return InputImpl.VideoFromComponents(components)
-
-
-def base64_to_video(base64_str: str) -> VideoInput:
-    if base64_str.startswith("data:video"):
-        _, base64_str = base64_str.split(",", 1)
-
-    try:
-        video_bytes = base64.b64decode(base64_str)
-    except Exception as e:
-        raise ValueError(f"Invalid base64 string: {e}")
-
-    return bytes_to_video(video_bytes)
-
-
 def audio_to_bytes(audio: AudioInput, filename: str = "audio.mp3", quality: str = "128k") -> BytesIO:
     waveform = audio["waveform"][0]  # Shape: (C, T)
     sample_rate = audio["sample_rate"]

apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/types.py

@@ -23,10 +23,6 @@ class QwenTTSModelSpecificParams(dict):
     pass
 
 
-class WanModelSpecificParams(dict):
-    pass
-
-
 class ModelMode(Enum):
     IMAGE_GENERATION = auto()
     VIDEO_GENERATION = auto()

apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/validators.py

@@ -16,7 +16,7 @@ def validate_model_and_sampling_params_types(
     # Skip if no spec or no sampling params
     pipeline_spec, _ = lookup_model_spec(model_name)
     if pipeline_spec is None:
-        logger.info(f"skipping sampling params check because spec for {model_name} is not found")
+        logger.info("skipping sampling params check because spec is not found")
         return
     if sampling_param_list is None:
         return
@@ -70,17 +70,13 @@ def add_sampling_parameters_to_stage(
 
     stages = pipeline_spec["stages"]
     if isinstance(sampling_param_list, dict):
-        sampling_param_list = sampling_param_list.__class__(sampling_param_list)
         sampling_param_list.update(params_to_add)
     elif sampling_param_list is None:
         sampling_param_list = params_to_add.copy()
     else:
         for i, stage in enumerate(stages):
             if stage == stage_type:
-                stage_param = sampling_param_list[i]
-                stage_param = stage_param.__class__(stage_param)
-                stage_param.update(params_to_add)
-                sampling_param_list[i] = stage_param
+                sampling_param_list[i].update(params_to_add)
 
     return sampling_param_list
 

apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni Video Generation.json

@@ -1,513 +0,0 @@
-{
-    "id": "3aea1d08-7de5-4997-ba71-b2f66ad72a21",
-    "revision": 0,
-    "last_node_id": 14,
-    "last_link_id": 28,
-    "nodes": [
-        {
-            "id": 10,
-            "type": "MarkdownNote",
-            "pos": [
-                227.99306819575278,
-                -231.24143306069843
-            ],
-            "size": [
-                240.3326478922474,
-                136.3505820791881
-            ],
-            "flags": {},
-            "order": 0,
-            "mode": 0,
-            "inputs": [],
-            "outputs": [],
-            "title": "Note: Task and Input",
-            "properties": {},
-            "widgets_values": [
-                "vLLM-Omni nodes are categorized based on the output modality. The \"Generate Video\" node supports both text-to-video generation and image-to-video generation. The node will create corresponding payload depending on whether an input image is present or not."
-            ],
-            "color": "#432",
-            "bgcolor": "#000"
-        },
-        {
-            "id": 4,
-            "type": "LoadImage",
-            "pos": [
-                496.31859627609606,
-                -229.71277089860084
-            ],
-            "size": [
-                270,
-                314
-            ],
-            "flags": {},
-            "order": 1,
-            "mode": 0,
-            "inputs": [
-                {
-                    "localized_name": "image",
-                    "name": "image",
-                    "type": "COMBO",
-                    "widget": {
-                        "name": "image"
-                    },
-                    "link": null
-                },
-                {
-                    "localized_name": "choose file to upload",
-                    "name": "upload",
-                    "type": "IMAGEUPLOAD",
-                    "widget": {
-                        "name": "upload"
-                    },
-                    "link": null
-                }
-            ],
-            "outputs": [
-                {
-                    "localized_name": "IMAGE",
-                    "name": "IMAGE",
-                    "type": "IMAGE",
-                    "links": [
-                        25
-                    ]
-                },
-                {
-                    "localized_name": "MASK",
-                    "name": "MASK",
-                    "type": "MASK",
-                    "links": null
-                }
-            ],
-            "properties": {
-                "Node name for S&R": "LoadImage"
-            },
-            "widgets_values": [
-                "cute-cat.jpg",
-                "image"
-            ]
-        },
-        {
-            "id": 8,
-            "type": "VLLMOmniDiffusionSampling",
-            "pos": [
-                465.5140218220927,
-                148.80072646828987
-            ],
-            "size": [
-                284.205078125,
-                226
-            ],
-            "flags": {},
-            "order": 2,
-            "mode": 0,
-            "inputs": [
-                {
-                    "localized_name": "n",
-                    "name": "n",
-                    "type": "INT",
-                    "widget": {
-                        "name": "n"
-                    },
-                    "link": null
-                },
-                {
-                    "localized_name": "num_inference_steps",
-                    "name": "num_inference_steps",
-                    "type": "INT",
-                    "widget": {
-                        "name": "num_inference_steps"
-                    },
-                    "link": null
-                },
-                {
-                    "localized_name": "guidance_scale",
-                    "name": "guidance_scale",
-                    "type": "FLOAT",
-                    "widget": {
-                        "name": "guidance_scale"
-                    },
-                    "link": null
-                },
-                {
-                    "localized_name": "true_cfg_scale",
-                    "name": "true_cfg_scale",
-                    "type": "FLOAT",
-                    "widget": {
-                        "name": "true_cfg_scale"
-                    },
-                    "link": null
-                },
-                {
-                    "localized_name": "vae_use_slicing",
-                    "name": "vae_use_slicing",
-                    "type": "BOOLEAN",
-                    "widget": {
-                        "name": "vae_use_slicing"
-                    },
-                    "link": null
-                },
-                {
-                    "localized_name": "vae_use_tiling",
-                    "name": "vae_use_tiling",
-                    "type": "BOOLEAN",
-                    "widget": {
-                        "name": "vae_use_tiling"
-                    },
-                    "link": null
-                },
-                {
-                    "localized_name": "seed",
-                    "name": "seed",
-                    "type": "INT",
-                    "widget": {
-                        "name": "seed"
-                    },
-                    "link": null
-                }
-            ],
-            "outputs": [
-                {
-                    "localized_name": "diffusion sampling params",
-                    "name": "diffusion sampling params",
-                    "type": "SAMPLING_PARAMS",
-                    "links": [
-                        27
-                    ]
-                }
-            ],
-            "properties": {
-                "Node name for S&R": "VLLMOmniDiffusionSampling"
-            },
-            "widgets_values": [
-                1,
-                50,
-                1,
-                1,
-                false,
-                false,
-                1992,
-                "randomize"
-            ]
-        },
-        {
-            "id": 14,
-            "type": "VLLMOmniWanParams",
-            "pos": [
-                518.6005521845747,
-                428.4455663719664
-            ],
-            "size": [
-                270,
-                106
-            ],
-            "flags": {},
-            "order": 3,
-            "mode": 0,
-            "inputs": [
-                {
-                    "localized_name": "guidance_scale_2",
-                    "name": "guidance_scale_2",
-                    "type": "FLOAT",
-                    "widget": {
-                        "name": "guidance_scale_2"
-                    },
-                    "link": null
-                },
-                {
-                    "localized_name": "boundary_ratio",
-                    "name": "boundary_ratio",
-                    "type": "FLOAT",
-                    "widget": {
-                        "name": "boundary_ratio"
-                    },
-                    "link": null
-                },
-                {
-                    "localized_name": "flow_shift",
-                    "name": "flow_shift",
-                    "type": "FLOAT",
-                    "widget": {
-                        "name": "flow_shift"
-                    },
-                    "link": null
-                }
-            ],
-            "outputs": [
-                {
-                    "localized_name": "Wan video params",
-                    "name": "Wan video params",
-                    "type": "VIDEO_PARAMS",
-                    "links": [
-                        28
-                    ]
-                }
-            ],
-            "properties": {
-                "Node name for S&R": "VLLMOmniWanParams"
-            },
-            "widgets_values": [
-                4,
-                0.875,
-                5
-            ]
-        },
-        {
-            "id": 12,
-            "type": "VLLMOmniGenerateVideo",
-            "pos": [
-                827.2566336087867,
-                -73.77449831827545
-            ],
-            "size": [
-                400,
-                346
-            ],
-            "flags": {},
-            "order": 4,
-            "mode": 0,
-            "inputs": [
-                {
-                    "localized_name": "image",
-                    "name": "image",
-                    "shape": 7,
-                    "type": "IMAGE",
-                    "link": 25
-                },
-                {
-                    "localized_name": "sampling_params",
-                    "name": "sampling_params",
-                    "shape": 7,
-                    "type": "SAMPLING_PARAMS",
-                    "link": 27
-                },
-                {
-                    "localized_name": "lora",
-                    "name": "lora",
-                    "shape": 7,
-                    "type": "REMOTE_LORA",
-                    "link": null
-                },
-                {
-                    "localized_name": "model_params",
-                    "name": "model_params",
-                    "shape": 7,
-                    "type": "VIDEO_PARAMS",
-                    "link": 28
-                },
-                {
-                    "localized_name": "url",
-                    "name": "url",
-                    "type": "STRING",
-                    "widget": {
-                        "name": "url"
-                    },
-                    "link": null
-                },
-                {
-                    "localized_name": "model",
-                    "name": "model",
-                    "type": "STRING",
-                    "widget": {
-                        "name": "model"
-                    },
-                    "link": null
-                },
-                {
-                    "localized_name": "prompt",
-                    "name": "prompt",
-                    "type": "STRING",
-                    "widget": {
-                        "name": "prompt"
-                    },
-                    "link": null
-                },
-                {
-                    "localized_name": "negative_prompt",
-                    "name": "negative_prompt",
-                    "type": "STRING",
-                    "widget": {
-                        "name": "negative_prompt"
-                    },
-                    "link": null
-                },
-                {
-                    "localized_name": "width",
-                    "name": "width",
-                    "type": "INT",
-                    "widget": {
-                        "name": "width"
-                    },
-                    "link": null
-                },
-                {
-                    "localized_name": "height",
-                    "name": "height",
-                    "type": "INT",
-                    "widget": {
-                        "name": "height"
-                    },
-                    "link": null
-                },
-                {
-                    "localized_name": "fps",
-                    "name": "fps",
-                    "type": "INT",
-                    "widget": {
-                        "name": "fps"
-                    },
-                    "link": null
-                },
-                {
-                    "localized_name": "num_frames",
-                    "name": "num_frames",
-                    "type": "INT",
-                    "widget": {
-                        "name": "num_frames"
-                    },
-                    "link": null
-                }
-            ],
-            "outputs": [
-                {
-                    "localized_name": "video",
-                    "name": "video",
-                    "type": "VIDEO",
-                    "links": [
-                        26
-                    ]
-                }
-            ],
-            "properties": {
-                "Node name for S&R": "VLLMOmniGenerateVideo"
-            },
-            "widgets_values": [
-                "http://localhost:8000/v1",
-                "Wan-AI/Wan2.2-TI2V-5B-Diffusers",
-                "Make the cat in the attached video yawning. Eye closed and mouth wide open. Keep the original realistic style.",
-                "Cartoonish",
-                900,
-                1200,
-                12,
-                41
-            ]
-        },
-        {
-            "id": 13,
-            "type": "SaveVideo",
-            "pos": [
-                1257.5441369501657,
-                -44.12956394419863
-            ],
-            "size": [
-                374,
-                372.34896826553796
-            ],
-            "flags": {},
-            "order": 5,
-            "mode": 0,
-            "inputs": [
-                {
-                    "localized_name": "video",
-                    "name": "video",
-                    "type": "VIDEO",
-                    "link": 26
-                },
-                {
-                    "localized_name": "filename_prefix",
-                    "name": "filename_prefix",
-                    "type": "STRING",
-                    "widget": {
-                        "name": "filename_prefix"
-                    },
-                    "link": null
-                },
-                {
-                    "localized_name": "format",
-                    "name": "format",
-                    "type": "COMBO",
-                    "widget": {
-                        "name": "format"
-                    },
-                    "link": null
-                },
-                {
-                    "localized_name": "codec",
-                    "name": "codec",
-                    "type": "COMBO",
-                    "widget": {
-                        "name": "codec"
-                    },
-                    "link": null
-                }
-            ],
-            "outputs": [],
-            "properties": {},
-            "widgets_values": [
-                "video/ComfyUI",
-                "auto",
-                "auto"
-            ]
-        }
-    ],
-    "links": [
-        [
-            25,
-            4,
-            0,
-            12,
-            0,
-            "IMAGE"
-        ],
-        [
-            26,
-            12,
-            0,
-            13,
-            0,
-            "VIDEO"
-        ],
-        [
-            27,
-            8,
-            0,
-            12,
-            1,
-            "SAMPLING_PARAMS"
-        ],
-        [
-            28,
-            14,
-            0,
-            12,
-            3,
-            "VIDEO_PARAMS"
-        ]
-    ],
-    "groups": [
-        {
-            "id": 1,
-            "title": "Input",
-            "bounding": [
-                213.1706010087147,
-                -313.20095750667554,
-                560.8246471437027,
-                407.4246532472182
-            ],
-            "color": "#3f789e",
-            "font_size": 24,
-            "flags": {}
-        }
-    ],
-    "config": {},
-    "extra": {
-        "workflowRendererVersion": "LG",
-        "ds": {
-            "scale": 1.1469075819486894,
-            "offset": [
-                101.7632950847089,
-                492.29121889347806
-            ]
-        }
-    },
-    "version": 0.4
-}

benchmarks/README.md

@@ -1,60 +1,42 @@
-# Benchmarks
+# Benchmarks Overview and Architecture
 
-This directory contains benchmark suites for evaluating different model families and infrastructure components in vLLM-Omni. Each subfolder targets a different benchmark family with its own scripts, configs, and metrics. See the per-subfolder READMEs for detailed usage.
+This document explains the benchmark architecture across all benchmark assets in this repo. It describes what we measure, and where to find or plug in new scenarios. Per-task details remain in subfolder READMEs (e.g., `benchmarks/<model>/README.md`).
 
-## Benchmark families
+## Scope and goals
+- Establish repeatable latency/throughput measurements for multimodal LLM pipelines.
+- Provide both HF Transformers (offline) and vLLM-Omni (multi-stage/pipeline) baselines.
+- Make it easy to plug in new datasets and models with minimal changes to the runner scripts.
 
-### [TTS](tts/README.md) — Text-to-Speech
+## Dataset and inputs
+- Default example: SeedTTS top-100 prompts (`benchmarks/build_dataset/top100.txt`) via `benchmarks/build_dataset/`.
+- Extensible: drop in new prompt files or modality-aligned payloads; keep the expected format for the consuming scripts (e.g., one prompt per line).
+- If you add a new dataset, document it under `benchmarks/<model>/README.md` and point scripts to your data path.
 
-Model-agnostic serving benchmarks for TTS models, including Qwen3-TTS and VoxCPM2.
+## 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.
 
-- **Layout**: `tts/bench_tts.py` (serving benchmark driver), `tts/model_configs.yaml` (model registry), `tts/plot_results.py` (result plotting)
-- **Dataset**: Seed-TTS full or text-only datasets, plus bundled smoke/design prompts under `build_dataset/`
-- **Key metrics**: TTFP (time to first audio packet), E2E latency, RTF (real-time factor), throughput (audio seconds / wall-clock second)
+## Reference workflows
+- **HF Transformers (offline, single process)**  
+  Script (example): `benchmarks/<model>/transformers/eval_qwen3_moe_omni_transformers.sh`  
+  Outputs: `benchmark_results/perf_stats.json`, `benchmark_results/results.json`, `benchmark_results/audio/` (if audio is produced).
 
-### [Diffusion](diffusion/README.md) — Image and Video Generation
+- **vLLM-Omni end-to-end pipeline**  
+  Script (example): `benchmarks/<model>/vllm_omni/eval_qwen3_moe_omni.sh`  
+  Outputs: `vllm_omni/logs/*.stats.jsonl` (per-stage/overall latency & TPS), `vllm_omni/logs/stage*.log`, `vllm_omni/outputs/` (text/audio artifacts).
 
-Online-serving benchmark for diffusion image/video models, sending requests to the configured vLLM serving endpoint (`/v1/chat/completions`, `/v1/images/generations`, or `/v1/videos`, depending on backend/task).
+- **Adding a new task/model**  
+  1) Create `benchmarks/<model>/transformers/` and/or `benchmarks/<model>/vllm_omni/` with scripts referencing your model and dataset.  
+  2) Add a task README describing dataset, configs, and expected outputs.  
+  3) Keep the output/log structure similar for easy comparison (perf_stats/results/audio or text outputs; stats.jsonl/logs for pipeline).
 
-- **Tasks**: text-to-image, text-to-video, image-to-image, image-to-video, text+image-to-image, text+image-to-video
-- **Datasets**: `vbench`, `trace`, `random`
-- **Key metrics**: request throughput, latency percentiles, optional SLO attainment
+## Metrics to watch
+- **Throughput**: `overall_tps`, `*_tps_avg` per stage.
+- **Latency distribution**: look for long tails in `*.stats.jsonl`.
+- **Quality/completeness**: missing outputs or errors in stage logs indicate pipeline failures or misconfigurations.
 
-### [GLM-Image](glm_image/README.md) — Text-to-Image and Image-to-Image
-
-Benchmarks for GLM-Image performance across HuggingFace baseline, vLLM-Omni offline inference, and vLLM-Omni online serving.
-
-- **Layout**: `glm_image/huggingface/` (HF baseline), `glm_image/vllm-omni/` (offline inference), `glm_image/benchmark_glm_image.py` (online serving)
-- **Tasks**: text-to-image and image-to-image
-- **Key metrics**: request/image throughput, latency percentiles, optional per-stage pipeline timings
-
-### [Distributed](distributed/omni_connectors/README.md) — RDMA Connector Testing
-
-RDMA environment setup and transfer tests for `MooncakeTransferEngineConnector`, including pytest-based single-node checks and manual cross-node benchmarks.
-
-- **Transfer modes**: `copy`, `zerocopy`, `gpu` (GPUDirect)
-- **Supports**: single-node pytest suites and manual multi-node/cross-node transfer testing
-
-### [Accuracy](accuracy/README.md) — Image Generation and Editing Quality
-
-Accuracy benchmarks for image generation/editing models, adapting external suites to vLLM-Omni serving and local judge-evaluation flows.
-
-- **Layout**: `accuracy/text_to_image/` (GEBench), `accuracy/image_to_image/` (GEdit-Bench)
-- **Method**: generation and judge scoring both run through local `vllm-omni serve` endpoints
-
-### Common serving metrics framework
-
-`vllm_omni/benchmarks/` extends `vllm bench serve --omni` with Omni-specific datasets, backends, and multimodal metrics. Key metrics include:
-
-- **Text output**: TTFT (time to first token), TPOT (time per output token), ITL (inter-token latency)
-- **Audio output**: TTFP (time to first audio packet), E2E latency, RTF (real-time factor)
-- **Throughput**: request throughput, output token throughput, total token throughput, audio throughput
-
-See `vllm_omni/benchmarks/serve.py` for the `vllm bench serve --omni` runner wrapper and `vllm_omni/benchmarks/metrics/` for Omni metric definitions.
-
-## Adding a new benchmark
-
-1. Create a subfolder under `benchmarks/<name>/` with scripts, configs if needed, and a `README.md`.
-2. If comparing against another runtime, use clear backend subfolders where applicable, such as `huggingface/` and `vllm-omni/`, or follow the shared TTS serving benchmark pattern in `tts/`.
-3. Add reusable dataset or prompt-building utilities to `build_dataset/` if applicable.
-4. Update this README with a link to the new benchmark family.
+## Troubleshooting
+- Verify GPU/driver/FlashAttention2 requirements for your chosen model/config.
+- Ensure network access for dataset/model downloads (Google Drive, Hugging Face, etc.).
+- If outputs are missing or slow, inspect per-stage logs and `*.stats.jsonl` for errors, stragglers, or contention.

benchmarks/__init__.py

@@ -1 +0,0 @@
-"""Benchmark helpers and runnable benchmark entrypoints."""

benchmarks/accuracy/README.md

@@ -1,27 +0,0 @@
-# Accuracy Benchmarks
-
-This directory hosts accuracy benchmark integrations that run entirely through a
-local `vllm-omni serve` deployment.
-
-Current integrations:
-
-- `text_to_image/`: GEBench generation + local judge scoring flow.
-- `image_to_image/`: GEdit-Bench generation + local VIEScore-style scoring flow.
-
-Design notes:
-
-- Generation is executed through the OpenAI-compatible endpoints exposed by
-  `vllm-omni serve`.
-- Evaluation is also executed through a local OpenAI-compatible judge model
-  served by `vllm-omni`.
-- Both generation and judge requests accept either `http://host:port` or
-  `http://host:port/v1`.
-- Output directory layout intentionally stays close to the upstream repos.
-
-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 `full_model +
-  benchmark + L4` for nightly when GPU capacity is available.

benchmarks/accuracy/__init__.py

@@ -1 +0,0 @@
-"""Accuracy benchmark integrations for vLLM-Omni."""

benchmarks/accuracy/common.py

@@ -1,208 +0,0 @@
-from __future__ import annotations
-
-import base64
-import io
-import json
-from pathlib import Path
-from typing import Any
-
-import requests
-from PIL import Image
-
-
-def ensure_dir(path: Path) -> Path:
-    path.mkdir(parents=True, exist_ok=True)
-    return path
-
-
-def load_json(path: Path) -> dict[str, Any]:
-    with path.open("r", encoding="utf-8") as handle:
-        return json.load(handle)
-
-
-def write_json(path: Path, payload: dict[str, Any]) -> None:
-    ensure_dir(path.parent)
-    with path.open("w", encoding="utf-8") as handle:
-        json.dump(payload, handle, indent=2, ensure_ascii=False)
-
-
-def save_image(path: Path, image: Image.Image) -> None:
-    ensure_dir(path.parent)
-    image.save(path)
-
-
-def find_first_image(folder: Path, stem: str | None = None) -> Path | None:
-    patterns = [f"{stem}.*"] if stem else ["*.png", "*.jpg", "*.jpeg", "*.webp"]
-    for pattern in patterns:
-        for candidate in sorted(folder.glob(pattern)):
-            if candidate.suffix.lower() in {".png", ".jpg", ".jpeg", ".webp"}:
-                return candidate
-    return None
-
-
-def extract_json_object(raw_text: str) -> dict[str, Any]:
-    raw_text = raw_text.strip()
-    delimiter = "||V^=^V||"
-    if raw_text.count(delimiter) >= 2:
-        start = raw_text.find(delimiter) + len(delimiter)
-        end = raw_text.rfind(delimiter)
-        raw_text = raw_text[start:end].strip()
-
-    start = raw_text.find("{")
-    end = raw_text.rfind("}")
-    if start == -1 or end == -1 or end < start:
-        raise ValueError(f"Could not find JSON object in: {raw_text[:200]}")
-    return json.loads(raw_text[start : end + 1])
-
-
-def build_openai_url(base_url: str, api_path: str) -> str:
-    base = base_url.rstrip("/")
-    normalized_path = api_path if api_path.startswith("/") else f"/{api_path}"
-    if base.endswith(normalized_path):
-        return base
-    if base.endswith("/v1"):
-        return f"{base}{normalized_path}"
-    return f"{base}/v1{normalized_path}"
-
-
-def pil_to_base64(image: Image.Image, image_format: str = "PNG") -> str:
-    buffer = io.BytesIO()
-    image.save(buffer, format=image_format)
-    return base64.b64encode(buffer.getvalue()).decode("utf-8")
-
-
-def pil_to_data_url(image: Image.Image, image_format: str = "PNG") -> str:
-    return f"data:image/{image_format.lower()};base64,{pil_to_base64(image, image_format=image_format)}"
-
-
-def decode_base64_image(encoded: str) -> Image.Image:
-    image = Image.open(io.BytesIO(base64.b64decode(encoded)))
-    image.load()
-    return image.convert("RGB")
-
-
-def pil_to_png_bytes(image: Image.Image) -> bytes:
-    buffer = io.BytesIO()
-    image.save(buffer, format="PNG")
-    return buffer.getvalue()
-
-
-class VllmOmniImageClient:
-    """Thin OpenAI-compatible image client for vLLM-Omni serving."""
-
-    def __init__(self, base_url: str, api_key: str = "EMPTY", timeout: int = 600):
-        self.base_url = base_url.rstrip("/")
-        self.api_key = api_key
-        self.timeout = timeout
-
-    @property
-    def _headers(self) -> dict[str, str]:
-        return {
-            "Authorization": f"Bearer {self.api_key}",
-            "Content-Type": "application/json",
-        }
-
-    def generate_text_to_image(
-        self,
-        *,
-        model: str,
-        prompt: str,
-        width: int,
-        height: int,
-        num_inference_steps: int = 20,
-        guidance_scale: float | None = None,
-        seed: int | None = None,
-        output_compression: int | None = None,
-    ) -> Image.Image:
-        payload: dict[str, Any] = {
-            "model": model,
-            "prompt": prompt,
-            "n": 1,
-            "size": f"{width}x{height}",
-            "response_format": "b64_json",
-            "num_inference_steps": num_inference_steps,
-        }
-        if guidance_scale is not None:
-            payload["guidance_scale"] = guidance_scale
-        if seed is not None:
-            payload["seed"] = seed
-        if output_compression is not None:
-            payload["output_compression"] = output_compression
-
-        response = requests.post(
-            build_openai_url(self.base_url, "/images/generations"),
-            json=payload,
-            headers=self._headers,
-            timeout=self.timeout,
-        )
-        response.raise_for_status()
-        return decode_base64_image(response.json()["data"][0]["b64_json"])
-
-    def generate_image_edit(
-        self,
-        *,
-        model: str,
-        prompt: str,
-        images: Image.Image | list[Image.Image],
-        width: int,
-        height: int,
-        num_inference_steps: int = 20,
-        guidance_scale: float | None = None,
-        seed: int | None = None,
-        negative_prompt: str | None = None,
-        output_compression: int | None = None,
-        bot_task: str | None = None,
-        sys_type: str | None = None,
-        system_prompt: str | None = None,
-    ) -> Image.Image:
-        if not isinstance(images, list):
-            images = [images]
-        data: dict[str, Any] = {
-            "model": model,
-            "prompt": prompt,
-            "n": 1,
-            "size": f"{width}x{height}",
-            "response_format": "b64_json",
-            "num_inference_steps": str(num_inference_steps),
-        }
-        if guidance_scale is not None:
-            data["guidance_scale"] = str(guidance_scale)
-        if seed is not None:
-            data["seed"] = str(seed)
-        if negative_prompt:
-            data["negative_prompt"] = negative_prompt
-        if output_compression is not None:
-            data["output_compression"] = str(output_compression)
-        if bot_task is not None:
-            data["bot_task"] = bot_task
-        if sys_type is not None:
-            data["sys_type"] = sys_type
-        if system_prompt is not None:
-            data["system_prompt"] = system_prompt
-
-        files = [
-            (
-                "image[]" if len(images) > 1 else "image",
-                (f"image_{index}.png", pil_to_png_bytes(image), "image/png"),
-            )
-            for index, image in enumerate(images)
-        ]
-
-        edit_paths = ["/images/edits", "/images/edit"]
-        last_response: requests.Response | None = None
-        for api_path in edit_paths:
-            response = requests.post(
-                build_openai_url(self.base_url, api_path),
-                data=data,
-                files=files,
-                headers={"Authorization": f"Bearer {self.api_key}"},
-                timeout=self.timeout,
-            )
-            last_response = response
-            if response.status_code != 404:
-                response.raise_for_status()
-                return decode_base64_image(response.json()["data"][0]["b64_json"])
-
-        assert last_response is not None
-        last_response.raise_for_status()
-        raise ValueError("No image payload returned from image edit endpoint")

benchmarks/accuracy/image_to_image/README.md

@@ -1,103 +0,0 @@
-# GEdit-Bench on vLLM-Omni
-
-This integration adapts the upstream `stepfun-ai/Step1X-Edit/GEdit-Bench`
-evaluation flow into `vllm-omni/benchmarks/accuracy/image_to_image`.
-
-Upstream mapping:
-
-- `run_gedit_score.py` -> `run_gedit_bench.py evaluate`
-- `calculate_statistics.py` -> `run_gedit_bench.py summarize`
-- upstream output layout under `results/<method>/fullset/...` is preserved
-
-What changed:
-
-- The upstream repo mainly ships evaluation scripts. This integration adds a
-  generation runner that uses the local `vllm-omni` OpenAI-compatible
-  `/v1/images/edits` endpoint to produce benchmark outputs in the expected
-  directory structure.
-- The evaluator keeps the same VIEScore-style decomposition:
-  - `semantics_score`
-  - `quality_score`
-  - `overall_score = sqrt(semantics_score * quality_score)`
-- Judge calls are routed to a local OpenAI-compatible model served by
-  `vllm-omni`, not a remote provider.
-
-Dataset:
-
-- Default `--dataset-ref` is `stepfun-ai/GEdit-Bench`
-- You can also pass a local dataset directory previously saved with
-  Hugging Face `datasets`
-
-Example usage:
-
-```bash
-python benchmarks/accuracy/image_to_image/run_gedit_bench.py generate \
-  --output-root benchmarks/accuracy/image_to_image/results \
-  --base-url http://127.0.0.1:8000 \
-  --model Qwen/Qwen-Image-Edit \
-  --model-name qwen_image_edit \
-  --dataset-ref stepfun-ai/GEdit-Bench \
-  --task-type all \
-  --instruction-language en
-```
-
-```bash
-python benchmarks/accuracy/image_to_image/run_gedit_bench.py evaluate \
-  --output-root benchmarks/accuracy/image_to_image/results \
-  --model-name qwen_image_edit \
-  --save-dir benchmarks/accuracy/image_to_image/scores \
-  --dataset-ref stepfun-ai/GEdit-Bench \
-  --judge-base-url http://127.0.0.1:8000 \
-  --judge-model Qwen/Qwen2.5-VL-7B-Instruct \
-  --judge-api-key EMPTY
-```
-
-```bash
-python benchmarks/accuracy/image_to_image/run_gedit_bench.py summarize \
-  --csv-path benchmarks/accuracy/image_to_image/scores/qwen_image_edit_all_all_vie_score.csv \
-  --language en
-```
-
-Example summary output:
-
-```json
-{
-  "language": "all",
-  "languages": {
-    "en": {
-      "overall": {"count": 110, "Q_SC": 6.58, "Q_PQ": 5.89, "Q_O": 5.86},
-      "intersection": {"count": 78, "Q_SC": 6.50, "Q_PQ": 5.66, "Q_O": 5.65}
-    },
-    "cn": {
-      "overall": {"count": 110, "Q_SC": 6.90, "Q_PQ": 5.78, "Q_O": 6.11},
-      "intersection": {"count": 63, "Q_SC": 7.22, "Q_PQ": 5.59, "Q_O": 6.28}
-    }
-  }
-}
-```
-
-Example generated images to inspect:
-
-- `benchmarks/accuracy/image_to_image/results/qwen_image_edit/fullset/background_change/en/<sample>.png`
-- `benchmarks/accuracy/image_to_image/results/qwen_image_edit/fullset/text_change/en/<sample>.png`
-- `benchmarks/accuracy/image_to_image/results/qwen_image_edit/fullset/subject-replace/cn/<sample>.png`
-
-Example score artifacts to inspect together with the images:
-
-- `benchmarks/accuracy/image_to_image/scores/qwen_image_edit_all_all_vie_score.csv`
-- `benchmarks/accuracy/image_to_image/scores/qwen_image_edit_all_all_summary.json`
-
-What to expect:
-
-- `Q_SC` measures instruction following and content preservation.
-- `Q_PQ` measures image naturalness and artifact quality.
-- `Q_O` is the combined overall score; higher is better.
-- `overall.count` is the number of evaluated samples for that language, while `intersection.count` is the subset with `intersection_exist == True`.
-
-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 `full_model`, `benchmark`,
-  and `L4` markers for nightly (or `advanced_model` for merge) or introduce a new repo-wide marker explicitly first.

benchmarks/accuracy/image_to_image/__init__.py

@@ -1 +0,0 @@
-"""GEdit-Bench integration for vLLM-Omni."""

benchmarks/accuracy/image_to_image/gedit_bench.py

@@ -1,787 +0,0 @@
-from __future__ import annotations
-
-import argparse
-import csv
-import json
-import logging
-import math
-import statistics
-from collections import defaultdict
-from concurrent.futures import ThreadPoolExecutor, as_completed
-from datetime import datetime, timezone
-from io import BytesIO
-from pathlib import Path
-from typing import Any
-
-import requests
-from PIL import Image
-from tqdm.auto import tqdm
-
-from benchmarks.accuracy.common import (
-    VllmOmniImageClient,
-    build_openai_url,
-    extract_json_object,
-    write_json,
-)
-
-GROUPS = [
-    "background_change",
-    "color_alter",
-    "material_alter",
-    "motion_change",
-    "ps_human",
-    "style_change",
-    "subject-add",
-    "subject-remove",
-    "subject-replace",
-    "text_change",
-    "tone_transfer",
-]
-DEFAULT_SAMPLES_PER_GROUP = 10
-logger = logging.getLogger(__name__)
-
-
-def infer_model_name(model: str) -> str:
-    normalized = model.rstrip("/\\")
-    name = Path(normalized).name
-    return name or normalized
-
-
-def resolve_model_name(*, model_name: str | None, model: str | None = None, output_root: Path | None = None) -> str:
-    if model_name:
-        return model_name
-    if model:
-        return infer_model_name(model)
-    if output_root is not None:
-        candidates = sorted(path.name for path in output_root.iterdir() if path.is_dir())
-        if len(candidates) == 1:
-            return candidates[0]
-        if not candidates:
-            raise ValueError(f"Could not infer model-name from empty output root: {output_root}")
-        raise ValueError(
-            f"Could not infer model-name from output root {output_root}; multiple candidates found: {candidates}"
-        )
-    raise ValueError("model-name is required when it cannot be inferred from model or output-root")
-
-
-def _coerce_score_int(value: Any) -> int:
-    if isinstance(value, bool):
-        raise TypeError(f"Expected numeric score, got bool: {value!r}")
-    if isinstance(value, (int, float)):
-        return int(value)
-    if isinstance(value, str):
-        stripped = value.strip()
-        if stripped.lstrip("-").isdigit():
-            return int(stripped)
-    raise TypeError(f"Expected numeric score, got {type(value).__name__}: {value!r}")
-
-
-def _expand_score_entry(entry: Any) -> list[int]:
-    if isinstance(entry, dict):
-        if "score" in entry:
-            return [_coerce_score_int(entry["score"])]
-        if "value" in entry:
-            return [_coerce_score_int(entry["value"])]
-        numeric_values = [_coerce_score_int(value) for value in entry.values()]
-        if numeric_values:
-            return numeric_values
-        raise TypeError(f"Score dict has no numeric values: {entry!r}")
-    return [_coerce_score_int(entry)]
-
-
-def _normalize_score_list(score: Any) -> list[int]:
-    if isinstance(score, dict):
-        entries: list[Any] = [score]
-    elif isinstance(score, list):
-        entries = score
-    else:
-        entries = [score]
-
-    normalized: list[int] = []
-    for entry in entries:
-        normalized.extend(_expand_score_entry(entry))
-    if not normalized:
-        raise ValueError(f"Empty score list after normalization: {score!r}")
-    return normalized
-
-
-def parse_score_payload(raw_text: str) -> dict[str, Any]:
-    try:
-        parsed = extract_json_object(raw_text)
-    except Exception:
-        stripped = raw_text.strip()
-        if stripped.startswith("[") and stripped.endswith("]"):
-            parsed = {"score": json.loads(stripped), "reasoning": ""}
-        else:
-            raise
-    parsed["score"] = _normalize_score_list(parsed.get("score", []))
-    return parsed
-
-
-def _utc_timestamp() -> str:
-    return datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
-
-
-def summarize_generated_records(records: list[dict[str, Any]]) -> dict[str, Any]:
-    by_task: dict[str, list[dict[str, Any]]] = defaultdict(list)
-    by_language: dict[str, list[dict[str, Any]]] = defaultdict(list)
-    for record in records:
-        by_task[record["task_type"]].append(record)
-        by_language[record["instruction_language"]].append(record)
-
-    return {
-        "count": len(records),
-        "by_task": {
-            task_type: {
-                "count": len(rows),
-                "samples": sorted(row["key"] for row in rows),
-            }
-            for task_type, rows in sorted(by_task.items())
-        },
-        "by_language": {
-            language: {
-                "count": len(rows),
-                "samples": sorted(row["key"] for row in rows),
-            }
-            for language, rows in sorted(by_language.items())
-        },
-    }
-
-
-def select_balanced_gedit_rows(
-    rows: list[dict[str, Any]],
-    *,
-    task_type: str = "all",
-    instruction_language: str = "all",
-    samples_per_group: int | None,
-) -> list[dict[str, Any]]:
-    filtered_rows = []
-    for row in rows:
-        if task_type != "all" and row["task_type"] != task_type:
-            continue
-        if instruction_language != "all" and row["instruction_language"] != instruction_language:
-            continue
-        filtered_rows.append(row)
-
-    if samples_per_group is None:
-        return filtered_rows
-
-    def _select_rows_for_group(group_rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
-        if instruction_language != "all":
-            return group_rows[:samples_per_group]
-
-        rows_by_language: dict[str, list[dict[str, Any]]] = defaultdict(list)
-        for row in group_rows:
-            rows_by_language[str(row.get("instruction_language", "")).strip()].append(row)
-
-        ordered_languages = ["en", "cn"]
-        per_language_quota = samples_per_group // max(len(ordered_languages), 1)
-        remainder = samples_per_group % max(len(ordered_languages), 1)
-
-        selected_rows: list[dict[str, Any]] = []
-        leftovers: list[dict[str, Any]] = []
-        for index, language in enumerate(ordered_languages):
-            language_rows = rows_by_language.get(language, [])
-            quota = per_language_quota + (1 if index < remainder else 0)
-            selected_rows.extend(language_rows[:quota])
-            leftovers.extend(language_rows[quota:])
-
-        if len(selected_rows) < samples_per_group:
-            selected_rows.extend(leftovers[: samples_per_group - len(selected_rows)])
-
-        return selected_rows
-
-    if task_type != "all":
-        return _select_rows_for_group(filtered_rows)
-
-    grouped_rows: dict[str, list[dict[str, Any]]] = defaultdict(list)
-    for row in filtered_rows:
-        grouped_rows[row["task_type"]].append(row)
-
-    selected: list[dict[str, Any]] = []
-    for group in GROUPS:
-        selected.extend(_select_rows_for_group(grouped_rows.get(group, [])))
-    return selected
-
-
-def summarize_gedit_rows(rows: list[dict[str, Any]], language: str = "all") -> dict[str, Any]:
-    return summarize_gedit_rows_with_backbone(rows, language=language)
-
-
-def _mean_or_none(values: list[float]) -> float | None:
-    return statistics.fmean(values) if values else None
-
-
-def _to_q_metrics(section: dict[str, float | None]) -> dict[str, float | None]:
-    return {
-        "count": section.get("count"),
-        "Q_SC": section.get("avg_semantics"),
-        "Q_PQ": section.get("avg_quality"),
-        "Q_O": section.get("avg_overall"),
-    }
-
-
-def _summarize_gedit_rows_single_language(
-    rows: list[dict[str, Any]],
-    *,
-    language: str,
-) -> dict[str, Any]:
-    filtered_rows = [row for row in rows if str(row.get("instruction_language", "")).strip() == language]
-
-    def _to_bool(value: Any) -> bool:
-        if isinstance(value, bool):
-            return value
-        return str(value).strip().lower() in {"1", "true", "yes"}
-
-    per_group: dict[str, dict[str, float | None]] = {}
-    intersection_group: dict[str, dict[str, float | None]] = {}
-
-    for group in GROUPS:
-        group_rows = [row for row in filtered_rows if row.get("task_type") == group]
-        semantics = [float(row["semantics_score"]) for row in group_rows]
-        quality = [float(row["quality_score"]) for row in group_rows]
-        overall = [math.sqrt(float(row["semantics_score"]) * float(row["quality_score"])) for row in group_rows]
-
-        per_group[group] = {
-            "count": len(group_rows),
-            "avg_semantics": _mean_or_none(semantics),
-            "avg_quality": _mean_or_none(quality),
-            "avg_overall": _mean_or_none(overall),
-        }
-
-        intersection_rows = [row for row in group_rows if _to_bool(row.get("intersection_exist", False))]
-        intersection_semantics = [float(row["semantics_score"]) for row in intersection_rows]
-        intersection_quality = [float(row["quality_score"]) for row in intersection_rows]
-        intersection_overall = [
-            math.sqrt(float(row["semantics_score"]) * float(row["quality_score"])) for row in intersection_rows
-        ]
-        intersection_group[group] = {
-            "count": len(intersection_rows),
-            "avg_semantics": _mean_or_none(intersection_semantics),
-            "avg_quality": _mean_or_none(intersection_quality),
-            "avg_overall": _mean_or_none(intersection_overall),
-        }
-
-    overall_section = {
-        "count": len(filtered_rows),
-        "avg_semantics": _mean_or_none(
-            [score["avg_semantics"] for score in per_group.values() if score["avg_semantics"] is not None]
-        ),
-        "avg_quality": _mean_or_none(
-            [score["avg_quality"] for score in per_group.values() if score["avg_quality"] is not None]
-        ),
-        "avg_overall": _mean_or_none(
-            [score["avg_overall"] for score in per_group.values() if score["avg_overall"] is not None]
-        ),
-    }
-    intersection_section = {
-        "count": sum(score["count"] for score in intersection_group.values()),
-        "avg_semantics": _mean_or_none(
-            [score["avg_semantics"] for score in intersection_group.values() if score["avg_semantics"] is not None]
-        ),
-        "avg_quality": _mean_or_none(
-            [score["avg_quality"] for score in intersection_group.values() if score["avg_quality"] is not None]
-        ),
-        "avg_overall": _mean_or_none(
-            [score["avg_overall"] for score in intersection_group.values() if score["avg_overall"] is not None]
-        ),
-    }
-
-    return {
-        "language": language,
-        "by_group": {group: _to_q_metrics(section) for group, section in per_group.items()},
-        "overall": _to_q_metrics(overall_section),
-        "intersection": _to_q_metrics(intersection_section),
-    }
-
-
-def summarize_gedit_rows_with_backbone(
-    rows: list[dict[str, Any]],
-    *,
-    language: str = "all",
-) -> dict[str, Any]:
-    if language == "all":
-        return {
-            "language": "all",
-            "languages": {
-                single_language: _summarize_gedit_rows_single_language(
-                    rows,
-                    language=single_language,
-                )
-                for single_language in ["en", "cn"]
-            },
-        }
-
-    return _summarize_gedit_rows_single_language(rows, language=language)
-
-
-def _require_datasets():
-    try:
-        from datasets import load_dataset, load_from_disk
-    except ImportError as exc:
-        raise ImportError("GEdit-Bench requires the optional `datasets` package.") from exc
-    return load_dataset, load_from_disk
-
-
-def _load_gedit_dataset(dataset_ref: str):
-    load_dataset, load_from_disk = _require_datasets()
-    dataset_path = Path(dataset_ref)
-    if dataset_path.exists():
-        if (dataset_path / "state.json").exists() and (dataset_path / "dataset_info.json").exists():
-            return load_from_disk(str(dataset_path))
-        return load_dataset(str(dataset_path))
-    return load_dataset(dataset_ref)
-
-
-def _resolve_gedit_split(dataset_obj: Any) -> Any:
-    if isinstance(dataset_obj, dict):
-        if "train" in dataset_obj:
-            return dataset_obj["train"]
-        return dataset_obj
-    try:
-        return dataset_obj["train"]
-    except Exception:
-        return dataset_obj
-
-
-def _to_pil_image(value: Any) -> Image.Image:
-    if isinstance(value, Image.Image):
-        return value.convert("RGB")
-    if isinstance(value, dict) and "bytes" in value:
-        image = Image.open(BytesIO(value["bytes"]))
-        image.load()
-        return image.convert("RGB")
-    raise TypeError(f"Unsupported image payload type: {type(value)!r}")
-
-
-class LocalVIEScorer:
-    def __init__(self, *, base_url: str, api_key: str, model: str, timeout: int = 600):
-        self.base_url = base_url.rstrip("/")
-        self.api_key = api_key
-        self.model = model
-        self.timeout = timeout
-        self.sc_prompt = (
-            "You are evaluating image editing quality.\n"
-            "Two images are provided: the source image and the edited image.\n"
-            'Return JSON only in the format {"score": [edit_success, content_preservation], '
-            '"reasoning": "..."}.\n'
-            "Each score must be an integer from 0 to 10.\n"
-            "Editing instruction: <instruction>"
-        )
-        self.pq_prompt = (
-            "You are evaluating image quality.\n"
-            'Return JSON only in the format {"score": [naturalness, artifact_free], "reasoning": "..."}.\n'
-            "Each score must be an integer from 0 to 10."
-        )
-
-    def _request(self, prompt: str, images: list[Image.Image]) -> dict[str, Any]:
-        from benchmarks.accuracy.common import pil_to_data_url
-
-        content: list[dict[str, Any]] = [{"type": "text", "text": prompt}]
-        for image in images:
-            content.append({"type": "image_url", "image_url": {"url": pil_to_data_url(image)}})
-
-        response = requests.post(
-            build_openai_url(self.base_url, "/chat/completions"),
-            json={
-                "model": self.model,
-                "messages": [{"role": "user", "content": content}],
-                "temperature": 0,
-            },
-            headers={
-                "Authorization": f"Bearer {self.api_key}",
-                "Content-Type": "application/json",
-            },
-            timeout=self.timeout,
-        )
-        response.raise_for_status()
-        message_content = response.json()["choices"][0]["message"]["content"]
-        if isinstance(message_content, list):
-            text = "\n".join(part.get("text", "") for part in message_content if part.get("type") == "text")
-        else:
-            text = str(message_content)
-        return parse_score_payload(text)
-
-    def evaluate(self, source_image: Image.Image, edited_image: Image.Image, instruction: str) -> dict[str, float]:
-        sc_payload = self._request(self.sc_prompt.replace("<instruction>", instruction), [source_image, edited_image])
-        pq_payload = self._request(self.pq_prompt, [edited_image])
-        semantics = float(min(sc_payload["score"])) if sc_payload["score"] else 0.0
-        quality = float(min(pq_payload["score"])) if pq_payload["score"] else 0.0
-        overall = math.sqrt(semantics * quality)
-        return {
-            "semantics_score": semantics,
-            "quality_score": quality,
-            "overall_score": overall,
-        }
-
-
-class GEditBenchRunner:
-    def __init__(
-        self,
-        *,
-        dataset_ref: str,
-        output_root: Path,
-        base_url: str,
-        model: str,
-        api_key: str = "EMPTY",
-        width: int = 512,
-        height: int = 512,
-        num_inference_steps: int = 20,
-        guidance_scale: float | None = None,
-        seed: int | None = 42,
-        bot_task: str | None = None,
-        sys_type: str | None = None,
-        system_prompt: str | None = None,
-    ):
-        self.dataset_ref = dataset_ref
-        self.output_root = output_root
-        self.model = model
-        self.width = width
-        self.height = height
-        self.num_inference_steps = num_inference_steps
-        self.guidance_scale = guidance_scale
-        self.seed = seed
-        self.bot_task = bot_task
-        self.sys_type = sys_type
-        self.system_prompt = system_prompt
-        self.client = VllmOmniImageClient(base_url=base_url, api_key=api_key)
-
-    def generate(
-        self,
-        *,
-        model_name: str,
-        task_type: str = "all",
-        instruction_language: str = "all",
-        workers: int = 1,
-        max_samples: int | None = None,
-        samples_per_group: int | None = None,
-    ) -> list[dict[str, Any]]:
-        dataset = _resolve_gedit_split(_load_gedit_dataset(self.dataset_ref))
-        rows = select_balanced_gedit_rows(
-            list(dataset),
-            task_type=task_type,
-            instruction_language=instruction_language,
-            samples_per_group=samples_per_group,
-        )
-        if max_samples is not None:
-            rows = rows[:max_samples]
-
-        outputs: list[dict[str, Any]] = []
-        total = len(rows)
-        if workers <= 1:
-            with tqdm(total=total, desc="GEdit generate", unit="sample") as progress:
-                for item in rows:
-                    result = self._safe_generate_one(model_name, item)
-                    if result:
-                        outputs.append(result)
-                    progress.update(1)
-            return outputs
-
-        with tqdm(total=total, desc="GEdit generate", unit="sample") as progress:
-            with ThreadPoolExecutor(max_workers=workers) as executor:
-                futures = [executor.submit(self._safe_generate_one, model_name, item) for item in rows]
-                for future in as_completed(futures):
-                    result = future.result()
-                    if result:
-                        outputs.append(result)
-                    progress.update(1)
-        return outputs
-
-    def _safe_generate_one(self, model_name: str, item: dict[str, Any]) -> dict[str, Any] | None:
-        try:
-            return self._generate_one(model_name, item)
-        except Exception:
-            logger.exception("Failed to generate GEdit-Bench sample %s", item.get("key", "<unknown>"))
-            return None
-
-    def _generate_one(self, model_name: str, item: dict[str, Any]) -> dict[str, Any] | None:
-        output_path = (
-            self.output_root
-            / model_name
-            / "fullset"
-            / item["task_type"]
-            / item["instruction_language"]
-            / f"{item['key']}.png"
-        )
-        output_path.parent.mkdir(parents=True, exist_ok=True)
-        if output_path.exists():
-            return {
-                "task_type": item["task_type"],
-                "instruction_language": item["instruction_language"],
-                "key": item["key"],
-                "output_path": str(output_path),
-            }
-
-        source_image = _to_pil_image(item["input_image_raw"])
-        edited_image = self.client.generate_image_edit(
-            model=self.model,
-            prompt=item["instruction"],
-            images=source_image,
-            width=self.width,
-            height=self.height,
-            num_inference_steps=self.num_inference_steps,
-            guidance_scale=self.guidance_scale,
-            seed=self.seed,
-            bot_task=self.bot_task,
-            sys_type=self.sys_type,
-            system_prompt=self.system_prompt,
-        )
-        edited_image.save(output_path)
-        return {
-            "task_type": item["task_type"],
-            "instruction_language": item["instruction_language"],
-            "key": item["key"],
-            "output_path": str(output_path),
-        }
-
-
-class GEditBenchEvaluator:
-    def __init__(self, *, dataset_ref: str, output_root: Path, scorer: LocalVIEScorer):
-        self.dataset_ref = dataset_ref
-        self.output_root = output_root
-        self.scorer = scorer
-
-    def evaluate(
-        self,
-        *,
-        model_name: str,
-        save_dir: Path,
-        task_type: str = "all",
-        instruction_language: str = "all",
-        workers: int = 1,
-        max_samples: int | None = None,
-        samples_per_group: int | None = None,
-    ) -> dict[str, Any]:
-        dataset = _resolve_gedit_split(_load_gedit_dataset(self.dataset_ref))
-        rows = select_balanced_gedit_rows(
-            list(dataset),
-            task_type=task_type,
-            instruction_language=instruction_language,
-            samples_per_group=samples_per_group,
-        )
-        if max_samples is not None:
-            rows = rows[:max_samples]
-
-        results: list[dict[str, Any]] = []
-        total = len(rows)
-        if workers <= 1:
-            with tqdm(total=total, desc="GEdit evaluate", unit="sample") as progress:
-                for item in rows:
-                    result = self._safe_evaluate_one(model_name, item)
-                    if result:
-                        results.append(result)
-                    progress.update(1)
-        else:
-            with tqdm(total=total, desc="GEdit evaluate", unit="sample") as progress:
-                with ThreadPoolExecutor(max_workers=workers) as executor:
-                    futures = [executor.submit(self._safe_evaluate_one, model_name, item) for item in rows]
-                    for future in as_completed(futures):
-                        result = future.result()
-                        if result:
-                            results.append(result)
-                        progress.update(1)
-
-        save_dir.mkdir(parents=True, exist_ok=True)
-        base_name = f"{model_name}_{task_type}_{instruction_language}"
-        timestamp = _utc_timestamp()
-        csv_path = save_dir / f"{base_name}_vie_score.csv"
-        timestamped_csv_path = save_dir / f"{base_name}_vie_score_{timestamp}.csv"
-
-        def _write_csv(path: Path) -> None:
-            with path.open("w", encoding="utf-8", newline="") as handle:
-                writer = csv.DictWriter(
-                    handle,
-                    fieldnames=[
-                        "key",
-                        "task_type",
-                        "edited_image",
-                        "instruction",
-                        "semantics_score",
-                        "quality_score",
-                        "overall_score",
-                        "intersection_exist",
-                        "instruction_language",
-                    ],
-                )
-                writer.writeheader()
-                for row in results:
-                    writer.writerow(row)
-
-        _write_csv(csv_path)
-        _write_csv(timestamped_csv_path)
-
-        summary = summarize_gedit_rows_with_backbone(
-            results,
-            language=instruction_language,
-        )
-        summary_path = save_dir / f"{base_name}_summary.json"
-        timestamped_summary_path = save_dir / f"{base_name}_summary_{timestamp}.json"
-        write_json(summary_path, summary)
-        write_json(timestamped_summary_path, summary)
-        return {
-            "results": results,
-            "summary": summary,
-            "csv_path": str(csv_path),
-            "summary_path": str(summary_path),
-            "timestamped_csv_path": str(timestamped_csv_path),
-            "timestamped_summary_path": str(timestamped_summary_path),
-        }
-
-    def _safe_evaluate_one(self, model_name: str, item: dict[str, Any]) -> dict[str, Any] | None:
-        try:
-            return self._evaluate_one(model_name, item)
-        except Exception:
-            logger.exception("Failed to evaluate GEdit-Bench sample %s", item.get("key", "<unknown>"))
-            return None
-
-    def _evaluate_one(self, model_name: str, item: dict[str, Any]) -> dict[str, Any] | None:
-        edited_image_path = (
-            self.output_root
-            / model_name
-            / "fullset"
-            / item["task_type"]
-            / item["instruction_language"]
-            / f"{item['key']}.png"
-        )
-        if not edited_image_path.exists():
-            return None
-
-        source_image = _to_pil_image(item["input_image_raw"])
-        edited_image = Image.open(edited_image_path).convert("RGB")
-        scores = self.scorer.evaluate(source_image, edited_image, item["instruction"])
-        return {
-            "key": item["key"],
-            "task_type": item["task_type"],
-            "edited_image": str(edited_image_path),
-            "instruction": item["instruction"],
-            "semantics_score": scores["semantics_score"],
-            "quality_score": scores["quality_score"],
-            "overall_score": scores["overall_score"],
-            "intersection_exist": item.get("Intersection_exist", False),
-            "instruction_language": item["instruction_language"],
-        }
-
-
-def build_parser() -> argparse.ArgumentParser:
-    parser = argparse.ArgumentParser(description="Run the GEdit-Bench integration against a local vLLM-Omni server.")
-    subparsers = parser.add_subparsers(dest="command", required=True)
-
-    generate = subparsers.add_parser("generate")
-    generate.add_argument("--dataset-ref", type=str, default="stepfun-ai/GEdit-Bench")
-    generate.add_argument("--output-root", type=Path, required=True)
-    generate.add_argument("--base-url", type=str, required=True)
-    generate.add_argument("--model", type=str, required=True)
-    generate.add_argument("--model-name", type=str, default=None)
-    generate.add_argument("--api-key", type=str, default="EMPTY")
-    generate.add_argument("--task-type", choices=["all", *GROUPS], default="all")
-    generate.add_argument("--instruction-language", choices=["all", "en", "cn"], default="all")
-    generate.add_argument("--width", type=int, default=512)
-    generate.add_argument("--height", type=int, default=512)
-    generate.add_argument("--num-inference-steps", type=int, default=20)
-    generate.add_argument("--guidance-scale", type=float, default=None)
-    generate.add_argument("--seed", type=int, default=42)
-    generate.add_argument(
-        "--bot-task",
-        type=str,
-        default=None,
-        choices=["think", "recaption", "think_recaption", "vanilla"],
-        help="HunyuanImage-3 prompting mode for /v1/images/edits (optional).",
-    )
-    generate.add_argument(
-        "--sys-type",
-        type=str,
-        default=None,
-        help="Hunyuan system-prompt type (e.g. en_unified). Defaults from --bot-task when omitted.",
-    )
-    generate.add_argument(
-        "--system-prompt",
-        type=str,
-        default=None,
-        help="Custom Hunyuan system prompt body (sys_type=custom).",
-    )
-    generate.add_argument("--workers", type=int, default=1)
-    generate.add_argument("--max-samples", type=int, default=None)
-    generate.add_argument("--samples-per-group", type=int, default=None)
-
-    evaluate = subparsers.add_parser("evaluate")
-    evaluate.add_argument("--dataset-ref", type=str, default="stepfun-ai/GEdit-Bench")
-    evaluate.add_argument("--output-root", type=Path, required=True)
-    evaluate.add_argument("--model-name", type=str, default=None)
-    evaluate.add_argument("--save-dir", type=Path, required=True)
-    evaluate.add_argument("--task-type", choices=["all", *GROUPS], default="all")
-    evaluate.add_argument("--instruction-language", choices=["all", "en", "cn"], default="all")
-    evaluate.add_argument("--judge-base-url", type=str, required=True)
-    evaluate.add_argument("--judge-model", type=str, required=True)
-    evaluate.add_argument("--judge-api-key", type=str, default="EMPTY")
-    evaluate.add_argument("--workers", type=int, default=1)
-    evaluate.add_argument("--max-samples", type=int, default=None)
-    evaluate.add_argument("--samples-per-group", type=int, default=None)
-
-    summarize = subparsers.add_parser("summarize")
-    summarize.add_argument("--csv-path", type=Path, required=True)
-    summarize.add_argument("--language", choices=["all", "en", "cn"], default="all")
-
-    return parser
-
-
-def main(argv: list[str] | None = None) -> int:
-    parser = build_parser()
-    args = parser.parse_args(argv)
-
-    if args.command == "generate":
-        model_name = resolve_model_name(model_name=args.model_name, model=args.model)
-        runner = GEditBenchRunner(
-            dataset_ref=args.dataset_ref,
-            output_root=args.output_root,
-            base_url=args.base_url,
-            model=args.model,
-            api_key=args.api_key,
-            width=args.width,
-            height=args.height,
-            num_inference_steps=args.num_inference_steps,
-            guidance_scale=args.guidance_scale,
-            seed=args.seed,
-            bot_task=args.bot_task,
-            sys_type=args.sys_type,
-            system_prompt=args.system_prompt,
-        )
-        records = runner.generate(
-            model_name=model_name,
-            task_type=args.task_type,
-            instruction_language=args.instruction_language,
-            workers=args.workers,
-            max_samples=args.max_samples,
-            samples_per_group=args.samples_per_group,
-        )
-        payload = {"records": records, "summary": summarize_generated_records(records)}
-        write_json(args.output_root / model_name / "generation_manifest.json", payload)
-        return 0
-
-    if args.command == "evaluate":
-        model_name = resolve_model_name(model_name=args.model_name, output_root=args.output_root)
-        scorer = LocalVIEScorer(
-            base_url=args.judge_base_url,
-            api_key=args.judge_api_key,
-            model=args.judge_model,
-        )
-        evaluator = GEditBenchEvaluator(dataset_ref=args.dataset_ref, output_root=args.output_root, scorer=scorer)
-        evaluator.evaluate(
-            model_name=model_name,
-            save_dir=args.save_dir,
-            task_type=args.task_type,
-            instruction_language=args.instruction_language,
-            workers=args.workers,
-            max_samples=args.max_samples,
-            samples_per_group=args.samples_per_group,
-        )
-        return 0
-
-    if args.command == "summarize":
-        with args.csv_path.open("r", encoding="utf-8", newline="") as handle:
-            rows = list(csv.DictReader(handle))
-        summary = summarize_gedit_rows_with_backbone(rows, language=args.language)
-        print(json.dumps(summary, indent=2, ensure_ascii=False))
-        return 0
-
-    parser.error(f"Unknown command: {args.command}")
-    return 1

benchmarks/accuracy/image_to_image/run_gedit_bench.py

@@ -1,13 +0,0 @@
-# ruff: noqa: E402, I001
-import sys
-from pathlib import Path
-
-REPO_ROOT = Path(__file__).resolve().parents[3]
-if str(REPO_ROOT) not in sys.path:
-    sys.path.insert(0, str(REPO_ROOT))
-
-from benchmarks.accuracy.image_to_image.gedit_bench import main
-
-
-if __name__ == "__main__":
-    raise SystemExit(main())

benchmarks/accuracy/text_to_image/README.md

@@ -1,103 +0,0 @@
-# GEBench on vLLM-Omni
-
-This integration adapts the upstream `stepfun-ai/GEBench` scripts into
-`vllm-omni/benchmarks/accuracy/text_to_image`.
-
-Upstream mapping:
-
-- `scripts/generate.py` -> `run_gebench.py generate`
-- `scripts/evaluate.py` -> `run_gebench.py evaluate`
-- upstream prompt / judge logic -> `gbench.py`
-
-What changed:
-
-- Generation calls the local OpenAI-compatible `vllm-omni` endpoints:
-  - `/v1/images/generations` for text-only frame generation
-  - `/v1/images/edits` for image-conditioned GUI transition generation
-- Evaluation still keeps the GEBench scoring dimensions:
-  - `goal`
-  - `logic`
-  - `cons`
-  - `ui`
-  - `qual`
-- Judge calls are also routed to a local OpenAI-compatible model served by
-  `vllm-omni` instead of a remote service.
-
-Expected dataset layout:
-
-- Clone the benchmark dataset from Hugging Face into a local directory:
-
-```bash
-git clone https://huggingface.co/datasets/stepfun-ai/GEBench /path/to/GEBench
-```
-
-Example usage:
-
-```bash
-python benchmarks/accuracy/text_to_image/run_gebench.py generate \
-  --dataset-root /path/to/GEBench \
-  --output-root benchmarks/accuracy/text_to_image/outputs \
-  --base-url http://127.0.0.1:8000 \
-  --model Tongyi-MAI/Z-Image-Turbo \
-  --data-type type3
-```
-
-```bash
-python benchmarks/accuracy/text_to_image/run_gebench.py evaluate \
-  --dataset-root /path/to/GEBench \
-  --output-root benchmarks/accuracy/text_to_image/outputs \
-  --data-type type3 \
-  --judge-base-url http://127.0.0.1:8000 \
-  --judge-model Qwen/Qwen2.5-VL-7B-Instruct \
-  --judge-api-key EMPTY
-```
-
-```bash
-python benchmarks/accuracy/text_to_image/run_gebench.py summarize \
-  --output-root benchmarks/accuracy/text_to_image/outputs
-```
-
-Example summary output:
-
-```json
-{
-  "generation": {
-    "count": 20,
-    "by_type": {
-      "type3": {"count": 10},
-      "type4": {"count": 10}
-    }
-  },
-  "evaluation": {
-    "count": 20,
-    "overall_mean": 0.52,
-    "by_type": {
-      "type3": {"count": 10, "overall_mean": 0.50, "overall_mean_100": 50.0},
-      "type4": {"count": 10, "overall_mean": 0.54, "overall_mean_100": 54.0}
-    }
-  }
-}
-```
-
-Example generated images to inspect:
-
-- `benchmarks/accuracy/text_to_image/outputs/03_trajectory_text_fictionalapp/<lang_device>/<sample>/frame0.png`
-- `benchmarks/accuracy/text_to_image/outputs/03_trajectory_text_fictionalapp/<lang_device>/<sample>/frame5.png`
-- `benchmarks/accuracy/text_to_image/outputs/04_trajectory_text_realapp/<lang_device>/<sample>/frame0.png`
-- `benchmarks/accuracy/text_to_image/outputs/04_trajectory_text_realapp/<lang_device>/<sample>/frame5.png`
-
-What to expect:
-
-- `overall_mean` is normalized to `0.0 ~ 1.0`; higher is better.
-- `frame0.png` is the initial GUI frame; `frame5.png` is the final trajectory frame most often used for quick inspection.
-- For full debugging, inspect the whole `frame0.png` ... `frame5.png` sequence for one sample directory.
-
-Notes:
-
-- GEBench upstream leaves type3/type4 generation unfinished. This integration
-  fills that gap with a trajectory runner that generates `frame0.png` followed
-  by `frame1.png` ... `frame5.png`.
-- Type1/2/5 require an image-edit capable model exposed through
-  `vllm-omni serve`.
-- `summarize` will report both generated coverage and any existing evaluation
-  summary files.

benchmarks/accuracy/text_to_image/__init__.py

@@ -1 +0,0 @@
-"""GEBench integration for vLLM-Omni."""

benchmarks/accuracy/text_to_image/gbench.py

@@ -1,927 +0,0 @@
-from __future__ import annotations
-
-import argparse
-import json
-import statistics
-from collections import defaultdict
-from concurrent.futures import ThreadPoolExecutor, as_completed
-from dataclasses import dataclass
-from datetime import datetime, timezone
-from pathlib import Path
-from typing import Any
-
-import requests
-from PIL import Image
-
-from benchmarks.accuracy.common import (
-    VllmOmniImageClient,
-    build_openai_url,
-    ensure_dir,
-    extract_json_object,
-    find_first_image,
-    load_json,
-    pil_to_data_url,
-    save_image,
-    write_json,
-)
-
-TYPE_TO_FOLDER = {
-    "type1": "01_single_step",
-    "type2": "02_multi_step",
-    "type3": "03_trajectory_text_fictionalapp",
-    "type4": "04_trajectory_text_realapp",
-    "type5": "05_grounding_data",
-}
-SCORE_KEYS = ("goal", "logic", "cons", "ui", "qual")
-DEFAULT_SAMPLES_PER_TYPE = 10
-
-
-def _utc_timestamp() -> str:
-    return datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
-
-
-def _write_json_with_timestamp(path: Path, payload: dict[str, Any]) -> Path:
-    write_json(path, payload)
-    timestamped_path = path.with_name(f"{path.stem}_{_utc_timestamp()}{path.suffix}")
-    write_json(timestamped_path, payload)
-    return timestamped_path
-
-
-@dataclass(frozen=True)
-class GEBenchSampleSpec:
-    sample_path: Path
-    metadata: dict[str, Any]
-    sample_name: str
-    lang_device: str
-
-
-def summarize_generated_records(records: list[dict[str, Any]]) -> dict[str, Any]:
-    by_type: dict[str, list[dict[str, Any]]] = defaultdict(list)
-    for record in records:
-        by_type[record["data_type"]].append(record)
-
-    return {
-        "count": len(records),
-        "by_type": {
-            data_type: {
-                "count": len(rows),
-            }
-            for data_type, rows in sorted(by_type.items())
-        },
-    }
-
-
-def summarize_gebench_results(results: list[dict[str, Any]]) -> dict[str, Any]:
-    by_type: dict[str, list[dict[str, Any]]] = defaultdict(list)
-    for result in results:
-        by_type[result["data_type"]].append(result)
-
-    summary: dict[str, Any] = {
-        "count": len(results),
-        "overall_mean": statistics.fmean(r["overall"] for r in results) if results else 0.0,
-        "by_type": {},
-    }
-    for data_type, rows in by_type.items():
-        score_means: dict[str, float] = {}
-        all_score_keys = {key for row in rows for key in row.get("scores", {}).keys()}
-        for score_key in all_score_keys:
-            values = [row["scores"][score_key] for row in rows if score_key in row.get("scores", {})]
-            score_means[score_key] = statistics.fmean(values) if values else 0.0
-        overall_mean = statistics.fmean(row["overall"] for row in rows)
-        summary["by_type"][data_type] = {
-            "count": len(rows),
-            "overall_mean": overall_mean,
-            "overall_mean_100": overall_mean * 100.0,
-            "score_means": score_means,
-        }
-    return summary
-
-
-def select_balanced_gebench_samples(
-    sample_paths_by_type: dict[str, list[Any]],
-    *,
-    samples_per_type: int | None,
-) -> dict[str, list[Any]]:
-    if samples_per_type is None:
-        return {data_type: list(paths) for data_type, paths in sample_paths_by_type.items()}
-    return {data_type: list(paths)[:samples_per_type] for data_type, paths in sample_paths_by_type.items()}
-
-
-def collect_gebench_generation_summary(output_root: Path) -> dict[str, Any]:
-    records: list[dict[str, Any]] = []
-    for data_type, folder_name in TYPE_TO_FOLDER.items():
-        type_root = output_root / folder_name
-        if not type_root.exists():
-            continue
-        for lang_dir in sorted(path for path in type_root.iterdir() if path.is_dir()):
-            for sample_dir in sorted(path for path in lang_dir.iterdir() if path.is_dir()):
-                expected = sample_dir / "frame5.png" if data_type in {"type2", "type3", "type4"} else None
-                if expected is None:
-                    expected = find_first_image(sample_dir)
-                elif not expected.exists():
-                    expected = None
-                if expected is None:
-                    continue
-                records.append(
-                    {
-                        "data_type": data_type,
-                        "sample_name": f"{lang_dir.name}/{sample_dir.name}",
-                        "output_path": str(expected),
-                    }
-                )
-    return summarize_generated_records(records)
-
-
-def _normalize_score_key(key: str) -> str:
-    mapping = {
-        "goal": "goal",
-        "logic": "logic",
-        "cons": "cons",
-        "consistency": "cons",
-        "ui": "ui",
-        "qual": "qual",
-        "quality": "qual",
-    }
-    return mapping.get(key.lower(), key.lower())
-
-
-def _normalize_scores(raw_scores: dict[str, Any]) -> dict[str, int]:
-    scores: dict[str, int] = {}
-    for key, value in raw_scores.items():
-        normalized = _normalize_score_key(key)
-        if normalized not in SCORE_KEYS:
-            continue
-        scalar = value.get("s", 0) if isinstance(value, dict) else value
-        try:
-            scores[normalized] = int(scalar)
-        except (TypeError, ValueError):
-            scores[normalized] = 0
-    for key in SCORE_KEYS:
-        scores.setdefault(key, 0)
-    return scores
-
-
-def _compute_overall(scores: dict[str, int]) -> float:
-    return sum(scores.values()) / (len(SCORE_KEYS) * 5.0)
-
-
-def _iter_sample_paths(dataset_root: Path, data_type: str) -> list[Path]:
-    data_dir = dataset_root / TYPE_TO_FOLDER[data_type]
-    if not data_dir.exists():
-        raise FileNotFoundError(f"Dataset path not found: {data_dir}")
-
-    samples: list[Path] = []
-    for lang_dir in sorted(path for path in data_dir.iterdir() if path.is_dir()):
-        for child in sorted(lang_dir.iterdir()):
-            if child.is_dir():
-                samples.append(child)
-            elif child.suffix.lower() == ".json":
-                samples.append(child)
-    return samples
-
-
-def _sample_name_from_metadata(metadata: dict[str, Any], sample_path: Path, item_index: int | None = None) -> str:
-    sample_id = metadata.get("id") or metadata.get("sample_id") or metadata.get("name")
-    if sample_id:
-        return str(sample_id)
-    if item_index is not None:
-        return f"{sample_path.stem}_{item_index:04d}"
-    return sample_path.stem if sample_path.is_file() else sample_path.name
-
-
-def _expand_sample_path(sample_path: Path) -> list[GEBenchSampleSpec]:
-    if sample_path.is_dir():
-        metadata = _load_metadata(sample_path)
-        return [
-            GEBenchSampleSpec(
-                sample_path=sample_path,
-                metadata=metadata,
-                sample_name=_sample_name(sample_path),
-                lang_device=_lang_device(sample_path, metadata),
-            )
-        ]
-
-    payload = load_json(sample_path)
-    if isinstance(payload, dict):
-        return [
-            GEBenchSampleSpec(
-                sample_path=sample_path,
-                metadata=payload,
-                sample_name=_sample_name_from_metadata(payload, sample_path),
-                lang_device=_lang_device(sample_path, payload),
-            )
-        ]
-
-    if isinstance(payload, list):
-        specs: list[GEBenchSampleSpec] = []
-        for item_index, item in enumerate(payload):
-            if not isinstance(item, dict):
-                continue
-            specs.append(
-                GEBenchSampleSpec(
-                    sample_path=sample_path,
-                    metadata=item,
-                    sample_name=_sample_name_from_metadata(item, sample_path, item_index),
-                    lang_device=_lang_device(sample_path, item),
-                )
-            )
-        return specs
-
-    raise TypeError(f"Unsupported metadata payload type for sample {sample_path}: {type(payload)!r}")
-
-
-def _iter_sample_specs(dataset_root: Path, data_type: str) -> list[GEBenchSampleSpec]:
-    sample_specs: list[GEBenchSampleSpec] = []
-    for sample_path in _iter_sample_paths(dataset_root, data_type):
-        sample_specs.extend(_expand_sample_path(sample_path))
-    return sample_specs
-
-
-def _load_metadata(sample_path: Path) -> dict[str, Any]:
-    if sample_path.is_file():
-        return load_json(sample_path)
-    for candidate in ("meta_data.json", "metadata.json"):
-        meta_path = sample_path / candidate
-        if meta_path.exists():
-            return load_json(meta_path)
-    raise FileNotFoundError(f"Metadata not found for sample: {sample_path}")
-
-
-def _sample_name(sample_path: Path) -> str:
-    return sample_path.stem if sample_path.is_file() else sample_path.name
-
-
-def _lang_device(sample_path: Path, metadata: dict[str, Any]) -> str:
-    return str(metadata.get("lang_device") or sample_path.parent.name)
-
-
-def _resolve_referenced_image(
-    *,
-    metadata: dict[str, Any],
-    sample_path: Path,
-    dataset_root: Path,
-    data_type: str,
-) -> Image.Image | None:
-    for key in ("image", "input_image", "initial_image", "reference_image"):
-        image_ref = metadata.get(key)
-        if not image_ref:
-            continue
-        candidate = dataset_root / TYPE_TO_FOLDER[data_type] / str(image_ref)
-        if candidate.exists():
-            image = Image.open(candidate)
-            image.load()
-            return image.convert("RGB")
-    if sample_path.is_dir():
-        local_image = find_first_image(sample_path)
-        if local_image:
-            image = Image.open(local_image)
-            image.load()
-            return image.convert("RGB")
-    return None
-
-
-def _trajectory_steps(metadata: dict[str, Any]) -> list[dict[str, Any]]:
-    for key in ("trajectory", "steps", "frames"):
-        value = metadata.get(key)
-        if isinstance(value, list):
-            return [step for step in value if isinstance(step, dict)]
-    extracted: list[dict[str, Any]] = []
-    for index in range(1, 6):
-        value = metadata.get(f"step{index}") or metadata.get(str(index))
-        if isinstance(value, dict):
-            extracted.append(value)
-    return extracted
-
-
-def _text_or_default(value: Any, default: str = "") -> str:
-    return str(value).strip() if value is not None else default
-
-
-def _type1_prompt(metadata: dict[str, Any]) -> str:
-    caption = _text_or_default(metadata.get("caption") or metadata.get("instruction"), "Transform the reference GUI.")
-    return (
-        "Using the reference GUI screenshot, generate the next GUI state after the requested interaction.\n\n"
-        f"Requested change:\n{caption}\n\n"
-        "Requirements:\n"
-        "- Preserve layout, visual identity, and unrelated regions.\n"
-        "- Only apply the requested state change.\n"
-        "- Keep all text and controls readable.\n"
-    )
-
-
-def _type2_prompt(goal: str, step_num: int) -> str:
-    return (
-        "Generate the next GUI state for a multi-step task.\n\n"
-        f"Overall goal: {goal}\n"
-        f"Current progress step: {step_num}/5\n\n"
-        "Requirements:\n"
-        "- The change should be incremental and plausible.\n"
-        "- Preserve layout and visual identity.\n"
-        "- Make text/buttons readable.\n"
-    )
-
-
-def _type34_initial_prompt(metadata: dict[str, Any], first_step: dict[str, Any]) -> str:
-    app_name = _text_or_default(metadata.get("app_name"), "App")
-    final_goal = _text_or_default(metadata.get("final_goal") or metadata.get("instruction"), "Complete the task.")
-    visual_description = _text_or_default(
-        metadata.get("visual_description") or first_step.get("visual_description") or first_step.get("description"),
-        "A clean product-quality app home screen.",
-    )
-    return (
-        "Generate the first GUI frame for a task trajectory.\n\n"
-        f"App name: {app_name}\n"
-        f"Final goal: {final_goal}\n"
-        f"Visual description:\n{visual_description}\n\n"
-        "Requirements:\n"
-        "- Generate a production-looking UI screenshot only.\n"
-        "- Keep the layout coherent and readable.\n"
-    )
-
-
-def _type34_next_prompt(step_num: int, step_info: dict[str, Any]) -> str:
-    action = _text_or_default(step_info.get("action") or step_info.get("instruction"), "Continue the task.")
-    visual_description = _text_or_default(
-        step_info.get("visual_description") or step_info.get("description"),
-        "Reflect the expected next GUI state.",
-    )
-    return (
-        "Using the previous frame as reference, generate the next GUI frame.\n\n"
-        f"Step {step_num} action: {action}\n"
-        f"Expected visual state:\n{visual_description}\n\n"
-        "Requirements:\n"
-        "- Only change UI regions affected by this action.\n"
-        "- Preserve persistent bars, layout, and style.\n"
-        "- Keep text and icons readable.\n"
-    )
-
-
-def _type5_prompt(metadata: dict[str, Any]) -> str:
-    grounding = metadata.get("grounding") or {}
-    explanation = _text_or_default(
-        metadata.get("grounding_explanation") or grounding.get("effect") or grounding.get("description"),
-        "Predict the immediate GUI reaction to the indicated target.",
-    )
-    return (
-        "Using the reference GUI screenshot, predict the immediate GUI state after the grounded interaction.\n\n"
-        f"Expected effect: {explanation}\n"
-        f"Grounding metadata: {json.dumps(grounding, ensure_ascii=False)}\n\n"
-        "Requirements:\n"
-        "- Apply only the interaction-triggered change.\n"
-        "- Preserve unrelated regions.\n"
-        "- Keep the UI realistic and readable.\n"
-    )
-
-
-def _make_storyboard_image(
-    frames: list[Image.Image],
-    *,
-    columns: int = 3,
-    background_color: tuple[int, int, int] = (255, 255, 255),
-) -> Image.Image:
-    if not frames:
-        raise ValueError("Expected at least one frame to build a storyboard image.")
-
-    normalized = [frame.convert("RGB") for frame in frames]
-    frame_width = max(frame.width for frame in normalized)
-    frame_height = max(frame.height for frame in normalized)
-    rows = (len(normalized) + columns - 1) // columns
-    storyboard = Image.new("RGB", (frame_width * columns, frame_height * rows), color=background_color)
-
-    for index, frame in enumerate(normalized):
-        x_offset = (index % columns) * frame_width
-        y_offset = (index // columns) * frame_height
-        if frame.size != (frame_width, frame_height):
-            frame = frame.resize((frame_width, frame_height))
-        storyboard.paste(frame, (x_offset, y_offset))
-    return storyboard
-
-
-def _trajectory_judge_payload(frames: list[Image.Image]) -> tuple[str, list[Image.Image]]:
-    storyboard = _make_storyboard_image(frames, columns=3)
-    prompt_suffix = (
-        "The attached image is a storyboard containing six frames arranged left-to-right, "
-        "top-to-bottom as frame0, frame1, frame2, frame3, frame4, frame5."
-    )
-    return prompt_suffix, [storyboard]
-
-
-class LocalJudgeClient:
-    def __init__(self, base_url: str, api_key: str, model: str, timeout: int = 600):
-        self.base_url = base_url.rstrip("/")
-        self.api_key = api_key
-        self.model = model
-        self.timeout = timeout
-
-    def _build_scoring_prompt(self, task_prompt: str) -> str:
-        return (
-            "You are an expert evaluator for GUI image editing and GUI trajectory generation.\n"
-            "Evaluate whether the generated image(s) satisfy the task.\n\n"
-            "Score these five dimensions from 0 to 5:\n"
-            "- goal: whether the user goal is completed correctly\n"
-            "- logic: whether the transition/state change is logically correct\n"
-            "- cons: whether unrelated regions remain consistent\n"
-            "- ui: whether the UI layout/components remain realistic and coherent\n"
-            "- qual: whether the images are visually clear and artifact-free\n\n"
-            "Return JSON only. Do not add any prose outside JSON.\n"
-            "Use exactly this schema:\n"
-            "{\n"
-            '  "goal": 0,\n'
-            '  "logic": 0,\n'
-            '  "cons": 0,\n'
-            '  "ui": 0,\n'
-            '  "qual": 0,\n'
-            '  "reasoning": "short explanation"\n'
-            "}\n\n"
-            "Scoring task:\n"
-            f"{task_prompt}"
-        )
-
-    def _request_text(self, prompt: str, images: list[Image.Image]) -> str:
-        content: list[dict[str, Any]] = [{"type": "text", "text": prompt}]
-        for image in images:
-            content.append({"type": "image_url", "image_url": {"url": pil_to_data_url(image)}})
-
-        response = requests.post(
-            build_openai_url(self.base_url, "/chat/completions"),
-            json={
-                "model": self.model,
-                "messages": [{"role": "user", "content": content}],
-                "temperature": 0,
-            },
-            headers={
-                "Authorization": f"Bearer {self.api_key}",
-                "Content-Type": "application/json",
-            },
-            timeout=self.timeout,
-        )
-        response.raise_for_status()
-        message_content = response.json()["choices"][0]["message"]["content"]
-        if isinstance(message_content, list):
-            return "\n".join(part.get("text", "") for part in message_content if part.get("type") == "text")
-        return str(message_content)
-
-    def evaluate(self, *, prompt: str, images: list[Image.Image]) -> dict[str, Any]:
-        primary_prompt = self._build_scoring_prompt(prompt)
-        text = self._request_text(primary_prompt, images)
-        try:
-            return extract_json_object(text)
-        except ValueError:
-            retry_prompt = (
-                self._build_scoring_prompt(prompt) + "\n\nYour previous response was not valid JSON. "
-                "Return only the JSON object with integer scores."
-            )
-            retry_text = self._request_text(retry_prompt, images)
-            try:
-                return extract_json_object(retry_text)
-            except ValueError:
-                return {
-                    "goal": 0,
-                    "logic": 0,
-                    "cons": 0,
-                    "ui": 0,
-                    "qual": 0,
-                    "reasoning": retry_text.strip() or text.strip() or "Judge response was not valid JSON.",
-                }
-
-
-class GEBenchRunner:
-    def __init__(
-        self,
-        *,
-        dataset_root: Path,
-        output_root: Path,
-        base_url: str,
-        model: str,
-        api_key: str = "EMPTY",
-        width: int = 768,
-        height: int = 576,
-        num_inference_steps: int = 8,
-        output_compression: int | None = 98,
-        guidance_scale: float | None = None,
-        seed: int | None = 42,
-    ):
-        self.dataset_root = dataset_root
-        self.output_root = output_root
-        self.model = model
-        self.width = width
-        self.height = height
-        self.num_inference_steps = num_inference_steps
-        self.output_compression = output_compression
-        self.guidance_scale = guidance_scale
-        self.seed = seed
-        self.client = VllmOmniImageClient(base_url=base_url, api_key=api_key)
-
-    def generate(
-        self,
-        *,
-        data_type: str,
-        workers: int = 1,
-        max_samples: int | None = None,
-        samples_per_type: int | None = None,
-    ) -> list[dict[str, Any]]:
-        sample_specs = select_balanced_gebench_samples(
-            {data_type: _iter_sample_specs(self.dataset_root, data_type)},
-            samples_per_type=samples_per_type,
-        )[data_type]
-        if max_samples is not None:
-            sample_specs = sample_specs[:max_samples]
-
-        results: list[dict[str, Any]] = []
-        if workers <= 1:
-            for sample_spec in sample_specs:
-                result = self._generate_one(data_type, sample_spec)
-                if result:
-                    results.append(result)
-            return results
-
-        with ThreadPoolExecutor(max_workers=workers) as executor:
-            futures = [executor.submit(self._generate_one, data_type, sample_spec) for sample_spec in sample_specs]
-            for future in as_completed(futures):
-                result = future.result()
-                if result:
-                    results.append(result)
-        return results
-
-    def _generate_one(self, data_type: str, sample_spec: GEBenchSampleSpec) -> dict[str, Any] | None:
-        sample_path = sample_spec.sample_path
-        metadata = sample_spec.metadata
-        lang_device = sample_spec.lang_device
-        sample_name = sample_spec.sample_name
-        output_dir = ensure_dir(self.output_root / TYPE_TO_FOLDER[data_type] / lang_device / sample_name)
-
-        if data_type == "type1":
-            output_path = output_dir / "generated.png"
-            if output_path.exists():
-                return {
-                    "data_type": data_type,
-                    "sample_name": f"{lang_device}/{sample_name}",
-                    "output_path": str(output_path),
-                }
-            source = _resolve_referenced_image(
-                metadata=metadata, sample_path=sample_path, dataset_root=self.dataset_root, data_type=data_type
-            )
-            if source is None:
-                return None
-            generated = self.client.generate_image_edit(
-                model=self.model,
-                prompt=_type1_prompt(metadata),
-                images=source,
-                width=self.width,
-                height=self.height,
-                num_inference_steps=self.num_inference_steps,
-                output_compression=self.output_compression,
-                guidance_scale=self.guidance_scale,
-                seed=self.seed,
-            )
-            save_image(output_path, generated)
-            return {
-                "data_type": data_type,
-                "sample_name": f"{lang_device}/{sample_name}",
-                "output_path": str(output_path),
-            }
-
-        if data_type == "type2":
-            goal = _text_or_default(metadata.get("question") or metadata.get("caption"), "Complete the task.")
-            source = _resolve_referenced_image(
-                metadata=metadata, sample_path=sample_path, dataset_root=self.dataset_root, data_type=data_type
-            )
-            if source is None:
-                return None
-            frame0_path = output_dir / "frame0.png"
-            if not frame0_path.exists():
-                save_image(frame0_path, source)
-            previous = source
-            for step_num in range(1, 6):
-                frame_path = output_dir / f"frame{step_num}.png"
-                if frame_path.exists():
-                    previous = Image.open(frame_path).convert("RGB")
-                    continue
-                generated = self.client.generate_image_edit(
-                    model=self.model,
-                    prompt=_type2_prompt(goal, step_num),
-                    images=previous,
-                    width=self.width,
-                    height=self.height,
-                    num_inference_steps=self.num_inference_steps,
-                    output_compression=self.output_compression,
-                    guidance_scale=self.guidance_scale,
-                    seed=self.seed,
-                )
-                save_image(frame_path, generated)
-                previous = generated
-            output_path = output_dir / "frame5.png"
-            return {
-                "data_type": data_type,
-                "sample_name": f"{lang_device}/{sample_name}",
-                "output_path": str(output_path),
-            }
-
-        if data_type in {"type3", "type4"}:
-            steps = _trajectory_steps(metadata)
-            frame0_path = output_dir / "frame0.png"
-            if frame0_path.exists():
-                previous = Image.open(frame0_path).convert("RGB")
-            else:
-                previous = self.client.generate_text_to_image(
-                    model=self.model,
-                    prompt=_type34_initial_prompt(metadata, steps[0] if steps else {}),
-                    width=self.width,
-                    height=self.height,
-                    num_inference_steps=self.num_inference_steps,
-                    output_compression=self.output_compression,
-                    guidance_scale=self.guidance_scale,
-                    seed=self.seed,
-                )
-                save_image(frame0_path, previous)
-
-            for step_num in range(1, 6):
-                frame_path = output_dir / f"frame{step_num}.png"
-                if frame_path.exists():
-                    previous = Image.open(frame_path).convert("RGB")
-                    continue
-                step_info = steps[step_num - 1] if step_num - 1 < len(steps) else {}
-                generated = self.client.generate_image_edit(
-                    model=self.model,
-                    prompt=_type34_next_prompt(step_num, step_info),
-                    images=previous,
-                    width=self.width,
-                    height=self.height,
-                    num_inference_steps=self.num_inference_steps,
-                    output_compression=self.output_compression,
-                    guidance_scale=self.guidance_scale,
-                    seed=self.seed,
-                )
-                save_image(frame_path, generated)
-                previous = generated
-            output_path = output_dir / "frame5.png"
-            return {
-                "data_type": data_type,
-                "sample_name": f"{lang_device}/{sample_name}",
-                "output_path": str(output_path),
-            }
-
-        if data_type == "type5":
-            output_path = output_dir / "generated.png"
-            if output_path.exists():
-                return {
-                    "data_type": data_type,
-                    "sample_name": f"{lang_device}/{sample_name}",
-                    "output_path": str(output_path),
-                }
-            source = _resolve_referenced_image(
-                metadata=metadata, sample_path=sample_path, dataset_root=self.dataset_root, data_type=data_type
-            )
-            if source is None:
-                return None
-            generated = self.client.generate_image_edit(
-                model=self.model,
-                prompt=_type5_prompt(metadata),
-                images=source,
-                width=self.width,
-                height=self.height,
-                num_inference_steps=self.num_inference_steps,
-                output_compression=self.output_compression,
-                guidance_scale=self.guidance_scale,
-                seed=self.seed,
-            )
-            save_image(output_path, generated)
-            return {
-                "data_type": data_type,
-                "sample_name": f"{lang_device}/{sample_name}",
-                "output_path": str(output_path),
-            }
-
-        raise ValueError(f"Unsupported data type: {data_type}")
-
-
-class GEBenchEvaluator:
-    def __init__(self, *, dataset_root: Path, output_root: Path, judge: LocalJudgeClient):
-        self.dataset_root = dataset_root
-        self.output_root = output_root
-        self.judge = judge
-
-    def evaluate(
-        self,
-        *,
-        data_type: str,
-        workers: int = 1,
-        max_samples: int | None = None,
-        samples_per_type: int | None = None,
-    ) -> dict[str, Any]:
-        output_type_dir = self.output_root / TYPE_TO_FOLDER[data_type]
-        sample_specs_by_name = {
-            (spec.lang_device, spec.sample_name): spec for spec in _iter_sample_specs(self.dataset_root, data_type)
-        }
-        if not output_type_dir.exists():
-            payload = {"data_type": data_type, "results": [], "summary": summarize_gebench_results([])}
-            write_json(self.output_root / "evaluations" / f"{data_type}.json", payload)
-            return payload
-        sample_dirs = [
-            sample_dir
-            for lang_dir in sorted(path for path in output_type_dir.iterdir() if path.is_dir())
-            for sample_dir in sorted(path for path in lang_dir.iterdir() if path.is_dir())
-            if (lang_dir.name, sample_dir.name) in sample_specs_by_name
-        ]
-        sample_dirs = select_balanced_gebench_samples(
-            {data_type: sample_dirs},
-            samples_per_type=samples_per_type,
-        )[data_type]
-        if max_samples is not None:
-            sample_dirs = sample_dirs[:max_samples]
-        results: list[dict[str, Any]] = []
-        if workers <= 1:
-            for sample_dir in sample_dirs:
-                result = self._evaluate_one(
-                    data_type,
-                    sample_dir,
-                    sample_specs_by_name[(sample_dir.parent.name, sample_dir.name)],
-                )
-                if result:
-                    results.append(result)
-        else:
-            with ThreadPoolExecutor(max_workers=workers) as executor:
-                futures = [
-                    executor.submit(
-                        self._evaluate_one,
-                        data_type,
-                        sample_dir,
-                        sample_specs_by_name[(sample_dir.parent.name, sample_dir.name)],
-                    )
-                    for sample_dir in sample_dirs
-                ]
-                for future in as_completed(futures):
-                    result = future.result()
-                    if result:
-                        results.append(result)
-
-        payload = {"data_type": data_type, "results": results, "summary": summarize_gebench_results(results)}
-        write_json(self.output_root / "evaluations" / f"{data_type}.json", payload)
-        return payload
-
-    def _evaluate_one(self, data_type: str, sample_dir: Path, sample_spec: GEBenchSampleSpec) -> dict[str, Any] | None:
-        lang_device = sample_dir.parent.name
-        sample_name = sample_dir.name
-        dataset_sample = sample_spec.sample_path
-        metadata = sample_spec.metadata
-
-        if data_type == "type1":
-            source = _resolve_referenced_image(
-                metadata=metadata, sample_path=dataset_sample, dataset_root=self.dataset_root, data_type=data_type
-            )
-            generated_path = find_first_image(sample_dir)
-            if source is None or generated_path is None:
-                return None
-            generated = Image.open(generated_path).convert("RGB")
-            raw_scores = self.judge.evaluate(prompt=_type1_prompt(metadata), images=[source, generated])
-        elif data_type == "type2":
-            frames = [Image.open(sample_dir / f"frame{i}.png").convert("RGB") for i in range(6)]
-            goal = _text_or_default(metadata.get("question") or metadata.get("caption"), "Complete the task.")
-            prompt_suffix, judge_images = _trajectory_judge_payload(frames)
-            raw_scores = self.judge.evaluate(
-                prompt=f"Evaluate a six-frame GUI trajectory.\nTask: {goal}\n{prompt_suffix}",
-                images=judge_images,
-            )
-        elif data_type in {"type3", "type4"}:
-            frames = [Image.open(sample_dir / f"frame{i}.png").convert("RGB") for i in range(6)]
-            instruction = _text_or_default(metadata.get("instruction") or metadata.get("caption"), "Complete the task.")
-            prompt_suffix, judge_images = _trajectory_judge_payload(frames)
-            raw_scores = self.judge.evaluate(
-                prompt=f"Evaluate a six-frame GUI trajectory.\nInstruction: {instruction}\n{prompt_suffix}",
-                images=judge_images,
-            )
-        elif data_type == "type5":
-            source = _resolve_referenced_image(
-                metadata=metadata, sample_path=dataset_sample, dataset_root=self.dataset_root, data_type=data_type
-            )
-            generated_path = find_first_image(sample_dir)
-            if source is None or generated_path is None:
-                return None
-            generated = Image.open(generated_path).convert("RGB")
-            raw_scores = self.judge.evaluate(prompt=_type5_prompt(metadata), images=[source, generated])
-        else:
-            raise ValueError(f"Unsupported data type: {data_type}")
-
-        scores = _normalize_scores(raw_scores)
-        return {
-            "sample_name": f"{lang_device}/{sample_name}",
-            "data_type": data_type,
-            "scores": scores,
-            "overall": _compute_overall(scores),
-            "raw_scores": raw_scores,
-        }
-
-
-def _data_types_arg(value: str) -> list[str]:
-    return list(TYPE_TO_FOLDER.keys()) if value == "all" else [value]
-
-
-def build_parser() -> argparse.ArgumentParser:
-    parser = argparse.ArgumentParser(description="Run local GEBench generation and scoring against vLLM-Omni.")
-    subparsers = parser.add_subparsers(dest="command", required=True)
-
-    generate = subparsers.add_parser("generate")
-    generate.add_argument("--dataset-root", type=Path, required=True)
-    generate.add_argument("--output-root", type=Path, required=True)
-    generate.add_argument("--base-url", type=str, required=True)
-    generate.add_argument("--model", type=str, required=True)
-    generate.add_argument("--data-type", choices=["all", *TYPE_TO_FOLDER.keys()], default="all")
-    generate.add_argument("--api-key", type=str, default="EMPTY")
-    generate.add_argument("--width", type=int, default=768)
-    generate.add_argument("--height", type=int, default=576)
-    generate.add_argument("--num-inference-steps", type=int, default=8)
-    generate.add_argument("--output-compression", type=int, default=98)
-    generate.add_argument("--guidance-scale", type=float, default=None)
-    generate.add_argument("--seed", type=int, default=42)
-    generate.add_argument("--workers", type=int, default=1)
-    generate.add_argument("--max-samples", type=int, default=None)
-    generate.add_argument("--samples-per-type", type=int, default=None)
-
-    evaluate = subparsers.add_parser("evaluate")
-    evaluate.add_argument("--dataset-root", type=Path, required=True)
-    evaluate.add_argument("--output-root", type=Path, required=True)
-    evaluate.add_argument("--data-type", choices=["all", *TYPE_TO_FOLDER.keys()], default="all")
-    evaluate.add_argument("--judge-base-url", type=str, required=True)
-    evaluate.add_argument("--judge-model", type=str, required=True)
-    evaluate.add_argument("--judge-api-key", type=str, default="EMPTY")
-    evaluate.add_argument("--workers", type=int, default=1)
-    evaluate.add_argument("--max-samples", type=int, default=None)
-    evaluate.add_argument("--samples-per-type", type=int, default=None)
-
-    summarize = subparsers.add_parser("summarize")
-    summarize.add_argument("--output-root", type=Path, required=True)
-
-    return parser
-
-
-def main(argv: list[str] | None = None) -> int:
-    parser = build_parser()
-    args = parser.parse_args(argv)
-
-    if args.command == "generate":
-        runner = GEBenchRunner(
-            dataset_root=args.dataset_root,
-            output_root=args.output_root,
-            base_url=args.base_url,
-            model=args.model,
-            api_key=args.api_key,
-            width=args.width,
-            height=args.height,
-            num_inference_steps=args.num_inference_steps,
-            output_compression=args.output_compression,
-            guidance_scale=args.guidance_scale,
-            seed=args.seed,
-        )
-        records: list[dict[str, Any]] = []
-        for data_type in _data_types_arg(args.data_type):
-            records.extend(
-                runner.generate(
-                    data_type=data_type,
-                    workers=args.workers,
-                    max_samples=args.max_samples,
-                    samples_per_type=args.samples_per_type,
-                )
-            )
-        payload = {"records": records, "summary": summarize_generated_records(records)}
-        write_json(args.output_root / "generation_manifest.json", payload)
-        return 0
-
-    if args.command == "evaluate":
-        judge = LocalJudgeClient(
-            base_url=args.judge_base_url,
-            api_key=args.judge_api_key,
-            model=args.judge_model,
-        )
-        evaluator = GEBenchEvaluator(dataset_root=args.dataset_root, output_root=args.output_root, judge=judge)
-        combined_results: list[dict[str, Any]] = []
-        for data_type in _data_types_arg(args.data_type):
-            payload = evaluator.evaluate(
-                data_type=data_type,
-                workers=args.workers,
-                max_samples=args.max_samples,
-                samples_per_type=args.samples_per_type,
-            )
-            combined_results.extend(payload["results"])
-        _write_json_with_timestamp(
-            args.output_root / "evaluations" / "summary.json",
-            {"summary": summarize_gebench_results(combined_results)},
-        )
-        return 0
-
-    if args.command == "summarize":
-        generation_summary = collect_gebench_generation_summary(args.output_root)
-        evaluation_dir = args.output_root / "evaluations"
-        result_records: list[dict[str, Any]] = []
-        if evaluation_dir.exists():
-            for file_path in sorted(evaluation_dir.glob("type*.json")):
-                payload = load_json(file_path)
-                result_records.extend(payload.get("results", []))
-        payload: dict[str, Any] = {"generation": generation_summary}
-        if result_records:
-            payload["evaluation"] = summarize_gebench_results(result_records)
-        _write_json_with_timestamp(args.output_root / "summary.json", payload)
-        print(json.dumps(payload, indent=2, ensure_ascii=False))
-        return 0
-
-    parser.error(f"Unknown command: {args.command}")
-    return 1

benchmarks/accuracy/text_to_image/run_gebench.py

@@ -1,13 +0,0 @@
-# ruff: noqa: E402, I001
-import sys
-from pathlib import Path
-
-REPO_ROOT = Path(__file__).resolve().parents[3]
-if str(REPO_ROOT) not in sys.path:
-    sys.path.insert(0, str(REPO_ROOT))
-
-from benchmarks.accuracy.text_to_image.gbench import main
-
-
-if __name__ == "__main__":
-    raise SystemExit(main())

benchmarks/build_dataset/download_process_data_seedtts.md

@@ -27,7 +27,7 @@ pip install gdown
 Download the dataset from Google Drive:
 
 ```bash
-gdown 1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP
+gdown --id 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  1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP
+gdown --id 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

@@ -1,30 +0,0 @@
-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.
-vd021|||Please fasten your seatbelts; we will be experiencing some turbulence shortly.|A calm, professional female flight attendant voice, reassuring and clear.
-vd022|||The fossil record suggests these creatures lived more than sixty million years ago.|A measured male documentary narrator voice, scientific and curious.
-vd023|||Hey kiddo, did you finish your homework before dinner today?|A gentle, caring father voice, warm and slightly tired after a long day.
-vd024|||Live from courtside, the home team has just taken a commanding lead.|A high-energy male sports broadcaster, urgent and engaged.
-vd025|||Please remain on the line while I transfer you to the next available agent.|A polite, neutral female call-center voice, professional and patient.
-vd026|||I have a dream that one day this nation will rise up and live out the true meaning of its creed.|A powerful, resonant male orator voice, slow and impassioned.
-vd027|||Mama always said life was like a box of chocolates; you never know what you're gonna get.|A folksy male voice with a Southern American accent, warm and reflective.
-vd028|||Ahoy there, matey! Hoist the colors and prepare to set sail.|A theatrical pirate male voice, gruff and boisterous.
-vd029|||Ladies and gentlemen, please welcome to the stage our keynote speaker.|A polished female event MC voice, gracious and articulate.
-vd030|||Once you've completed the form, click submit and we'll review it within two business days.|A friendly female product-onboarding voice, encouraging and clear.

benchmarks/build_dataset/seed_tts_smoke/en/meta.lst

@@ -1,30 +0,0 @@
-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.
-smoke021|||A long line of cars stretched along the coastal highway after the lighthouse opened.
-smoke022|||The chef recommended the seared scallops with a lemon butter sauce as the special.
-smoke023|||Our quarterly engineering review begins promptly at two in the afternoon on Thursday.
-smoke024|||Despite the heavy rain, the marathon runners kept a steady pace through downtown.
-smoke025|||She finished the crossword puzzle in under fifteen minutes while sipping her tea.
-smoke026|||The new release adds streaming output, faster startup, and better error reporting.
-smoke027|||His grandfather taught him how to whittle small wooden birds from cedar offcuts.
-smoke028|||Travelers were advised to check the timetable since trains were running on a holiday schedule.
-smoke029|||The startup launched a public beta after eighteen months of closed development.
-smoke030|||At night the desert sky filled with stars so bright the Milky Way was clearly visible.

benchmarks/build_dataset/sound_effect_smoke/en/meta.lst

@@ -1,10 +0,0 @@
-se001|ocean waves crashing on a rocky beach|10.0
-se002|heavy rain on a metal roof at night|8.0
-se003|a dog barking in the distance|6.0
-se004|a crackling campfire in a quiet forest|10.0
-se005|footsteps on dry leaves|6.0
-se006|thunder rolling across a wide valley|8.0
-se007|a busy coffee shop with chatter and espresso machines|10.0
-se008|wind howling through tall pine trees|8.0
-se009|a small stream flowing over pebbles|10.0
-se010|a creaking wooden door slowly opening|6.0

benchmarks/build_dataset/ttsd_smoke/en/meta.lst

@@ -1,10 +0,0 @@
-ttsd001|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] Hello and welcome to the show today. [S2] Thanks for having me, it's a pleasure to be here. [S1] Let's dive right in.
-ttsd002|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] Did you finish the report? [S2] Almost — I'm just polishing the conclusion. [S1] Great, send it over when you're done.
-ttsd003|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] Have you tried that new restaurant downtown? [S2] Yes, the pasta is amazing. [S1] I'll book a table for Friday.
-ttsd004|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] What did you think of yesterday's meeting? [S2] Productive, but a bit long. [S1] Agreed, we should set an hour cap next time.
-ttsd005|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] The weather is beautiful today. [S2] Perfect for a hike. [S1] Want to drive out to the trails this afternoon?
-ttsd006|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] How did the demo go? [S2] Better than expected — the client signed off on the prototype. [S1] Fantastic news!
-ttsd007|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] Coffee or tea this morning? [S2] Coffee please, with a splash of milk. [S1] Coming right up.
-ttsd008|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] Did you read the email about the deadline change? [S2] I did, we now have until next Friday. [S1] That should give us enough buffer.
-ttsd009|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] Are you joining the team dinner tonight? [S2] I'll be there around seven. [S1] See you then.
-ttsd010|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] Have we backed up the database? [S2] Last backup ran two hours ago. [S1] Good, I'll trigger a fresh snapshot before the migration.

benchmarks/diffusion/README.md

@@ -30,7 +30,7 @@ python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
 
 **Notes**
 
-- By default, image tasks talk to `http://<host>:<port>/v1/chat/completions`; video tasks talk to `/v1/videos`.
+- The benchmark talks to `http://<host>:<port>/v1/chat/completions`.
 - If you run the server on another host or port, pass `--base-url` accordingly.
 
 ## 2. Supported Datasets
@@ -43,27 +43,6 @@ The benchmark supports three dataset modes via `--dataset`:
 
 ### VBench dataset
 
-`vbench` only provides prompt data (and image paths for i2v/i2i); it does not carry
-per-request generation fields. In this mode, all requests share CLI values:
-`--width --height --num-frames --fps --num-inference-steps`
-(pass `--width` and `--height` together).
-
-Example (`t2v`):
-
-```bash
-python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
-	--base-url http://localhost:8099 \
-	--model Wan-AI/Wan2.2-T2V-A14B-Diffusers \
-	--task t2v \
-	--dataset vbench \
-	--num-prompts 50 \
-	--width 640 --height 480 \
-	--num-frames 81 --fps 16 \
-	--num-inference-steps 40
-```
-
-Note: `vbench` can also be used for other tasks such as `t2i` / `i2v` (and `i2i`). For `t2i`, the loader reuses VBench t2v text prompts; for `i2v` / `i2i`, it loads the VBench i2v dataset (with image paths).
-
 If you use i2v/i2i bench datasets and need auto-download support, you may need:
 
 ```bash
@@ -95,9 +74,8 @@ You can point to your own trace using `--dataset-path`.
 
 ### Basic flags
 
-- `--base-url`: Server address; `--endpoint` selects the path appended to this base URL.
+- `--base-url`: Server address (the script calls `.../v1/chat/completions`).
 - `--model`: The OpenAI-compatible `model` field.
-- `--endpoint`: API endpoint path. Leading `/` is optional, e.g. `/v1/videos` or `v1/videos`.
 - `--task`: Task type (e.g., `t2i`, `t2v`, `i2i`, `i2v`).
 - `--dataset`: Dataset mode (`vbench` / `trace` / `random`).
 - `--num-prompts`: Number of requests to send.
@@ -112,7 +90,7 @@ Common optional flags:
 Related flags: `--width`, `--height`, `--num-frames`, `--fps`, `--num-inference-steps`.
 
 - For `vbench` / `random`: these CLI flags act as global defaults for all generated requests.
-- For `trace`: requests can carry their own fields (e.g., `width/height/num_frames/num_inference_steps`), with overrides/fallbacks as below.
+- For `trace`: each request can carry its own fields (e.g., `width/height/num_frames/num_inference_steps`).
 
 Precedence rules for `trace` (i.e., what actually gets sent):
 
@@ -131,21 +109,9 @@ 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

@@ -1,4 +1,3 @@
-import asyncio
 import base64
 import mimetypes
 import os
@@ -10,8 +9,6 @@ from typing import Any
 import aiohttp
 from tqdm import tqdm
 
-DEFAULT_EDITS_BOT_TASK = "think"
-
 
 @dataclass
 class RequestFuncInput:
@@ -29,7 +26,6 @@ class RequestFuncInput:
     extra_body: dict[str, Any] = field(default_factory=dict)
     image_paths: list[str] | None = None
     request_id: str = field(default_factory=lambda: str(uuid.uuid4()))
-    default_bot_task: str | None = DEFAULT_EDITS_BOT_TASK
 
 
 @dataclass
@@ -39,7 +35,6 @@ class RequestFuncOutput:
     error: str = ""
     start_time: float = 0.0
     response_body: dict[str, Any] = field(default_factory=dict)
-    stage_durations: dict[str, float] = field(default_factory=dict)
     peak_memory_mb: float = 0.0
     slo_achieved: bool | None = None
 
@@ -56,100 +51,10 @@ def _encode_image_as_data_url(path: str) -> str:
     return f"data:{mime};base64,{encoded}"
 
 
-async def async_request_image_edits(
-    input: RequestFuncInput,
-    session: aiohttp.ClientSession,
-    pbar: tqdm | None = None,
-    enable_diffusion_pipeline_profiler: bool = False,
-) -> RequestFuncOutput:
-    """POST /v1/images/edits (multipart)."""
-    del enable_diffusion_pipeline_profiler
-    output = RequestFuncOutput()
-    output.start_time = time.perf_counter()
-
-    extra_body = dict(input.extra_body)
-    width = input.width or extra_body.get("width") or 1024
-    height = input.height or extra_body.get("height") or 1024
-    edits_url = input.api_url
-
-    form = aiohttp.FormData()
-    form.add_field("model", input.model)
-    form.add_field("prompt", input.prompt)
-    form.add_field("size", f"{width}x{height}")
-    form.add_field("response_format", "b64_json")
-
-    if input.num_inference_steps is not None:
-        form.add_field("num_inference_steps", str(input.num_inference_steps))
-    elif extra_body.get("num_inference_steps") is not None:
-        form.add_field("num_inference_steps", str(extra_body["num_inference_steps"]))
-
-    if input.seed is not None:
-        form.add_field("seed", str(input.seed))
-    elif extra_body.get("seed") is not None:
-        form.add_field("seed", str(extra_body["seed"]))
-
-    if extra_body.get("guidance_scale") is not None:
-        form.add_field("guidance_scale", str(extra_body["guidance_scale"]))
-    if extra_body.get("negative_prompt") is not None:
-        form.add_field("negative_prompt", str(extra_body["negative_prompt"]))
-    if extra_body.get("true_cfg_scale") is not None:
-        form.add_field("true_cfg_scale", str(extra_body["true_cfg_scale"]))
-    if extra_body.get("sys_type") is not None:
-        form.add_field("sys_type", str(extra_body["sys_type"]))
-    if extra_body.get("system_prompt") is not None:
-        form.add_field("system_prompt", str(extra_body["system_prompt"]))
-
-    bot_task = extra_body.get("bot_task")
-    if bot_task is None and input.default_bot_task is not None:
-        bot_task = input.default_bot_task
-    if bot_task is not None:
-        form.add_field("bot_task", str(bot_task))
-
-    assert input.image_paths is not None
-    for img_path in input.image_paths:
-        if not os.path.exists(img_path):
-            output.error = f"Image file not found: {img_path}"
-            output.success = False
-            if pbar:
-                pbar.update(1)
-            return output
-        with open(img_path, "rb") as img_f:
-            image_bytes = img_f.read()
-        form.add_field(
-            "image",
-            image_bytes,
-            filename=os.path.basename(img_path),
-            content_type=_guess_mime_type(img_path),
-        )
-
-    try:
-        async with session.post(edits_url, data=form) as response:
-            if response.status == 200:
-                resp_json = await response.json()
-                output.response_body = resp_json
-                output.success = True
-            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 output.success and input.slo_ms is not None:
-        output.slo_achieved = (output.latency * 1000.0) <= float(input.slo_ms)
-
-    if pbar:
-        pbar.update(1)
-    return output
-
-
 async def async_request_chat_completions(
     input: RequestFuncInput,
     session: aiohttp.ClientSession,
     pbar: tqdm | None = None,
-    enable_diffusion_pipeline_profiler: bool = False,
 ) -> RequestFuncOutput:
     output = RequestFuncOutput()
     output.start_time = time.perf_counter()
@@ -201,31 +106,8 @@ async def async_request_chat_completions(
                 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):
-                            content = msg.get("content", [])
-                            if content and isinstance(content, list) and len(content) > 0:
-                                first_item = 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
-
-                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
+                if "peak_memory_mb" in resp_json:
+                    output.peak_memory_mb = resp_json["peak_memory_mb"]
             else:
                 output.error = f"HTTP {response.status}: {await response.text()}"
                 output.success = False
@@ -243,7 +125,7 @@ async def async_request_chat_completions(
     return output
 
 
-async def async_request_openai_image_generations(
+async def async_request_openai_images(
     input: RequestFuncInput,
     session: aiohttp.ClientSession,
     pbar: tqdm | None = None,
@@ -310,169 +192,7 @@ async def async_request_openai_image_generations(
     return output
 
 
-async def async_request_v1_videos(
-    input: RequestFuncInput,
-    session: aiohttp.ClientSession,
-    pbar: tqdm | None = None,
-) -> RequestFuncOutput:
-    output = RequestFuncOutput()
-    output.start_time = time.perf_counter()
-
-    files = dict(input.extra_body)
-    if input.prompt:
-        files.setdefault("prompt", input.prompt)
-    if input.width and input.height:
-        files.setdefault("height", input.height)
-        files.setdefault("width", input.width)
-    if input.num_frames:
-        files.setdefault("num_frames", input.num_frames)
-    if input.num_inference_steps:
-        files.setdefault("num_inference_steps", input.num_inference_steps)
-    if input.seed is not None:
-        files.setdefault("seed", input.seed)
-    if input.fps:
-        files.setdefault("fps", input.fps)
-
-    form = aiohttp.FormData()
-    for k, v in files.items():
-        form.add_field(k, str(v))
-
-    image_file = None
-    if input.image_paths and len(input.image_paths) > 0:
-        image_path = input.image_paths[0]
-        image_file = open(image_path, "rb")
-        form.add_field(
-            "input_reference",
-            image_file,
-            filename=os.path.basename(image_path),
-            content_type="application/octet-stream",
-        )
-
-    job_id = None
-    job_status = None
-    poll_json = {}
-    resp_json = {}
-
-    try:
-        # invoke a post request (POST /v1/videos)
-        async with session.post(input.api_url, data=form) as response:
-            if response.status == 200:
-                resp_json = await response.json()
-                job_id = resp_json.get("id")
-                job_status = resp_json.get("status")
-                if not job_id or not job_status:
-                    output.error = "API response missing job 'id' or 'status' field."
-                    output.success = False
-                    return output
-            else:
-                output.error = f"HTTP {response.status}: {await response.text()}"
-                output.success = False
-                return output
-
-        # invoke a poll request (GET /v1/videos/{video_id})
-        poll_interval = 2.0  # Unit(s)
-        timeout_seconds = 600.0
-        deadline = time.perf_counter() + timeout_seconds
-        job_url = f"{input.api_url}/{job_id}"
-
-        while job_status not in {"completed", "failed"}:
-            await asyncio.sleep(poll_interval)
-
-            async with session.get(job_url) as poll_response:
-                if poll_response.status != 200:
-                    output.error = f"Polling failed HTTP {poll_response.status}: {await poll_response.text()}"
-                    output.success = False
-                    return output
-
-                poll_json = await poll_response.json()
-                job_status = poll_json.get("status")
-
-                if time.perf_counter() >= deadline:
-                    output.error = f"Timed out waiting for video job {job_id} to complete."
-                    output.success = False
-                    return output
-
-        if job_status == "failed":
-            output.error = f"Video job failed: {poll_json}"
-            output.success = False
-            return output
-
-        # invoke a get request (GET /v1/videos/{video_id}/content)
-        content_url = f"{job_url}/content"
-        async with session.get(content_url) as content_response:
-            if content_response.status != 200:
-                output.error = (
-                    f"Content retrieval failed HTTP {content_response.status}: {await content_response.text()}"
-                )
-                output.success = False
-                return output
-
-            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:
-                output.peak_memory_mb = resp_json["peak_memory_mb"]
-    except Exception as e:
-        output.error = str(e)
-        output.success = False
-    finally:
-        if image_file is not None:
-            image_file.close()
-
-        if job_id is not None:
-            try:
-                async with session.delete(f"{input.api_url}/{job_id}") as _:
-                    pass
-            except Exception as e:
-                print(f"Failed to clean up video job {job_id}: {e}")
-
-    output.latency = time.perf_counter() - output.start_time
-
-    if output.success and input.slo_ms is not None:
-        output.slo_achieved = (output.latency * 1000.0) <= float(input.slo_ms)
-
-    if pbar:
-        pbar.update(1)
-    return output
-
-
-LEGACY_BACKEND_ENDPOINT_ALIASES = {
-    "vllm-omni": "/v1/chat/completions",
-    "openai": "/v1/images/generations",
-}
-
-
-def normalize_endpoint(value: str) -> str:
-    endpoint = str(value).strip()
-    if not endpoint:
-        raise ValueError("endpoint must not be empty.")
-    endpoint = LEGACY_BACKEND_ENDPOINT_ALIASES.get(
-        endpoint,
-        LEGACY_BACKEND_ENDPOINT_ALIASES.get(endpoint.lstrip("/"), endpoint),
-    )
-    if not endpoint.startswith("/"):
-        endpoint = f"/{endpoint}"
-    return endpoint
-
-
-def endpoint_filename_token(value: str) -> str:
-    token = normalize_endpoint(value).lstrip("/")
-    for bad in ("/", "\\", ":", "*", "?", '"', "<", ">", "|"):
-        token = token.replace(bad, "_")
-    return token or "endpoint"
-
-
 backends_function_mapping = {
-    "2i": {
-        "/v1/chat/completions": (async_request_chat_completions, "/v1/chat/completions"),
-        "/v1/images/generations": (async_request_openai_image_generations, "/v1/images/generations"),
-        "/v1/images/edits": (async_request_image_edits, "/v1/images/edits"),
-    },
-    "2v": {
-        "/v1/videos": (async_request_v1_videos, "/v1/videos"),
-    },
+    "vllm-omni": (async_request_chat_completions, "/v1/chat/completions"),
+    "openai": (async_request_openai_images, "/v1/images/generations"),
 }

benchmarks/diffusion/bench_attention_backends.py

@@ -1,328 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Standalone diagnostic for the diffusion attention backends added in #3079.
-
-Exercises the same synthetic attention shape through each backend and each
-torch SDPA sub-kernel, with and without an attn_mask, so we can pinpoint why
-CUDNN_ATTN underperforms on some SKUs. Two typical causes:
-
-  (a) cuDNN lacks a tuned kernel for this (SM, head_dim, seq) combination
-      and silently falls back to MATH.
-  (b) attn_mask is non-None and the cuDNN SDPA dispatch rejects it,
-      walking CUDNN -> FLASH -> MATH.
-
-Run:
-    python benchmarks/diffusion/bench_attention_backends.py --preset hv15
-    python benchmarks/diffusion/bench_attention_backends.py --preset wan22
-    python benchmarks/diffusion/bench_attention_backends.py \
-        --batch 1 --heads 24 --seq 14336 --head-dim 128
-
-Optional extras:
-    pip install --pre flash-attn-4        # FA4 is currently pre-release only
-    pip install -U flashinfer              # latest FlashInfer (0.6.9)
-
-The table at the end is the data we want on the PR — surface the row where a
-backend is >1.5x the SDPA baseline, that's the one to gate off the auto-route.
-"""
-
-from __future__ import annotations
-
-import argparse
-import time
-
-import torch
-from torch.nn.attention import SDPBackend, sdpa_kernel
-
-# Presets mirror the DiT hot-path attention call (single-stream block) at the
-# shapes we run in the PR's validation configs.
-_PRESETS = {
-    # HunyuanVideo-1.5 480p/33f: (30 * 52 * 9) latent + ~256 text tokens.
-    "hv15": {"batch": 1, "heads": 24, "seq": 14336, "head_dim": 128},
-    # Wan 2.2 480p/33f rough estimate; adjust when we measure real shapes.
-    "wan22": {"batch": 1, "heads": 40, "seq": 16384, "head_dim": 128},
-    # Smaller image-gen shape for quick smoke tests.
-    "flux": {"batch": 1, "heads": 24, "seq": 4096, "head_dim": 128},
-}
-
-_SDPA_BACKENDS = [
-    ("CUDNN_ATTENTION", [SDPBackend.CUDNN_ATTENTION]),
-    ("FLASH_ATTENTION", [SDPBackend.FLASH_ATTENTION]),
-    ("EFFICIENT_ATTENTION", [SDPBackend.EFFICIENT_ATTENTION]),
-    ("MATH", [SDPBackend.MATH]),
-    # The PR's CUDNN_ATTN impl uses this priority chain; keep it here to show
-    # what torch actually picks when multiple are allowed.
-    ("CUDNN_ATTN_CHAIN", [SDPBackend.CUDNN_ATTENTION, SDPBackend.FLASH_ATTENTION, SDPBackend.MATH]),
-]
-
-
-def _env_report() -> None:
-    if not torch.cuda.is_available():
-        raise SystemExit("CUDA not available; this diagnostic needs a GPU.")
-    cc = torch.cuda.get_device_capability()
-    print("=" * 72)
-    print(f"GPU           : {torch.cuda.get_device_name()}")
-    print(f"SM capability : sm_{cc[0]}{cc[1]}")
-    print(f"torch         : {torch.__version__}")
-    print(f"cuDNN version : {torch.backends.cudnn.version()}")
-    print("flashinfer    : ", end="")
-    try:
-        import flashinfer
-
-        print(getattr(flashinfer, "__version__", "present"))
-    except Exception as e:
-        print(f"not installed ({type(e).__name__})")
-    print("=" * 72)
-
-
-def _make_qkv(batch: int, heads: int, seq: int, head_dim: int, device: str, dtype: torch.dtype):
-    # Layout: (B, S, H, D) — what the PR's backends accept as input.
-    q = torch.randn(batch, seq, heads, head_dim, device=device, dtype=dtype)
-    k = torch.randn(batch, seq, heads, head_dim, device=device, dtype=dtype)
-    v = torch.randn(batch, seq, heads, head_dim, device=device, dtype=dtype)
-    return q, k, v
-
-
-def _make_mask(batch: int, seq: int, device: str, dtype: torch.dtype, pad_tokens: int = 128) -> torch.Tensor:
-    # Mirrors a text-encoder padding mask: last `pad_tokens` positions masked.
-    mask = torch.zeros(batch, 1, seq, seq, device=device, dtype=dtype)
-    if pad_tokens > 0:
-        mask[..., -pad_tokens:] = float("-inf")
-    return mask
-
-
-def _time_call(fn, *args, warmup: int = 3, iters: int = 10) -> tuple[float, str]:
-    """Return (median ms, error string). Skips on any exception so that a
-    backend that rejects our args (wrong dtype, missing JIT module, unsupported
-    kwarg value) doesn't abort the whole sweep."""
-    try:
-        for _ in range(warmup):
-            fn(*args)
-        torch.accelerator.synchronize()
-        times = []
-        for _ in range(iters):
-            torch.accelerator.synchronize()
-            t0 = time.perf_counter()
-            fn(*args)
-            torch.accelerator.synchronize()
-            times.append((time.perf_counter() - t0) * 1000.0)
-        times.sort()
-        return times[len(times) // 2], ""
-    except Exception as e:  # noqa: BLE001 — probe script, keep the whole table going
-        msg = str(e).split("\n", 1)[0][:60]
-        return float("nan"), f"{type(e).__name__}: {msg}" if msg else type(e).__name__
-
-
-def _run_sdpa_variants(q, k, v, attn_mask, scale: float) -> list[tuple[str, float, str]]:
-    # q/k/v come in (B, S, H, D); F.sdpa wants (B, H, S, D).
-    q_t = q.permute(0, 2, 1, 3)
-    k_t = k.permute(0, 2, 1, 3)
-    v_t = v.permute(0, 2, 1, 3)
-    rows: list[tuple[str, float, str]] = []
-    for name, backends in _SDPA_BACKENDS:
-
-        def _call(mask=attn_mask):
-            with sdpa_kernel(backends):
-                return torch.nn.functional.scaled_dot_product_attention(
-                    q_t, k_t, v_t, attn_mask=mask, dropout_p=0.0, is_causal=False, scale=scale
-                )
-
-        ms, err = _time_call(_call)
-        rows.append((name, ms, err))
-    return rows
-
-
-def _run_flashinfer(q, k, v, scale: float, backend: str | None = None, attn_mask=None) -> tuple[float, str]:
-    """Call FlashInfer's dense single-prefill.
-
-    ``backend`` hints at cutlass/fa3/trtllm if the installed version exposes
-    that kwarg. ``attn_mask`` is the same additive 4D mask we pass to SDPA —
-    we convert it to the 2D boolean ``custom_mask`` form FlashInfer accepts.
-    Per FlashInfer docs, ``custom_mask`` only applies when ``causal=False``.
-    """
-    try:
-        import inspect
-
-        from flashinfer.prefill import single_prefill_with_kv_cache
-    except Exception as e:
-        return float("nan"), f"import-{type(e).__name__}"
-
-    kwargs: dict = {"sm_scale": scale, "causal": False, "return_lse": False}
-    if backend is not None:
-        sig = inspect.signature(single_prefill_with_kv_cache)
-        if "backend" not in sig.parameters:
-            return float("nan"), "no-backend-kwarg"
-        kwargs["backend"] = backend
-
-    if attn_mask is not None:
-        # Collapse (B, 1, S, S) additive float mask to (S, S) boolean.
-        # FlashInfer expects True = keep, False = masked out.
-        mask_2d = attn_mask[0, 0]
-        kwargs["custom_mask"] = mask_2d != float("-inf")
-
-    def _call():
-        out = single_prefill_with_kv_cache(q[0], k[0], v[0], **kwargs)
-        return out.unsqueeze(0)
-
-    return _time_call(_call)
-
-
-def _run_fa4(q, k, v, scale: float) -> tuple[float, str]:
-    """Call FlashAttention-4 directly (``pip install flash-attn-4``). FA4
-    ships a Blackwell-native kernel via CuTe-DSL; on sm_120 it should beat
-    cuDNN by ~20%. API lives under ``flash_attn.cute``; the older
-    ``flash_attn.flash_attn_func`` path is FA2/FA3 only."""
-    try:
-        from flash_attn.cute import flash_attn_func
-    except Exception as e:
-        return float("nan"), f"import-{type(e).__name__}"
-
-    def _call():
-        # FA4 accepts (B, S, H, D) directly — same layout the PR's backends use.
-        return flash_attn_func(q, k, v, softmax_scale=scale, causal=False)
-
-    return _time_call(_call)
-
-
-def _run_flashinfer_cudnn_batch(q, k, v, scale: float) -> tuple[float, str]:
-    """Call FlashInfer's direct cuDNN wrapper. Bypasses PyTorch SDPA dispatch
-    (which has a few hundred ns of overhead per call) and talks to cuDNN FMHA
-    straight. Useful as a ceiling for 'pure cuDNN, no SDPA' on Blackwell."""
-    try:
-        from flashinfer.prefill import cudnn_batch_prefill_with_kv_cache
-    except Exception as e:
-        return float("nan"), f"import-{type(e).__name__}"
-
-    b, s, h, d = q.shape
-    qo_indptr = torch.tensor([0, s], dtype=torch.int32, device=q.device)
-    kv_indptr = torch.tensor([0, s], dtype=torch.int32, device=q.device)
-
-    def _call():
-        return cudnn_batch_prefill_with_kv_cache(
-            q.reshape(b * s, h, d),
-            k.reshape(b * s, h, d),
-            v.reshape(b * s, h, d),
-            qo_indptr=qo_indptr,
-            kv_indptr=kv_indptr,
-            max_qo_len=s,
-            max_kv_len=s,
-            sm_scale=scale,
-            causal=False,
-        )
-
-    return _time_call(_call)
-
-
-def _print_table(title: str, rows: list[tuple[str, float, str]], baseline_name: str | None = None) -> None:
-    print(f"\n{title}")
-    print("-" * 88)
-    header = f"{'backend':<24} {'median (ms)':>14}    {'vs baseline':>12}    status"
-    print(header)
-    print("-" * 88)
-    baseline_ms = None
-    if baseline_name is not None:
-        for name, ms, _ in rows:
-            if name == baseline_name and ms == ms:  # not NaN
-                baseline_ms = ms
-                break
-    for name, ms, err in rows:
-        ms_str = f"{ms:>14.3f}" if ms == ms else f"{'n/a':>14}"
-        if baseline_ms is not None and ms == ms:
-            ratio = baseline_ms / ms
-            ratio_str = f"{ratio:>11.2f}x"
-        else:
-            ratio_str = f"{'—':>12}"
-        status = "ok" if not err else f"FAILED ({err})"
-        print(f"{name:<24} {ms_str}    {ratio_str}    {status}")
-
-
-def _pick_winner(rows: list[tuple[str, float, str]]) -> tuple[str, float] | None:
-    """Return (backend, ms) of the fastest non-failing row, or None."""
-    ok_rows = [(n, m) for n, m, err in rows if not err and m == m]
-    if not ok_rows:
-        return None
-    return min(ok_rows, key=lambda x: x[1])
-
-
-def _bench_one_shape(shape: dict, dtype: torch.dtype, device: str) -> tuple[list, list]:
-    torch.manual_seed(0)
-    q, k, v = _make_qkv(**shape, device=device, dtype=dtype)
-    scale = 1.0 / (shape["head_dim"] ** 0.5)
-
-    rows_nomask = _run_sdpa_variants(q, k, v, attn_mask=None, scale=scale)
-    rows_nomask.append(("FLASHINFER (default)", *_run_flashinfer(q, k, v, scale)))
-    # `trtllm-gen` has no sm_120 cubins (NVIDIA/TensorRT-LLM#11799) — skipped.
-    for fi_backend in ("fa2", "fa3", "cutlass", "auto"):
-        rows_nomask.append((f"FLASHINFER ({fi_backend})", *_run_flashinfer(q, k, v, scale, backend=fi_backend)))
-    rows_nomask.append(("FLASHINFER (cudnn-batch)", *_run_flashinfer_cudnn_batch(q, k, v, scale)))
-    rows_nomask.append(("FA4 (direct)", *_run_fa4(q, k, v, scale)))
-
-    mask = _make_mask(shape["batch"], shape["seq"], device=device, dtype=dtype, pad_tokens=256)
-    rows_mask = _run_sdpa_variants(q, k, v, attn_mask=mask, scale=scale)
-    rows_mask.append(("FLASHINFER (dense)", *_run_flashinfer(q, k, v, scale, attn_mask=mask)))
-    rows_mask.append(("FA4 (direct)", float("nan"), "mask-not-supported"))
-
-    return rows_nomask, rows_mask
-
-
-def main() -> None:
-    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
-    parser.add_argument("--preset", choices=list(_PRESETS.keys()), default="hv15")
-    parser.add_argument("--sweep", action="store_true", help="Run all presets and print a ranking")
-    parser.add_argument("--batch", type=int, default=None)
-    parser.add_argument("--heads", type=int, default=None)
-    parser.add_argument("--seq", type=int, default=None)
-    parser.add_argument("--head-dim", type=int, default=None)
-    parser.add_argument("--dtype", choices=["bf16", "fp16"], default="bf16")
-    parser.add_argument("--device", default="cuda:0")
-    args = parser.parse_args()
-
-    dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float16
-    _env_report()
-    print(f"dtype         : {args.dtype}")
-
-    presets = list(_PRESETS.keys()) if args.sweep else [args.preset]
-    summary: dict[str, tuple[str, float] | None] = {}
-
-    for preset_name in presets:
-        shape = dict(_PRESETS[preset_name])
-        if not args.sweep:
-            for k in ("batch", "heads", "seq", "head_dim"):
-                v = getattr(args, k if k != "head_dim" else "head_dim")
-                if v is not None:
-                    shape[k] = v
-
-        print("\n" + "=" * 88)
-        print(
-            f"Preset: {preset_name}  |  batch={shape['batch']} heads={shape['heads']} "
-            f"seq={shape['seq']} head_dim={shape['head_dim']}"
-        )
-        print("=" * 88)
-
-        rows_nomask, rows_mask = _bench_one_shape(shape, dtype, args.device)
-        _print_table(f"[{preset_name}] No attention mask", rows_nomask, baseline_name="CUDNN_ATTENTION")
-        _print_table(
-            f"[{preset_name}] With attention mask (pad 256 tokens)", rows_mask, baseline_name="CUDNN_ATTENTION"
-        )
-
-        summary[preset_name] = _pick_winner(rows_nomask)
-
-    if args.sweep or len(presets) > 1:
-        print("\n" + "=" * 88)
-        print("Winners per preset (no-mask path)")
-        print("=" * 88)
-        for preset_name, winner in summary.items():
-            if winner is None:
-                print(f"  {preset_name:<10} — no successful backend")
-            else:
-                name, ms = winner
-                print(f"  {preset_name:<10} {name:<24} {ms:>8.3f} ms")
-
-    print("\nNotes:")
-    print("  * Ratios are relative to CUDNN_ATTENTION. >1.0x means faster than cuDNN.")
-    print("  * Mask-path winner inherits CUDNN_ATTN's fallback in the PR's backends.")
-    print("  * `trtllm-gen` and FA4 4.0.0b10 are known-broken on sm_120 as of Apr 2026.")
-    print("  * For e2e timings, run text_to_video.py / text_to_image.py while varying DIFFUSION_ATTENTION_BACKEND.")
-
-
-if __name__ == "__main__":
-    main()

benchmarks/diffusion/bench_sage_attn3_matrix.sh

@@ -1,108 +0,0 @@
-#!/usr/bin/env bash
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-#
-# 2x2 timing matrix (backend x compile mode) for SAGE_ATTN_3 vs TORCH_SDPA,
-# on HunyuanVideo-1.5 480p T2V and Wan 2.2 TI2V-5B. Run from the repo root:
-#
-#     bash benchmarks/diffusion/bench_sage_attn3_matrix.sh
-#
-# Requires: sageattn3 installed (for SAGE_ATTN_3 rows) and a Blackwell GPU.
-
-set -u
-
-STEPS=${STEPS:-50}
-SEED=${SEED:-42}
-PROMPT=${PROMPT:-"A dog running across a field of golden wheat."}
-OUT_DIR=${OUT_DIR:-/tmp/bench_sage3}
-mkdir -p "$OUT_DIR"
-
-declare -A TOTAL PER_STEP TEXT_ENC TRANSFORMER VAE
-
-run() {
-    local name="$1"; shift
-    local model="$1"; shift
-    local h="$1"; shift
-    local w="$1"; shift
-    local f="$1"; shift
-    local g="$1"; shift
-    local log="$OUT_DIR/${name}.log"
-    echo "[$(date +%T)] Running $name ..."
-    env "$@" python examples/offline_inference/text_to_video/text_to_video.py \
-        --model "$model" \
-        --prompt "$PROMPT" \
-        --height "$h" --width "$w" --num-frames "$f" \
-        --num-inference-steps "$STEPS" --seed "$SEED" --guidance-scale "$g" \
-        --enable-diffusion-pipeline-profiler \
-        --output "$OUT_DIR/${name}.mp4" > "$log" 2>&1 || {
-            echo "  FAILED — see $log"
-            TOTAL[$name]="FAIL"; PER_STEP[$name]="FAIL"
-            TEXT_ENC[$name]="-"; TRANSFORMER[$name]="-"; VAE[$name]="-"
-            return
-        }
-
-    local wait_ms
-    wait_ms=$(grep "add_req_and_wait=" "$log" | tail -1 | sed -nE 's/.*add_req_and_wait=([0-9.]+) ms.*/\1/p')
-    local total_s
-    total_s=$(grep -oE "Total generation time: [0-9.]+ seconds" "$log" | tail -1 | sed -nE 's/.*: ([0-9.]+) seconds.*/\1/p')
-    local text_enc_s
-    text_enc_s=$(grep -oE "text_encoder.forward took [0-9.]+s" "$log" | awk '{sub(/s$/,"",$NF); sum+=$NF} END {printf "%.3f", sum+0}')
-    local vae_s
-    vae_s=$(grep -oE "vae.decode took [0-9.]+s" "$log" | awk '{sub(/s$/,"",$NF); sum+=$NF} END {printf "%.3f", sum+0}')
-    local pipeline_s
-    pipeline_s=$(grep -oE "Pipeline\.forward took [0-9.]+s" "$log" | tail -1 | sed -nE 's/.*took ([0-9.]+)s.*/\1/p')
-    local transformer_s
-    transformer_s=$(awk "BEGIN {printf \"%.3f\", ${pipeline_s:-0} - ${vae_s:-0} - ${text_enc_s:-0}}")
-
-    TOTAL[$name]="${total_s:-0}"
-    PER_STEP[$name]=$(awk "BEGIN {printf \"%.3f\", ${wait_ms:-0} / 1000 / $STEPS}")
-    TEXT_ENC[$name]="$text_enc_s"
-    TRANSFORMER[$name]="$transformer_s"
-    VAE[$name]="$vae_s"
-
-    echo "  total=${TOTAL[$name]}s  per-step=${PER_STEP[$name]}s/it  transformer=${TRANSFORMER[$name]}s  vae=${VAE[$name]}s"
-}
-
-# ========== HunyuanVideo-1.5 480p T2V (480x832x33, CFG=6.0) ==========
-HV_MODEL=hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v
-run hv15_sdpa_eager     "$HV_MODEL" 480 832 33 6.0 TORCH_COMPILE_DISABLE=1
-run hv15_sdpa_compiled  "$HV_MODEL" 480 832 33 6.0
-run hv15_sage3_eager    "$HV_MODEL" 480 832 33 6.0 TORCH_COMPILE_DISABLE=1 DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN_3
-run hv15_sage3_compiled "$HV_MODEL" 480 832 33 6.0                         DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN_3
-
-# ========== Wan 2.2 TI2V-5B (704x1280x33, CFG=5.0) ==========
-WAN_MODEL=Wan-AI/Wan2.2-TI2V-5B-Diffusers
-run wan22_sdpa_eager     "$WAN_MODEL" 704 1280 33 5.0 TORCH_COMPILE_DISABLE=1
-run wan22_sdpa_compiled  "$WAN_MODEL" 704 1280 33 5.0
-run wan22_sage3_eager    "$WAN_MODEL" 704 1280 33 5.0 TORCH_COMPILE_DISABLE=1 DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN_3
-run wan22_sage3_compiled "$WAN_MODEL" 704 1280 33 5.0                         DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN_3
-
-print_table() {
-    local title="$1"; shift
-    echo
-    echo "$title (${STEPS} steps, seed=${SEED})"
-    echo
-    printf "| %-23s | %-9s | %-15s | %-12s | %-15s | %-14s |\n" \
-        "Config" "Total (s)" "Per-step (s/it)" "Text Enc (s)" "Transformer (s)" "VAE Decode (s)"
-    printf "| %-23s | %-9s | %-15s | %-12s | %-15s | %-14s |\n" \
-        "-----------------------" "---------" "---------------" "------------" "---------------" "--------------"
-    for row in "$@"; do
-        case "$row" in
-            *_sdpa_eager)     label="SDPA + Eager" ;;
-            *_sdpa_compiled)  label="SDPA + Compiled" ;;
-            *_sage3_eager)    label="SAGE_ATTN_3 + Eager" ;;
-            *_sage3_compiled) label="SAGE_ATTN_3 + Compiled" ;;
-        esac
-        printf "| %-23s | %-9s | %-15s | %-12s | %-15s | %-14s |\n" \
-            "$label" "${TOTAL[$row]}" "${PER_STEP[$row]}" "${TEXT_ENC[$row]}" "${TRANSFORMER[$row]}" "${VAE[$row]}"
-    done
-}
-
-print_table "HunyuanVideo-1.5 480p T2V, 832x480x33" \
-    hv15_sdpa_eager hv15_sdpa_compiled hv15_sage3_eager hv15_sage3_compiled
-
-print_table "Wan 2.2 TI2V-5B, 1280x704x33" \
-    wan22_sdpa_eager wan22_sdpa_compiled wan22_sage3_eager wan22_sage3_compiled
-
-echo
-echo "Logs: $OUT_DIR/*.log"

benchmarks/diffusion/diffusion_benchmark_serving.py

@@ -1,4 +1,4 @@
-# adapted from fastvideo
+# adapted from sglang and fastvideo
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
@@ -6,74 +6,38 @@
 Benchmark online serving for diffusion models (Image/Video Generation).
 If you want to use i2v, i2i dataset, you should `uv pip install gdown` first
 
-Supports multiple endpoints:
-    - /v1/chat/completions: OpenAI chat-compatible image requests (e.g. t2i, Qwen i2i)
-    - /v1/images/edits: OpenAI image edit / IT2I (multipart; e.g. Hunyuan --bot-task think)
-    - /v1/images/generations: OpenAI image generation requests
-    - /v1/videos: Async video jobs
-
-Legacy --backend vllm-omni and openai are aliases for chat/completions and images/generations.
+Supports multiple backends:
+    - vllm-omni: Uses /v1/chat/completions endpoint (default)
+    - openai: Uses /v1/images/generations endpoint
 
 Usage:
-    # Video (/v1/videos endpoint)
+    # Video (vllm-omni backend)
     t2v:
     python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
-        --endpoint /v1/videos --dataset vbench --task t2v --num-prompts 10 \
+        --backend vllm-omni --dataset vbench --task t2v --num-prompts 10 \
         --height 480 --width 640 --fps 16 --num-frames 80
 
     i2v:
     python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
-        --endpoint /v1/videos --dataset vbench --task i2v --num-prompts 10
+        --backend vllm-omni --dataset vbench --task i2v --num-prompts 10
 
 
-    # Image (/v1/chat/completions endpoint)
+    # Image (vllm-omni backend)
     t2i:
     python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
-        --endpoint /v1/chat/completions --dataset vbench --task t2i --num-prompts 10 \
+        --backend vllm-omni --dataset vbench --task t2i --num-prompts 10 \
         --height 1024 --width 1024
 
+    i2i:
     python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
-        --endpoint /v1/chat/completions --dataset random --task t2i --num-prompts 1 \
-        --max-concurrency 1 --enable-negative-prompt \
-        --random-request-config '[
-            {"width":512,"height":512,"num_inference_steps":20,"weight":0.15},
-            {"width":768,"height":768,"num_inference_steps":20,"weight":0.25},
-            {"width":1024,"height":1024,"num_inference_steps":25,"weight":0.45},
-            {"width":1536,"height":1536,"num_inference_steps":35,"weight":0.15}
-        ]'
+        --backend vllm-omni --dataset vbench --task i2i --num-prompts 10
 
-    ti2i (Hunyuan / OpenAI image edit API):
-    python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
-        --endpoint /v1/images/edits --dataset random --task ti2i --num-prompts 10 \
-        --bot-task think
-
-    python benchmarks/diffusion/diffusion_benchmark_serving.py \
-      --endpoint /v1/images/edits \
-      --dataset custom \
-      --dataset-path custom_requests.jsonl \
-      --task ti2i \
-      --bot-task think
-
-    i2i (chat-based models such as Qwen-Image-Edit):
-    python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
-        --endpoint /v1/chat/completions --dataset vbench --task i2i --num-prompts 10
-
-    # Image (/v1/images/generations endpoint)
+    # Image (openai backend)
     t2i:
     python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
-        --endpoint /v1/images/generations --dataset vbench --task t2i --num-prompts 10 \
+        --backend openai --dataset vbench --task t2i --num-prompts 10 \
         --height 1024 --width 1024 --port 3000
 
-    # Video (v1/videos)
-    t2v:
-    python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
-        --endpoint /v1/videos --dataset random --task t2v --num-prompts 1 \
-        --max-concurrency 1 --enable-negative-prompt \
-        --random-request-config '[
-            {"width":854,"height":480,"num_inference_steps":18,"num_frames":120,"fps":24,"weight":1}
-        ]'
-
-
 """
 
 import argparse
@@ -81,10 +45,7 @@ import ast
 import asyncio
 import glob
 import json
-import logging
 import os
-import random
-import tempfile
 import time
 import uuid
 from abc import ABC, abstractmethod
@@ -95,17 +56,9 @@ from typing import Any
 import aiohttp
 import numpy as np
 import requests
-from backends import (
-    RequestFuncInput,
-    RequestFuncOutput,
-    backends_function_mapping,
-    normalize_endpoint,
-)
-from PIL import Image
+from backends import RequestFuncInput, RequestFuncOutput, backends_function_mapping
 from tqdm.asyncio import tqdm
 
-logger = logging.getLogger(__name__)
-
 
 class BaseDataset(ABC):
     def __init__(self, args, api_url: str, model: str):
@@ -147,7 +100,7 @@ class VBenchDataset(BaseDataset):
     def _load_data(self) -> list[dict[str, Any]]:
         if self.args.task == "t2v":
             return self._load_t2v_prompts()
-        elif self.args.task in ["i2v", "ti2v", "ti2i", "i2i", "it2i"]:
+        elif self.args.task in ["i2v", "ti2v", "ti2i", "i2i"]:
             return self._load_i2v_data()
         else:
             return self._load_t2v_prompts()
@@ -309,9 +262,6 @@ class VBenchDataset(BaseDataset):
 
     def _resize_data(self, data: list[dict[str, Any]]) -> list[dict[str, Any]]:
         """Resize data to match num_prompts."""
-        if not data:
-            raise ValueError("No benchmark data available. Install Pillow or provide --dataset-path.")
-
         if not self.args.num_prompts:
             return data
 
@@ -535,14 +485,6 @@ class TraceDataset(BaseDataset):
         timestamp = self._coerce_float(row.get("timestamp"))
         slo_ms = self._coerce_float(row.get("slo_ms"))
         image_paths = row.get("image_paths")
-        if not image_paths:
-            single = row.get("image_path")
-            image_paths = [single] if single else None
-
-        if not image_paths and self.args.task in ["i2v", "i2i", "ti2v", "ti2i", "it2i"]:
-            raise ValueError(
-                f"Task {self.args.task} requires image input, but no image_path or image_paths found in trace row."
-            )
 
         override_w = self.args.width
         override_h = self.args.height
@@ -573,210 +515,32 @@ class TraceDataset(BaseDataset):
         return [self[i] for i in range(len(self))]
 
 
-class CustomDataset(BaseDataset):
-    """
-    Custom dataset that loads requests from a JSONL file.
-
-    Each line in the JSONL file should be a JSON object with the following fields:
-    - prompt (required): The text prompt for the request
-    - width (optional): Image/video width
-    - height (optional): Image height
-    - num_inference_steps (optional): Number of diffusion steps
-    - seed (optional): Random seed
-    - image_paths (optional): List of input image paths for i2i/i2v/ti2i tasks
-    - image_urls (optional): List of input image URLs (alternative to image_paths)
-
-    Example JSONL for ti2i:
-    {"prompt": "Add sunset lighting", "width": 1024, "height": 1024, "image_urls": ["https://example.com/image.jpg"]}
-    """
-
-    def __init__(self, args, api_url: str, model: str):
-        super().__init__(args, api_url, model)
-        self.dataset_path = args.dataset_path
-        if not self.dataset_path:
-            raise ValueError("--dataset-path must be provided when using 'custom' dataset")
-        self.load_data()
-
-    def load_data(self) -> None:
-        """Load data from JSONL file."""
-        import pandas as pd
-
-        if not self.dataset_path.endswith(".jsonl"):
-            raise ValueError("Custom dataset must be a JSONL file")
-
-        try:
-            df = pd.read_json(path_or_buf=self.dataset_path, lines=True)
-        except Exception as e:
-            raise ValueError(f"Failed to load JSONL file {self.dataset_path}: {e}")
-
-        if "prompt" not in df.columns:
-            raise ValueError("JSONL file must contain a 'prompt' column")
-
-        self.data = df.to_dict("records")
-        print(f"Loaded {len(self.data)} requests from {self.dataset_path}")
-
-    def _resolve_image_paths(self, item: dict) -> list[str] | None:
-        # Handle local file paths
-        if "image_paths" in item and item["image_paths"]:
-            paths = item["image_paths"]
-            if isinstance(paths, str):
-                paths = [paths]
-
-            # Verify all paths exist
-            valid_paths = []
-            for path in paths:
-                if os.path.exists(path):
-                    valid_paths.append(path)
-                else:
-                    raise ValueError(f"Image file not found: {path}")
-
-            return valid_paths if valid_paths else None
-
-        # Handle URLs - download to temp files
-        if "image_urls" in item and item["image_urls"]:
-            urls = item["image_urls"]
-            if isinstance(urls, str):
-                urls = [urls]
-
-            downloaded_paths = []
-            for url in urls:
-                try:
-                    response = requests.get(url, timeout=30)
-                    response.raise_for_status()
-
-                    # Create temp file with appropriate extension
-                    suffix = os.path.splitext(url)[1] or ".png"
-                    with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
-                        tmp.write(response.content)
-                        downloaded_paths.append(tmp.name)
-                except Exception as e:
-                    raise ValueError(f"Failed to download image from {url}: {e}")
-
-            return downloaded_paths if downloaded_paths else None
-
-        return None
-
-    def __len__(self) -> int:
-        return len(self.data)
-
-    def __getitem__(self, idx: int) -> RequestFuncInput:
-        item = self.data[idx]
-
-        # Extract fields with fallback to args defaults
-        prompt = item["prompt"]
-        width = item.get("width", self.args.width)
-        height = item.get("height", self.args.height)
-        num_inference_steps = item.get("num_inference_steps", self.args.num_inference_steps)
-        seed = item.get("seed", self.args.seed)
-        reserved_keys = {"prompt", "width", "height", "num_inference_steps", "seed", "image_paths", "image_urls"}
-        extra_body = {k: v for k, v in item.items() if k not in reserved_keys}
-
-        # Handle image paths/URLs
-        image_paths = self._resolve_image_paths(item)
-
-        return RequestFuncInput(
-            prompt=prompt,
-            api_url=self.api_url,
-            model=self.model,
-            seed=seed,
-            image_paths=image_paths,
-            extra_body=extra_body,
-            width=width,
-            height=height,
-            num_inference_steps=num_inference_steps,
-        )
-
-    def get_requests(self) -> list[RequestFuncInput]:
-        """Get all requests, cycling through data if num_prompts > dataset size."""
-        num_requests = getattr(self.args, "num_prompts", len(self.data))
-
-        if num_requests <= len(self.data):
-            return [self[i] for i in range(num_requests)]
-
-        # Cycle through the dataset to reach num_prompts
-        requests = []
-        for i in range(num_requests):
-            idx = i % len(self.data)
-            requests.append(self[idx])
-
-        return requests
-
-
 class RandomDataset(BaseDataset):
-    def __init__(self, args, api_url: str, model: str, enable_negative_prompt: bool = False):
-        super().__init__(args, api_url, model)
+    def __init__(self, args, api_url: str, model: str):
+        self.args = args
+        self.api_url = api_url
+        self.model = 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)
-            self._weights = [p["weight"] for p in self.random_request_config]
-
-            self.random_request_config = [
-                {k: v for k, v in p.items() if k != "weight"} for p in self.random_request_config
-            ]
-
-            seed = getattr(args, "random_request_seed", 42)
-            self._rng = random.Random(seed)
-
-            self._sampled_requests = self._rng.choices(
-                self.random_request_config,
-                weights=self._weights,
-                k=self.num_prompts,
-            )
-        else:
-            self._sampled_requests = None
-
-        # Random image generate
-        if self.args.task in ["i2v", "ti2v", "ti2i", "i2i", "it2i"]:
-            self._random_image_path = self._generate_random_image_paths()
-        else:
-            self._random_image_path = None
 
     def __len__(self) -> int:
         return self.num_prompts
 
     def __getitem__(self, idx: int) -> RequestFuncInput:
-        extra_body = {}
-        if self.enable_negative_prompt:
-            extra_body["negative_prompt"] = f"Negative prompt {idx} for benchmarking diffusion models"
-
-        params = {
-            "width": self.args.width,
-            "height": self.args.height,
-            "num_frames": self.args.num_frames,
-            "num_inference_steps": self.args.num_inference_steps,
-            "fps": self.args.fps,
-        }
-        if self._sampled_requests:
-            profile = self._sampled_requests[idx]
-            params.update(profile)
         return RequestFuncInput(
             prompt=f"Random prompt {idx} for benchmarking diffusion models",
             api_url=self.api_url,
             model=self.model,
+            width=self.args.width,
+            height=self.args.height,
+            num_frames=self.args.num_frames,
+            num_inference_steps=self.args.num_inference_steps,
             seed=self.args.seed,
-            extra_body=extra_body,
-            image_paths=self._random_image_path,
-            **params,
+            fps=self.args.fps,
         )
 
     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.
@@ -883,69 +647,23 @@ async def iter_requests(
     requests_list: list[RequestFuncInput],
     request_rate: float,
 ) -> AsyncGenerator[RequestFuncInput, None]:
-    """Yield requests using a Poisson process if request_rate is set.
+    """Yield requests using a fixed interval if request_rate is set.
 
     - If request_rate is inf, all requests are yielded immediately (no sleep).
-    - Otherwise, inter-arrival times follow an exponential distribution.
+    - Otherwise, requests are emitted at a fixed cadence of 1 / request_rate seconds.
     """
 
     if request_rate != float("inf"):
         if request_rate <= 0:
             raise ValueError(f"request_rate must be positive or inf, got {request_rate}.")
+        interval_s = 1.0 / float(request_rate)
 
     for i, req in enumerate(requests_list):
         if request_rate != float("inf") and i > 0:
-            interval_s = random.expovariate(request_rate)
             await asyncio.sleep(interval_s)
         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,
@@ -960,15 +678,6 @@ def calculate_metrics(
     latencies = [o.latency for o in success_outputs]
     peak_memories = [o.peak_memory_mb for o in success_outputs if o.peak_memory_mb > 0]
 
-    # Aggregate per-stage durations across all successful requests that reported them.
-    stage_duration_lists: dict[str, list[float]] = {}
-    for o in success_outputs:
-        for stage, duration in (o.stage_durations or {}).items():
-            stage_duration_lists.setdefault(stage, []).append(duration)
-    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()}
-    stage_durations_p99 = {s: float(np.percentile(v, 99)) for s, v in stage_duration_lists.items()}
-
     metrics = {
         "duration": total_duration,
         "completed_requests": num_success,
@@ -977,14 +686,10 @@ def calculate_metrics(
         "latency_mean": np.mean(latencies) if latencies else 0,
         "latency_median": np.median(latencies) if latencies else 0,
         "latency_p99": np.percentile(latencies, 99) if latencies else 0,
-        "latency_p95": np.percentile(latencies, 95) if latencies else 0,
         "latency_p50": np.percentile(latencies, 50) if latencies else 0,
         "peak_memory_mb_max": max(peak_memories) if peak_memories else 0,
         "peak_memory_mb_mean": np.mean(peak_memories) if peak_memories else 0,
         "peak_memory_mb_median": np.median(peak_memories) if peak_memories else 0,
-        "stage_durations_mean": stage_durations_mean,
-        "stage_durations_p50": stage_durations_p50,
-        "stage_durations_p99": stage_durations_p99,
     }
 
     if slo_enabled:
@@ -1032,52 +737,13 @@ def wait_for_service(base_url: str, timeout: int = 120) -> None:
         time.sleep(1)
 
 
-def _default_endpoint_for_task(task: str) -> str:
-    if task in {"t2v", "i2v", "ti2v"}:
-        return "/v1/videos"
-    if task in {"i2i", "ti2i", "it2i"}:
-        return "/v1/images/edits"
-    if task == "t2i":
-        return "/v1/chat/completions"
-    raise ValueError(f"Unsupported task for endpoint resolution: {task}")
-
-
 async def benchmark(args):
     # Construct base_url if not provided
     if args.base_url is None:
         args.base_url = f"http://{args.host}:{args.port}"
 
-    VIDEO_TASKS = {"t2v", "i2v", "ti2v"}
-    IMAGE_TASKS = {"t2i", "i2i", "ti2i", "it2i"}
-
-    if args.task in VIDEO_TASKS:
-        task_type = "2v"
-    elif args.task in IMAGE_TASKS:
-        task_type = "2i"
-    else:
-        raise ValueError(
-            f"Unsupported task: '{args.task}'. "
-            f"Valid video tasks: {sorted(VIDEO_TASKS)}, "
-            f"Valid image tasks: {sorted(IMAGE_TASKS)}"
-        )
-
-    raw_endpoint = args.endpoint if args.endpoint is not None else args.backend
-    if raw_endpoint is None:
-        raw_endpoint = _default_endpoint_for_task(args.task)
-    args.endpoint = normalize_endpoint(raw_endpoint)
-
-    valid_endpoints = sorted(backends_function_mapping[task_type].keys())
-
-    if args.endpoint not in valid_endpoints:
-        logger.error(
-            f"Invalid endpoint '{args.endpoint}' for task '{args.task}' (task type: '{task_type}').\n"
-            f"Valid endpoints for this task type: {valid_endpoints}\n"
-            f"Example usage: --task {args.task} --endpoint {valid_endpoints[0]}"
-        )
-        raise ValueError("Endpoint validation failed. See log above for valid options.")
-
-    # Setup API URL and request function based on endpoint.
-    request_func, api_url = backends_function_mapping[task_type][args.endpoint]
+    # Setup API URL and request function based on backend
+    request_func, api_url = backends_function_mapping[args.backend]
     api_url = f"{args.base_url}{api_url}"
 
     if args.dataset == "vbench":
@@ -1085,9 +751,7 @@ async def benchmark(args):
     elif args.dataset == "trace":
         dataset = TraceDataset(args, api_url, args.model)
     elif args.dataset == "random":
-        dataset = RandomDataset(args, api_url, args.model, args.enable_negative_prompt)
-    elif args.dataset == "custom":
-        dataset = CustomDataset(args, api_url, args.model)
+        dataset = RandomDataset(args, api_url, args.model)
     else:
         raise ValueError(f"Unknown dataset: {args.dataset}")
 
@@ -1095,10 +759,6 @@ async def benchmark(args):
     requests_list = dataset.get_requests()
     print(f"Prepared {len(requests_list)} requests from {args.dataset} dataset.")
 
-    if args.endpoint == "/v1/images/edits":
-        for req in requests_list:
-            req.default_bot_task = args.bot_task
-
     # Limit concurrency
     if args.max_concurrency is not None:
         semaphore = asyncio.Semaphore(args.max_concurrency)
@@ -1116,12 +776,21 @@ async def benchmark(args):
     pbar = tqdm(total=len(requests_list), disable=args.disable_tqdm)
 
     async with aiohttp.ClientSession() as session:
-        warmup_pairs = await _run_warmups(
-            requests_list=requests_list,
-            args=args,
-            session=session,
-            request_func=request_func,
-        )
+        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,
+                    )
+                warm_out = await limited_request_func(warm_req, session, None)
+                warmup_pairs.append((warm_req, warm_out))
 
         if args.slo:
             # Prefer trace-provided per-request slo_ms. Only populate when missing.
@@ -1146,17 +815,15 @@ async def benchmark(args):
     metrics = calculate_metrics(outputs, total_duration, requests_list, args, args.slo)
 
     # Add configuration info to metrics for JSON output
-    metrics["endpoint"] = args.endpoint
+    metrics["backend"] = args.backend
     metrics["model"] = args.model
     metrics["dataset"] = args.dataset
     metrics["task"] = args.task
-    if args.endpoint == "/v1/images/edits":
-        metrics["bot_task"] = args.bot_task
 
-    print("\n{s:{c}^{n}}".format(s=" Serving Benchmark Result ", n=50, c="="))
+    print("\n{s:{c}^{n}}".format(s=" Serving Benchmark Result ", n=60, c="="))
 
     # Section 1: Configuration
-    print("{:<40} {:<15}".format("Endpoint:", args.endpoint))
+    print("{:<40} {:<15}".format("Backend:", args.backend))
     print("{:<40} {:<15}".format("Model:", args.model))
     print("{:<40} {:<15}".format("Dataset:", args.dataset))
     print("{:<40} {:<15}".format("Task:", args.task))
@@ -1180,7 +847,6 @@ async def benchmark(args):
     print("{:<40} {:<15.4f}".format("Latency Mean (s):", metrics["latency_mean"]))
     print("{:<40} {:<15.4f}".format("Latency Median (s):", metrics["latency_median"]))
     print("{:<40} {:<15.4f}".format("Latency P99 (s):", metrics["latency_p99"]))
-    print("{:<40} {:<15.4f}".format("Latency P95 (s):", metrics["latency_p95"]))
 
     if args.slo:
         print(f"{'-' * 50}")
@@ -1194,12 +860,6 @@ async def benchmark(args):
         print("{:<40} {:<15.2f}".format("Peak Memory Mean (MB):", metrics["peak_memory_mb_mean"]))
         print("{:<40} {:<15.2f}".format("Peak Memory Median (MB):", metrics["peak_memory_mb_median"]))
 
-    if metrics["stage_durations_mean"]:
-        print(f"{'-' * 50}")
-        print("Stage Durations Mean (s):")
-        for stage, val in metrics["stage_durations_mean"].items():
-            print("{:<40} {:<15.4f}".format(f"  {stage}:", val))
-
     print("\n" + "=" * 60)
 
     if args.output_file:
@@ -1219,30 +879,25 @@ if __name__ == "__main__":
     parser.add_argument("--host", type=str, default="localhost", help="Server host.")
     parser.add_argument("--port", type=int, default=8091, help="Server port.")
     parser.add_argument("--model", type=str, default="default", help="Model name.")
-    parser.add_argument(
-        "--endpoint",
-        type=str,
-        default=None,
-        help=("Endpoint path to target. Leading '/' is optional, e.g. /v1/videos or v1/videos."),
-    )
     parser.add_argument(
         "--backend",
         type=str,
-        default=None,
-        help=argparse.SUPPRESS,
+        default="vllm-omni",
+        choices=["vllm-omni", "openai"],
+        help="Backend to target the benchmark to.",
     )
     parser.add_argument(
         "--dataset",
         type=str,
         default="vbench",
-        choices=["vbench", "trace", "random", "custom"],
+        choices=["vbench", "trace", "random"],
         help="Dataset to use.",
     )
     parser.add_argument(
         "--task",
         type=str,
         default="t2v",
-        choices=["t2v", "i2v", "ti2v", "ti2i", "i2i", "it2i", "t2i"],
+        choices=["t2v", "i2v", "ti2v", "ti2i", "i2i", "t2i"],
         help="Task type.",
     )
     parser.add_argument(
@@ -1284,13 +939,6 @@ 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).")
@@ -1324,40 +972,7 @@ if __name__ == "__main__":
         help="SLO target multiplier: slo_ms = estimated_exec_time_ms * slo_scale (default: 3).",
     )
     parser.add_argument("--disable-tqdm", action="store_true", help="Disable progress bar.")
-    parser.add_argument(
-        "--enable-negative-prompt",
-        action="store_true",
-        default=False,
-        help="Generate negative prompts when using the random dataset.",
-    )
-    parser.add_argument(
-        "--random-request-config",
-        type=str,
-        default=None,
-        help=(
-            "JSON string defining random request profiles. "
-            "Each profile may contain: width, height, num_inference_steps, etc. "
-            "The 'weight' field controls sampling probability (relative weight). "
-            "Example: "
-            '[{"width":512,"height":512,"num_inference_steps":20,"weight":0.15},'
-            '{"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, it2i) when using random dataset."
-        ),
-    )
-    parser.add_argument(
-        "--bot-task",
-        type=str,
-        default="think",
-        help=("bot_task form field for --endpoint /v1/images/edits (think, recaption, think_recaption, vanilla)."),
-    )
 
     args = parser.parse_args()
+
     asyncio.run(benchmark(args))

benchmarks/diffusion/performance_dashboard/qwen_image_serving_performance.md

@@ -1,183 +0,0 @@
-# Qwen-Image Serving Performance Dashboard
-
-This document describes how to deploy and benchmark **Qwen-Image** using vLLM-Omni. It includes service startup configuration, acceleration-related options, benchmark methodology, dataset settings, and performance results.
-
----
-
-# 1. Overview
-
-Qwen-Image is a multimodal text-to-image generation model served through the vLLM-Omni infrastructure.
-
-This document covers:
-
-* Service launch configuration (including acceleration options)
-* Benchmark scripts and usage
-* Dataset and workload settings
-* Performance measurement results
-* Reproducibility guidelines
-
----
-
-# 2. Test Environment
-| Component | Specification |
-|------------|----------------|
-| GPU | NVIDIA A100-SXM4-80GB |
-| Diffusion Attention Backend | FlashAttention |
-
-# 3. Service Launch Configuration
-
-## 3.1 Basic Serving Command
-
-```bash
-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              |
-| --------------------- | ------------------------ |
-| `--cfg-parallel-size` | CFG parallelism degree   |
-| `--ulysses-degree`    | Ulysses parallel degree  |
-| `--vae-patch-parallel-size`    | VAE parallel degree  |
-| `--tensor-parallel-size` | Tensor parallelism degree |
-
-Record these parameters when reporting performance results.
-
----
-
-# 4. Benchmark Script
-
-## 4.1 Benchmark Entry
-
-```bash
-python benchmarks/diffusion/diffusion_benchmark_serving.py \
-    --endpoint /v1/chat/completions \
-    --dataset <DATASET_NAME> \
-    --task t2i \
-    --num-prompts <N> \
-    --max-concurrency <C> \
-    --enable-negative-prompt \
-    --random-request-config <CFG>
-```
-
-## 4.2 Key Benchmark Arguments
-
-| Parameter              | Description                       |
-| ---------------------- | --------------------------------- |
-| `--endpoint`           | API endpoint (use `/v1/chat/completions`; leading `/` is optional) |
-| `--dataset`            | Dataset name (`random` or custom) |
-| `--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 |
-
----
-
-# 5. Dataset & Workload Settings
-
-## 5.1 Recommended Evaluation Configurations
-
-### Dataset A ( 512 Resolution)
-
-* Dataset: `random`
-* Task: t2i
-* Concurrency: 1
-* Mix Resolution
-```
-[
- {"width":512,"height":512,"num_inference_steps":20,"weight":1}
-]
-```
-
-### Dataset B (1536 Resolution)
-
-* Dataset: `random`
-* Task: t2i
-* Concurrency: 1
-* Mix Resolution
-```
-[
- {"width":1536,"height":1536,"num_inference_steps":35,"weight":1}
-]
-```
-
-### Dataset C (Mix Resolution)
-
-* Dataset: `random`
-* Task: t2i
-* Concurrency: 1
-* Mix Resolution
-```
-[
- {"width":512,"height":512,"num_inference_steps":20,"weight":0.15},
- {"width":768,"height":768,"num_inference_steps":20,"weight":0.25},
- {"width":1024,"height":1024,"num_inference_steps":25,"weight":0.45},
- {"width":1536,"height":1536,"num_inference_steps":35,"weight":0.15}
-]
-```
----
-
-## 5.2 Example Benchmark Command
-
-```bash
-python benchmarks/diffusion/diffusion_benchmark_serving.py \
-    --endpoint /v1/chat/completions \
-    --dataset random \
-    --task t2i \
-    --num-prompts 1 \
-    --max-concurrency 1 \
-    --enable-negative-prompt \
-    --random-request-config '[
-        {"width":512,"height":512,"num_inference_steps":20,"weight":1}
-    ]'
-```
-
-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
-
-The following metrics are collected during benchmarking:
-
-| Metric             | Description                   | Unit    |
-| ------------------ | ----------------------------- | ------- |
-| Mean Latency        | Mean of latency       | seconds |
-| P99 Latency        | P99 of latency             | seconds |
-
----
-
-# 7. Performance Results
-
-| Dataset Configuration | Max Concur. | CFG | Usp | Tp | VAE Parallel | Mean Latency (s) | P99 Latency (s) |
-|-----------------------|-----|-----|-----|----|--------------|------------------|------------------|
-| Dataset A | 1 | 2 | 2 | Off | Off          | 2.2087          | 2.2087          |
-| Dataset B | 1 | 2 | 2 | Off | Off          | 19.6739          | 19.6739          |
-| Dataset C | 1 | 2 | 2 | Off | Off          | 5.67259        | 18.6234 |
----
-
-# 8. Reproducibility Checklist
-
-To ensure consistent and comparable benchmark results:
-
-* Record GPU type
-* 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
-
----
-
-This document serves as the official Qwen-Image serving performance reference under vLLM-Omni.

benchmarks/diffusion/performance_dashboard/wan_2_2_serving_performance.md

@@ -1,170 +0,0 @@
-# Wan2.2 Serving Performance Dashboard
-
-This document describes how to deploy and benchmark **Wan-AI/Wan2.2-T2V-A14B-Diffusers** using vLLM-Omni. It includes service startup configuration, acceleration-related options, benchmark methodology, dataset settings, and performance results.
-
----
-
-# 1. Overview
-
-Wan-AI/Wan2.2-T2V-A14B-Diffusers is a multimodal text-to-video generation model served through the vLLM-Omni infrastructure.
-
-This document covers:
-
-* Service launch configuration (including acceleration options)
-* Benchmark scripts and usage
-* Dataset and workload settings
-* Performance measurement results
-* Reproducibility guidelines
-
----
-
-# 2. Test Environment
-| Component | Specification |
-|------------|----------------|
-| GPU | NVIDIA A100-SXM4-80GB |
-| Diffusion Attention Backend | FlashAttention |
-
-# 3. Service Launch Configuration
-
-## 3.1 Basic Serving Command
-
-```bash
-vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni \
-    --port 8091
-```
-
-## 3.2 Key Parameters
-
-| Parameter             | Description              |
-| --------------------- | ------------------------ |
-| `--cfg-parallel-size` | CFG parallelism degree   |
-| `--ulysses-degree`    | Ulysses parallel degree  |
-| `--vae-patch-parallel-size`    | VAE parallel degree  |
-| `--tensor-parallel-size` | Tensor parallelism degree |
-| `--use-hsdp` | Enable Hybrid Sharded Data Parallel to shard model weights across GPUs |
-
-Record these parameters when reporting performance results.
-
----
-
-# 4. Benchmark Script
-
-## 4.1 Benchmark Entry
-
-```bash
-python benchmarks/diffusion/diffusion_benchmark_serving.py \
-    --endpoint /v1/videos \
-    --dataset <DATASET_NAME> \
-    --task t2v \
-    --num-prompts <N> \
-    --max-concurrency <C> \
-    --enable-negative-prompt \
-    --random-request-config <CFG>
-```
-
-## 4.2 Key Benchmark Arguments
-
-| Parameter              | Description                       |
-| ---------------------- | --------------------------------- |
-| `--endpoint`           | API endpoint (use `/v1/videos`; leading `/` is optional) |
-| `--dataset`            | Dataset name (`random` or custom) |
-| `--task`               | Task type (e.g., `t2v`)           |
-| `--num-prompts`        | Total number of requests          |
-| `--max-concurrency`    | Client-side concurrency           |
-| `--random-request-config`| JSON string defining random request |
-
----
-
-# 5. Dataset & Workload Settings
-
-## 5.1 Recommended Evaluation Configurations
-
-### Dataset A (480p)
-
-* Dataset: `random`
-* Task: t2v
-* Concurrency: 1
-* Mix Resolution
-```
-[
-    {"width":854,"height":480,"num_inference_steps":3,"num_frames":80,"fps":16,"weight":1}
-]
-```
-### Dataset B (720p)
-
-* Dataset: `random`
-* Task: t2v
-* Concurrency: 1
-* Mix Resolution
-```
-[
-    {"width":1280,"height":720,"num_inference_steps":6,"num_frames":80,"fps":16,"weight":1}
-]
-```
-### Dataset C (Mix Resolution)
-
-* Dataset: `random`
-* Task: t2v
-* Concurrency: 1
-* Mix Resolution
-```
-[
- {"width":854,"height":480,"num_inference_steps":3,"num_frames":80,"fps":16,"weight":0.15},
- {"width":854,"height":480,"num_inference_steps":4,"num_frames":120,"fps":24,"weight":0.25},
- {"width":1280,"height":720,"num_inference_steps":6,"num_frames":80,"fps":16,"weight":0.6}
-]
-```
----
-
-## 5.2 Example Benchmark Command
-
-```bash
-python benchmarks/diffusion/diffusion_benchmark_serving.py \
-    --endpoint /v1/videos \
-    --dataset random \
-    --task t2v \
-    --num-prompts 1 \
-    --max-concurrency 1 \
-    --enable-negative-prompt \
-    --random-request-config '[
-        {"width":854,"height":480,"num_inference_steps":18,"num_frames": 33,"fps":16",weight":1}
-    ]'
-```
-
----
-
-# 6. Performance Metrics
-
-The following metrics are collected during benchmarking:
-
-| Metric             | Description                   | Unit    |
-| ------------------ | ----------------------------- | ------- |
-| Mean Latency        | Mean of latency       | seconds |
-| P99 Latency        | P99 of latency             | seconds |
-
----
-
-# 7. Performance Results
-
-| Dataset Configuration | Max Concur. | CFG | Usp | Tp | Hsdp | VAE Parallel | Mean Latency (s) | P99 Latency (s) |
-|-----------------------|-----|-----|-----|-----|----|--------------|------------------|------------------|
-| Dataset A | 1 | 2 | 2 | 1 | On | 1          | 24.6766          | 24.6766          |
-| Dataset A | 1 | 2 | 2 | 1 | On | 4          | 21.6810          | 21.6810          |
-| Dataset B | 1 | 2 | 2 | 1 | On | 1          | 124.6639         | 124.6639          |
-| Dataset B | 1 | 2 | 2 | 1 | On | 4          | 117.44          | 117.44          |
-| Dataset C | 1 | 2 | 2 | 1 | On | 1          | 79.2175        | 124.2565 |
-| Dataset C | 1 | 2 | 2 | 1 | On | 4          | 74.4977        | 117.710 |
----
-
-# 8. Reproducibility Checklist
-
-To ensure consistent and comparable benchmark results:
-
-* Record GPU type
-* Record parallel configuration
-* Record benchmark parameters (resolution, concurrency, number of prompts)
-* Ensure no background workload on GPUs during testing
-
----
-
-This document serves as the official Wan2.2 serving performance reference under vLLM-Omni.

benchmarks/diffusion/quantization_quality.py

@@ -1,487 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""
-Benchmark quantization quality loss for diffusion models (image & video).
-
-Generates outputs with BF16 (baseline) and a quantized config using the same
-seed, then computes LPIPS perceptual distance between them. Results are printed
-as a Markdown table ready to paste into a PR description.
-
-Requirements:
-    pip install lpips Pillow numpy
-
-Image example (text-to-image):
-    python benchmarks/diffusion/quantization_quality.py \
-        --model Tongyi-MAI/Z-Image-Turbo \
-        --task t2i \
-        --quantization fp8 \
-        --prompts \
-            "an aerial view of a coral reef with crystal clear turquoise water" \
-            "a campfire in a dark forest with sparks rising into a starry sky" \
-            "a gourmet dessert plate with chocolate mousse and gold leaf" \
-        --height 1024 --width 1024 \
-        --num-inference-steps 50 --seed 42
-
-Video example (text-to-video):
-    python benchmarks/diffusion/quantization_quality.py \
-        --model Wan-AI/Wan2.2-T2V-A14B-Diffusers \
-        --task t2v \
-        --quantization fp8 \
-        --prompts \
-            "A serene lakeside sunrise with mist over the water" \
-            "A cat walking across a wooden bridge in autumn" \
-        --height 720 --width 1280 \
-        --num-frames 81 --num-inference-steps 40 --seed 42
-
-LTX-2 example (text-to-video; audio output is generated but not scored by LPIPS):
-    python benchmarks/diffusion/quantization_quality.py \
-        --model Lightricks/LTX-Video-2 \
-        --task t2v \
-        --quantization fp8 int8 \
-        --prompts \
-            "A serene lakeside sunrise with mist over the water" \
-        --height 704 --width 1216 \
-        --num-frames 121 --num-inference-steps 40 --seed 42
-
-Multiple quantization methods:
-    python benchmarks/diffusion/quantization_quality.py \
-        --model Tongyi-MAI/Z-Image-Turbo \
-        --task t2i \
-        --quantization fp8 int8 bitsandbytes \
-        --prompts "a cup of coffee on the table" \
-        --height 1024 --width 1024 \
-        --num-inference-steps 50 --seed 42
-
-Output directory structure (--output-dir, default: ./quant_bench_output):
-    quant_bench_output/
-        baseline/           # BF16 outputs
-        <method>/           # Quantized outputs per method
-        results.md          # Markdown table
-"""
-
-import argparse
-import gc
-import time
-from pathlib import Path
-
-import numpy as np
-import torch
-
-
-def compute_lpips_images(
-    baseline_images: list,
-    quantized_images: list,
-    net: str = "alex",
-) -> list[float]:
-    """Compute LPIPS between paired lists of PIL images."""
-    import lpips
-    from torchvision import transforms
-
-    loss_fn = lpips.LPIPS(net=net).eval()
-    if torch.cuda.is_available():
-        loss_fn = loss_fn.cuda()
-
-    transform = transforms.Compose(
-        [
-            transforms.Resize((256, 256)),
-            transforms.ToTensor(),
-            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
-        ]
-    )
-
-    scores = []
-    for img_bl, img_qt in zip(baseline_images, quantized_images):
-        t_bl = transform(img_bl.convert("RGB")).unsqueeze(0)
-        t_qt = transform(img_qt.convert("RGB")).unsqueeze(0)
-        if torch.cuda.is_available():
-            t_bl, t_qt = t_bl.cuda(), t_qt.cuda()
-        with torch.no_grad():
-            score = loss_fn(t_bl, t_qt).item()
-        scores.append(score)
-    return scores
-
-
-def compute_lpips_video(
-    baseline_frames: np.ndarray,
-    quantized_frames: np.ndarray,
-    net: str = "alex",
-) -> float:
-    """Compute mean per-frame LPIPS for a video pair.
-
-    Args:
-        baseline_frames: (F, H, W, C) float array in [0, 1].
-        quantized_frames: same shape.
-
-    Returns:
-        Mean LPIPS across all frames.
-    """
-    import lpips
-
-    loss_fn = lpips.LPIPS(net=net).eval()
-    if torch.cuda.is_available():
-        loss_fn = loss_fn.cuda()
-
-    num_frames = min(len(baseline_frames), len(quantized_frames))
-    scores = []
-    for i in range(num_frames):
-        # Convert (H, W, C) float [0,1] -> (1, C, H, W) float [-1, 1]
-        f_bl = torch.from_numpy(baseline_frames[i]).permute(2, 0, 1).unsqueeze(0).float() * 2 - 1
-        f_qt = torch.from_numpy(quantized_frames[i]).permute(2, 0, 1).unsqueeze(0).float() * 2 - 1
-        if torch.cuda.is_available():
-            f_bl, f_qt = f_bl.cuda(), f_qt.cuda()
-        with torch.no_grad():
-            score = loss_fn(f_bl, f_qt).item()
-        scores.append(score)
-    return float(np.mean(scores))
-
-
-def _build_omni_kwargs(args, quantization=None):
-    """Build kwargs dict for Omni() constructor."""
-    from vllm_omni.diffusion.data import DiffusionParallelConfig
-
-    parallel_config = DiffusionParallelConfig(
-        ulysses_degree=args.ulysses_degree,
-        ring_degree=args.ring_degree,
-        tensor_parallel_size=args.tensor_parallel_size,
-    )
-    kwargs = {
-        "model": args.model,
-        "parallel_config": parallel_config,
-        "enforce_eager": args.enforce_eager,
-    }
-    if quantization:
-        kwargs["quantization_config"] = quantization
-    return kwargs
-
-
-def _generate_image(omni, args, prompt, seed):
-    """Generate a single image and return (PIL.Image, time_seconds, memory_gib)."""
-    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-    from vllm_omni.outputs import OmniRequestOutput
-    from vllm_omni.platforms import current_omni_platform
-
-    generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(seed)
-    torch.accelerator.reset_peak_memory_stats()
-    start = time.perf_counter()
-    outputs = omni.generate(
-        {"prompt": prompt},
-        OmniDiffusionSamplingParams(
-            height=args.height,
-            width=args.width,
-            generator=generator,
-            num_inference_steps=args.num_inference_steps,
-        ),
-    )
-    elapsed = time.perf_counter() - start
-    peak_mem = torch.accelerator.max_memory_allocated() / (1024**3)
-
-    req_out = OmniRequestOutput.unwrap_result(outputs)
-    if not req_out.images:
-        raise ValueError("Could not extract image output from result.")
-    img = req_out.images[0]
-    return img, elapsed, peak_mem
-
-
-def _generate_video(omni, args, prompt, seed):
-    """Generate a video and return (np.ndarray [F,H,W,C], time_seconds, memory_gib)."""
-    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-    from vllm_omni.outputs import OmniRequestOutput
-    from vllm_omni.platforms import current_omni_platform
-
-    generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(seed)
-    torch.accelerator.reset_peak_memory_stats()
-    start = time.perf_counter()
-    outputs = omni.generate(
-        {"prompt": prompt, "negative_prompt": ""},
-        OmniDiffusionSamplingParams(
-            height=args.height,
-            width=args.width,
-            generator=generator,
-            guidance_scale=args.guidance_scale,
-            num_inference_steps=args.num_inference_steps,
-            num_frames=args.num_frames,
-        ),
-    )
-    elapsed = time.perf_counter() - start
-    peak_mem = torch.accelerator.max_memory_allocated() / (1024**3)
-
-    first = outputs[0]
-    if hasattr(first, "request_output") and isinstance(first.request_output, list):
-        inner = first.request_output[0]
-        if isinstance(inner, OmniRequestOutput) and hasattr(inner, "images"):
-            frames = inner.images[0] if inner.images else None
-        else:
-            frames = inner
-    elif hasattr(first, "images") and first.images:
-        frames = first.images[0]
-    else:
-        raise ValueError("Could not extract video frames from output.")
-
-    # LTX-2 (and similar audio+video models) may surface a dict or (video, audio) tuple
-    if isinstance(frames, dict):
-        frames = frames.get("video") or frames.get("frames")
-    elif isinstance(frames, tuple) and len(frames) == 2:
-        frames = frames[0]
-
-    if frames is None:
-        raise ValueError("Could not extract video frames from output.")
-
-    if isinstance(frames, torch.Tensor):
-        video = frames.detach().cpu()
-        if video.dim() == 5:
-            video = video[0].permute(1, 2, 3, 0) if video.shape[1] in (3, 4) else video[0]
-        elif video.dim() == 4 and video.shape[0] in (3, 4):
-            video = video.permute(1, 2, 3, 0)
-        if video.is_floating_point():
-            video = video.clamp(-1, 1) * 0.5 + 0.5
-        frames_array = video.float().numpy()
-    else:
-        frames_array = np.asarray(frames)
-        if frames_array.ndim == 5:
-            frames_array = frames_array[0]
-
-    return frames_array, elapsed, peak_mem
-
-
-def _free_gpu_memory():
-    """Force GC and release cached GPU memory.
-
-    Must be called AFTER the caller has dropped (i.e., via `del`)
-    every reference to the Omni instance"""
-    gc.collect()
-    if torch.cuda.is_available():
-        torch.accelerator.empty_cache()
-        torch.accelerator.synchronize()
-
-
-def run_benchmark(args):
-    from vllm_omni.entrypoints.omni import Omni
-
-    output_dir = Path(args.output_dir)
-    output_dir.mkdir(parents=True, exist_ok=True)
-
-    is_video = args.task == "t2v"
-    prompts = args.prompts
-    seed = args.seed
-
-    # Determine configs to benchmark
-    configs = []  # list of (label, quantization_method)
-    for method in args.quantization:
-        configs.append((method, method))
-
-    # --- Baseline run ---
-    print("\n" + "=" * 60)
-    print("Running BF16 baseline...")
-    print("=" * 60)
-    bl_kwargs = _build_omni_kwargs(args, quantization=None)
-    omni_bl = Omni(**bl_kwargs)
-
-    baseline_outputs = {}  # prompt -> (output, time, mem)
-    for prompt in prompts:
-        print(f"  Generating: {prompt[:60]}...")
-        if is_video:
-            out, t, mem = _generate_video(omni_bl, args, prompt, seed)
-        else:
-            out, t, mem = _generate_image(omni_bl, args, prompt, seed)
-        baseline_outputs[prompt] = (out, t, mem)
-
-    bl_avg_time = np.mean([v[1] for v in baseline_outputs.values()])
-    bl_mem = baseline_outputs[prompts[0]][2]  # use first prompt's memory
-    omni_bl.shutdown()
-    del omni_bl
-    _free_gpu_memory()
-
-    # Save baseline outputs
-    bl_dir = output_dir / "baseline"
-    bl_dir.mkdir(parents=True, exist_ok=True)
-    for i, prompt in enumerate(prompts):
-        out = baseline_outputs[prompt][0]
-        if is_video:
-            try:
-                from diffusers.utils import export_to_video
-
-                frames_list = list(out) if isinstance(out, np.ndarray) and out.ndim == 4 else out
-                export_to_video(frames_list, str(bl_dir / f"prompt_{i}.mp4"), fps=args.fps)
-            except ImportError:
-                np.save(bl_dir / f"prompt_{i}.npy", out)
-        else:
-            out.save(bl_dir / f"prompt_{i}.png")
-
-    # --- Quantized runs ---
-    all_results = []  # list of dicts
-
-    for config_label, quant_method in configs:
-        print(f"\n{'=' * 60}")
-        print(f"Running: {config_label}...")
-        print("=" * 60)
-
-        qt_kwargs = _build_omni_kwargs(args, quantization=quant_method)
-        omni_qt = Omni(**qt_kwargs)
-
-        qt_outputs = {}
-        for prompt in prompts:
-            print(f"  Generating: {prompt[:60]}...")
-            if is_video:
-                out, t, mem = _generate_video(omni_qt, args, prompt, seed)
-            else:
-                out, t, mem = _generate_image(omni_qt, args, prompt, seed)
-            qt_outputs[prompt] = (out, t, mem)
-
-        qt_avg_time = np.mean([v[1] for v in qt_outputs.values()])
-        qt_mem = qt_outputs[prompts[0]][2]
-        omni_qt.shutdown()
-        del omni_qt
-        _free_gpu_memory()
-
-        # Save quantized outputs
-        qt_dir = output_dir / config_label.replace(" ", "_")
-        qt_dir.mkdir(parents=True, exist_ok=True)
-
-        # Compute LPIPS per prompt
-        per_prompt = []
-        for i, prompt in enumerate(prompts):
-            bl_out = baseline_outputs[prompt][0]
-            qt_out = qt_outputs[prompt][0]
-            if is_video:
-                lpips_score = compute_lpips_video(bl_out, qt_out, net=args.lpips_net)
-                try:
-                    from diffusers.utils import export_to_video
-
-                    frames_list = list(qt_out) if isinstance(qt_out, np.ndarray) and qt_out.ndim == 4 else qt_out
-                    export_to_video(frames_list, str(qt_dir / f"prompt_{i}.mp4"), fps=args.fps)
-                except ImportError:
-                    np.save(qt_dir / f"prompt_{i}.npy", qt_out)
-            else:
-                lpips_score = compute_lpips_images([bl_out], [qt_out], net=args.lpips_net)[0]
-                qt_out.save(qt_dir / f"prompt_{i}.png")
-            per_prompt.append({"prompt": prompt, "lpips": lpips_score})
-
-        mean_lpips = np.mean([p["lpips"] for p in per_prompt])
-        speedup = bl_avg_time / qt_avg_time if qt_avg_time > 0 else float("inf")
-        mem_reduction = (bl_mem - qt_mem) / bl_mem * 100
-
-        all_results.append(
-            {
-                "config": config_label,
-                "avg_time": qt_avg_time,
-                "speedup": speedup,
-                "memory_gib": qt_mem,
-                "mem_reduction_pct": mem_reduction,
-                "mean_lpips": mean_lpips,
-                "per_prompt": per_prompt,
-            }
-        )
-
-    # --- Print results ---
-    print("\n\n")
-    print("=" * 80)
-    print("RESULTS")
-    print("=" * 80)
-
-    # Summary table
-    lines = []
-    lines.append(f"## Quantization Quality Benchmark — {args.model.split('/')[-1]}")
-    lines.append(
-        f"Setup: {args.height}x{args.width}, {args.num_inference_steps} steps, "
-        f"seed={args.seed}, LPIPS ({args.lpips_net})"
-    )
-    if is_video:
-        lines.append(f"Video: {args.num_frames} frames")
-    lines.append("")
-    lines.append("### Summary")
-    lines.append("")
-    lines.append("| Config | Avg Time | Speedup | Memory (GiB) | Mem Reduction | Mean LPIPS |")
-    lines.append("|--------|----------|---------|--------------|---------------|------------|")
-    lines.append(f"| BF16 baseline | {bl_avg_time:.2f}s | 1.00x | {bl_mem:.2f} | — | (ref) |")
-    for r in all_results:
-        lines.append(
-            f"| {r['config']} | {r['avg_time']:.2f}s | {r['speedup']:.2f}x "
-            f"| {r['memory_gib']:.2f} | {r['mem_reduction_pct']:.0f}% "
-            f"| {r['mean_lpips']:.4f} |"
-        )
-    lines.append("")
-    lines.append("> LPIPS < 0.01 = imperceptible, > 0.1 = clearly noticeable.")
-    lines.append("")
-
-    # Per-prompt table
-    if len(prompts) > 1:
-        lines.append("### Per-Prompt LPIPS")
-        lines.append("")
-        header = "| Prompt |"
-        sep = "|--------|"
-        for r in all_results:
-            header += f" {r['config']} |"
-            sep += "--------|"
-        lines.append(header)
-        lines.append(sep)
-        for i, prompt in enumerate(prompts):
-            short = prompt[:50] + "..." if len(prompt) > 50 else prompt
-            row = f"| {short} |"
-            for r in all_results:
-                row += f" {r['per_prompt'][i]['lpips']:.4f} |"
-            lines.append(row)
-        lines.append("")
-
-    md = "\n".join(lines)
-    print(md)
-
-    # Save markdown
-    results_path = output_dir / "results.md"
-    results_path.write_text(md, encoding="utf-8")
-    print(f"\nResults saved to {results_path}")
-    print(f"Baseline outputs in {bl_dir}")
-    for r in all_results:
-        qt_dir = output_dir / r["config"].replace(" ", "_")
-        print(f"Quantized outputs in {qt_dir}")
-
-
-def parse_args():
-    parser = argparse.ArgumentParser(
-        description="Benchmark quantization quality loss for diffusion models.",
-        formatter_class=argparse.RawDescriptionHelpFormatter,
-    )
-    parser.add_argument("--model", required=True, help="Model name or local path.")
-    parser.add_argument(
-        "--task",
-        default="t2i",
-        choices=["t2i", "t2v"],
-        help="Task type: t2i (text-to-image) or t2v (text-to-video).",
-    )
-    parser.add_argument(
-        "--quantization",
-        nargs="+",
-        required=True,
-        help="One or more quantization methods to benchmark (e.g. fp8 int8 bitsandbytes).",
-    )
-    parser.add_argument(
-        "--prompts",
-        nargs="+",
-        default=["a cup of coffee on the table"],
-        help="One or more prompts to generate.",
-    )
-    parser.add_argument("--seed", type=int, default=42)
-    parser.add_argument("--height", type=int, default=1024)
-    parser.add_argument("--width", type=int, default=1024)
-    parser.add_argument("--num-inference-steps", type=int, default=50)
-    parser.add_argument("--num-frames", type=int, default=81, help="Number of video frames (t2v only).")
-    parser.add_argument("--fps", type=int, default=24, help="Video FPS for saving (t2v only).")
-    parser.add_argument("--guidance-scale", type=float, default=4.0, help="CFG scale (used for video).")
-    parser.add_argument("--output-dir", type=str, default="./quant_bench_output", help="Directory to save outputs.")
-    parser.add_argument(
-        "--lpips-net",
-        type=str,
-        default="alex",
-        choices=["alex", "vgg", "squeeze"],
-        help="LPIPS backbone network.",
-    )
-    parser.add_argument("--ulysses-degree", type=int, default=1)
-    parser.add_argument("--ring-degree", type=int, default=1)
-    parser.add_argument("--tensor-parallel-size", type=int, default=1)
-    parser.add_argument("--enforce-eager", action="store_true")
-    return parser.parse_args()
-
-
-if __name__ == "__main__":
-    args = parse_args()
-    run_benchmark(args)

benchmarks/distributed/omni_connectors/cross_node_mooncake_transfer_engine.py

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

benchmarks/fish-speech/bench_speaker_cache.py

@@ -1,290 +0,0 @@
-"""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())