Merge branch 'main' into fix/rebase-pipeline-upload-all-e2e

Signed-off-by: Weiming Liao <liaowm5@gmail.com>

.buildkite/bootstrap-amd-omni.sh

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

.buildkite/bootstrap-intel-omni.sh

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

@@ -0,0 +1,16 @@
+# .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.23.0"
+        priority: 100
+        timeout_in_minutes: 60
+        soft_fail: false

.buildkite/pipeline-npu-a3.yaml

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

.buildkite/pipeline-npu.yaml

@@ -0,0 +1,112 @@
+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,256 +1,77 @@
+# Document 1: Buildkite loads only this block on first parse. The next step resolves docs-only,
+# pytest skip-mark-only, or combined skip-ci from git diff, then uploads document 2. When skip-ci
+# applies, image-build still runs for the nightly exception in upload_pipeline.py (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.py (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 (docs / skip marks) & upload pipeline"
+    key: upload-ci-pipeline
+    commands:
+      - "python3 .buildkite/scripts/upload_pipeline.py --upload .buildkite/pipeline.yml"
+    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 --file docker/Dockerfile.ci -t vllm-omni-ci ."
+      - "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"
 
-  # - label: "Test on NPU"
-  #   depends_on: ~
-  #   key: npu-test
-  #   commands:
-  #     - ".buildkite/scripts/hardware_ci/run_npu_test.sh"
-  #   agents:
-  #     queue: "ascend"
-
-  - label: "Simple Unit Test"
-    depends_on: ~
+  # L2 Test
+  - label: "Upload Ready Pipeline"
+    depends_on: image-build
+    key: upload-ready-pipeline
+    if: __UPLOAD_READY_IF__
     commands:
-      - ".buildkite/scripts/simple_test.sh"
+      - python3 .buildkite/scripts/upload_pipeline.py --upload .buildkite/test-ready.yml
     agents:
       queue: "cpu_queue_premerge"
 
-  - label: "Diffusion Model Test"
-    timeout_in_minutes: 20
+  # L3 Test
+  - label: "Upload Merge Pipeline"
     depends_on: image-build
+    key: upload-merge-pipeline
+    if: __UPLOAD_MERGE_IF__
     commands:
-      - pytest -s -v tests/e2e/offline_inference/test_t2i_model.py
+      - python3 .buildkite/scripts/upload_pipeline.py --upload .buildkite/test-merge.yml
     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"
+      queue: "cpu_queue_premerge"
 
-  - label: "Diffusion Images API LoRA E2E"
-    timeout_in_minutes: 20
+  # 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:
-      - pytest -s -v tests/e2e/online_serving/test_images_generations_lora.py
+      - buildkite-agent pipeline upload .buildkite/test-nightly.yml
     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"
+      queue: "cpu_queue_premerge"
 
-  - label: "Diffusion Model CPU offloading Test"
-    timeout_in_minutes: 20
+  # 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:
-      - pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
+      - buildkite-agent pipeline upload .buildkite/test-weekly.yml
     agents:
-      queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
-    plugins:
-      - docker#v5.2.0:
-          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-          always-pull: true
-          propagate-environment: true
-          environment:
-            - "HF_HOME=/fsx/hf_cache"
-          volumes:
-            - "/fsx/hf_cache:/fsx/hf_cache"
-
-  - label: "Audio Generation Model Test"
-    timeout_in_minutes: 20
-    depends_on: image-build
-    commands:
-      - pytest -s -v tests/e2e/offline_inference/test_stable_audio_model.py
-    agents:
-      queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
-    plugins:
-      - docker#v5.2.0:
-          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-          always-pull: true
-          propagate-environment: true
-          environment:
-            - "HF_HOME=/fsx/hf_cache"
-          volumes:
-            - "/fsx/hf_cache:/fsx/hf_cache"
-
-  - label: "Diffusion Cache Backend Test"
-    timeout_in_minutes: 15
-    depends_on: image-build
-    commands:
-      - pytest -s -v tests/e2e/offline_inference/test_cache_dit.py tests/e2e/offline_inference/test_teacache.py
-    agents:
-      queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
-    plugins:
-      - docker#v5.2.0:
-          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-          always-pull: true
-          propagate-environment: true
-          environment:
-            - "HF_HOME=/fsx/hf_cache"
-          volumes:
-            - "/fsx/hf_cache:/fsx/hf_cache"
-
-  - label: "Diffusion Sequence Parallelism Test"
-    timeout_in_minutes: 20
-    depends_on: image-build
-    commands:
-      - pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py
-    agents:
-      queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
-    plugins:
-      - docker#v5.2.0:
-          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-          always-pull: true
-          propagate-environment: true
-          shm-size: "8gb"
-          environment:
-            - "HF_HOME=/fsx/hf_cache"
-          volumes:
-            - "/fsx/hf_cache:/fsx/hf_cache"
-
-  - label: "Diffusion Tensor Parallelism Test"
-    timeout_in_minutes: 20
-    depends_on: image-build
-    commands:
-      - pytest -s -v tests/e2e/offline_inference/test_zimage_tensor_parallel.py
-    agents:
-      queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
-    plugins:
-      - docker#v5.2.0:
-          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-          always-pull: true
-          propagate-environment: true
-          shm-size: "8gb"
-          environment:
-            - "HF_HOME=/fsx/hf_cache"
-          volumes:
-            - "/fsx/hf_cache:/fsx/hf_cache"
-
-  - label: "Diffusion GPU Worker Test"
-    timeout_in_minutes: 20
-    depends_on: image-build
-    commands:
-      - pytest -s -v tests/diffusion/test_gpu_diffusion_worker.py
-    agents:
-      queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
-    plugins:
-      - docker#v5.2.0:
-          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-          always-pull: true
-          propagate-environment: true
-          shm-size: "8gb"
-          environment:
-            - "HF_HOME=/fsx/hf_cache"
-          volumes:
-            - "/fsx/hf_cache:/fsx/hf_cache"
-
-  - label: "Omni Model Test"
-    timeout_in_minutes: 15
-    depends_on: image-build
-    commands:
-      - export VLLM_LOGGING_LEVEL=DEBUG
-      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
-    agents:
-      queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
-    plugins:
-      - docker#v5.2.0:
-          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-          always-pull: true
-          propagate-environment: true
-          environment:
-            - "HF_HOME=/fsx/hf_cache"
-          volumes:
-            - "/fsx/hf_cache:/fsx/hf_cache"
-
-  - label: "Omni Model Test with H100"
-    timeout_in_minutes: 30
-    depends_on: image-build
-    commands:
-      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - export VLLM_TEST_CLEAN_GPU_MEMORY="1"
-      - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
-      - pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py
-      - pytest -s -v tests/e2e/online_serving/test_async_omni.py
-    agents:
-      queue: "mithril-h100-pool"
-    plugins:
-      - kubernetes:
-          podSpec:
-            containers:
-              - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-                resources:
-                  limits:
-                    nvidia.com/gpu: 2
-                volumeMounts:
-                  - name: devshm
-                    mountPath: /dev/shm
-                  - name: hf-cache
-                    mountPath: /root/.cache/huggingface
-                env:
-                  - name: HF_HOME
-                    value: /root/.cache/huggingface
-            nodeSelector:
-              node.kubernetes.io/instance-type: gpu-h100-sxm
-            volumes:
-              - name: devshm
-                emptyDir:
-                  medium: Memory
-              - name: hf-cache
-                hostPath:
-                  path: /mnt/hf-cache
-                  type: DirectoryOrCreate
-
-  - label: "Diffusion Image Edit Test with H100 (1 GPU)"
-    timeout_in_minutes: 20
-    depends_on: image-build
-    commands:
-      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - pytest -s -v tests/e2e/online_serving/test_i2i_multi_image_input.py
-    agents:
-      queue: "mithril-h100-pool"
-    plugins:
-      - kubernetes:
-          podSpec:
-            containers:
-              - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-                resources:
-                  limits:
-                    nvidia.com/gpu: 1
-                volumeMounts:
-                  - name: devshm
-                    mountPath: /dev/shm
-                  - name: hf-cache
-                    mountPath: /root/.cache/huggingface
-                env:
-                  - name: HF_HOME
-                    value: /root/.cache/huggingface
-            nodeSelector:
-              node.kubernetes.io/instance-type: gpu-h100-sxm
-            volumes:
-              - name: devshm
-                emptyDir:
-                  medium: Memory
-              - name: hf-cache
-                hostPath:
-                  path: /mnt/hf-cache
-                  type: DirectoryOrCreate
+      queue: "cpu_queue_premerge"

.buildkite/rebase-pipeline.yaml

@@ -0,0 +1,38 @@
+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:
+      - python3 .buildkite/scripts/upload_pipeline.py --upload --all .buildkite/test-ready.yml
+    agents:
+      queue: "cpu_queue_premerge"
+
+  # L3 Test
+  - label: "Upload Merge Pipeline"
+    depends_on: image-build
+    key: upload-merge-pipeline
+    commands:
+      - python3 .buildkite/scripts/upload_pipeline.py --upload --all .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:
+      - python3 .buildkite/scripts/upload_pipeline.py --upload --all .buildkite/test-nightly.yml
+    agents:
+      queue: "cpu_queue_premerge"

.buildkite/release-pipeline.yaml

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

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

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

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

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

@@ -11,15 +11,6 @@ set -o pipefail
 export PYTHONPATH=".."
 
 # Print ROCm version
-echo "--- Confirming Clean Initial State"
-while true; do
-        sleep 3
-        if grep -q clean /opt/amdgpu/etc/gpu_state; then
-                echo "GPUs state is \"clean\""
-                break
-        fi
-done
-
 echo "--- ROCm info"
 rocminfo
 
@@ -51,33 +42,30 @@ 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"
-image_name="public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:${BUILDKITE_COMMIT}-rocm-omni"
+## 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}"
 container_name="rocm_${BUILDKITE_COMMIT}_$(tr -dc A-Za-z0-9 < /dev/urandom | head -c 10; echo)"
 
-# Install AWS CLI to authenticate to ECR Public Gallery to get higher rate limit for pulling images
-sudo apt-get update && sudo apt-get install -y awscli
-# Use safe docker login helper to prevent race conditions
-source "$(dirname "${BASH_SOURCE[0]}")/../docker_login_ecr_public.sh"
-safe_docker_login_ecr_public
-# Pull the container from ECR Public Gallery
+# 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
 
 docker pull "${image_name}"
 
 remove_docker_container() {
-   docker rm -f "${container_name}" || docker image rm -f "${image_name}" || true
+   docker ps -aq --filter "name=${container_name}" | xargs -r docker rm -f || true
+   docker image rm -f "${image_name}" || true
 }
 trap remove_docker_container EXIT
 
@@ -87,7 +75,11 @@ HF_CACHE="$(realpath ~)/huggingface"
 mkdir -p "${HF_CACHE}"
 HF_MOUNT="/root/.cache/huggingface"
 
-commands=$@
+if [[ -n "${TEST_COMMAND:-}" ]]; then
+    commands="$TEST_COMMAND"
+else
+    commands="$@"
+fi
 echo "Commands:$commands"
 
 PARALLEL_JOB_COUNT=8
@@ -102,6 +94,7 @@ 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')
@@ -115,11 +108,11 @@ if [[ $commands == *"--shard-id="* ]]; then
         --network=host \
         --shm-size=16gb \
         --group-add "$render_gid" \
-        --rm \
+        -e MIOPEN_DEBUG_CONV_DIRECT=0 \
+        -e MIOPEN_DEBUG_CONV_GEMM=0 \
+        -e VLLM_ROCM_USE_AITER=0 \
         -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}" \
@@ -147,10 +140,10 @@ else
           --network=host \
           --shm-size=16gb \
           --group-add "$render_gid" \
-          --rm \
+          -e MIOPEN_DEBUG_CONV_DIRECT=0 \
+          -e MIOPEN_DEBUG_CONV_GEMM=0 \
+          -e VLLM_ROCM_USE_AITER=0 \
           -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

@@ -0,0 +1,58 @@
+#!/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.py
+    VLLM_USE_MODELSCOPE=True pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni_expansion.py
 '

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

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

@@ -0,0 +1,52 @@
+#!/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/simple_test.sh

@@ -1,58 +0,0 @@
-#!/usr/bin/env bash
-
-set -euo pipefail
-
-SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
-REPO_ROOT="$(cd -- "${SCRIPT_DIR}/../.." && pwd)"
-UV_BIN="uv"
-PYTHON_BOOTSTRAP="${PYTHON:-python3}"
-VENV_DIR="${VENV_DIR:-${REPO_ROOT}/.venv-simple-test}"
-TARGET_PY_VERSION="3.12"
-
-if ! command -v "${PYTHON_BOOTSTRAP}" >/dev/null 2>&1; then
-  PYTHON_BOOTSTRAP="python"
-fi
-
-cd "${REPO_ROOT}"
-
-# Ensure pip-installed scripts (including uv) are on PATH
-PYTHON_SCRIPTS="$(${PYTHON_BOOTSTRAP} -c 'import sysconfig; print(sysconfig.get_path("scripts"))')"
-USER_BASE="$(${PYTHON_BOOTSTRAP} -m site --user-base 2>/dev/null || true)"
-PATH="${PYTHON_SCRIPTS}:${PATH}"
-if [[ -n "${USER_BASE}" ]]; then
-  PATH="${USER_BASE}/bin:${PATH}"
-fi
-
-if ! command -v "${UV_BIN}" >/dev/null 2>&1; then
-  "${PYTHON_BOOTSTRAP}" -m pip install --upgrade pip
-  "${PYTHON_BOOTSTRAP}" -m pip install uv
-  hash -r
-fi
-
-TARGET_PY_PATH="$(${UV_BIN} python find "${TARGET_PY_VERSION}" 2>/dev/null | head -n1 || true)"
-if [[ -z "${TARGET_PY_PATH}" ]]; then
-  "${UV_BIN}" python install "${TARGET_PY_VERSION}"
-  TARGET_PY_PATH="$(${UV_BIN} python find "${TARGET_PY_VERSION}" 2>/dev/null | head -n1 || true)"
-fi
-
-if [[ -z "${TARGET_PY_PATH}" ]]; then
-  echo "Failed to locate python ${TARGET_PY_VERSION}" >&2
-  exit 1
-fi
-
-if [[ ! -d "${VENV_DIR}" ]]; then
-  "${UV_BIN}" venv --python "${TARGET_PY_PATH}" "${VENV_DIR}"
-fi
-
-VENV_PYTHON="${VENV_DIR}/bin/python"
-[[ -x "${VENV_PYTHON}" ]] || { echo "Python not found in ${VENV_DIR}"; exit 1; }
-
-"${UV_BIN}" pip install --python "${VENV_PYTHON}" vllm==0.14.0
-"${UV_BIN}" pip install --python "${VENV_PYTHON}" -e ".[dev]"
-
-"${VENV_PYTHON}" -m pytest -v -s tests/entrypoints/
-"${VENV_PYTHON}" -m pytest -v -s tests/diffusion/cache/
-"${VENV_PYTHON}" -m pytest -v -s tests/diffusion/lora/
-"${VENV_PYTHON}" -m pytest -v -s tests/model_executor/models/qwen2_5_omni/test_audio_length.py
-"${VENV_PYTHON}" -m pytest -v -s tests/worker/
-"${VENV_PYTHON}" -m pytest -v -s tests/distributed/omni_connectors/test_kv_flow.py

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

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

.buildkite/scripts/upload_pipeline.py

@@ -0,0 +1,513 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+"""Render and optionally upload Buildkite pipeline YAML with diff-aware logic.
+
+Bootstrap mode (pipeline.yml with __IMAGE_BUILD_IF__ placeholders):
+  - Detect docs-only, pytest skip-mark-only, or combined docs+skip-mark PR/main changes and
+    substitute skip-ci ``if`` expressions.
+
+Test pipeline mode (e.g. test-merge.yml):
+  - Drop steps or groups whose ``source_file_dependencies`` do not match changed files.
+  - ``source_file_dependencies`` may be set on a leaf step or on a group step.
+  - Always strip ``source_file_dependencies`` from uploaded YAML (Buildkite does not
+    recognize this uploader-only key).
+
+Usage:
+  python3 upload_pipeline.py [--upload] [--all] <pipeline.yml>
+
+  # Bootstrap (replaces upload_pipeline_with_skip_ci.sh):
+  python3 upload_pipeline.py --upload .buildkite/pipeline.yml
+
+  # Test pipeline (replaces upload_test_pipeline_with_diff_skip.py):
+  python3 upload_pipeline.py --upload .buildkite/test-merge.yml
+
+  # Full suite (rebase pipeline): keep all steps, still strip the uploader-only key:
+  python3 upload_pipeline.py --upload --all .buildkite/test-merge.yml
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import re
+import subprocess
+import sys
+from pathlib import Path
+from typing import Any
+
+try:
+    import yaml
+except ImportError:
+    subprocess.check_call(
+        [sys.executable, "-m", "pip", "install", "-q", "pyyaml"],
+    )
+    import yaml
+
+LOG = "upload_pipeline"
+ROOT = Path(__file__).resolve().parent.parent.parent
+DOC_SEP = "\n---\n"
+BOOTSTRAP_MARKER = "__IMAGE_BUILD_IF__"
+
+
+def _log(message: str) -> None:
+    print(f"{LOG}: {message}", file=sys.stderr)
+
+
+def _git(*args: str) -> subprocess.CompletedProcess[str]:
+    return subprocess.run(
+        ["git", *args],
+        check=False,
+        capture_output=True,
+        text=True,
+    )
+
+
+def resolve_diff_range() -> str | None:
+    """Return a git diff range for PR or main builds, or None when unavailable."""
+    is_pr = os.environ.get("BUILDKITE_PULL_REQUEST", "false") != "false" and os.environ.get(
+        "BUILDKITE_PULL_REQUEST", ""
+    )
+    commit = os.environ.get("BUILDKITE_COMMIT", "")
+
+    if is_pr:
+        base_branch = os.environ.get("BUILDKITE_PULL_REQUEST_BASE_BRANCH", "main")
+        base_ref = f"origin/{base_branch}"
+        if _git("rev-parse", "--verify", base_ref).returncode != 0:
+            _log(f"origin/{base_branch} not found locally; trying fetch")
+            _git("fetch", "--depth=200", "origin", base_branch)
+        if _git("rev-parse", "--verify", base_ref).returncode != 0:
+            if _git("rev-parse", "--verify", base_branch).returncode == 0:
+                base_ref = base_branch
+            else:
+                _log(f"cannot resolve PR base {base_branch}; using safe defaults")
+                return None
+        return f"{base_ref}...{commit}"
+
+    if os.environ.get("BUILDKITE_BRANCH", "") == "main":
+        if _git("rev-parse", "--verify", f"{commit}^").returncode != 0:
+            _log("main commit has no parent; using safe defaults")
+            return None
+        return f"{commit}^..{commit}"
+
+    _log("not PR/main build; using safe defaults")
+    return None
+
+
+def resolve_changed_files() -> list[str] | None:
+    """Return changed file paths, or None when diff cannot be resolved."""
+    diff_range = resolve_diff_range()
+    if diff_range is None:
+        return None
+
+    result = _git("diff", "--name-only", diff_range)
+    if result.returncode != 0:
+        _log(f"git diff failed ({diff_range}); using safe defaults")
+        return None
+
+    files = [line for line in result.stdout.splitlines() if line.strip()]
+    _log(f"{len(files)} changed file(s)")
+    return files
+
+
+def _is_doc_file(file_path: str) -> bool:
+    if not file_path:
+        return False
+    if file_path.startswith("docs/"):
+        return True
+    if file_path.endswith(".md"):
+        return True
+    return file_path == "mkdocs.yml"
+
+
+def is_docs_only_change(changed_files: list[str]) -> bool:
+    has_any = False
+    for file_path in changed_files:
+        if not file_path:
+            continue
+        has_any = True
+        if not _is_doc_file(file_path):
+            return False
+    return has_any
+
+
+def _is_test_python_file(file_path: str) -> bool:
+    path = Path(file_path)
+    return path.suffix == ".py" and bool(path.parts) and path.parts[0] == "tests"
+
+
+_SKIP_MARK_RE = re.compile(r"pytest\.mark\.skip(?:if)?\b|pytest\.skip\s*\(")
+_PYTESTMARK_SKIP_RE = re.compile(r"pytest\.mark\.skip\b")
+_PYTEST_MARK_ONLY_RE = re.compile(r"pytest\.mark\.\w+")
+
+
+def _paren_balance(line: str) -> int:
+    return line.count("(") - line.count(")")
+
+
+def _is_skip_mark_related_content(line: str) -> bool:
+    stripped = line.strip()
+    if not stripped:
+        return True
+    if stripped.startswith("#"):
+        return True
+    return _SKIP_MARK_RE.search(stripped) is not None
+
+
+def _is_pytestmark_adjacent_content(line: str) -> bool:
+    """Allow pytestmark list refactors that only add skip/skipif alongside existing marks."""
+    stripped = line.strip().rstrip(",")
+    if not stripped:
+        return True
+    if stripped in {"[", "]"}:
+        return True
+    if stripped.startswith("#"):
+        return True
+    if stripped.startswith("pytestmark"):
+        return True
+    return _PYTEST_MARK_ONLY_RE.search(stripped) is not None
+
+
+def diff_only_contains_skip_mark_changes(diff_text: str) -> bool:
+    """True when diff only edits skip marks or reformats pytestmark to add them."""
+    pending_depth = 0
+    saw_change = False
+    has_skip_mark_edit = False
+    for raw_line in diff_text.splitlines():
+        if raw_line.startswith("@@"):
+            pending_depth = 0
+            continue
+        if not (raw_line.startswith("+") or raw_line.startswith("-")):
+            continue
+        if raw_line.startswith("+++") or raw_line.startswith("---"):
+            continue
+
+        saw_change = True
+        content = raw_line[1:]
+
+        if pending_depth > 0:
+            pending_depth += _paren_balance(content)
+            if pending_depth < 0:
+                pending_depth = 0
+            continue
+
+        if _is_skip_mark_related_content(content):
+            has_skip_mark_edit = True
+            pending_depth = max(0, _paren_balance(content))
+            continue
+
+        if _is_pytestmark_adjacent_content(content):
+            continue
+
+        return False
+
+    return saw_change and has_skip_mark_edit
+
+
+def _git_file_diff(file_path: str, diff_range: str) -> str | None:
+    result = _git("diff", diff_range, "--", file_path)
+    if result.returncode != 0:
+        _log(f"skip-mark-only: git diff failed for {file_path}")
+        return None
+    return result.stdout
+
+
+def _is_new_file_diff(diff_text: str) -> bool:
+    return "new file mode" in diff_text or "--- /dev/null" in diff_text
+
+
+def _file_has_module_level_pytest_skip(content: str) -> bool:
+    """True when ``pytestmark`` applies unconditional ``pytest.mark.skip`` to the module."""
+    match = re.search(r"^pytestmark\s*=\s*\[(.*?)^\s*\]", content, re.MULTILINE | re.DOTALL)
+    if match is not None:
+        return _PYTESTMARK_SKIP_RE.search(match.group(1)) is not None
+    match = re.search(r"^pytestmark\s*=\s*pytest\.mark\.skip\b", content, re.MULTILINE)
+    return match is not None
+
+
+def _file_content_at_commit(file_path: str, commit: str) -> str | None:
+    result = _git("show", f"{commit}:{file_path}")
+    if result.returncode != 0:
+        _log(f"skip-mark-only: cannot read {file_path} at {commit}")
+        return None
+    return result.stdout
+
+
+def _qualifies_as_skip_mark_test_change(file_path: str, diff_range: str) -> bool:
+    diff_text = _git_file_diff(file_path, diff_range)
+    if diff_text is None:
+        return False
+    if not diff_text.strip():
+        _log(f"skip-mark-only: empty diff for {file_path}")
+        return False
+
+    if diff_only_contains_skip_mark_changes(diff_text):
+        return True
+
+    if not _is_new_file_diff(diff_text):
+        _log(f"skip-mark-only: non-skip changes in {file_path}")
+        return False
+
+    commit = os.environ.get("BUILDKITE_COMMIT", "")
+    if not commit:
+        return False
+    content = _file_content_at_commit(file_path, commit)
+    if content is None:
+        return False
+    if _file_has_module_level_pytest_skip(content):
+        _log(f"skip-mark-only: new test file with module-level pytest.mark.skip: {file_path}")
+        return True
+
+    _log(f"skip-mark-only: new test file without module-level pytest.mark.skip: {file_path}")
+    return False
+
+
+def is_skip_mark_only_change(changed_files: list[str], *, diff_range: str) -> bool:
+    """True when every changed file is a tests/ module with only skip-mark edits."""
+    has_any = False
+    for file_path in changed_files:
+        if not file_path:
+            continue
+        has_any = True
+        if not _is_test_python_file(file_path):
+            return False
+        if not _qualifies_as_skip_mark_test_change(file_path, diff_range):
+            return False
+    return has_any
+
+
+def is_docs_or_skip_mark_only_change(changed_files: list[str], *, diff_range: str) -> bool:
+    """True when each changed file is a doc path or a qualifying skip-mark test change."""
+    has_any = False
+    for file_path in changed_files:
+        if not file_path:
+            continue
+        has_any = True
+        if _is_doc_file(file_path):
+            continue
+        if _is_test_python_file(file_path) and _qualifies_as_skip_mark_test_change(file_path, diff_range):
+            continue
+        _log(f"docs/skip-mark-only: rejecting {file_path}")
+        return False
+    return has_any
+
+
+def resolve_skip_ci(changed_files: list[str] | None, *, diff_range: str | None = None) -> bool:
+    if changed_files is None:
+        _log("skip-ci=0 (could not resolve changed files)")
+        return False
+
+    if diff_range is not None and is_docs_or_skip_mark_only_change(changed_files, diff_range=diff_range):
+        if is_docs_only_change(changed_files):
+            _log("docs-only change detected; skip-ci=1")
+        elif is_skip_mark_only_change(changed_files, diff_range=diff_range):
+            _log("pytest skip-mark-only change detected; skip-ci=1")
+        else:
+            _log("docs + pytest skip-mark-only change detected; skip-ci=1")
+        return True
+
+    if diff_range is None and is_docs_only_change(changed_files):
+        _log("docs-only change detected; skip-ci=1")
+        return True
+
+    _log("non-doc/non-skip-mark changes detected; skip-ci=0")
+    return False
+
+
+def render_bootstrap_pipeline(text: str, *, skip_ci: bool) -> str:
+    if DOC_SEP in text:
+        _, continuation = text.split(DOC_SEP, 1)
+    else:
+        continuation = text
+
+    nightly_only = (
+        '(build.pull_request.labels includes "nightly-test") || (build.branch == "main" && build.env("NIGHTLY") == "1")'
+    )
+    if skip_ci:
+        image_if = f"'{nightly_only}'"
+        ready_if = "'false'"
+        merge_if = "'false'"
+    else:
+        image_if = "'true'"
+        ready_if = '\'build.branch != "main" && build.pull_request.labels includes "ready"\''
+        merge_if = (
+            '\'(build.branch == "main" && build.env("NIGHTLY") != "1" '
+            '&& build.env("WEEKLY") != "1") || '
+            '(build.branch != "main" && build.pull_request.labels includes "merge-test")\''
+        )
+
+    return (
+        continuation.replace("__IMAGE_BUILD_IF__", image_if)
+        .replace("__UPLOAD_READY_IF__", ready_if)
+        .replace("__UPLOAD_MERGE_IF__", merge_if)
+    )
+
+
+def _matches_dependencies(changed_files: list[str], prefixes: list[str]) -> bool:
+    for path in changed_files:
+        for prefix in prefixes:
+            normalized = prefix.rstrip("/")
+            if path == normalized or path.startswith(f"{normalized}/"):
+                return True
+    return False
+
+
+def _strip_source_file_dependencies(step: dict[str, Any]) -> dict[str, Any]:
+    return {key: value for key, value in step.items() if key != "source_file_dependencies"}
+
+
+def _step_label(step: dict[str, Any]) -> str:
+    return str(step.get("group") or step.get("label") or "<step>")
+
+
+def _filter_steps(steps: list[Any], changed_files: list[str]) -> list[Any]:
+    filtered: list[Any] = []
+    for step in steps:
+        if not isinstance(step, dict):
+            filtered.append(step)
+            continue
+
+        deps = step.get("source_file_dependencies")
+        if deps is not None and not isinstance(deps, list):
+            raise ValueError(
+                f"source_file_dependencies must be a list in step {_step_label(step)!r}",
+            )
+        if deps is not None and not _matches_dependencies(changed_files, deps):
+            _log(
+                f"skip {_step_label(step)!r} (no changes under {deps})",
+            )
+            continue
+
+        nested = step.get("steps")
+        if nested is not None:
+            kept_nested = _filter_steps(nested, changed_files)
+            if not kept_nested:
+                _log(f"omit empty group {_step_label(step)!r}")
+                continue
+            new_step = _strip_source_file_dependencies(step)
+            new_step["steps"] = kept_nested
+            filtered.append(new_step)
+            continue
+
+        if deps is not None:
+            filtered.append(_strip_source_file_dependencies(step))
+        else:
+            filtered.append(step)
+
+    return filtered
+
+
+def _strip_uploader_metadata_from_steps(steps: list[Any]) -> list[Any]:
+    """Remove uploader-only keys while keeping all steps (no diff filtering)."""
+    stripped: list[Any] = []
+    for step in steps:
+        if not isinstance(step, dict):
+            stripped.append(step)
+            continue
+
+        deps = step.get("source_file_dependencies")
+        if deps is not None and not isinstance(deps, list):
+            raise ValueError(
+                f"source_file_dependencies must be a list in step {_step_label(step)!r}",
+            )
+
+        nested = step.get("steps")
+        new_step = _strip_source_file_dependencies(step)
+        if nested is not None:
+            new_step["steps"] = _strip_uploader_metadata_from_steps(nested)
+        stripped.append(new_step)
+
+    return stripped
+
+
+def render_test_pipeline(
+    doc: dict[str, Any],
+    changed_files: list[str] | None,
+) -> dict[str, Any]:
+    steps = doc.get("steps")
+    if not isinstance(steps, list):
+        return doc
+    if changed_files is not None:
+        steps = _filter_steps(steps, changed_files)
+    else:
+        steps = _strip_uploader_metadata_from_steps(steps)
+    return {**doc, "steps": steps}
+
+
+def resolve_pipeline_path(arg: str) -> Path:
+    path = Path(arg)
+    if path.is_absolute():
+        return path
+    return ROOT / path
+
+
+def render_pipeline(path: Path, *, force_all: bool = False) -> str:
+    text = path.read_text(encoding="utf-8")
+    diff_range = resolve_diff_range()
+    changed_files = resolve_changed_files()
+
+    # ``--all`` forces the keep-all-steps path (no diff-aware skipping) while still
+    # stripping ``source_file_dependencies``. Used by the rebase pipeline so main builds
+    # run the full e2e suite (see .buildkite/rebase-pipeline.yaml).
+    if force_all:
+        changed_files = None
+
+    if BOOTSTRAP_MARKER in text:
+        rendered = render_bootstrap_pipeline(
+            text,
+            skip_ci=resolve_skip_ci(changed_files, diff_range=diff_range),
+        )
+        return rendered
+
+    doc = yaml.safe_load(text)
+    if not isinstance(doc, dict):
+        raise ValueError(f"invalid pipeline YAML: {path}")
+
+    doc = render_test_pipeline(doc, changed_files)
+
+    return yaml.safe_dump(doc, sort_keys=False)
+
+
+def upload_to_buildkite(content: str) -> None:
+    subprocess.run(
+        ["buildkite-agent", "pipeline", "upload"],
+        input=content,
+        text=True,
+        check=True,
+    )
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "pipeline",
+        nargs="?",
+        default=".buildkite/pipeline.yml",
+        help="Pipeline YAML path (default: .buildkite/pipeline.yml)",
+    )
+    parser.add_argument(
+        "--upload",
+        action="store_true",
+        help="Pipe rendered YAML to buildkite-agent pipeline upload",
+    )
+    parser.add_argument(
+        "--all",
+        action="store_true",
+        help="Keep all steps (disable diff-aware skipping); still strips source_file_dependencies",
+    )
+    args = parser.parse_args()
+
+    path = resolve_pipeline_path(args.pipeline)
+    if not path.is_file():
+        _log(f"missing pipeline file: {path}")
+        return 1
+
+    rendered = render_pipeline(path, force_all=args.all)
+    if args.upload:
+        upload_to_buildkite(rendered)
+    else:
+        sys.stdout.write(rendered)
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

.buildkite/test-amd-merge.yml

@@ -0,0 +1,246 @@
+env:
+  VLLM_WORKER_MULTIPROC_METHOD: spawn
+  HF_HUB_DOWNLOAD_TIMEOUT: 300
+  HF_HUB_ETAG_TIMEOUT: 60
+
+steps:
+
+- group: ":card_index_dividers: Simple Test"
+  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"
+
+- group: ":card_index_dividers: Diffusion Test"
+  steps:
+    - label: "Diffusion · LoRA Test"
+      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))
+          '
+
+    - 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: "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'
+          "
+
+## 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: "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.
+- group: ":card_index_dividers: E2E Test"
+  steps:
+    - 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"
+          '
+
+    # It is the whisper ASR validation step that is slow, not the Voxtral TTS generation.
+    # whisper ASR runs on CPU.
+    - label: "Voxtral-TTS E2E Test (Online Serving)"
+      agent_pool: mi325_1
+      depends_on: amd-build
+      mirror_hardwares: [amdproduction]
+      grade: Blocking
+      commands:
+        - |
+          timeout 40m bash -c '
+            export VLLM_WORKER_MULTIPROC_METHOD=spawn
+            pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
+          '
+
+    - label: "Voxtral-TTS E2E Test (Offline Inference)"
+      agent_pool: mi325_1
+      depends_on: amd-build
+      mirror_hardwares: [amdproduction]
+      grade: Blocking
+      commands:
+        - |
+          timeout 30m bash -c '
+            export VLLM_WORKER_MULTIPROC_METHOD=spawn
+            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
+# # 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"

.buildkite/test-amd-ready.yaml

@@ -0,0 +1,236 @@
+env:
+  VLLM_WORKER_MULTIPROC_METHOD: spawn
+  HF_HUB_DOWNLOAD_TIMEOUT: 300
+  HF_HUB_ETAG_TIMEOUT: 60
+
+steps:
+
+- group: ":card_index_dividers: Simple Test"
+  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: "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: "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
+
+- group: ":card_index_dividers: Diffusion Test"
+  steps:
+    - 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: "Diffusion · Offloader 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))
+          '
+
+    - label: "Diffusion · Cache 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'
+          "
+
+## 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: "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"
+
+- group: ":card_index_dividers: E2E Test"
+  steps:
+    - 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: "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"
+          '
+
+# 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"

.buildkite/test-amd.yaml

@@ -1,20 +1,39 @@
 steps:
 
 - label: "Diffusion Model Test"
-  timeout_in_minutes: 20
+  timeout_in_minutes: 30
   agent_pool: mi325_2
   depends_on: amd-build
   mirror_hardwares: [amdproduction]
   grade: Blocking
   commands:
     - export GPU_ARCHS=gfx942
-    - export MIOPEN_DEBUG_CONV_DIRECT=0
-    - export MIOPEN_DEBUG_CONV_GEMM=0
-    - export VLLM_ROCM_USE_AITER=1
-    - export VLLM_ROCM_USE_AITER_MHA=1
-    - export VLLM_ROCM_USE_AITER_LINEAR=0
-    - export VLLM_ROCM_USE_AITER_RMSNORM=0
-    - pytest -s -v tests/e2e/offline_inference/test_t2i_model.py
+    - pytest -s -v tests/e2e/offline_inference/test_z_image.py
+
+- label: "Diffusion Images API LoRA E2E"
+  timeout_in_minutes: 30
+  agent_pool: mi325_1
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export GPU_ARCHS=gfx942
+    - export VLLM_LOGGING_LEVEL=DEBUG
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - pytest -s -v tests/diffusion/lora/test_images_generations_lora.py
+
+- label: "Diffusion Model CPU offloading Test"
+  timeout_in_minutes: 20
+  agent_pool: mi325_1
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export GPU_ARCHS=gfx942
+    - export VLLM_LOGGING_LEVEL=DEBUG
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py
+    - pytest -s -v tests/diffusion/offloader/test_diffusion_layerwise_offload.py
 
 - label: "Diffusion Cache Backend Test"
   timeout_in_minutes: 15
@@ -26,24 +45,29 @@ steps:
     - export GPU_ARCHS=gfx942
     - export VLLM_LOGGING_LEVEL=DEBUG
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - export MIOPEN_DEBUG_CONV_DIRECT=0
-    - export MIOPEN_DEBUG_CONV_GEMM=0
-    - export VLLM_ROCM_USE_AITER=1
-    - export VLLM_ROCM_USE_AITER_MHA=1
-    - export VLLM_ROCM_USE_AITER_LINEAR=0
-    - export VLLM_ROCM_USE_AITER_RMSNORM=0
-    - pytest -s -v tests/e2e/offline_inference/test_cache_dit.py tests/e2e/offline_inference/test_teacache.py
+    - pytest -s -v -m 'core_model and cache and diffusion and not distributed_rocm and MI325'
 
-- label: "Diffusion Parallelism Test"
-  timeout_in_minutes: 15
+- label: "Diffusion Sequence Parallelism Test"
+  timeout_in_minutes: 20
   agent_pool: mi325_2
   depends_on: amd-build
   mirror_hardwares: [amdproduction]
   grade: Blocking
   commands:
-    - export MIOPEN_DEBUG_CONV_DIRECT=0
-    - export MIOPEN_DEBUG_CONV_GEMM=0
-    - pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py
+    - 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
+
+- label: "Diffusion Tensor Parallelism Test"
+  timeout_in_minutes: 20
+  agent_pool: mi325_2
+  depends_on: amd-build
+  commands:
+    - export GPU_ARCHS=gfx942
+    - export VLLM_LOGGING_LEVEL=DEBUG
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - pytest -s -v tests/diffusion/distributed/test_zimage_parallelism.py
 
 - label: "Diffusion GPU Worker Test"
   timeout_in_minutes: 20
@@ -52,9 +76,7 @@ steps:
   mirror_hardwares: [amdproduction]
   grade: Blocking
   commands:
-    - export MIOPEN_DEBUG_CONV_DIRECT=0
-    - export MIOPEN_DEBUG_CONV_GEMM=0
-    - pytest -s -v tests/diffusion/test_gpu_diffusion_worker.py
+    - pytest -s -v tests/diffusion/test_diffusion_worker.py
 
 - label: "Omni Model Test Qwen2-5-Omni"
   timeout_in_minutes: 15
@@ -66,13 +88,8 @@ steps:
     - export GPU_ARCHS=gfx942
     - export VLLM_LOGGING_LEVEL=DEBUG
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - export MIOPEN_DEBUG_CONV_DIRECT=0
-    - export MIOPEN_DEBUG_CONV_GEMM=0
-    - export VLLM_ROCM_USE_AITER=1
-    - export VLLM_ROCM_USE_AITER_MHA=1
-    - export VLLM_ROCM_USE_AITER_LINEAR=0
-    - export VLLM_ROCM_USE_AITER_RMSNORM=0
-    - pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
+    - pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni_expansion.py
+    - pytest -s -v tests/engine/test_async_omni_engine_abort.py
 
 - label: "Omni Model Test Qwen3-Omni"
   timeout_in_minutes: 15
@@ -83,13 +100,14 @@ steps:
   commands:
     - export VLLM_LOGGING_LEVEL=DEBUG
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - export MIOPEN_DEBUG_CONV_DIRECT=0
-    - export MIOPEN_DEBUG_CONV_GEMM=0
-    - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py
+    - export VLLM_ROCM_USE_AITER=0
+    - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
+    - pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py
 
-- label: "Diffusion Image Edit Test"
-  timeout_in_minutes: 15
-  agent_pool: mi325_1
+
+- label: "Omni Sleep Mode Test"
+  timeout_in_minutes: 40
+  agent_pool: mi325_2
   depends_on: amd-build
   mirror_hardwares: [amdproduction]
   grade: Blocking
@@ -97,10 +115,4 @@ steps:
     - export GPU_ARCHS=gfx942
     - export VLLM_LOGGING_LEVEL=DEBUG
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - export MIOPEN_DEBUG_CONV_DIRECT=0
-    - export MIOPEN_DEBUG_CONV_GEMM=0
-    - export VLLM_ROCM_USE_AITER=1
-    - export VLLM_ROCM_USE_AITER_MHA=1
-    - export VLLM_ROCM_USE_AITER_LINEAR=0
-    - export VLLM_ROCM_USE_AITER_RMSNORM=0
-    - pytest -s -v tests/e2e/online_serving/test_i2i_multi_image_input.py
+    - pytest -s -v tests/entrypoints/test_omni_sleep_mode.py -m "advanced_model and omni and MI325" --run-level "advanced_model"

.buildkite/test-merge.yml

@@ -0,0 +1,943 @@
+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-merge-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"
+
+  - group: ":card_index_dividers: Diffusion Test"
+    depends_on: upload-merge-pipeline
+    steps:
+      - label: "Diffusion · LoRA Test"
+        timeout_in_minutes: 30
+        commands:
+          - pytest -sv tests/diffusion/lora/ -m 'advanced_model and cuda' --run-level "advanced_model"
+        agents:
+          queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
+        plugins:
+          - docker#v5.2.0:
+              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+              always-pull: true
+              propagate-environment: true
+              environment:
+                - "HF_HOME=/fsx/hf_cache"
+                - "HF_TOKEN"
+              volumes:
+                - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Distributed Test"
+    timeout_in_minutes: 60
+    commands:
+      - |
+        timeout 55m bash -c "
+          set -e
+          export VLLM_IMAGE_FETCH_TIMEOUT=60
+          pytest -s -v tests/distributed -m 'advanced_model and cuda' --run-level 'advanced_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
+
+  - group: ":card_index_dividers: Entrypoints Test"
+    depends_on: upload-merge-pipeline
+    steps:
+      - label: "Entrypoint Test with H100"
+        timeout_in_minutes: 60
+        commands:
+          - pytest -s -v tests/entrypoints/ -m "advanced_model and cuda and H100" --run-level "advanced_model"
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 2
+                    volumeMounts:
+                      - name: devshm
+                        mountPath: /dev/shm
+                      - name: hf-cache
+                        mountPath: /root/.cache/huggingface
+                    env:
+                      - name: HF_HOME
+                        value: /root/.cache/huggingface
+                      - name: HF_TOKEN
+                        valueFrom:
+                          secretKeyRef:
+                            name: hf-token-secret
+                            key: token
+                nodeSelector:
+                  node.kubernetes.io/instance-type: gpu-h100-sxm
+                volumes:
+                  - name: devshm
+                    emptyDir:
+                      medium: Memory
+                  - name: hf-cache
+                    hostPath:
+                      path: /mnt/hf-cache
+                      type: DirectoryOrCreate
+
+      - label: "Entrypoint Test with L4"
+        timeout_in_minutes: 30
+        commands:
+          - pytest -s -v tests/entrypoints/ -m "advanced_model and cuda and L4" --run-level "advanced_model"
+        agents:
+          queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
+        plugins:
+          - docker#v5.2.0:
+              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+              always-pull: true
+              propagate-environment: true
+              environment:
+                - "HF_HOME=/fsx/hf_cache"
+                - "HF_TOKEN"
+              volumes:
+                - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - group: ":card_index_dividers: E2E Test"
+    depends_on: upload-merge-pipeline
+    steps:
+      - label: "TTS · Step-Audio2 E2E 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_step_audio2.py -m 'advanced_model' --run-level 'advanced_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"
+              volumes:
+                - "/fsx/hf_cache:/fsx/hf_cache"
+
+      - label: "TTS · Qwen3-TTS CustomVoice Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_qwen3_tts_customvoice.py
+          - tests/e2e/offline_inference/test_qwen3_tts_customvoice.py
+          - vllm_omni/model_executor/models/qwen3_tts/
+          - vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
+          - vllm_omni/deploy/qwen3_tts.yaml
+        commands:
+          - |
+            timeout 20m bash -c "
+              export VLLM_LOGGING_LEVEL=DEBUG
+              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 and cuda' --run-level 'advanced_model'
+            "
+        agents:
+          queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
+        plugins:
+          - docker#v5.2.0:
+              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+              always-pull: true
+              propagate-environment: true
+              shm-size: "8gb"
+              environment:
+                - "HF_HOME=/fsx/hf_cache"
+              volumes:
+                - "/fsx/hf_cache:/fsx/hf_cache"
+
+      - label: "TTS · Qwen3-TTS Base Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_qwen3_tts_base.py
+          - tests/e2e/offline_inference/test_qwen3_tts_base.py
+          - vllm_omni/model_executor/models/qwen3_tts/
+          - vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
+          - vllm_omni/deploy/qwen3_tts.yaml
+        commands:
+          - |
+            timeout 20m bash -c "
+              export VLLM_LOGGING_LEVEL=DEBUG
+              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 and cuda' --run-level 'advanced_model'
+            "
+        agents:
+          queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
+        plugins:
+          - docker#v5.2.0:
+              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+              always-pull: true
+              propagate-environment: true
+              shm-size: "8gb"
+              environment:
+                - "HF_HOME=/fsx/hf_cache"
+              volumes:
+                - "/fsx/hf_cache:/fsx/hf_cache"
+
+      - label: "TTS · Ming-Omni-TTS Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_ming_tts.py
+          - tests/e2e/offline_inference/test_ming_tts.py
+          - vllm_omni/model_executor/models/ming_tts/
+          - vllm_omni/model_executor/stage_input_processors/ming_tts.py
+          - vllm_omni/deploy/ming_tts.yaml
+        timeout_in_minutes: 20
+        commands:
+          - |
+            timeout 20m bash -c "
+              export VLLM_LOGGING_LEVEL=DEBUG
+              pytest -s -v tests/e2e/online_serving/test_ming_tts.py tests/e2e/offline_inference/test_ming_tts.py -m 'advanced_model and cuda' --run-level 'advanced_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"
+              volumes:
+                - "/fsx/hf_cache:/fsx/hf_cache"
+
+      - label: "TTS · Voxtral-TTS Test"
+        source_file_dependencies:
+          - tests/e2e/offline_inference/test_voxtral_tts.py
+          - tests/e2e/online_serving/test_voxtral_tts.py
+          - vllm_omni/model_executor/models/voxtral_tts/
+          - vllm_omni/model_executor/stage_input_processors/voxtral_tts.py
+          - vllm_omni/deploy/voxtral_tts.yaml
+        timeout_in_minutes: 20
+        commands:
+          - |
+            timeout 20m bash -c "
+              export VLLM_LOGGING_LEVEL=DEBUG
+              pytest -s -v tests/e2e/offline_inference/test_voxtral_tts.py tests/e2e/online_serving/test_voxtral_tts.py -m 'advanced_model and cuda' --run-level 'advanced_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: "TTS · GLM-TTS Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_glm_tts.py
+          - tests/e2e/offline_inference/test_glm_tts.py
+          - vllm_omni/model_executor/models/glm_tts/
+          - vllm_omni/model_executor/stage_input_processors/glm_tts.py
+          - vllm_omni/deploy/glm_tts.yaml
+        timeout_in_minutes: 20
+        commands:
+          - |
+            timeout 20m bash -c "
+              export VLLM_LOGGING_LEVEL=DEBUG
+              pytest -s -v tests/e2e/online_serving/test_glm_tts.py tests/e2e/offline_inference/test_glm_tts.py -m 'advanced_model and cuda' --run-level 'advanced_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: "Omni · Qwen3-Omni Test"
+        source_file_dependencies:
+          - tests/e2e/offline_inference/test_qwen3_omni.py
+          - tests/e2e/online_serving/test_qwen3_omni.py
+          - vllm_omni/model_executor/models/qwen3_omni/
+          - vllm_omni/model_executor/stage_input_processors/qwen3_omni.py
+          - vllm_omni/deploy/qwen3_omni_moe.yaml
+        timeout_in_minutes: 30
+        commands:
+          - 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'
+        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: "Omni · Mimo Audio Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_mimo_audio.py
+          - vllm_omni/model_executor/models/mimo_audio/
+          - vllm_omni/model_executor/stage_input_processors/mimo_audio.py
+          - vllm_omni/deploy/mimo_audio.yaml
+        timeout_in_minutes: 30
+        commands:
+          - pytest -s -v tests/e2e/online_serving/test_mimo_audio.py -m 'advanced_model and cuda' --run-level 'advanced_model'
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 2
+                    volumeMounts:
+                      - name: devshm
+                        mountPath: /dev/shm
+                      - name: hf-cache
+                        mountPath: /root/.cache/huggingface
+                    env:
+                      - name: HF_HOME
+                        value: /root/.cache/huggingface
+                      - name: HF_TOKEN
+                        valueFrom:
+                          secretKeyRef:
+                            name: hf-token-secret
+                            key: token
+                nodeSelector:
+                  node.kubernetes.io/instance-type: gpu-h100-sxm
+                volumes:
+                  - name: devshm
+                    emptyDir:
+                      medium: Memory
+                  - name: hf-cache
+                    hostPath:
+                      path: /mnt/hf-cache
+                      type: DirectoryOrCreate
+
+      - label: "TTS · VoxCPM2 Test"
+        source_file_dependencies:
+          - tests/e2e/offline_inference/test_voxcpm2_tts.py
+          - tests/e2e/online_serving/test_voxcpm2_tts.py
+          - vllm_omni/model_executor/models/voxcpm2/
+          - vllm_omni/deploy/voxcpm2.yaml
+        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/offline_inference/test_voxcpm2_tts.py tests/e2e/online_serving/test_voxcpm2_tts.py -m 'advanced_model and cuda' --run-level 'advanced_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: "Omni · Higgs-Audio V2 Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_higgs_audio_v2.py
+          - tests/e2e/offline_inference/test_higgs_audio_v2.py
+          - vllm_omni/model_executor/models/higgs_audio_v2/
+          - vllm_omni/model_executor/stage_input_processors/higgs_audio_v2.py
+          - vllm_omni/deploy/higgs_audio_v2.yaml
+        timeout_in_minutes: 30
+        commands:
+          - |
+            timeout 30m bash -c "
+              export VLLM_LOGGING_LEVEL=DEBUG
+              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_v2.py tests/e2e/offline_inference/test_higgs_audio_v2.py -m 'advanced_model and cuda' --run-level 'advanced_model'
+            "
+        agents:
+          queue: "gpu_1_queue"
+        plugins:
+          - docker#v5.2.0:
+              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+              always-pull: true
+              propagate-environment: true
+              shm-size: "8gb"
+              environment:
+                - "HF_HOME=/fsx/hf_cache"
+                - "HF_TOKEN"
+              volumes:
+                - "/fsx/hf_cache:/fsx/hf_cache"
+
+      - label: "Diffusion · Qwen Image Edit Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_qwen_image_edit.py
+          - vllm_omni/diffusion/models/qwen_image/
+        timeout_in_minutes: 20
+        commands:
+          - pytest -s -v tests/e2e/online_serving/test_qwen_image_edit.py -m 'advanced_model and cuda' --run-level 'advanced_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 · Longcat Image Edit Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_longcat_image_edit.py
+          - vllm_omni/diffusion/models/longcat_image/
+        timeout_in_minutes: 20
+        commands:
+          - pytest -s -v tests/e2e/online_serving/test_longcat_image_edit.py -m 'advanced_model and cuda' --run-level 'advanced_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 · Longcat Image Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_longcat_image.py
+          - vllm_omni/diffusion/models/longcat_image/
+        timeout_in_minutes: 20
+        commands:
+          - pytest -s -v tests/e2e/online_serving/test_longcat_image.py -m 'advanced_model and cuda' --run-level 'advanced_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 · Bagel Test"
+        source_file_dependencies:
+          - tests/e2e/offline_inference/test_bagel.py
+          - tests/e2e/online_serving/test_bagel.py
+          - vllm_omni/model_executor/models/bagel/
+          - vllm_omni/diffusion/models/bagel/
+          - vllm_omni/model_executor/stage_input_processors/bagel.py
+          - vllm_omni/deploy/bagel.yaml
+        timeout_in_minutes: 20
+        commands:
+          - pytest -s -v tests/e2e/offline_inference/test_bagel.py tests/e2e/online_serving/test_bagel.py -m 'advanced_model and cuda' --run-level 'advanced_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"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_qwen_image.py
+          - vllm_omni/diffusion/models/qwen_image/
+        commands:
+          - |
+            timeout 40m bash -c '
+              export VLLM_IMAGE_FETCH_TIMEOUT=60
+              pytest -s -v tests/e2e/online_serving/test_qwen_image.py -m 'advanced_model and cuda' --run-level 'advanced_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"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_qwen_image_layered.py
+          - vllm_omni/diffusion/models/qwen_image/
+        commands:
+          - |
+            timeout 40m bash -c "
+              export VLLM_IMAGE_FETCH_TIMEOUT=60
+              pytest -s -v tests/e2e/online_serving/test_qwen_image_layered.py -m 'advanced_model and cuda' --run-level 'advanced_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"
+        source_file_dependencies:
+          - tests/e2e/offline_inference/test_wan22_t2v.py
+          - tests/e2e/online_serving/test_wan22_t2v.py
+          - vllm_omni/diffusion/models/wan2_2/
+        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'
+            "
+        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 · HunyuanVideo 1.5 Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_hunyuan_video_15.py
+          - vllm_omni/diffusion/models/hunyuan_video/
+        timeout_in_minutes: 40
+        commands:
+          - |
+            timeout 40m bash -c "
+              pytest -s -v tests/e2e/online_serving/test_hunyuan_video_15.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_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 · Z Image Test"
+        source_file_dependencies:
+          - tests/e2e/offline_inference/test_z_image.py
+          - tests/e2e/online_serving/test_z_image.py
+          - vllm_omni/diffusion/models/z_image/
+        timeout_in_minutes: 30
+        commands:
+          - pytest -s -v tests/e2e/offline_inference/test_z_image.py tests/e2e/online_serving/test_z_image.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
+        agents:
+          queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
+        plugins:
+          - docker#v5.2.0:
+              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+              always-pull: true
+              propagate-environment: true
+              environment:
+                - "HF_HOME=/fsx/hf_cache"
+                - "HF_TOKEN"
+              volumes:
+                - "/fsx/hf_cache:/fsx/hf_cache"
+
+      - label: "Diffusion · FLUX.2 Klein Test"
+        source_file_dependencies:
+          - tests/e2e/offline_inference/test_flux2_klein.py
+          - tests/e2e/online_serving/test_flux2_klein.py
+          - vllm_omni/diffusion/models/flux2_klein/
+        timeout_in_minutes: 30
+        commands:
+          - pytest -s -v tests/e2e/offline_inference/test_flux2_klein.py tests/e2e/online_serving/test_flux2_klein.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
+        agents:
+          queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
+        plugins:
+          - docker#v5.2.0:
+              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+              always-pull: true
+              propagate-environment: true
+              environment:
+                - "HF_HOME=/fsx/hf_cache"
+                - "HF_TOKEN"
+              volumes:
+                - "/fsx/hf_cache:/fsx/hf_cache"
+
+      - label: "Diffusion · SD3 Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_sd3.py
+          - vllm_omni/diffusion/models/sd3/
+        timeout_in_minutes: 45
+        commands:
+          - |
+            timeout 45m bash -c "
+              export VLLM_IMAGE_FETCH_TIMEOUT=60
+              pytest -s -v tests/e2e/online_serving/test_sd3.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
+            "
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 4
+                    volumeMounts:
+                      - name: devshm
+                        mountPath: /dev/shm
+                      - name: hf-cache
+                        mountPath: /root/.cache/huggingface
+                    env:
+                      - name: HF_HOME
+                        value: /root/.cache/huggingface
+                      - name: HF_TOKEN
+                        valueFrom:
+                          secretKeyRef:
+                            name: hf-token-secret
+                            key: token
+                nodeSelector:
+                  node.kubernetes.io/instance-type: gpu-h100-sxm
+                volumes:
+                  - name: devshm
+                    emptyDir:
+                      medium: Memory
+                  - name: hf-cache
+                    hostPath:
+                      path: /mnt/hf-cache
+                      type: DirectoryOrCreate

.buildkite/test-npu-nightly.yaml

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

.buildkite/test-npu-ready.yaml

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

@@ -0,0 +1,674 @@
+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"
+    depends_on: upload-ready-pipeline
+    commands:
+      - timeout 20m pytest -s -v tests/e2e/offline_inference/custom_pipeline/ -m "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"
+
+  - group: ":card_index_dividers: Diffusion Test"
+    depends_on: upload-ready-pipeline
+    steps:
+      - label: "Diffusion · Offloader Test"
+        commands:
+          - timeout 40m pytest -s -v tests/diffusion/offloader -m "core_model and cuda" --run-level "core_model"
+        agents:
+          queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
+        plugins:
+          - docker#v5.2.0:
+              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+              always-pull: true
+              propagate-environment: true
+              environment:
+                - "HF_HOME=/fsx/hf_cache"
+                - "HF_TOKEN"
+              volumes:
+                - "/fsx/hf_cache:/fsx/hf_cache"
+
+      - label: "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"
+    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'
+        "
+    agents:
+      queue: "gpu_1_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          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: "Distributed Test"
+    depends_on: upload-ready-pipeline
+    commands:
+      - |
+        timeout 20m bash -c "
+                  pytest -sv tests/distributed/ -m 'core_model and cuda' --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
+
+  - group: ":card_index_dividers: Entrypoints Test"
+    depends_on: upload-ready-pipeline
+    steps:
+      - label: "Entrypoints Test with H100"
+        commands:
+          - |
+            timeout 30m bash -c "
+              export VLLM_WORKER_MULTIPROC_METHOD=spawn
+              pytest -sv tests/entrypoints/ -m 'core_model and cuda and H100' --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: "Entrypoints Test with L4"
+        commands:
+          - |
+            timeout 30m bash -c "
+              export VLLM_WORKER_MULTIPROC_METHOD=spawn
+              pytest -sv tests/entrypoints/ -m 'core_model and cuda and L4' --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"
+
+  - group: ":card_index_dividers: E2E Test"
+    depends_on: upload-ready-pipeline
+    steps:
+      - label: "RLHF · VeRL-Omni E2E Test"
+        source_file_dependencies:
+          - tests/e2e/offline_inference/rlhf_test/test_verl_omni_e2e.py
+          - tests/e2e/offline_inference/custom_pipeline/
+          - vllm_omni/diffusion/
+          - vllm_omni/engine/
+          - vllm_omni/entrypoints/
+          - vllm_omni/inputs/
+          - vllm_omni/outputs/
+        commands:
+          - |
+            timeout 20m bash -c "
+              pytest -s -v tests/e2e/offline_inference/rlhf_test/test_verl_omni_e2e.py
+            "
+        agents:
+          queue: "gpu_1_queue"
+        plugins:
+          - docker#v5.2.0:
+              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+              always-pull: true
+              propagate-environment: true
+              shm-size: "8gb"
+              environment:
+                - "HF_HOME=/fsx/hf_cache"
+                - "HF_TOKEN"
+              volumes:
+                - "/fsx/hf_cache:/fsx/hf_cache"
+
+      - label: "Omni · Qwen3-Omni Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_qwen3_omni.py
+          - vllm_omni/model_executor/models/qwen3_omni/
+          - vllm_omni/model_executor/stage_input_processors/qwen3_omni.py
+          - vllm_omni/deploy/qwen3_omni_moe.yaml
+        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: "TTS · VoxCPM2 Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_voxcpm2_tts.py
+          - vllm_omni/model_executor/models/voxcpm2/
+          - vllm_omni/deploy/voxcpm2.yaml
+        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' --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"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_qwen3_tts_customvoice.py
+          - vllm_omni/model_executor/models/qwen3_tts/
+          - vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
+          - vllm_omni/deploy/qwen3_tts.yaml
+        commands:
+          - |
+            timeout 20m bash -c "
+              export VLLM_ALLOW_LONG_MAX_MODEL_LEN=1
+              pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m 'core_model' --run-level 'core_model'
+            "
+        agents:
+          queue: "gpu_1_queue"
+        plugins:
+          - docker#v5.2.0:
+              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+              always-pull: true
+              propagate-environment: true
+              shm-size: "8gb"
+              environment:
+                - "HF_HOME=/fsx/hf_cache"
+                - "HF_TOKEN"
+              volumes:
+                - "/fsx/hf_cache:/fsx/hf_cache"
+
+      - label: "TTS · Qwen3-TTS Base Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_qwen3_tts_base.py
+          - vllm_omni/model_executor/models/qwen3_tts/
+          - vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
+          - vllm_omni/deploy/qwen3_tts.yaml
+        commands:
+          - |
+            timeout 20m bash -c "
+              export VLLM_WORKER_MULTIPROC_METHOD=spawn
+              export VLLM_ALLOW_LONG_MAX_MODEL_LEN=1
+              export VLLM_OMNI_USE_V2_RUNNER=1
+              pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py -m 'core_model' --run-level 'core_model'
+            "
+        agents:
+          queue: "gpu_1_queue"
+        plugins:
+          - docker#v5.2.0:
+              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+              always-pull: true
+              propagate-environment: true
+              shm-size: "8gb"
+              environment:
+                - "HF_HOME=/fsx/hf_cache"
+                - "HF_TOKEN"
+              volumes:
+                - "/fsx/hf_cache:/fsx/hf_cache"
+
+      - label: "TTS · Higgs-Audio-V3 Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_higgs_audio_v3.py
+          - vllm_omni/model_executor/models/higgs_audio_v3/
+          - vllm_omni/model_executor/stage_input_processors/higgs_audio_v3.py
+          - vllm_omni/deploy/higgs_multimodal_qwen3.yaml
+        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_1_queue"
+        plugins:
+          - docker#v5.2.0:
+              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+              always-pull: true
+              propagate-environment: true
+              shm-size: "8gb"
+              environment:
+                - "HF_HOME=/fsx/hf_cache"
+                - "HF_TOKEN"
+              volumes:
+                - "/fsx/hf_cache:/fsx/hf_cache"
+
+      - label: "Diffusion · Qwen Image Edit Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_qwen_image_edit.py
+          - vllm_omni/diffusion/models/qwen_image/
+        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"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_bagel.py
+          - vllm_omni/model_executor/models/bagel/
+          - vllm_omni/diffusion/models/bagel/
+          - vllm_omni/model_executor/stage_input_processors/bagel.py
+          - vllm_omni/deploy/bagel.yaml
+        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"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_wan22_t2v.py
+          - vllm_omni/diffusion/models/wan2_2/
+        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"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_qwen_image.py
+          - vllm_omni/diffusion/models/qwen_image/
+        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"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_qwen_image_layered.py
+          - vllm_omni/diffusion/models/qwen_image/
+        commands:
+          - |
+            timeout 40m bash -c "
+              export VLLM_IMAGE_FETCH_TIMEOUT=60
+              pytest -s -v tests/e2e/online_serving/test_qwen_image_layered.py -m 'core_model' --run-level 'core_model'
+            "
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 1
+                    volumeMounts:
+                      - name: devshm
+                        mountPath: /dev/shm
+                      - name: hf-cache
+                        mountPath: /root/.cache/huggingface
+                    env:
+                      - name: HF_HOME
+                        value: /root/.cache/huggingface
+                      - name: HF_TOKEN
+                        valueFrom:
+                          secretKeyRef:
+                            name: hf-token-secret
+                            key: token
+                nodeSelector:
+                  node.kubernetes.io/instance-type: gpu-h100-sxm
+                volumes:
+                  - name: devshm
+                    emptyDir:
+                      medium: Memory
+                  - name: hf-cache
+                    hostPath:
+                      path: /mnt/hf-cache
+                      type: DirectoryOrCreate

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

@@ -3,9 +3,16 @@
    Last synced: 2025-12-15
    Modifications: Removed unused CUDA/NVIDIA logic, keeping only AMD tests
 #}
-{% set docker_image_amd = "public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT-rocm-omni" %}
+{% set docker_image_amd = "rocm/vllm-omni:$BUILDKITE_COMMIT" %}
 {% set default_working_dir = "/app/vllm-omni" %}
 
+{% if env %}
+env:
+{% for name in env %}
+  {{ name }}: "{{ env[name] }}"
+{% endfor %}
+{% endif %}
+steps:
   - group: "AMD Tests"
     depends_on: ~
     steps:
@@ -13,8 +20,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 final --progress plain ."
+          - "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"
         env:
@@ -22,16 +28,25 @@
         retry:
           automatic:
             - exit_status: -1  # Agent was lost
-              limit: 1
+              limit: 2
             - exit_status: -10  # Agent was lost
-              limit: 1
+              limit: 2
+            - exit_status: 128  # Git connectivity issues
+              limit: 2
             - exit_status: 1  # Machine occasionally fail
               limit: 1
         agents:
-          queue: cpu_queue_premerge
+          queue: amd-cpu
 
-    {% for step in steps %}
-    {% if step.mirror_hardwares and mirror_hw in step.mirror_hardwares %}
+{% for step_group in steps %}
+{% if step_group.group %}
+  - group: "{{ step_group.group }}"
+{% if step_group.depends_on %}
+    depends_on: {{ step_group.depends_on }}
+{% endif %}
+    steps:
+{% for step in step_group.steps %}
+{% if step.mirror_hardwares and mirror_hw in step.mirror_hardwares %}
       - label: "{{ step.agent_pool }}: {{ step.label }}"
         depends_on: amd-build
         agents:
@@ -40,14 +55,87 @@
           {% else %}
           queue: amd_mi325_1
           {% endif %}
-        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 }}"
+{% 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
         env:
           DOCKER_BUILDKIT: "1"
+{% if step.env %}
+{% for name in step.env %}
+          {{ name }}: "{{ step.env[name] }}"
+{% endfor %}
+{% endif %}
+          TEST_COMMAND: |-
+            (command rocm-smi || true) && cd {{ (step.working_dir or default_working_dir) | safe }}
+{% if env %}
+{% for name in env %}
+            export {{ name }}="{{ env[name] }}"
+{% endfor %}
+{% endif %}
+{% if step.env %}
+{% for name in step.env %}
+            export {{ name }}="{{ step.env[name] }}"
+{% endfor %}
+{% endif %}
+{% 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 %}
         soft_fail: true
         {% endif%}
+{% endif %}
+{% endfor %}
+{% elif step_group.mirror_hardwares and mirror_hw in step_group.mirror_hardwares %}
+{% set step = step_group %}
+  - label: "{{ step.agent_pool }}: {{ step.label }}"
+    depends_on: amd-build
+    agents:
+      {% if step.agent_pool %}
+      queue: amd_{{ step.agent_pool }}
+      {% 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
+    env:
+      DOCKER_BUILDKIT: "1"
+{% if step.env %}
+{% for name in step.env %}
+      {{ name }}: "{{ step.env[name] }}"
+{% endfor %}
+{% endif %}
+      TEST_COMMAND: |-
+        (command rocm-smi || true) && cd {{ (step.working_dir or default_working_dir) | safe }}
+{% if env %}
+{% for name in env %}
+        export {{ name }}="{{ env[name] }}"
+{% endfor %}
+{% endif %}
+{% if step.env %}
+{% for name in step.env %}
+        export {{ name }}="{{ step.env[name] }}"
+{% endfor %}
+{% endif %}
+{% 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 %}
-    {% endfor %}
+    {% if step.grade and step.grade == "Blocking" %}
+    soft_fail: false
+    {% else %}
+    soft_fail: true
+    {% endif%}
+{% endif %}
+{% endfor %}

.buildkite/test-weekly.yml

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

@@ -0,0 +1,630 @@
+---
+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 implementation, then validate model-specific examples, tests, and docs as needed |
+| **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 for acceleration features such as `_repeated_blocks` and `_layerwise_offload_blocks_attrs` (see [references/transformer-adaptation.md](references/transformer-adaptation.md) for examples).
+
+### 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`.
+
+### A4.1 Add progress bar support (recommended)
+
+For pipelines with a standard denoising loop, prefer the existing progress bar pattern instead of hand-rolled logging.
+
+```python
+from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
+
+class YourPipeline(nn.Module, ProgressBarMixin):
+    def forward(self, req: OmniDiffusionRequest) -> DiffusionOutput:
+        # ... prepare timesteps / latents ...
+        with self.progress_bar(total=len(timesteps)) as progress_bar:
+            for i, t in enumerate(timesteps):
+                # predict noise / scheduler step
+                latents = ...
+                progress_bar.update()
+
+        return DiffusionOutput(output=output)
+```
+
+For custom loop structures, follow `vllm_omni/diffusion/models/progress_bar.py` and existing pipelines using `ProgressBarMixin`.
+
+### 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. See [references/custom-model-patterns.md](references/custom-model-patterns.md) for concrete custom-model integration patterns.
+
+### 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 common 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
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+
+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:
+                param = params[name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, 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/overview.md` — parallelism support overview/table
+2. `docs/user_guide/diffusion/cpu_offload_diffusion.md` — if CPU offload supported (add to supported models table)
+3. `docs/user_guide/diffusion/cache_acceleration/teacache.md` — if TeaCache supported
+4. `docs/user_guide/diffusion/cache_acceleration/cache_dit.md` — if Cache-DiT supported
+5. Offline example docs under `examples/offline_inference/<name>/` (`README.md` or category-specific `.md`)
+6. `examples/online_serving/<name>/README.md` — online serving docs
+
+### Step 8: Add E2E Tests
+
+**Follow the [vllm-omni-test skill](../vllm-omni-test/SKILL.md)** for markers, file naming, Buildkite wiring, and run commands. Also read [l4_functionality_tests.inc.md](https://github.com/vllm-project/vllm-omni/blob/main/docs/contributing/ci/test_examples/l4_functionality_tests.inc.md) and [CI_5levels.md](https://github.com/vllm-project/vllm-omni/blob/main/docs/contributing/ci/CI_5levels.md).
+
+Classify the model's **CI priority** first:
+
+| Priority | Required test levels | Files & markers |
+|----------|---------------------|-----------------|
+| **High** (listed in [#1832](https://github.com/vllm-project/vllm-omni/issues/1832) or on the diffusion hot path) | **L1** · **L2** online · **L3** online + offline · **L4** feature + performance | See table below |
+| **Medium** (*normal priority* in L4 docs) | **L3** online + offline · **L4** feature only | Fewer L4 parametrized rows |
+| **Low** | **L4** feature only | One or two `*_expansion.py` cases |
+
+**Per-level deliverables (diffusion / `pytest.mark.diffusion`):**
+
+| Level | Location | Marker | CI pipeline | Notes |
+|-------|----------|--------|-------------|-------|
+| **L1** | `tests/diffusion/models/{slug}/`, `tests/diffusion/cache/`, transformer unit tests | `core_model` + `cpu` | `test-ready.yml` | Weight remap, `_sp_plan`, cache enabler registration, shape contracts |
+| **L2** | `tests/e2e/online_serving/test_{slug}.py` (and offline if the category is offline-first) | **`core_model` + `advanced_model`** (both on baseline smoke) + `diffusion` + `@hardware_test` / `hardware_marks` | `test-ready.yml` | Default deploy smoke — minimal `num_inference_steps`, single prompt |
+| **L3** | `tests/e2e/online_serving/test_{slug}.py` **and** `tests/e2e/offline_inference/test_{slug}.py` when offline matters | Baseline smoke: **`core_model` + `advanced_model`**; heavier cases: `advanced_model` only (+ `diffusion`) | `test-merge.yml` **or** merged into nightly diffusion function job | Real weights, streaming/API paths, LoRA/offload smoke |
+| **L4** | `tests/e2e/online_serving/test_{slug}_expansion.py` (+ offline expansion if needed) | `full_model` + `diffusion` | `test-nightly.yml` (X2I / X2V / X2A function groups) | Feature combos per [#1832](https://github.com/vllm-project/vllm-omni/issues/1832); perf → `tests/dfx/perf/tests/test_{model}_vllm_omni.json` |
+
+**L2 & L3 online — same file, dual marks on the baseline smoke:** The **first / simplest** case in `test_{slug}.py` (default deploy, minimal steps, single prompt) should carry **both** `@pytest.mark.core_model` and `@pytest.mark.advanced_model` on the **same** function so L2 (`test-ready.yml`) and L3 (`test-merge.yml`) share one smoke test. Heavier deploy variants or API paths in the same file use **`advanced_model` only**. When L3 moves to nightly, migrate those heavier cases into `test_{slug}_expansion.py` with `full_model` and remove the dedicated `test-merge.yml` job (see `test_longcat_image_expansion.py`, `test_qwen_image_expansion.py`).
+
+**L4 design (high priority):** Combine multiple supported features (Cache-DiT, TP, USP, CFG, HSDP, CPU offload, quantization) into **few parametrized** `OmniServerParams` rows so each feature appears in at least one case without exploding GPU jobs. Shard single-GPU vs multi-GPU cases across the nightly X2I/X2V function steps (1-GPU vs `distributed_cuda`).
+
+**L4 design (medium / low):** One or two parametrized rows covering the best quality/perf trade-off; skip perf JSON unless the model is high priority.
+
+**Reference implementations:** `tests/e2e/online_serving/test_qwen_image_edit_expansion.py`, `tests/e2e/online_serving/test_longcat_image_expansion.py`, `tests/e2e/online_serving/test_hunyuan_video_15_expansion.py`.
+
+**Keep model-specific code inside test modules — not `tests/helpers/{slug}.py`:** deploy constants, prompts, sampling dicts, and inline `request_config` / `form_data` belong in each `test_{slug}.py` and `test_{slug}_expansion.py`. Do not add per-model files under `tests/helpers/`; reuse only repo-wide harness (`mark`, `media`, `runtime`, `stage_config`, `assertions`). **L2+ online/offline e2e:** reuse or add `send_*_request` in `tests/helpers/runtime.py` — tests call the handler, not raw `omni.generate` / HTTP. See [vllm-omni-test skill](../vllm-omni-test/SKILL.md) § **Runtime send helpers**.
+
+### 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/overview.md` — add your model to the support overview/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
+
+- [vllm-omni-test skill](../vllm-omni-test/SKILL.md) — L1–L4 markers, naming, Buildkite wiring, run commands
+- [Transformer Adaptation](references/transformer-adaptation.md) — porting transformers from diffusers
+- [Custom Model Patterns](references/custom-model-patterns.md) — patterns for non-diffusers models
+- [Parallelism Patterns](references/parallelism-patterns.md) — TP, SP/USP, CFG parallel, HSDP implementation details
+- [Cache-DiT Patterns](references/cache-dit-patterns.md) — cache-dit acceleration for standard and custom architectures
+- [Troubleshooting](references/troubleshooting.md) — common errors and fixes

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

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

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

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

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

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

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

@@ -0,0 +1,239 @@
+# 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_attrs = ["blocks"]
+
+    # LoRA: mapping of fused param names to original param names
+    packed_modules_mapping = {"to_qkv": ["to_q", "to_k", "to_v"]}
+
+    # Sequence parallelism plan (advanced — add after basic impl works)
+    _sp_plan = {
+        "blocks.0": SequenceParallelInput(split_dim=1),
+        "proj_out": SequenceParallelOutput(gather_dim=1),
+    }
+```

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

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

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

@@ -0,0 +1,564 @@
+---
+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},
+    ]
+```
+
+### Test Case Writing (CI Levels)
+
+**Follow the [vllm-omni-test skill](../vllm-omni-test/SKILL.md)** for markers, file naming (`test_{slug}.py` / `test_{slug}_expansion.py`), Buildkite wiring, and copy-paste run commands. Also read [CI_5levels.md](https://github.com/vllm-project/vllm-omni/blob/main/docs/contributing/ci/CI_5levels.md) and [tests_style.md](https://github.com/vllm-project/vllm-omni/blob/main/docs/contributing/ci/tests_style.md).
+
+Classify the model's **CI priority** first (high / medium / low). High-priority TTS models are typically those on the integration hot path or listed in tracking issues such as [#1832](https://github.com/vllm-project/vllm-omni/issues/1832); medium and low tiers cover the long tail. When unsure, ask the reviewer which tier applies.
+
+| Priority | Required test levels | Files & markers |
+|----------|---------------------|-----------------|
+| **High** | **L1** unit/logic · **L2** online smoke · **L3** online + offline integration · **L4** feature + performance | See table below |
+| **Medium** | **L3** online + offline · **L4** feature only | Skip dedicated L1/L2 unless fixing a logic bug |
+| **Low** | **L4** feature only | One or two `*_expansion.py` parametrized cases |
+
+**Per-level deliverables (TTS / `pytest.mark.tts`):**
+
+| Level | Location | Marker | CI pipeline | Notes |
+|-------|----------|--------|-------------|-------|
+| **L1** | `tests/model_executor/…`, `tests/entrypoints/openai_api/…`, stage-processor tests | `core_model` + `cpu` | `test-ready.yml` | Prompt assembly, async_chunk helpers, `serving_speech` validation — no GPU |
+| **L2** | `tests/e2e/online_serving/test_{slug}.py` | **`core_model` + `advanced_model`** (both on baseline smoke) + `tts` + `@hardware_test(...)` | `test-ready.yml` (`ready` label) | Default deploy smoke: single `/v1/audio/speech` or offline `OmniRunner` path |
+| **L3** | `tests/e2e/online_serving/test_{slug}.py` **and** `tests/e2e/offline_inference/test_{slug}.py` | Baseline smoke: **`core_model` + `advanced_model`**; heavier cases: `advanced_model` only (+ `tts`) | `test-merge.yml` **or** merged into nightly TTS function job | Streaming, voice clone, batch/queue, async_chunk |
+| **L4** | `tests/e2e/online_serving/test_{slug}_expansion.py`, optional offline expansion | `full_model` + `tts` | `test-nightly.yml` (`:full_moon: TTS · Function Test with L4`) | Feature matrix; perf → `tests/dfx/perf/tests/test_tts.json` |
+
+**L2 & L3 online — same file, dual marks on the baseline smoke:** The **first / simplest** case in `test_{slug}.py` (default deploy, single non-streaming `/v1/audio/speech` or equivalent offline path) should carry **both** `@pytest.mark.core_model` **and** `@pytest.mark.advanced_model` on the **same** function so it runs in L2 (`test-ready.yml`, `--run-level core_model`, basic validation) and L3 (`test-merge.yml`, `--run-level advanced_model`, deeper validation) without duplicating the test. In-tree examples: `test_voxcpm2_tts.py::test_text_to_audio_001`, `test_qwen3_tts_customvoice.py::test_text_to_audio_001`.
+
+Heavier scenarios in the same file use **`advanced_model` only** (streaming, extra languages, concurrency, async_chunk, batch). Example: `test_voice_clone_en_streaming_001` → `advanced_model` only. When migrating L3 to nightly, move those heavier cases into `test_{slug}_expansion.py` with `full_model` and drop the dedicated merge job (see `test_ming_tts_expansion.py`, `test_glm_tts_expansion.py`).
+
+```python
+@pytest.mark.core_model
+@pytest.mark.advanced_model
+@pytest.mark.tts
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
+def test_voice_clone_en_non_streaming_001(omni_server, openai_client) -> None:
+    openai_client.send_audio_speech_request({...})
+```
+
+**L4 consolidation:** Prefer parametrized `OmniServerParams` rows (`default`, `async_chunk`, feature flags) in one expansion module rather than many merge-only files ([#1832](https://github.com/vllm-project/vllm-omni/issues/1832)).
+
+**L4 performance (high-priority models):** Add latency / throughput / stress rows in `tests/dfx/perf/tests/test_tts.json`, or a dedicated `tests/dfx/perf/tests/test_{slug}.json` when the model must not join the shared nightly server matrix before integration lands (see VoxCPM2 / Coqui XTTS pattern). Register the model in `benchmarks/tts/model_configs.yaml` for local `bench_tts.py`. Wire a **separate** `test-nightly.yml` Perf Test step when the JSON is not merged into `test_tts.json` yet.
+
+**Keep model-specific code inside test modules — not `tests/helpers/{slug}.py`:**
+
+- Put `MODEL`, deploy path, vendored `REF_AUDIO_URL`, `get_prompt()`, and inline `request_config` dicts **in each** `test_{slug}.py`, `test_{slug}_expansion.py`, offline `test_{slug}.py`, and L1 `test_{slug}_*.py` as needed.
+- **Do not** add `tests/helpers/{slug}.py` (or `tests/helpers/{model_name}.py`) to deduplicate constants or request builders across those files. A little duplication is intentional; follow in-tree references such as `tests/e2e/online_serving/test_glm_tts.py` and `tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py`.
+- `tests/helpers/` is for **repo-wide** harness code only (`mark.py`, `media.py`, `runtime.py`, `stage_config.py`, `assertions.py`, `fixtures/`). Import those; do not extend the tree with per-model modules.
+
+**Runtime send helpers (`tests/helpers/runtime.py`) — online and offline e2e:**
+
+| Path | Fixture | Call |
+|------|---------|------|
+| Online `/v1/*` | `openai_client` | `openai_client.send_*_request(request_config)` |
+| Offline inference | `omni_runner_handler` | `omni_runner_handler.send_*_request(request_config)` |
+
+1. **Grep `runtime.py` first** — reuse `send_omni_request`, `send_diffusion_request`, `send_audio_speech_request` (online + offline Qwen-style TTS), `send_single_stage_tts_request` (Coqui XTTS / MOSS-TTS-Nano offline), etc.
+2. **No matching helper** → add `send_<feature>_request` (or `send_<route>_http_request` for negative/dfx) in **`runtime.py`** with general `assert_*` bundled inside, **then** call it from the test.
+3. **Test file** holds `request_config` dicts only — not `omni.generate`, not `_collect_audio()`, not raw HTTP/SDK.
+4. Wire `tests/helpers/runtime.py` into Buildkite `source_file_dependencies` when you add helpers.
+
+See [vllm-omni-test skill](../vllm-omni-test/SKILL.md) § **Runtime send helpers** for full tables and exceptions.
+
+### 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 tests per **Test Case Writing (CI Levels)** above (priority tier determines L1–L4 scope)
+- Buildkite wired per level: `test-ready.yml` (L1/L2), `test-merge.yml` or nightly function job (L3), `test-nightly.yml` (L4) — see [vllm-omni-test skill](../vllm-omni-test/SKILL.md)
+- 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 tests per the CI Levels table** — baseline smoke: `core_model` + `advanced_model`; heavier cases: `advanced_model` only; L4 expansion: `full_model`
+- **No per-model helper modules** — do not create `tests/helpers/{slug}.py`; keep constants and `request_config` payloads in the test file
+- **Online and offline e2e go through `runtime.py`** — `openai_client.send_audio_speech_request` (online); `omni_runner_handler.send_audio_speech_request` (Qwen-style offline) or `send_single_stage_tts_request` (single-stage offline). Add a new `send_*_request` in `runtime.py` when none fits; do not embed `omni.generate` or HTTP in tests
+
+## 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 tests added per model priority tier (see **Test Case Writing (CI Levels)**)
+- [ ] Buildkite entries match level: `test-ready.yml` / `test-merge.yml` or nightly TTS job / `test-nightly.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
+- [vllm-omni-test skill](../vllm-omni-test/SKILL.md) — L1–L4 markers, naming, Buildkite wiring, run commands
+- [Adding a TTS model (developer guide)](https://github.com/vllm-project/vllm-omni/blob/main/docs/contributing/model/adding_tts_model.md)
+- `plan/voxcpm2_native_ar_design.md` — VoxCPM2's vLLM-native AR + side-computation pattern (distinct from the generator-based single-stage described above)

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

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

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

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

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

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

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

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

.claude/skills/diffusion-perf-opt/SKILL.md

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

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

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

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

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

@@ -0,0 +1,135 @@
+# 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/quantization/SKILL.md

@@ -0,0 +1,103 @@
+---
+name: quantization
+description: Work on vLLM-Omni quantization for diffusion, autoregressive, omni, or multi-stage models. Use when choosing or adding methods such as fp8, int8, gguf, mxfp8, mxfp4, mxfp4_dualscale, ModelOpt, AutoRound, INC, msModelSlim, awq, or gptq; debugging quantized loading; or validating memory, speed, and output quality.
+---
+
+# vLLM-Omni Quantization
+
+Use this skill for quantization work in `vllm-omni`. Start from the local
+code and docs, then pick the smallest validated path for the model, method,
+hardware, and quality target.
+
+## First Checks
+
+Before changing code or recommending a command, identify:
+
+- Model and task: AR text, omni/TTS, text-to-image, text-to-video, or multi-stage diffusion.
+- Hardware backend: CUDA generation, Ascend NPU, Intel XPU, or another platform.
+- Quantization mode: online from BF16/FP16 weights, pre-quantized checkpoint, runtime KV/FA quantization, or per-component routing.
+- Scope: transformer/DiT, thinker language model, text encoder, VAE, audio/vision encoder, talker, or a specific stage.
+- Validation target: memory reduction, latency/throughput, output quality, or all three.
+
+Do not assume a method is supported just because the CLI accepts the string.
+Check `docs/user_guide/quantization/` and the closest implementation first.
+
+## Quick Decision
+
+| Task | Start With |
+|------|------------|
+| Choose a method or command | `references/methods.md` and `references/modality-compat.md` |
+| Use `build_quant_config()` or per-component routing | `references/methods.md` |
+| Work on diffusion quantization | `references/diffusion.md` |
+| Add quantization to a new model | `references/adding-models.md` |
+| Convert or load ModelOpt FP8 checkpoints | `references/modelopt-fp8.md` |
+| Debug ModelOpt, GGUF, AutoRound, MXFP, or serialized Int8 loading | `references/diffusion.md` and method docs under `docs/user_guide/quantization/` |
+
+## Current Runtime Shape
+
+The unified entrypoint is:
+
+```python
+from vllm_omni.quantization import build_quant_config
+```
+
+It supports method strings, flat method dictionaries, per-component
+dictionaries, existing `QuantizationConfig` objects, and `None`. The factory
+delegates generic methods to upstream `vllm` and keeps vLLM-Omni overrides for
+diffusion or omni-specific routing:
+
+- `gguf`
+- `int8`
+- `mxfp8`
+- `mxfp4`
+- `mxfp4_dualscale`
+- `inc`, `auto-round`, `auto_round`
+- ModelOpt auto-detected configs: `modelopt`, `modelopt_fp4`, `modelopt_mixed`
+
+Use `ComponentQuantizationConfig` when only one stage or component should be
+quantized. Pre-quantized ModelOpt-style checkpoints should not spill into
+vision/audio encoders that have no corresponding scale tensors.
+
+## Ownership Boundary
+
+- Upstream `vllm` owns generic quantization configs, kernels, loader semantics,
+  hardware capability rules, and generic AR quantization methods.
+- `vllm-omni` owns unified config routing, diffusion-specific wrappers,
+  component scoping, GGUF or ModelOpt checkpoint adapters, model-specific
+  prefix mapping, docs, examples, and validation.
+
+If a new method needs missing generic kernels or loader behavior, fix upstream
+`vllm` first. In `vllm-omni`, add thin integration and model wiring.
+
+## Working Loop
+
+1. Read the local method doc under `docs/user_guide/quantization/`.
+2. Inspect the active factory and config class in `vllm_omni/quantization/`.
+3. Inspect the target model pipeline and transformer for stable prefixes and
+   whether every relevant vLLM linear layer receives `quant_config`.
+4. For pre-quantized checkpoints, inspect `transformer/config.json` and the
+   checkpoint tensor names before running full generation.
+5. Validate with the same prompt, seed, size, scheduler, steps, dtype, eager or
+   non-eager mode, and parallelism as the BF16 baseline.
+6. Report quality, latency, and memory separately. Do not call a method
+   supported until quality is checked.
+
+## Common Mistakes
+
+| Symptom | Likely Cause | Fix |
+|---------|--------------|-----|
+| `--quantization` has no visible effect | Wrong scope or unsupported model path | Check component routing and method docs |
+| Some layers stay BF16 unexpectedly | `quant_config` was not threaded into all vLLM linear layers | Audit transformer constructors and prefixes |
+| Quality collapses but loading succeeds | Too many sensitive layers were quantized | Add model-specific `ignored_layers` and compare to BF16 |
+| ModelOpt checkpoint loads but output is corrupted | Prefixes, packed-module mapping, scale routing, or BF16 fallback is wrong | Use `references/modelopt-fp8.md` |
+| GGUF shape or tensor mismatch | Missing architecture-specific adapter | Add explicit adapter mapping; avoid generic fallback |
+| Online and offline MXFP4 disagree | Wrong `mxfp4` vs `mxfp4_dualscale` mode | Check `docs/user_guide/quantization/mxfp4.md` |
+| Quantized path is slower than BF16 | Kernel path, compile mode, dtype, or shape mismatch | Compare eager-to-eager and non-eager-to-non-eager |
+
+## References
+
+- General methods and config routing: [references/methods.md](references/methods.md)
+- Modality and model scope: [references/modality-compat.md](references/modality-compat.md)
+- Diffusion implementation patterns: [references/diffusion.md](references/diffusion.md)
+- Adding model support: [references/adding-models.md](references/adding-models.md)
+- ModelOpt FP8 conversion and loading: [references/modelopt-fp8.md](references/modelopt-fp8.md)

.claude/skills/quantization/references/adding-models.md

@@ -0,0 +1,70 @@
+# Adding Quantization Support to a Model
+
+Use this checklist when extending quantization to a new architecture or stage.
+
+## Classify Ownership
+
+| Case | Owner |
+|------|-------|
+| New generic AR method, kernel, or loader semantics | upstream `vllm` first |
+| Existing upstream method needs diffusion wiring | `vllm-omni` model and config integration |
+| GGUF or ModelOpt tensor mapping for a diffusion architecture | `vllm-omni` adapter |
+| NPU MXFP or msModelSlim diffusion path | `vllm-omni` platform/method wrapper and docs |
+| Multi-stage omni checkpoint quantizes only the thinker LM | `vllm-omni` component scoping |
+
+Avoid building a private quantization stack in model code.
+
+## Model Wiring Checklist
+
+1. Find the closest supported model implementation.
+2. Confirm the model uses vLLM linear layers that accept `quant_config`.
+3. Add `quant_config` to the transformer constructor.
+4. Pass stable `prefix` values through every linear layer.
+5. Handle fused modules and packed mappings explicitly.
+6. Ensure `load_weights()` and any `WeightsMapper` remap `ignored_layers` or
+   `modules_to_not_convert`.
+7. Keep tokenizer, scheduler, VAE, text encoder, audio/vision encoders, and
+   talker BF16 unless method-specific docs validate them.
+8. Add or update a method doc under `docs/user_guide/quantization/`.
+
+## Checkpoint Adapter Checklist
+
+For pre-quantized checkpoints:
+
+- Inspect `transformer/config.json` first.
+- Check `quant_method`, `method`, `producer.name`, `quant_algo`, serialized
+  flags, `ignored_layers`, and any method-specific fields.
+- Verify tensor name mapping before generation.
+- Verify packed modules such as `to_qkv`, `add_kv_proj`, and `w13`.
+- Confirm scale tensor names and shapes.
+- Allow dequantization back to BF16 for intentionally skipped target layers.
+
+## Sensitive Layer Discovery
+
+Start conservative. Quantize the backbone first and keep output-side or routing
+layers BF16 until quality is proven.
+
+Frequently sensitive areas:
+
+- attention output projections such as `to_out`
+- FFN down projections such as `w2`
+- final projections such as `proj_out`, `final_layer`, and output heads
+- modulation, timestep, position, image, or text embedders
+- refiner blocks
+- MoE routers, experts selected by routing, and layernorm-adjacent modules
+- VAE, text/vision towers, audio encoders, and codec/talker stages
+
+Do not reuse one model family's ignore list blindly.
+
+## Tests and Evidence
+
+Minimum evidence for a PR:
+
+- loader/config test or smoke command
+- fixed-seed BF16 vs quantized output comparison
+- quality metrics and saved outputs for precision-changing work
+- memory and latency comparison under identical execution mode
+- failure behavior for unsupported models or missing checkpoint fields
+
+When the PR adds a method to a support table, include the exact model,
+hardware, command, quality threshold, and artifact path used for validation.

.claude/skills/quantization/references/diffusion.md

@@ -0,0 +1,113 @@
+# Diffusion Quantization Reference
+
+Use this file for DiT quantization, config routing, checkpoint adapters, and
+quality or loader debugging.
+
+## Supported Method Families
+
+The current local diffusion quantization families are:
+
+- CUDA/GPU online or checkpoint: `fp8`, `int8`, `gguf`, ModelOpt
+- Ascend NPU: `int8`, `mxfp8`, `mxfp4`, `mxfp4_dualscale`, msModelSlim paths
+- Intel XPU: AutoRound/INC and online MXFP8 paths where documented
+
+Always cross-check the method doc before claiming model support. The CLI
+accepting a method string is not enough.
+
+## Implementation Starting Points
+
+| Area | Files |
+|------|-------|
+| Unified factory | `vllm_omni/quantization/factory.py` |
+| Per-component routing | `vllm_omni/quantization/component_config.py` |
+| GGUF diffusion method | `vllm_omni/quantization/gguf_config.py` |
+| Int8 diffusion method | `vllm_omni/quantization/int8_config.py` |
+| MXFP8 diffusion method | `vllm_omni/quantization/mxfp8_config.py` |
+| MXFP4 and dual-scale method | `vllm_omni/quantization/mxfp4_config.py` |
+| AutoRound/INC omni extension | `vllm_omni/quantization/inc_config.py` |
+| ModelOpt checkpoint adapters | `vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py` |
+| FP8 scaled-mm runtime patches | `vllm_omni/patch.py` |
+| Similarity tool | `vllm_omni/quantization/tools/compare_diffusion_trajectory_similarity.py` |
+
+## Threading `quant_config`
+
+For a diffusion transformer:
+
+1. Accept `quant_config: QuantizationConfig | None` in the transformer.
+2. Pass `quant_config` and stable `prefix` values to every vLLM linear layer.
+3. Preserve packed-module mappings for fused QKV, FFN, or model-specific fused
+   projections.
+4. Call `apply_vllm_mapper()` on configs that need HF-to-vLLM name remapping.
+5. For component configs, keep prefixes specific enough to avoid quantizing VAE
+   or encoders by accident.
+
+Common skipped areas are `proj_out`, output projections, modulation layers,
+embedders, refiner blocks, MoE routers, and layernorm-adjacent paths.
+
+## Online vs Offline
+
+Online quantization:
+
+- starts from BF16/FP16 checkpoint weights
+- computes or stores quantized weights at load time
+- typically uses `--quantization <method>` or `quantization="<method>"`
+
+Offline or pre-quantized checkpoints:
+
+- include quantized weights and scales on disk
+- usually carry `quantization_config` in `transformer/config.json`
+- should auto-detect when the config is complete
+- may require model-specific adapter logic for tensor names and packed modules
+
+Do not mix online flags with incompatible offline checkpoint configs. If the
+active config and disk config differ, `resolve_quant_config_from_disk()` should
+raise or rebuild rather than silently loading corrupted weights.
+
+## GGUF Rules
+
+- Require `quantization_config.gguf_model`.
+- Use the normal base model for tokenizer, scheduler, text encoder, and VAE.
+- Add explicit architecture adapters for tensor-name mapping.
+- Use `named_parameters()` and `named_buffers()`, not `state_dict()`, to
+  discover loadable names.
+- Guard fused QKV/KV rewrites so tensors are not rewritten twice.
+- Flatten N-D activations around the GGUF matmul and restore shape.
+
+Unsupported architectures should fail clearly instead of falling back to a fake
+generic mapper.
+
+## MXFP Rules
+
+- `mxfp8` is W8A8 microscaling, currently centered on NPU and documented XPU
+  behavior.
+- `mxfp4` is online single-scale W4A4 on NPU.
+- `mxfp4_dualscale` is online/offline dual-scale W4A4 with BF16 fallback; the
+  offline path is the production path when calibrated `mul_scale` exists.
+- Offline dual-scale checkpoints must inject per-transformer `ignored_layers`
+  into config because cascade transformers can have different BF16 fallback
+  sets.
+
+Do not load an offline dual-scale checkpoint with the online single-scale path.
+
+## Common Failures
+
+| Symptom | Likely Cause | Fix |
+|---------|--------------|-----|
+| `Unknown quantization method` | Factory override or upstream registration is missing | Add the right override or upstream method |
+| Missing scale tensors | Wrong scope, wrong adapter, or BF16 layer not allowed to dequantize | Check config scope and adapter scale routing |
+| Output rank collapses from 3-D to 2-D | FP8 scaled-mm kernel ignored output shape | Check the patch in `vllm_omni/patch.py` |
+| Quality degrades silently | Sensitive layers quantized | Add `ignored_layers`; compare to BF16 |
+| Only one cascade transformer works | Disk config differs per transformer but active config is reused | Rebuild from each transformer's config |
+| Quantized path slower than BF16 | Kernel, dtype, compile, or benchmark mode mismatch | Compare eager-to-eager and non-eager-to-non-eager |
+
+## Verification
+
+For a new or changed diffusion quantization path, run:
+
+1. Config and loader smoke test.
+2. One fixed-seed generation against BF16.
+3. Numeric comparison with the similarity tool or equivalent metrics.
+4. Serving test if the path is user-facing online.
+5. Memory and latency check under the same execution mode as the baseline.
+
+Keep output artifacts when precision changes are part of the PR.

.claude/skills/quantization/references/methods.md

@@ -0,0 +1,82 @@
+# General Quantization Methods
+
+Use this file for method selection and unified `vllm-omni` config syntax. For
+diffusion-specific wiring, see `diffusion.md`. For ModelOpt checkpoint export
+and loading details, see `modelopt-fp8.md`.
+
+## Unified Entry Point
+
+`vllm_omni.quantization.build_quant_config()` accepts:
+
+- `None` or `"none"` to disable quantization
+- a method string such as `"fp8"`, `"int8"`, `"mxfp8"`, `"awq"`, or `"auto-round"`
+- a flat dict with `"method"` or `"quant_method"`
+- a per-component dict such as `{"transformer": {"method": "fp8"}, "vae": None}`
+- an already-built `QuantizationConfig`
+
+Per-component dicts are routed by `ComponentQuantizationConfig` using
+longest-prefix matching. Prefix stability matters: if a model remaps HF names
+to vLLM names, apply the mapper to `ignored_layers` or component prefixes.
+
+## Local Overrides
+
+The local factory delegates most methods to upstream `vllm`, but has vLLM-Omni
+overrides for methods that need diffusion, NPU, XPU, or component behavior:
+
+| Method | Local Config | Main Use |
+|--------|--------------|----------|
+| `gguf` | `DiffusionGGUFConfig` | Pre-quantized diffusion transformer weights |
+| `int8` | `DiffusionInt8Config` | Online or serialized W8A8 diffusion transformers |
+| `mxfp8` | `DiffusionMXFP8Config` | W8A8 MXFP8 diffusion on NPU and online XPU path |
+| `mxfp4` | `DiffusionMXFP4Config` | Online W4A4 MXFP4 diffusion on NPU |
+| `mxfp4_dualscale` | `DiffusionMXFP4DualScaleMixedConfig` | Online/offline W4A4 dual-scale plus BF16 fallback on NPU |
+| `inc`, `auto-round`, `auto_round` | `OmniINCConfig` | AutoRound/INC checkpoint loading, including MXFP8 checkpoint metadata |
+| ModelOpt config dicts | upstream ModelOpt configs through auto-detect | FP8, NVFP4, or mixed ModelOpt checkpoints |
+
+Check `SUPPORTED_QUANTIZATION_METHODS` in `vllm_omni/quantization/factory.py`
+before adding a new method name.
+
+## AR and Generic Quantization
+
+For AR-backed models, upstream `vllm` methods usually apply:
+
+- `awq` and `gptq`: 4-bit or 8-bit weight quantization with calibration,
+  commonly used for text and AR stages.
+- `fp8`: FP8 weight path or KV-cache FP8 depending on flags and hardware.
+- `inc` or `auto-round`: pre-quantized checkpoints with Intel Neural
+  Compressor / AutoRound metadata.
+
+Use these for the AR language model stage only when the checkpoint and model
+guide support it. Omni and TTS models often keep audio, vision, talker, and
+codec stages in BF16.
+
+## Diffusion and Offline Methods
+
+For diffusion transformers, prefer the method-specific docs:
+
+- `docs/user_guide/quantization/fp8.md`
+- `docs/user_guide/quantization/int8.md`
+- `docs/user_guide/quantization/gguf.md`
+- `docs/user_guide/quantization/modelopt.md`
+- `docs/user_guide/quantization/mxfp8.md`
+- `docs/user_guide/quantization/mxfp4.md`
+- `docs/user_guide/quantization/autoround.md`
+- `docs/user_guide/quantization/msmodelslim.md`
+
+Pre-quantized checkpoints should normally auto-detect from config. Runtime
+flags such as `--force-cutlass-fp8`, `--linear-backend cutlass`, and
+`--moe-backend cutlass` select kernels; they do not quantize a BF16 checkpoint.
+
+## Validation
+
+Use fixed-seed A/B comparisons and record:
+
+- model path and quantized checkpoint path
+- method config and `ignored_layers`
+- prompt, size, frame count, steps, scheduler, dtype, and seed
+- eager/non-eager mode, compile flags, and parallelism
+- output metrics such as PSNR, MAE, cosine similarity, SSIM, LPIPS, plus visual artifacts
+- latency and peak memory, reported separately from quality
+
+The helper `vllm_omni.quantization.tools.compare_diffusion_trajectory_similarity`
+is the default lightweight image/video comparison tool for diffusion paths.

.claude/skills/quantization/references/modality-compat.md

@@ -0,0 +1,47 @@
+# Quantization Compatibility by Modality
+
+Use this as a routing guide. The source of truth is the local method docs under
+`docs/user_guide/quantization/` and the current model implementation.
+
+## Summary
+
+- AR and omni language-model quantization mostly follows upstream `vllm`, with
+  vLLM-Omni scoping for multi-stage models.
+- Diffusion quantization is method-specific and model-specific. It usually
+  targets the diffusion transformer only.
+- Pre-quantized checkpoints should quantize only the component that has
+  quantized weights and scale tensors.
+
+## Model Type Matrix
+
+| Model Type | Typical Quantization Scope | Notes |
+|------------|----------------------------|-------|
+| AR text models | Language model | Upstream `vllm` methods usually apply |
+| Omni/TTS models | Thinker or AR language-model stage | Audio, vision, talker, and codec stages usually stay BF16 |
+| DiT image models | Diffusion transformer | Check FP8, Int8, GGUF, ModelOpt, AutoRound docs |
+| DiT video models | Diffusion transformer | HunyuanVideo FP8 and Wan2.2 MXFP paths are documented; validate other video paths |
+| Multi-stage diffusion | Selected stage only | Use per-component routing and per-stage validation |
+| Audio diffusion | Usually unsupported | Prefer dtype, offload, cache, or parallelism unless docs say otherwise |
+
+## Method Hints
+
+| Goal | Candidate |
+|------|-----------|
+| AR memory reduction | AWQ, GPTQ, AutoRound, ModelOpt where checkpoint supports it |
+| CUDA diffusion W8A8 | FP8 or Int8 |
+| Pre-quantized CUDA diffusion checkpoint | ModelOpt or GGUF |
+| Ascend NPU diffusion | Int8, MXFP8, MXFP4, MXFP4 DualScale, msModelSlim |
+| Intel XPU checkpoint | AutoRound/INC paths |
+| Conservative image quality | selective FP8/Int8 with `ignored_layers` |
+| Unsupported diffusion model | cache acceleration, CPU/layerwise offload, dtype, TP/SP/CFG parallelism |
+
+## Validation Rule
+
+Compatibility means more than successful loading. A method is compatible only
+after it has:
+
+- correct component scope
+- complete checkpoint/config metadata if offline
+- fixed-seed quality comparison against BF16
+- acceptable latency and memory behavior
+- docs or examples that state model, method, hardware, and caveats

.claude/skills/quantization/references/modelopt-fp8.md

@@ -0,0 +1,122 @@
+# ModelOpt FP8 Conversion and Loading
+
+Use this when exporting or loading ModelOpt FP8 checkpoints for diffusion
+transformers. This reference captures the practical conversion lessons from the
+2026-04-21 ModelOpt FP8 paste and aligns them with the current
+`vllm-omni` ModelOpt adapter.
+
+## Core Rule
+
+Treat ModelOpt FP8 as selective checkpoint quantization, not one-click
+whole-model FP8. Quantize the main transformer or DiT backbone, keep sensitive
+or peripheral modules BF16, serialize a complete ModelOpt checkpoint, and let
+vLLM-Omni auto-detect it from config.
+
+Do not pass `--quantization fp8` to force a ModelOpt checkpoint path. The
+checkpoint's `quantization_config` should select `modelopt`.
+
+## Conversion Strategy
+
+Prefer:
+
+- backbone attention Q/K/V and FFN expansion projections
+- conservative per-model ignore lists
+- MHA quantizers disabled until basic linear quantization is stable
+- full module-name prefixes in ignore/routing logic
+- BF16 fallback for output-side and routing-sensitive layers
+
+Avoid quantizing by default:
+
+- VAE
+- text, vision, or audio encoders
+- embedders and modulation layers
+- final heads and output projections
+- refiner modules
+- MoE routers and layernorm-sensitive blocks
+
+## Model-Specific Starting Points
+
+| Model Family | Quantize First | Keep BF16 First |
+|--------------|----------------|-----------------|
+| Z-Image | attention q/k/v, FFN w1/w3 | `to_out`, `w2`, `noise_refiner`, `context_refiner`, embedders, modulation, final/output |
+| Qwen-Image | transformer backbone linear layers | `proj_out`, `img_in`, `txt_in`, norms, modulation, position/time/embed layers |
+| FLUX2-dev / FLUX2-klein | backbone linear layers | `proj_out`, context/x embedders, `norm_out`, stream modulation, time embeddings |
+| HunyuanImage-3.0 | `model.model` backbone | VAE, vision tower, token/time/image embedders, final image head, LM head, routers, layernorm-sensitive blocks |
+
+These are starting points, not support claims. Validate with BF16 outputs.
+
+## Required Checkpoint Metadata
+
+`transformer/config.json` must include a complete `quantization_config`.
+At minimum it should identify ModelOpt and the algorithm:
+
+```json
+{
+  "quantization_config": {
+    "producer": {"name": "modelopt"},
+    "quant_method": "modelopt",
+    "quant_algo": "FP8"
+  }
+}
+```
+
+Current vLLM-Omni auto-detects:
+
+- `FP8` and `FP8_PER_CHANNEL_PER_TOKEN` -> `modelopt`
+- `NVFP4` -> `modelopt_fp4`
+- `MIXED_PRECISION` -> `modelopt_mixed`
+
+Packed modules and scales must map correctly:
+
+- `to_qkv` from `to_q`, `to_k`, `to_v`
+- `add_kv_proj` from `add_q_proj`, `add_k_proj`, `add_v_proj`
+- `w13` from `w1`, `w3`
+- `.input_scale`, `.weight_scale`, `.weight_scale_2`, `.weight_scale_inv`
+
+For selective FP8, target layers that remain BF16 must be allowed to
+dequantize from FP8 weights using the matching `weight_scale`, or the adapter
+must skip the quantized scale tensors for that BF16 target.
+
+## Runtime Signals
+
+Healthy loading should show ModelOpt auto-detection and a ModelOpt linear
+kernel path. Useful signals include:
+
+- auto-detected quantization from config
+- selected ModelOpt FP8 linear method
+- CUTLASS or FlashInfer FP8 scaled-mm backend when enabled and supported
+- non-eager path compiling the transformer when expected
+
+If online serving fails while offline inference works, inspect warmup and
+serving initialization before changing checkpoint conversion.
+
+## Validation Order
+
+1. Check config completeness.
+2. Run an export checker such as `examples/quantization/check_modelopt_fp8_export.py`.
+3. Confirm actual FP8 tensors and expected disk-size reduction.
+4. Run one offline generation with fixed prompt, seed, size, and steps.
+5. Run serving with `vllm serve --omni` and verify `/v1/images/generations`.
+6. Compare against BF16 using visual review plus metrics such as LPIPS, PSNR,
+   MAE, cosine similarity, and SSIM.
+7. Compare performance only under matched modes:
+   - BF16 eager vs FP8 eager
+   - BF16 non-eager vs FP8 non-eager
+
+FP8 is not automatically faster. Check whether execution reaches FP8 kernels,
+whether compile mode is comparable, and whether the workload shape fits the
+kernel path.
+
+## Triage
+
+| Symptom | First Check |
+|---------|-------------|
+| Startup failure | `config.json`, `model_index.json`, `quantization_config`, missing scales |
+| Loads but output is corrupted | Over-quantized sensitive layers or wrong full-name prefixes |
+| Offline works but online fails | Dummy warmup, serving init, compile path |
+| FP8 is not faster | Kernel selection, dtype, compile mode, workload shape |
+| BF16 fallback fails | Missing `weight_scale` for dequantization or wrong packed-module mapping |
+
+The success criterion is stable, correctly scoped checkpoint loading with
+quality preserved enough for the target use case, not maximizing the number of
+FP8 layers.

.claude/skills/readme.md

@@ -0,0 +1,50 @@
+# 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
+- `quantization`: guides quantization method selection, model integration,
+  checkpoint loading, and quality/performance validation for vLLM-Omni
+- `add-omni-model`: covers addition of new omni-modality model support
+- `add-tts-model`: covers integration of new TTS models and related serving
+  workflows
+- `generate-release-note`: helps prepare release notes for repository changes
+- `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
+- `vllm-omni-test`: guides generation and execution of CI-aligned tests (L1–L4),
+  pytest marker selection (`core_model` / `advanced_model` / `full_model`,
+  `omni` / `tts` / `diffusion`), Buildkite wiring (`test-ready.yml`,
+  `test-merge.yml`, `test-nightly.yml`, `test-weekly.yml`), and copy-paste
+  local plus CI-like `pytest` commands; see `references/test-routing.md` for
+  level-to-command mapping
+- `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

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

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

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

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

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

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

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

.claude/skills/vllm-omni-test/SKILL.md

@@ -0,0 +1,890 @@
+---
+name: vllm-omni-test
+description: Generate and run tests for vllm-project/vllm-omni with CI-aligned levels and markers; wire new tests into Buildkite (test-ready.yml for L1/L2, test-merge.yml for L3, test-nightly.yml for L4). On completion, always provide copy-paste local and CI-like pytest commands plus prerequisites. Use when creating regression tests, adding L1-L4 coverage, selecting pytest markers, or validating fixes from issues/PRs.
+---
+
+# vLLM-Omni Test Generator & Runner
+
+## Purpose
+
+Use this skill to generate minimal, stable test cases and run them with the correct marker/level strategy for [vllm-project/vllm-omni](https://github.com/vllm-project/vllm-omni).
+
+**Link convention:** Paths such as `.buildkite/` and `docs/contributing/` live at the **vllm-omni repo root**. Markdown links use repo-relative paths from this skill file (e.g. `../../../.buildkite/test-ready.yml`, `../../../docs/contributing/ci/CI_5levels.md`).
+
+Default priorities:
+
+1. Reproducible regression coverage for bug fixes
+2. Correct test level and marker selection
+3. Low flake, low dependency tests first
+4. CI-compatible run commands
+5. **Actionable run commands for the human**: whenever you add or change tests, always finish with **copy-paste-ready** `pytest` lines (local: single file and/or single test; CI-like: markers + `--run-level`), plus short **prerequisites** (GPU tier, HF cache, optional `model_prefix`). Do not assume the reader will infer commands from `test-routing.md` alone.
+
+## Inputs
+
+- Issue/PR link and summary
+- Changed files or suspected code path
+- Whether the user wants local quick validation or CI-equivalent validation
+- Hardware constraints (CPU only / CUDA / ROCm / NPU)
+
+## Workflow
+
+### Step 1: Classify Test Goal
+
+- **Bugfix regression**: start from a minimal failing scenario and add assertions that prevent recurrence. Before writing tests, output **`required` / `recommended` / `not_needed`**: **`required`** — stable logic/contract bug that should have been caught; **`recommended`** — environment-sensitive but a small regression still helps; **`not_needed`** — one-off external/config failure or existing tests already cover the path. Prefer the narrowest stable L1 (CPU) case; escalate to L2/L3 only when the bug needs real weights or serving.
+- **Feature coverage**: verify new behavior and one negative/boundary case.
+- **Perf/benchmark claim**: require benchmark-oriented tests and explicit metrics.
+
+### Step 2: Select Test Level
+
+- **L1**: unit/logic, deterministic, CPU-friendly, fastest feedback.
+- **L2**: basic e2e/integration and platform-dependent checks.
+- **L3/L4**: advanced model/integration/perf validation.
+
+Use [references/test-routing.md](references/test-routing.md) for level-to-marker and command mapping.
+
+### Step 3: Pick Markers
+
+Always attach markers deliberately:
+
+- **Level**: `core_model` (L1/L2) and/or `advanced_model` (L3) and/or `full_model` (L4 nightly)
+- **Model type** (required on model-centric e2e — pick exactly one):
+  - `omni` — end-to-end multimodal LLM pipelines (thinker/talker/stages; Qwen-Omni family)
+  - `tts` — speech synthesis / TTS-only models (`/v1/audio/speech`, voice clone, etc.)
+  - `diffusion` — generative diffusion models (image / audio / text / video from noise)
+- **Cross-cutting area** (when relevant): `parallel`, `cache`, `example`, `benchmark`
+- **Hardware**: `cpu`, `gpu`, `cuda`, `rocm`, `npu`, `L4`, `H100`, `distributed_cuda`, …
+- Optional: `slow`, distributed markers when multi-card is required
+
+**Baseline smoke (L2 + L3):** The simplest e2e case per model — default deploy, minimal request — should usually carry **both** `@pytest.mark.core_model` and `@pytest.mark.advanced_model` on the **same** test function so `test-ready.yml` and `test-merge.yml` share one test. `send_*_request` picks validation depth from `--run-level`. References: `test_voxcpm2_tts.py::test_text_to_audio_001`, `test_qwen3_tts_customvoice.py::test_text_to_audio_001`. Heavier scenarios use **`advanced_model` only**; L4 expansion uses **`full_model`**.
+
+For hardware-aware tests, prefer `@hardware_test(...)` or `hardware_marks(...)` in `tests/helpers/mark.py`.
+
+**Diffusion nightly split** (`test-nightly.yml`): all diffusion tests use `pytest.mark.diffusion`, but CI groups them by **output modality**, not by separate pytest markers:
+
+| Nightly group | Scope | Typical models / paths |
+|---------------|--------|-------------------------|
+| **Diffusion X2I(&A&T) Model Test** | x2**i** (image), x2**a** (audio), x2**t** (text) and other **non-video** diffusion | Qwen-Image*, BAGEL, FLUX, SD3, Z-Image, LongCat, DreamZero, … — `test_*_expansion.py` under X2I shards; L4 sweep uses `-k "not test_wan and not test_bagel_expansion and not hunyuan"` for L4 |
+| **Diffusion X2V Model Test** | x2**v** (video) only | Wan2.2, HunyuanVideo 1.5, Wan VACE, … — e.g. `test_wan22_expansion.py`, `test_hunyuan_video_15_expansion.py` |
+
+Wire new **video** diffusion expansion tests into the **X2V** group; wire **image/audio/text** diffusion expansion tests into **X2I(&A&T)**. Do not place x2v modules in X2I shards (see comment in `test-nightly.yml` above the X2I group).
+
+### Naming: generated test module files (L2–L4, model-centric e2e)
+
+When adding **new** pytest modules whose primary scope is a **specific model** (typical under `tests/e2e/offline_inference/` or `tests/e2e/online_serving/`), use this filename pattern:
+
+| Level | Filename pattern | Example (model `Qwen/Qwen2.5-Omni-7B`) |
+|-------|------------------|----------------------------------------|
+| **L2**, **L3** | `test_{lowercase_model_slug}.py` | `test_qwen2_5_omni.py` |
+| **L4** | `test_{lowercase_model_slug}_expansion.py` | `test_qwen2_5_omni_expansion.py` |
+
+**Slug rules for `{lowercase_model_slug}`:**
+
+1. Start from the HuggingFace-style id (e.g. `Qwen/Qwen2.5-Omni-7B`), but **do not** put the org into the filename: use the **repo segment only** (`Qwen2.5-Omni-7B`), not `Qwen_...` / `qwen_qwen2_5_...`.
+2. **Lowercase**; replace `.`, `-`, and whitespace with a single `_` (e.g. `Qwen2.5-Omni-7B` → `qwen2_5_omni`). **Omit** trailing size tokens such as `7b` / `30b` in the basename when a single file covers that model line in the directory (matches `test_qwen2_5_omni.py` in-tree).
+3. If two checkpoints in the same folder need separate modules, add a **minimal** disambiguator (e.g. `_7b` vs `_3b`) only then.
+4. **L1** unit tests are **not** bound to this pattern; use `tests/<area>/test_<feature>.py` as today.
+
+Routing tables and commands: [references/test-routing.md](references/test-routing.md) § *Model-centric e2e filename convention*.
+
+Existing references: `tests/e2e/offline_inference/test_qwen2_5_omni.py` (L2-style omni), `tests/e2e/offline_inference/test_qwen3_5_9b.py` (L2-style omni, single-stage VL), `tests/e2e/online_serving/test_qwen3_omni_expansion.py` (L4-style omni), `tests/e2e/online_serving/test_qwen_image_edit_expansion.py` / `test_qwen_image_expansion.py` (L4-style diffusion).
+
+### Step 4: Generate Test Case Skeleton
+
+**1. Pick the functional scenario** (then choose directory, fixtures, and markers):
+
+| Scenario | Typical location | Fixtures / runner pattern | Baseline markers & level |
+|----------|------------------|---------------------------|---------------------------|
+| **Offline inference e2e** | `tests/e2e/offline_inference/` | **Module (default):** `omni_runner` + `omni_runner_handler`. **Function (isolation only):** `omni_runner_function` + `omni_runner_handler_function`. Diffusion/TTS may use `Omni(...).generate` directly | L2: `core_model` + **one of** `omni` / `tts` / `diffusion`; `@hardware_test(...)` when GPU/NPU is required |
+| **Online serving e2e** | `tests/e2e/online_serving/` | **Module (default):** `omni_server` + `openai_client`. **Function (isolation only):** `omni_server_function` + `openai_client_function`. Clients: `send_omni_request` (omni), `send_audio_speech_request` (tts), `send_diffusion_request` / `send_video_diffusion_request` / `send_images_generations_request` (diffusion) | Baseline smoke: **`core_model` + `advanced_model`**; heavier paths: `advanced_model` only; L4 expansion: `full_model` |
+| **Documentation / runnable examples** | `tests/examples/offline_inference/`, `tests/examples/online_serving/` | **Offline docs (preferred):** extract Python/Bash blocks from the doc README (e.g. `ReadmeSnippet.extract_readme_snippets`), `pytest.mark.parametrize` each snippet, run via `example_runner.run` with a stable `output_subfolder`. **Online docs:** copy client/request scripts into dedicated tests and keep them in sync with the doc page. | Usually **L4**: `advanced_model`, often `example` plus hardware marks matching the nightly docs-example job (see `.buildkite/test-nightly.yml`). Full conventions: [docs/contributing/ci/test_examples/l4_doc_example_tests.inc.md](../../../docs/contributing/ci/test_examples/l4_doc_example_tests.inc.md) (introduced in [PR #1910](https://github.com/vllm-project/vllm-omni/pull/1910): naming, output directory layout, skip rules, avoid trimming `num_inference_steps` without a strong CI reason). |
+| **Performance / benchmark** | `tests/dfx/perf/tests/*.json` + `run_*_benchmark.py` | JSON or script-driven server + load config; assert explicit metrics / baselines | L4 Perf: `full_model` + `benchmark`; wire `test-nightly.yml` Perf steps |
+| **Invalid parameter / negative HTTP validation** | `tests/dfx/reliability/invalid_param_test/` | Live `omni_server` + low-level `send_*_http_request` with `err_code` / `err_message` | `pytest.mark.slow` + `omni` / `tts` / `diffusion` + `@hardware_marks` (`H100` or `L4`); CI in **`test-weekly.yml`** (not ready/merge/nightly) |
+
+**If the user’s test plan includes invalid parameter validation / invalid params / negative HTTP / 400 validation:** do **not** add those `test_*` functions to `tests/e2e/online_serving/test_*.py` or `*_expansion.py`. **Move or author them** under `tests/dfx/reliability/invalid_param_test/` in the **endpoint-matching script** (see **Invalid parameter validation** below). Success-path e2e and invalid-param dfx tests must stay in separate modules.
+
+**1b. Model type — `omni` vs `tts` vs `diffusion`**
+
+After choosing offline/online/docs/perf, classify the **product under test** and attach **exactly one** model-type marker. All three can live under the same `tests/e2e/` trees; conventions diverge:
+
+| Dimension | **Omni** (`pytest.mark.omni`) | **TTS** (`pytest.mark.tts`) | **Diffusion** (`pytest.mark.diffusion`) |
+|-----------|------------------------------|----------------------------|----------------------------------------|
+| **What it is** | Multimodal LLM pipeline (thinker/talker/stages; text + vision + audio I/O) | Speech synthesis / voice models | Generative diffusion (noise → image, audio, text, or video) |
+| **Examples** | Qwen2.5-Omni, Qwen3-Omni | Qwen3-TTS, VoxCPM2, Higgs-Audio, Step-Audio2 | Qwen-Image, BAGEL, Wan2.2, HunyuanVideo |
+| **Offline runner** | `omni_runner` + `omni_runner_handler`, `generate_multimodal` | `Omni(...).generate` or stage YAML + TTS params | `Omni(...).generate` + `OmniDiffusionSamplingParams` |
+| **Online client** | `openai_client.send_omni_request` (chat completions, modalities) | `openai_client.send_audio_speech_request` (`/v1/audio/speech`) | `send_diffusion_request` (chat/T2I), `send_video_diffusion_request` (`/v1/videos`, X2V), or `send_images_generations_http_request` / `send_images_edits_http_request` (DALL-E routes) |
+| **Typical assertions** | Stage outputs, text/audio keywords via `OmniRunnerHandler` / response handler | WAV bytes, stream chunks, speech endpoint contract | Image/video dimensions, `final_output_type`, `assert_diffusion_response` |
+| **Stage / deploy YAML** | Per-model omni stage configs (`ci/qwen3_omni_moe.yaml`, …) | `qwen3_tts.yaml`, `voxcpm2.yaml`, … | Often default serve; parallel/offload YAML for heavy DiT |
+| **Nightly group** (`test-nightly.yml`) | **Omni Model Test** — `-m "full_model and omni"` | **TTS Model Test** — `-m "full_model and tts"` | **Diffusion X2I(&A&T)** *or* **Diffusion X2V** (see Step 3 table; same `diffusion` marker, different YAML group / file shard) |
+| **L4 pressure** | Expansion per modality/model as needed | Expansion + accuracy/perf in TTS group | X2I: merge feature combos per [#1832](https://github.com/vllm-project/vllm-omni/issues/1832); X2V: separate nightly group |
+
+Do **not** mix fixtures across types (e.g. do not use `omni_runner` layout for a pure diffusion or TTS model without mirroring an in-tree test in that family).
+
+**Diffusion only — X2I(&A&T) vs X2V (nightly routing, not extra markers):**
+
+- **X2I(&A&T)**: image / audio / text generation — Qwen-Image*, FLUX, SD3, Z-Image, BAGEL (expansion), LongCat, audio diffusion, etc.
+- **X2V**: video generation only — Wan2.2, HunyuanVideo 1.5, Wan VACE, LTX video similarity paths.
+
+When adding a new `test_*_expansion.py`, place it in the matching **nightly group** step (explicit file list in `test-nightly.yml`), not only by marker expression.
+
+**2. Use the narrowest deterministic skeleton for the scenario**
+
+*L1 unit / logic (CPU-first):*
+
+**Mocking rule (L1 only):** use **pytest** integration — **`mocker`** (`pytest-mock`) or **`monkeypatch`** (built-in). **Do not** import or call **`unittest.mock`** (`patch`, `MagicMock`, `@patch`, `with patch(...)`, etc.) in L1 tests; patches must auto-revert with the test lifecycle.
+
+```python
+import pytest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_<scenario_name>(mocker):
+    # Prefer mocker.patch / mocker.spy / mocker.Mock — not unittest.mock.patch
+    fake_fn = mocker.patch("vllm_omni.some.module.expensive_call", return_value=...)
+    # Act
+    # Assert
+    fake_fn.assert_called_once()
+
+
+def test_<env_or_attr>(monkeypatch):
+    # For simple env / attribute substitution without a Mock object
+    monkeypatch.setenv("SOME_FLAG", "1")
+    monkeypatch.setattr("vllm_omni.some.module.CONST", 42)
+    ...
+```
+
+**Avoid in L1:**
+
+```python
+# BAD — do not use in L1 unit tests
+from unittest.mock import patch, MagicMock
+
+@patch("vllm_omni.some.module.fn")
+def test_bad(mock_fn): ...
+
+def test_also_bad():
+    with patch("...") as m: ...
+```
+
+See **L1 unit test constraints (mocking)** below for the full do/don't list.
+
+*Offline multimodal e2e — **Omni** (representative):*
+
+```python
+@pytest.mark.core_model
+@pytest.mark.omni
+@hardware_test(...)
+@pytest.mark.parametrize("omni_runner", test_params, indirect=True)
+def test_<scenario>(omni_runner, omni_runner_handler) -> None:
+    request_config = {"prompts": ..., "modalities": [...]}  # optional: images, videos, audios
+    omni_runner_handler.send_omni_request(request_config)
+```
+
+*Offline generative e2e — **Diffusion** (representative):*
+
+```python
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@hardware_test(...)
+@pytest.mark.parametrize("omni_runner", test_params, indirect=True)
+def test_text_to_image_001(omni_runner_handler) -> None:
+    omni_runner_handler.send_diffusion_request({"prompt": "...", "extra_body": {"num_inference_steps": 4, ...}})
+```
+
+*Offline TTS e2e — **Qwen3-TTS** (two-stage; representative):*
+
+```python
+@pytest.mark.advanced_model
+@pytest.mark.tts
+@hardware_test(...)
+@pytest.mark.parametrize("omni_runner", tts_server_params, indirect=True)
+def test_text_to_audio_001(omni_runner, omni_runner_handler) -> None:
+    omni_runner_handler.send_audio_speech_request({
+        "input": "...",
+        "task_type": "Base",
+        "ref_audio": REF_AUDIO_URL,
+        "ref_text": REF_TEXT,
+    })
+```
+
+*Offline TTS e2e — **single-stage** (Coqui XTTS, MOSS-TTS-Nano; representative):*
+
+```python
+@pytest.mark.advanced_model
+@pytest.mark.tts
+@hardware_test(...)
+@pytest.mark.parametrize("omni_runner", tts_server_params, indirect=True)
+def test_voice_clone_001(omni_runner_handler) -> None:
+    omni_runner_handler.send_single_stage_tts_request({
+        "input": "...",
+        "language": "en",
+        "prompt_audio_path": REF_AUDIO_PATH,
+        "response_format": "wav",
+        "run_level": "advanced_model",
+    })
+```
+
+*Online serving e2e — **Omni** (representative):*
+
+```python
+@pytest.mark.core_model
+@pytest.mark.omni
+@hardware_test(...)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_text_to_text_001(omni_server, openai_client) -> None:
+    request_config = {"model": omni_server.model, "messages": ..., "modalities": ["text"]}
+    openai_client.send_omni_request(request_config)
+```
+
+*Online serving e2e — **TTS** (representative):*
+
+```python
+@pytest.mark.core_model
+@pytest.mark.advanced_model
+@pytest.mark.tts
+@hardware_test(...)
+@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
+def test_text_to_audio_001(omni_server, openai_client) -> None:
+    request_config = {"model": omni_server.model, "input": "...", "response_format": "wav", ...}
+    openai_client.send_audio_speech_request(request_config)
+```
+
+*Online serving e2e — **Diffusion** X2I (representative):*
+
+```python
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@pytest.mark.parametrize("omni_server", _get_default_case(MODEL), indirect=True)
+def test_text_to_image_001(omni_server, openai_client) -> None:
+    openai_client.send_diffusion_request({...})  # chat completions + extra_body
+```
+
+*Online serving e2e — **Diffusion** X2V (representative):*
+
+```python
+@pytest.mark.core_model
+@pytest.mark.diffusion
+def test_text_to_video_001(omni_server, openai_client) -> None:
+    openai_client.send_video_diffusion_request({"model": ..., "form_data": {...}})  # /v1/videos
+```
+
+*Documentation example tests:* follow the **Preferred Test Strategy** in [l4_doc_example_tests.inc.md](../../../docs/contributing/ci/test_examples/l4_doc_example_tests.inc.md): dynamic extraction for offline READMEs; explicit copied client code for online pages until extraction is justified; use the documented **naming**, **output directory** (page folder + case id), and **skipping** rules (e.g. Gradio-only scripts).
+
+*Performance tests:* add or extend entries under `tests/perf/` (and JSON configs where the project uses them), with **explicit baselines** and the same marker/run-level pairing as the CI step that will execute them.
+
+**3. Cross-cutting rules**
+
+- Reuse existing fixtures for the chosen scenario; do not mix “online client” assumptions into offline `OmniRunner` tests without a clear reason.
+- Avoid external network dependency in assertions unless the scenario is explicitly “online serving” or doc examples that require a model hub (then align with CI secrets/cache).
+- Keep **one test function = one intent** (one modality combo, one endpoint contract, or one acceleration combo).
+- **E2E test function layout**: **one case → one `test_<scenario>` function** with a name that states what is validated (endpoint, size/`n`, server flag, or route). **Do not** merge multiple cases into a single test that branches on `request.node.callspec.id`, `param.id`, or `if case_id == ...`. Use `@pytest.mark.parametrize("omni_server", [...], indirect=True)` per function (usually **one** `OmniServerParams` per test). A loop inside one test is OK only when it serves **that function’s single intent** (e.g. three standard sizes in `test_*_sizes_256_512_1024`).
+- **Runtime fixture scope** (`tests/helpers/fixtures/runtime.py`): default **`omni_server` / `omni_runner`** (module) + matching client/handler; use **`omni_server_function` / `omni_runner_function`** only when each `test_*` must start a fresh instance (see below).
+- **L1 mocks**: never `unittest.mock`; use `mocker` or `monkeypatch` only (see below).
+- **API calls (L2+ e2e, online and offline)**: **reuse** `send_*_request` in `tests/helpers/runtime.py` when it exists; **otherwise add** the helper in `runtime.py` first, then call it from the test. General `assert_*` inside `send_*_request`; special `assert_*` only in the test. See **Runtime send helpers** below — do not call `omni.generate`, raw HTTP, or SDK clients from `test_*.py`.
+- **Response assertions**: reusable checks on API bodies / decoded media belong in **`tests/helpers/assertions.py`** — not as private `_assert_*` helpers inside `test_*.py` (see below).
+- **Model-specific payloads stay in test modules** — per-model `MODEL`, deploy path, `REF_AUDIO_URL`, `get_prompt()`, `_build_request()`, and inline `request_config` dicts live in `test_{slug}.py` / `test_{slug}_expansion.py` (and offline/L1 siblings). **Do not** create `tests/helpers/{slug}.py` to deduplicate them; a little copy across files is preferred (see `test_glm_tts.py`, `test_cosyvoice3_tts_expansion.py`). `tests/helpers/` is repo-wide harness only (`mark`, `media`, `runtime`, `stage_config`, `assertions`, `fixtures/`).
+
+#### L1 unit test constraints (mocking)
+
+L1 tests (`core_model and cpu`, under `tests/diffusion/`, `tests/engine/`, `tests/model_executor/`, etc.) must follow the repo’s **pytest-mock** convention (`pytest-mock>=3.10.0` in `pyproject.toml` `[project.optional-dependencies] dev`).
+
+| Do | Don't |
+|----|-------|
+| `def test_foo(mocker):` + `mocker.patch(...)`, `mocker.spy(...)`, `mocker.Mock()`, `mocker.MagicMock()`, `mocker.AsyncMock()` | `from unittest.mock import patch, MagicMock, Mock, AsyncMock` |
+| `def test_bar(monkeypatch):` + `monkeypatch.setattr` / `setenv` / `delenv` / `setitem` | `@patch(...)` decorator |
+| Let `mocker` auto-stop patches after the test | `with patch(...):` / `patch.object(...)` context managers |
+| Mirror neighboring L1 tests in the same directory | `unittest.mock.create_autospec` unless an existing file already documents an exception |
+
+**Rationale:** `mocker` ties patch lifecycle to pytest fixtures (no leaked patches across tests). `unittest.mock` decorators/context managers are easy to compose incorrectly with parametrized or async tests and are inconsistent with in-tree L1 style.
+
+**Minimal patterns:**
+
+```python
+def test_returns_cached_config(mocker):
+    loader = mocker.patch(
+        "vllm_omni.foo.load_yaml",
+        return_value={"stages": []},
+    )
+    result = get_deploy_config("ci/foo.yaml")
+    assert result["stages"] == []
+    loader.assert_called_once_with("ci/foo.yaml")
+
+
+def test_skips_when_env_unset(monkeypatch):
+    monkeypatch.delenv("VLLM_OMNI_FEATURE", raising=False)
+    assert should_enable_feature() is False
+```
+
+E2E levels (L2+) generally **avoid mocks**; if a rare L2 stub is unavoidable, still prefer `mocker` over `unittest.mock` for consistency.
+
+#### Runtime fixtures — scope (`tests/helpers/fixtures/runtime.py`)
+
+vLLM-Omni e2e tests start a real **OmniServer** (online) or **OmniRunner** (offline). Pick scope by **how often the process must be recreated**, not by test level alone.
+
+| Scope | Online fixtures | Offline fixtures | When to use |
+|-------|-----------------|------------------|-------------|
+| **Module** (default) | `omni_server` → `openai_client` | `omni_runner` → `omni_runner_handler` | **Default for L2/L3/L4 expansion** — amortize model/server init across `test_*` in the same module. Same `OmniServerParams` / runner config can be reused by multiple tests. |
+| **Function** | `omni_server_function` → `openai_client_function` | `omni_runner_function` → `omni_runner_handler_function` | **Only when required** — each `test_*` must get a **clean** server/runner (no shared engine/GPU state). Typical: `tests/dfx/reliability/`, sleep/wakeup, crash/restart, tests that mutate global server state. |
+
+**Rules:**
+
+1. **Default to module scope** (`omni_server` / `omni_runner`) unless the scenario **explicitly** needs a fresh instance per test function.
+2. **Indirect parametrize name must match the fixture name**: `@pytest.mark.parametrize("omni_server", ...)` with `omni_server` + `openai_client`; `@pytest.mark.parametrize("omni_server_function", ...)` with `omni_server_function` + `openai_client_function`. Do not mix module fixture with function client (or vice versa).
+3. Different `OmniServerParams` per `test_*` (e.g. default vs `--enable-cpu-offload`) is still OK with **module** `omni_server` — pytest parametrizes per test node; only switch to `_function` when isolation between tests matters, not merely because `server_args` differ.
+4. One `test_*` with **many** server configs (expansion matrix) → single function + `@pytest.mark.parametrize("omni_server", [...], indirect=True)` + module `omni_server` (see in-tree `test_qwen_image_expansion.py`).
+
+```python
+# Default — module-scoped server (L4 expansion)
+@pytest.mark.parametrize("omni_server", [pytest.param(OmniServerParams(model=MODEL), marks=H100)], indirect=True)
+def test_foo_images_generations_default_1024(omni_server, openai_client) -> None:
+    openai_client.send_images_generations_request({...})
+
+# Function-scoped — reliability / per-test clean state only
+@pytest.mark.parametrize(
+    "omni_server_function",
+    [pytest.param(OmniServerParams(model=MODEL), marks=H100)],
+    indirect=True,
+)
+def test_foo_sleep_wakeup_cycle(omni_server_function, openai_client_function) -> None:
+    openai_client_function.send_omni_sleep_http_request({...})
+    openai_client_function.send_omni_wakeup_http_request({...})
+```
+
+#### Runtime send helpers — online **and** offline (`tests/helpers/runtime.py`)
+
+**L2+ e2e (online serving and offline inference) must call APIs through `tests/helpers/runtime.py`.** Fixtures live in `tests/helpers/fixtures/runtime.py`; **send/assert implementation** lives in `tests/helpers/runtime.py` + `tests/helpers/assertions.py`.
+
+| Principle | Action |
+|-----------|--------|
+| **Reuse first** | Grep `runtime.py` for an existing **`send_*_request`** on `OpenAIClientHandler` (online) or `OmniRunnerHandler` (offline) that matches the endpoint / pipeline shape. |
+| **Extend when close** | If an existing `send_*_request` almost fits (e.g. missing one optional field), extend it in `runtime.py` — do not fork logic in the test file. |
+| **Add when missing** | No suitable helper → add **`send_<feature>_request`** (high-level: call + general `assert_*`) or **`send_<route>_<verb>_http_request`** (low-level HTTP for negative/dfx) in **`runtime.py` first**, then call it from tests. |
+| **Test module owns payload only** | In `test_*.py`: `MODEL`, deploy path, vendored media, `get_prompt()`, and inline **`request_config` dicts** only. **No** `omni.generate(...)`, raw `requests.post`, OpenAI SDK calls, or `_collect_audio()` / `_process_output()` in e2e tests. |
+
+**Online** (`openai_client` from `omni_server`): `OpenAIClientHandler.send_*_request`.
+
+**Offline** (`omni_runner_handler` from `omni_runner`): `OmniRunnerHandler.send_*_request`.
+
+| Do | Don't |
+|----|-------|
+| Online: `openai_client.send_omni_request`, `send_diffusion_request`, `send_audio_speech_request`, `send_video_diffusion_request`, `send_images_generations_request`, … | `requests.post(f"{base_url}/v1/...", json=…)` or `client.chat.completions.create(...)` inside a test |
+| Offline: `omni_runner_handler.send_omni_request`, `send_diffusion_request`, `send_audio_speech_request`, `send_single_stage_tts_request`, `send_single_stage_tts_batch_request`, … | `omni_runner.omni.generate(...)` + hand-rolled tensor/WAV extraction in `test_*.py` |
+| Add missing `send_*` to **`runtime.py`** first; bundle general `assert_*` inside the send helper | A one-off `def _post_*` / `def _collect_audio` at the bottom of a test module |
+| Mirror naming/style of neighboring `send_*` (docstring, `request_config` dict, optional `run_level`, `err_code` / `err_message` for negative cases) | Different parameter shapes per test file for the same endpoint |
+
+**Workflow when generating online/offline e2e tests:**
+
+1. Decide whether the needed check is **general** (every success call of this `send_*`) or **special** (one case / one parameter combo only). See **General vs special assert placement** below.
+2. **Search `runtime.py`** for a matching **`send_*_request`**; if missing, add it (with general `assert_*` inside) before writing the test body.
+3. In the test module: build `request_config` → call **`send_*_request` only** for the general contract; call **`assert_*` in the test file only** for special, case-specific checks (import from `assertions.py`).
+4. Reserve **low-level** `send_*_http_request` for negative/dfx tests (`err_code` / `err_message`) — not for ordinary L2+ success-path e2e.
+5. When wiring Buildkite `source_file_dependencies`, include **`tests/helpers/runtime.py`** and/or **`tests/helpers/assertions.py`** when new helpers were added.
+
+**Exceptions (document in the test docstring why):** models whose offline prompt path cannot go through existing handlers yet (e.g. `test_higgs_audio_v2.py`, `test_voxtral_tts.py` with custom tokenizer compose) may call `omni.generate` directly until a `send_*_request` is added to `runtime.py` — treat as **debt**, not the default for new TTS/diffusion/omni e2e.
+
+**General vs special assert placement:**
+
+| Kind | Definition | Where it lives | Test file calls |
+|------|------------|----------------|-----------------|
+| **General** | Default success contract for this endpoint/client method — e.g. HTTP 200, `data[]` shape, `n` count, `size` dimensions, decodable image/audio, bundled omni/diffusion fields | Implement in `assertions.py`, invoke **inside** `send_*_request` in `runtime.py` (after low-level send / SDK call, when `err_code` is not set) | **`send_*_request` only** — do **not** repeat the same `assert_*` |
+| **Special** | Extra check tied to **this test case only** — e.g. seed byte-identical replay, accuracy/CLIP threshold, perf ceiling, model-specific optional field | Add or reuse `assert_*` in `assertions.py` (never inline in test) | `send_*_request` **then** `assert_<special>(...)` **once** for that case |
+
+Changing `request_config` fields (`size`, `n`, `seed`, server flags) is **not** special validation — the general `assert_*` should read those from `request_config` inside `send_*`.
+
+**`send_*` ↔ `assert_*` pairing (online `OpenAIClientHandler`):**
+
+| High-level `send_*` (prefer in L2+ e2e) | Assert already invoked inside `runtime.py` | Low-level HTTP-only sibling |
+|----------------------------------------|--------------------------------------------|-----------------------------|
+| `send_omni_request` | `assert_omni_response` | `send_chat_completions_http_request` → `assert_http_error` only |
+| `send_diffusion_request` | `assert_diffusion_response` | — |
+| `send_audio_speech_request` | `assert_audio_speech_response` | `send_audio_speech_http_request` → `assert_http_error` only |
+| `send_video_diffusion_request` | `assert_diffusion_response` | `send_videos_*_http_request` → `assert_http_error` only |
+| `send_images_generations_request` *(add when needed)* | **`assert_images_generations_response`** *(general — inside send)* | `send_images_generations_http_request` → `assert_http_error` only |
+| `send_images_edits_request` *(add when needed)* | **`assert_images_edits_response`** *(general — inside send)* | `send_images_edits_http_request` → `assert_http_error` only |
+
+**Rule:** Tests call high-level `send_*_request` for the general contract. **Never** call the same bundled `assert_*` again. Call an extra `assert_*` in the test **only** for special, case-specific validation.
+
+**Common `OpenAIClientHandler` entry points** (non-exhaustive — grep `runtime.py` before adding):
+
+| Area | High-level (SDK + assert) | Low-level HTTP (`*_http_request`) |
+|------|---------------------------|-----------------------------------|
+| Omni chat | `send_omni_request` | `send_chat_completions_http_request` |
+| Diffusion T2I (chat route) | `send_diffusion_request` | — |
+| Diffusion X2V | `send_video_diffusion_request` | `send_videos_create_http_request`, `send_video_content_http_request`, … |
+| DALL-E T2I / edit | `send_images_generations_request`, `send_images_edits_request` | `send_images_generations_http_request`, `send_images_edits_http_request` |
+| TTS (online) | `send_audio_speech_request` | `send_audio_speech_http_request`, `send_audio_generate_http_request`, … |
+| Ops / meta | — | `send_health_http_request`, `send_models_http_request`, `send_omni_sleep_http_request`, … |
+
+**Common `OmniRunnerHandler` entry points** (offline — grep `runtime.py` before adding):
+
+| Area | High-level (`send_*_request` + assert) | Notes |
+|------|----------------------------------------|-------|
+| Omni multimodal | `send_omni_request` | `generate_multimodal` path |
+| Diffusion offline | `send_diffusion_request` | chat-route / `OmniTextPrompt` offline |
+| Qwen-style TTS | `send_audio_speech_request` | two-stage, `generate_multimodal` + `mm_processor_kwargs` |
+| Single-stage TTS (Coqui XTTS, MOSS-TTS-Nano, …) | `send_single_stage_tts_request`, `send_single_stage_tts_batch_request` | `prompt` + `additional_information` + `omni.generate` — **do not** duplicate in tests |
+| New model family | **Add** `send_<family>_request` here first | Then call from `tests/e2e/offline_inference/test_*.py` |
+
+L1 tests under `tests/entrypoints/` may use **FastAPI `TestClient`** or direct handler calls with mocks; they do **not** need `OpenAIClientHandler` / `OmniRunnerHandler`. **L2+ online and offline e2e** must use `runtime.py` `send_*` helpers.
+
+#### Invalid parameter validation (`tests/dfx/reliability/invalid_param_test/`)
+
+When the user asks for **invalid parameter validation**, **invalid request bodies**, **HTTP 4xx contract tests**, or any case that sends **malformed / out-of-range / mismatched** API payloads against a **live** server:
+
+1. **Do not** place these in `tests/e2e/online_serving/` or `*_expansion.py`. If already drafted there, **move** the `test_*` into the correct `invalid_param_test` script and delete the duplicate from e2e.
+2. **Pick the script by HTTP route** (extend an existing file; add a new `test_invalid_<area>.py` only when no in-tree script covers that route family):
+
+| API / area | Script |
+|------------|--------|
+| Omni chat completions, WebSocket video/realtime paths | `test_invalid_omni_chat.py` |
+| `POST /v1/audio/speech`, stream, batch, voices | `test_invalid_audio_speech.py` |
+| Audio diffusion endpoints | `test_invalid_audio_diffusion.py` |
+| `POST /v1/images/generations` | `test_invalid_image_generation.py` |
+| `POST /v1/images/edits` | `test_invalid_image_editing.py` |
+| `POST/GET/DELETE /v1/videos*` | `test_invalid_video_generation.py` |
+| Sleep / wakeup / server control | `test_invalid_server_control.py` |
+
+3. **Match in-tree style** in the chosen script:
+
+| Element | Convention |
+|---------|------------|
+| **Module markers** | `pytestmark = [pytest.mark.slow, pytest.mark.<omni or tts or diffusion>]` |
+| **Server fixture** | `_PARAMS` / `_QWEN3_TTS_SPEECH`-style list of `pytest.param(OmniServerParams(...), id="...", marks=hardware_marks(...))`; `@pytest.mark.parametrize("omni_server", _PARAMS, indirect=True)` |
+| **Hardware** | `hardware_marks(res={"cuda": "H100"})` for heavy diffusion/omni/video; `hardware_marks(res={"cuda": "L4"})` for smaller TTS models (must match weekly `-m "slow and L4"` step) |
+| **HTTP client** | **Low-level** `openai_client.send_*_http_request({..., "err_code": 400, "err_message": (...)} )` — **not** `send_*_request` (success path) |
+| **Case shape** | Prefer **one `test_*` per route family** + `@pytest.mark.parametrize("body_spec, err_message", [...])` with stable `id=` per case; or dedicated `test_<route>_malformed_json` when not parametrized |
+| **Body helpers** | `_minimal_<endpoint>_json(omni_server)` / `_minimal_*_form_data()`; merge overrides with `body.update(body_spec)` |
+| **Known gaps** | `pytest.mark.skip(reason="…#3649")` as `_SKIP_ISSUE_3649` when server validation is not yet strict (mirror neighboring cases) |
+| **Sections** | Route banner comments (`# ─── POST /v1/images/generations ───`) like existing files |
+| **Shared fixtures** | Reuse `tests/dfx/reliability/invalid_param_test/conftest.py` (`tiny_png_bytes`, env defaults) |
+
+4. **Adding a new model** to invalid-param coverage: append a `pytest.param(OmniServerParams(model="...", stage_config_path=..., server_args=...), ...)` entry to the script’s `_PARAMS` list — do **not** create `test_invalid_<model>.py` unless the route is new.
+
+5. **CI — `.buildkite/test-weekly.yml` only** (not `test-ready.yml` / `test-merge.yml` / `test-nightly.yml`):
+
+| Weekly step | Command | When your cases run |
+|-------------|---------|---------------------|
+| **Invalid parameters Test · H100** | `pytest -s -v tests/dfx/reliability/invalid_param_test/ -m "slow and H100"` | Diffusion / omni / video invalid-param tests with `H100` hardware mark |
+| **Invalid parameters Test · L4** | `pytest -s -v tests/dfx/reliability/invalid_param_test/ -m "slow and L4"` | TTS / lighter models marked `L4` |
+
+- **Trigger:** `build.env("WEEKLY") == "1"` or PR label `weekly-test`.
+- **Default:** extending an existing `invalid_param_test` script needs **no YAML edit** — the weekly steps already sweep the whole directory.
+- **Edit YAML only** when adding a **new hardware queue**, a **new top-level script** that must run in isolation, or a **model-specific weekly shard** (mirror neighboring reliability steps).
+- Weekly steps **do not** use `source_file_dependencies`.
+
+**Example — append to `test_invalid_image_generation.py`:**
+
+```python
+@pytest.mark.parametrize(
+    "body_spec, err_message",
+    [
+        pytest.param({"seed": -1}, ("seed", "greater_than_equal", "0"), id="seed_negative"),
+    ],
+)
+@pytest.mark.parametrize("omni_server", _PARAMS, indirect=True)
+def test_images_generations_invalid_requests(
+    omni_server: OmniServer,
+    openai_client: OpenAIClientHandler,
+    body_spec: dict[str, object],
+    err_message: str | tuple[str, ...],
+) -> None:
+    body = _minimal_images_gen_json(omni_server)
+    body.update(body_spec)
+    openai_client.send_images_generations_http_request(
+        {"json": body, "timeout": 300, "err_code": 400, "err_message": err_message}
+    )
+```
+
+See [references/test-routing.md](references/test-routing.md) **Invalid parameter / weekly CI**.
+
+#### Assertion helpers (`tests/helpers/assertions.py`)
+
+**Do not** add module-local helpers such as `_assert_images_generations_payload` or `_send_and_assert_*` in e2e test files. Response/media validation belongs in **`tests/helpers/assertions.py`**, grouped by **category**:
+
+| Category | Existing anchors | When to extend vs add |
+|----------|------------------|------------------------|
+| **Image** (chat diffusion + DALL-E JSON) | `assert_image_diffusion_response`, `assert_image_valid` | DALL-E `/v1/images/generations` JSON → add **`assert_images_generations_response`** beside image helpers (reuses `assert_image_valid`); do not fork decode logic into tests |
+| **Video** | `assert_video_diffusion_response`, `assert_video_valid` | Extend these for new video contracts |
+| **Audio** | `assert_audio_diffusion_response`, `assert_audio_speech_response`, `assert_audio_valid` | Extend for new TTS/audio endpoints |
+| **Omni multimodal** | `assert_omni_response` | Extend for new modality combos |
+| **HTTP errors** | `assert_http_error`, `assert_err_message_in_text` | Used inside low-level `send_*_http_request` and negative tests |
+
+| Do | Don't |
+|----|-------|
+| Put **general** `assert_*` inside `send_*_request` in `runtime.py`; tests call `send_*_request` only | Call `assert_diffusion_response` after `send_diffusion_request` (general assert already bundled) |
+| Put **special** `assert_*` in `assertions.py` and call it **in the test** after `send_*_request` when that case needs extra checks | Put general decode/count/size logic in the test because “this case uses `n=4`” (that belongs in general assert reading `request_config`) |
+| Extend the matching category function, or add `assert_<endpoint>_response` next to its category | `_send_and_assert_*` or PIL/base64 loops in `test_*.py` |
+| Use low-level `send_*_http_request` in dfx/negative tests (`err_code`) | Use low-level HTTP send + manual general assert in every L2+ expansion file |
+
+**Workflow when generating tests:**
+
+1. Implement or reuse **general** `assert_*` in `assertions.py` (by category).
+2. Wire it into **`send_*_request`** in `runtime.py` if not already bundled.
+3. In the test: `request_config` → `send_*_request`. Add a separate `assert_*` import/call **only** for special case validation.
+4. Update `source_file_dependencies` when `runtime.py` or `assertions.py` changes.
+
+**Example — DALL-E `/v1/images/generations`:**
+
+```python
+# tests/helpers/assertions.py — general contract (image category)
+def assert_images_generations_response(resp_body: dict, request_config: dict, *, run_level: str | None = None) -> None:
+    ...  # data[], n, size → assert_image_valid
+
+# tests/helpers/runtime.py — general assert INSIDE send
+def send_images_generations_request(self, request_config: dict[str, Any], ...) -> list[HttpResponse]:
+    responses = self.send_images_generations_http_request(request_config, ...)
+    if request_config.get("err_code") is None:
+        body = responses[0].json_body
+        assert isinstance(body, dict)
+        assert_images_generations_response(body, request_config, run_level=self.run_level)
+    return responses
+
+# tests/e2e/.../test_foo_expansion.py — one case per test_* (no case_id branching)
+@pytest.mark.parametrize("omni_server", [pytest.param(OmniServerParams(model=MODEL), marks=SINGLE_L4)], indirect=True)
+def test_foo_images_generations_default_1024(omni_server, openai_client) -> None:
+    openai_client.send_images_generations_request({"json": body, "timeout": 300})
+
+@pytest.mark.parametrize("omni_server", [pytest.param(OmniServerParams(model=MODEL), marks=SINGLE_L4)], indirect=True)
+def test_foo_images_generations_sizes_256_512_1024(omni_server, openai_client) -> None:
+    for size in ("256x256", "512x512", "1024x1024"):
+        openai_client.send_images_generations_request({"json": {**body, "size": size}, "timeout": 300})
+
+# chat diffusion — separate test; send_diffusion_request bundles assert_diffusion_response
+@pytest.mark.parametrize("omni_server", [pytest.param(OmniServerParams(model=MODEL), marks=SINGLE_L4)], indirect=True)
+def test_foo_chat_completions_t2i_fallback(omni_server, openai_client) -> None:
+    openai_client.send_diffusion_request({...})
+```
+
+L1 tests under `tests/entrypoints/` may keep **minimal** asserts next to the handler under test when validating a single branch; still prefer `assertions.py` when the same JSON/media check appears in more than one file or level (e.g. L1 protocol test + L4 e2e).
+
+#### Omni Test Writing Guidance (L1-L4 Layering)
+
+When the goal is a general Omni/multimodal test case, prioritize mapping the test to the correct purpose, directory, and resource assumptions aligned with the layers (see [CI_5levels.md](../../../docs/contributing/ci/CI_5levels.md)):
+
+- **L1**: Unit/logic validation on CPU (`core_model and cpu`). Cover input validation, branches, and exception paths (`tests/<component>/test_*.py`). **Mock with `mocker` / `monkeypatch` only — not `unittest.mock`.**
+- **L2**: Basic e2e (online/offline basic scenarios). Prefer dummy/lightweight models to validate the end-to-end request-to-output-structure/streaming chain (typically `tests/e2e/online_serving/` and `tests/e2e/offline_inference/`).
+- **L3/L4**: Important integration, performance, and accuracy validation. L4 emphasizes “full functional scenarios + performance/stress + runnable doc examples” (typically `*_expansion.py` plus related expansion cases).
+
+Also keep markers consistent with the run level: use `core_model` for L1/L2 and `advanced_model` for L3/L4, and pair with `--run-level` to select the intended CI strategy.
+
+#### Diffusion Test Writing Guidance (L4 Coverage Combinations)
+
+When the task involves diffusion models/features, organize L4 test cases following [`#1832`](https://github.com/vllm-project/vllm-omni/issues/1832): combine multiple diffusion features into as few test cases as possible to fit limited CI GPU resources.
+
+Implementation strategy:
+
+1. If “full L4 is too heavy”: first provide a **reduced local validation case in L1/L2** (so the key assertion and the contract fix point are covered deterministically).
+2. Then provide **CI/nightly-ready L3/L4 high-marked cases** (e.g. `advanced_model`) to broaden coverage under resource constraints.
+
+#### L4 nightly sub-pillars — Function vs Accuracy vs Perf
+
+**L4 is not only `test_*_expansion.py`.** In `test-nightly.yml`, each model-type group typically runs **separate jobs**:
+
+| Sub-pillar | Nightly step label pattern | What you add | Typical paths |
+|------------|----------------------------|--------------|---------------|
+| **Function** | `· Function Test with …` | E2e expansion / feature matrix | `tests/e2e/online_serving/test_<model>_expansion.py` (or offline) |
+| **Accuracy** | `· Accuracy Test` | Quality / similarity vs baseline | `tests/e2e/accuracy/test_<model>*.py` |
+| **Perf** | `· Perf Test · <Model>` | Throughput / latency / memory benchmarks | `tests/dfx/perf/tests/test_<model>_vllm_omni.json` + runner script |
+| **Doc** (optional) | `· Doc Test` | Runnable doc examples | `tests/examples/*/test_text_to_image.py`, … |
+
+**Default when the user asks for “L4 functional cases”:** deliver **Function** pillar only (`*_expansion.py` + Function Test shard in `test-nightly.yml`). **Do not** silently add Perf or Accuracy unless the user also asks for **performance** / **benchmark** / **accuracy** / **full L4** / **full L4 coverage**.
+
+**When the user explicitly asks for L4 perf (or full L4 including perf):**
+
+1. **Do not** put throughput/latency assertions inside `test_*_expansion.py` — perf uses the **dfx benchmark harness**, not `omni_server` e2e fixtures.
+2. Add or extend a **JSON config** under `tests/dfx/perf/tests/` (mirror in-tree names: `test_qwen_image_vllm_omni.json`, `test_bagel_vllm_omni.json`).
+3. Each JSON entry: `test_name`, `server_params` (`model`, `serve_args`), `benchmark_params[]` with explicit **`baseline`** (`throughput_qps`, `latency_mean`, `peak_memory_mb_mean`, …).
+4. Run locally via the matching script (model type → runner):
+
+| Model type | Runner script | Example config |
+|------------|---------------|----------------|
+| **Diffusion X2I/X2V** | `tests/dfx/perf/scripts/run_diffusion_benchmark.py` | `tests/dfx/perf/tests/test_qwen_image_vllm_omni.json` |
+| **TTS** | `tests/dfx/perf/scripts/run_benchmark.py` | `tests/dfx/perf/tests/test_tts.json` (shared) or `test_{slug}.json` (dedicated; use when the model must not join the shared nightly matrix before integration — e.g. `test_voxcpm2.json`, `test_coqui_tts.json`) |
+| **Omni** | (see existing Omni Perf steps — `run_benchmark.py` / dedicated omni configs in-tree) | `tests/dfx/perf/tests/test_omni*.json` |
+
+5. Wire **`test-nightly.yml` · Perf Test** step for that model (separate from Function Test): export `DIFFUSION_BENCHMARK_DIR` / `BENCHMARK_DIR`, run pytest on the script with `--test-config-file`, **upload artifacts** (`buildkite-agent artifact upload`), often **multi-GPU H100** for diffusion perf.
+6. **One benchmark scenario → one `test_name` block** in JSON (same spirit as one `test_*` per function case). Combine server `serve_args` + workload in one entry; use multiple `benchmark_params` rows for size/step sweeps under the same server config.
+
+**Perf local command (diffusion example):**
+
+```bash
+cd tests
+export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
+export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
+pytest -s -v dfx/perf/scripts/run_diffusion_benchmark.py \
+  --test-config-file dfx/perf/tests/test_<model>_vllm_omni.json
+```
+
+**Clarify in the test plan** which L4 pillars you deliver: `Function only` | `Function + Perf` | `Function + Accuracy + Perf`.
+
+See [references/test-routing.md](references/test-routing.md) **L4 nightly pillars** for Buildkite step patterns.
+
+### Step 5: Wire Buildkite (when CI must run the new test)
+
+If the test is not already collected by an existing pipeline command (for example, L1 tests marked `core_model and cpu` are already covered by the **Simple Unit Test** step in ready/merge), update the appropriate pipeline under `.buildkite/`:
+
+| Test level | Edit this file | Typical trigger / intent |
+|------------|----------------|---------------------------|
+| **L1** and **L2** | [`.buildkite/test-ready.yml`](../../../.buildkite/test-ready.yml) | PR **ready** label; L1 CPU + L2 GPU/basic e2e (steps labeled **Omni ·**, **TTS ·**, **Diffusion ·**) |
+| **L3** | [`.buildkite/test-merge.yml`](../../../.buildkite/test-merge.yml) | Post-merge; `advanced_model` integration per model type |
+| **L4** | [`.buildkite/test-nightly.yml`](../../../.buildkite/test-nightly.yml) | Nightly; grouped by model type (see below) |
+| **Invalid param / reliability (weekly)** | [`.buildkite/test-weekly.yml`](../../../.buildkite/test-weekly.yml) | Weekly; `tests/dfx/reliability/invalid_param_test/` — **not** L1–L4 e2e pipelines |
+
+**Level-specific Buildkite delivery (follow the user’s requested level):**
+
+| User asks for | Document / edit **only** | Do **not** add by default |
+|---------------|---------------------------|---------------------------|
+| **L4** (`*_expansion.py`, `full_model`) | `test-nightly.yml` — append file to the matching nightly shard (X2I / X2V / Omni / TTS) **or** note it is already collected by an existing `-m` / `-k` sweep | `test-merge.yml` (L3 / `advanced_model`) or `test-ready.yml` (L2) |
+| **L3** | `test-merge.yml` with `advanced_model` + `source_file_dependencies` | `test-nightly.yml` unless user also wants nightly |
+| **L2** | `test-ready.yml` with `core_model` + `source_file_dependencies` | merge / nightly pipelines |
+| **Invalid param** | `test-weekly.yml` — **Invalid parameters Test** group (`-m "slow and H100"` / `-m "slow and L4"`). Usually **no YAML edit** when appending cases to existing `invalid_param_test` scripts | `test-ready.yml`, `test-merge.yml`, `test-nightly.yml` |
+
+`test-nightly.yml` shards use **explicit pytest file paths** in `commands` (and PR labels like `diffusion-x2iat-test`); they generally **do not** use `source_file_dependencies`. Only suggest merge/ready YAML when the requested level is L3/L2 (or the user explicitly asks for multi-level CI).
+
+**`test-nightly.yml` top-level groups** (each uses `full_model` + type marker in pytest `-m`):
+
+| Group | Type marker | Notes |
+|-------|-------------|--------|
+| **Omni Model Test** | `omni` | Function / doc / accuracy / **perf** for Qwen-Omni family |
+| **TTS Model Test** | `tts` | `-m "full_model and L4 and tts"` (and H100 variants when enabled); **Perf Test** uses `run_benchmark.py` + `test_tts.json` |
+| **Diffusion X2I(&A&T) Model Test** | `diffusion` | Image/audio/text diffusion; explicit file shards for **Function**; **Perf Test · &lt;Model&gt;** uses `run_diffusion_benchmark.py` + JSON config |
+| **Diffusion X2V Model Test** | `diffusion` | Video-only (Wan, HunyuanVideo, …); PR label `diffusion-x2v-test`; separate **Perf** steps when in-tree |
+
+Guidelines when editing YAML:
+
+- **Match markers and `--run-level`** to the test: L1/L2 in `test-ready.yml` use `core_model` + `omni` / `tts` / `diffusion` as appropriate; L3 merge uses `advanced_model`; L4 nightly uses `full_model` with the same type marker.
+- **Diffusion L4**: add x2i/x2a/x2t expansion files to an **X2I(&A&T)** step; add x2v expansion files to **X2V** — never rely on a broad `tests/e2e/` sweep alone (shards exist to save GPU and avoid wrong queue).
+- **`source_file_dependencies` (required for new E2E jobs in `test-ready.yml` and `test-merge.yml`)**: every step under the **`:card_index_dividers: E2E Test`** group **must** declare `source_file_dependencies` so Buildkite only schedules the job when relevant paths change. Omitting it breaks path-based CI skipping for GPU e2e jobs.
+- **Reuse or extend an existing step** when the new test shares the same marker expression, queue, and timeout; otherwise add a new `steps:` entry with the correct `agents.queue`, `timeout_in_minutes`, and docker/kubernetes plugin blocks consistent with neighboring jobs.
+- **Job timeout — use `timeout_in_minutes` on the step, not `timeout` before pytest**: when adding or documenting a new E2E job in `test-ready.yml` / `test-merge.yml`, set `timeout_in_minutes` on the step and run `pytest` directly in `commands`. **Do not** wrap pytest in `timeout 40m bash -c "..."` / `timeout 20m bash -c "..."`; Buildkite already enforces the step deadline via `timeout_in_minutes`.
+- **Never omit `plugins` in skill examples or PR YAML snippets**: H100 E2E steps on `mithril-h100-pool` use the full `kubernetes` block (`resources`, `volumeMounts`, `env` with `HF_TOKEN` secret, `nodeSelector: gpu-h100-sxm`, `devshm` + `hf-cache` volumes). L4/docker steps on `gpu_1_queue` use the full `docker#v5.2.0` block (`shm-size`, `HF_HOME`, volumes). Do not write `# ... plugins unchanged` — paste the complete block from a sibling step.
+- **Platform forks** (e.g. AMD-ready / AMD-merge) live alongside these files; apply the same level → file mapping for those pipelines when the test is platform-specific.
+
+#### `source_file_dependencies` for E2E Test jobs (`test-ready.yml`, `test-merge.yml`)
+
+When adding or splitting a step inside the **E2E Test** group, **always** include `source_file_dependencies` on that step. List paths the job actually depends on — at minimum:
+
+1. **Test module(s)** — every `tests/e2e/...` file the `pytest` command runs (online and offline if both are in one step).
+2. **Model implementation** — `vllm_omni/model_executor/models/<family>/` and/or `vllm_omni/diffusion/models/<family>/` as applicable.
+3. **Stage / input plumbing** (omni & multi-stage TTS) — `vllm_omni/model_executor/stage_input_processors/<family>.py` when present.
+4. **Deploy / stage YAML** — `vllm_omni/deploy/<config>.yaml` or `vllm_omni/deploy/ci/<config>.yaml` referenced by the test’s `OmniServerParams` / `get_deploy_config_path(...)`.
+5. **Test helpers** (when changed in the same PR) — `tests/helpers/runtime.py` for new `send_*` client methods; `tests/helpers/assertions.py` for new shared `assert_*` helpers.
+
+Use **repo-relative paths** (no leading `./`). Prefer **directory** entries (`.../models/bagel/`) when the whole package matters; use **file** entries for single test modules and YAML configs. Mirror a neighboring step for the same model family (e.g. **Diffusion · Bagel Test** in `test-ready.yml`).
+
+**Example** (L2 ready — diffusion online smoke, H100 kubernetes):
+
+```yaml
+- label: "Diffusion · Bagel Test"
+  source_file_dependencies:
+    - tests/e2e/online_serving/test_bagel.py
+    - vllm_omni/model_executor/models/bagel/
+    - vllm_omni/diffusion/models/bagel/
+    - vllm_omni/model_executor/stage_input_processors/bagel.py
+    - vllm_omni/deploy/bagel.yaml
+  timeout_in_minutes: 40
+  commands:
+    - |
+      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
+```
+
+**Example** (L3 merge — TTS online + offline, L4 docker):
+
+```yaml
+- label: "TTS · Qwen3-TTS Base Test"
+  source_file_dependencies:
+    - tests/e2e/online_serving/test_qwen3_tts_base.py
+    - tests/e2e/offline_inference/test_qwen3_tts_base.py
+    - vllm_omni/model_executor/models/qwen3_tts/
+    - vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
+    - vllm_omni/deploy/qwen3_tts.yaml
+  timeout_in_minutes: 20
+  commands:
+    - |
+      export VLLM_LOGGING_LEVEL=DEBUG
+      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 and cuda' --run-level 'advanced_model'
+  agents:
+    queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
+  plugins:
+    - docker#v5.2.0:
+        image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+        always-pull: true
+        propagate-environment: true
+        shm-size: "8gb"
+        environment:
+          - "HF_HOME=/fsx/hf_cache"
+        volumes:
+          - "/fsx/hf_cache:/fsx/hf_cache"
+```
+
+**Example** (L3 merge — diffusion online + offline, H100 kubernetes):
+
+```yaml
+- label: "Diffusion · Qwen Image Test"
+  source_file_dependencies:
+    - tests/e2e/online_serving/test_qwen_image.py
+    - tests/e2e/offline_inference/test_qwen_image.py
+    - vllm_omni/diffusion/models/qwen_image/
+  timeout_in_minutes: 40
+  commands:
+    - |
+      export VLLM_IMAGE_FETCH_TIMEOUT=60
+      pytest -s -v \
+        tests/e2e/online_serving/test_qwen_image.py \
+        tests/e2e/offline_inference/test_qwen_image.py \
+        -m 'advanced_model and cuda' --run-level 'advanced_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
+```
+
+When extending an **existing** E2E step to run an additional test file, **append** that file (and any new model/deploy paths) to `source_file_dependencies` on the same step. When authoring a PR, mention the dependency list in the **Buildkite change** section of your test plan.
+
+#### L3 / L4: validating Buildkite from a feature branch (root `pipeline.yml`)
+
+Production behavior is driven by [`.buildkite/pipeline.yml`](../../../.buildkite/pipeline.yml): it builds the CI image, then uploads **L2** (`test-ready.yml` on non-`main`), **L3** (`test-merge.yml` when `build.branch == "main"`), and **L4** (`test-nightly.yml` when `build.env("NIGHTLY") == "1"`).
+
+When you add or debug **L3** or **L4** tests and need those child pipelines to run **off `main`** or **without** `NIGHTLY=1`, apply **temporary** edits (revert before merge):
+
+1. **Point the upload at the right file (optional shortcut)**  
+   In the **Upload Ready Pipeline** step, the command is `buildkite-agent pipeline upload .buildkite/test-ready.yml`. For a one-off run you may change that path to `.buildkite/test-merge.yml` (L3) or `.buildkite/test-nightly.yml` (L4), and adjust the step’s `if:` so it does not conflict with the other upload steps (avoid double-uploading unless intentional).
+
+2. **L3 — merge pipeline**  
+   The gate `if: build.branch == "main"` lives on the **Upload Merge Pipeline** step in `pipeline.yml` (not inside `test-merge.yml`). **Comment out that `if` line** so `test-merge.yml` is uploaded on your feature branch. Alternatively rely on step (1) instead of the dedicated merge upload step.
+
+3. **L4 — nightly pipeline**  
+   **Comment out** `if: build.env("NIGHTLY") == "1"` on the **Upload Nightly Pipeline** step in `pipeline.yml` so the nightly definition is uploaded without setting the env var.  
+   Child steps in [`test-nightly.yml`](../../../.buildkite/test-nightly.yml) often repeat `if: build.env("NIGHTLY") == "1"`; **comment out those lines on the steps you need to run**, otherwise they will still be skipped after upload.
+
+### Step 6: Run Tests
+
+Pick one command path:
+
+- **Quick local regression (preferred first)**: single file or `file.py::test_name`
+- **CI-like level run**: marker expression + `--run-level`
+
+Full templates live in [references/test-routing.md](references/test-routing.md). **After authoring tests, always emit concrete commands** (see **Output Format**).
+
+**Typical copy-paste examples** (run from repo `tests/` directory; requires matching vLLM/vllm-omni install and hardware):
+
+| Area | Example |
+|------|---------|
+| **Omni offline L2** | `pytest -s -v e2e/offline_inference/test_qwen2_5_omni.py -m "core_model and omni and not cpu" --run-level=core_model` |
+| **Omni online L2** | `pytest -s -v e2e/online_serving/test_qwen3_omni.py -m "core_model and omni" --run-level=core_model` |
+| **TTS online L2** | `pytest -s -v e2e/online_serving/test_qwen3_tts_base.py -m "core_model and tts" --run-level=core_model` |
+| **Diffusion X2I L2 online** | `pytest -s -v e2e/online_serving/test_qwen_image.py -m "core_model and diffusion" --run-level=core_model` |
+| **Diffusion X2V L2 online** | `pytest -s -v e2e/online_serving/test_wan22_t2v.py -m "core_model and diffusion" --run-level=core_model` |
+| **Diffusion X2I L4 nightly** | `pytest -s -v e2e/online_serving/test_qwen_image_expansion.py -m "full_model and diffusion and H100" --run-level=full_model` |
+| **Diffusion X2I L4 perf (nightly)** | `pytest -s -v dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file dfx/perf/tests/test_qwen_image_vllm_omni.json` |
+| **Diffusion X2V L4 nightly** | `pytest -s -v e2e/online_serving/test_wan22_expansion.py -m "full_model and cuda" --run-level=full_model` |
+| **Invalid param (weekly H100)** | `pytest -s -v dfx/reliability/invalid_param_test/ -m "slow and H100"` |
+| **Invalid param (weekly L4)** | `pytest -s -v dfx/reliability/invalid_param_test/ -m "slow and L4"` |
+| **L1 CPU** | `pytest -s -v -m "core_model and cpu"` |
+
+**Prerequisites to mention when relevant**: GPU model (e.g. L4 vs H100), `HF_HOME` / token for hub weights, module-level `skipif` (NPU/XPU-only gaps), and whether CI already collects the path (e.g. `test_*_expansion.py` glob in `test-nightly.yml`).
+
+### Step 7: Validate Result Quality
+
+Before finishing:
+
+- Is the new assertion directly tied to the bug/feature contract?
+- Are API calls made via **`runtime.py` `send_*_request`** (not inline HTTP/SDK in the test file)?
+- Is the **general** `assert_*` bundled inside `send_*_request` (not duplicated in the test)?
+- Is each scenario a **separate `test_*` function** (no `case_id` / `if param.id` branching across cases)?
+- **Fixture scope**: default **`omni_server` / `omni_runner`** (module); **`_function`** variants only when each test must restart the instance?
+- Is the test deterministic (no fragile timing/network coupling)?
+- Is runtime appropriate for the selected level?
+- Are markers and `--run-level` consistent?
+- **Invalid-param cases** in `tests/dfx/reliability/invalid_param_test/` (not e2e), using `send_*_http_request` + `err_code`, with `pytest.mark.slow` and correct `H100` / `L4` hardware mark?
+
+## Output Format
+
+When completing a request, return:
+
+1. **Test plan** (level, markers, file target, **L4 pillar(s)** if applicable: Function / Accuracy / Perf / Doc; **or** **Invalid param** pillar with target `invalid_param_test/test_invalid_<area>.py`; module basename for e2e: `test_{slug}.py` / `test_{slug}_expansion.py`)
+2. **Generated/updated test file(s)** — model-specific constants and request payloads stay **inside** those test modules (**never** `tests/helpers/{slug}.py`). Extend **`tests/helpers/runtime.py`** / **`tests/helpers/assertions.py`** only when new **repo-wide** `send_*` or shared assert helpers are required (do not leave ad-hoc HTTP or `_assert_*` in test modules). **If L4 Perf was requested**, also list **`tests/dfx/perf/tests/*.json`**. **If invalid-param cases were requested**, list the **`invalid_param_test/` script** and parametrized `test_*` / new `id=` rows — **not** e2e paths.
+3. **Buildkite change** — **match the requested level and pillar**:
+   - **L4 Function**: `test-nightly.yml` **Function Test** shard (file list / note existing sweep). **No** `test-merge.yml` unless user also asked for L3.
+   - **L4 Perf**: `test-nightly.yml` **Perf Test · &lt;Model&gt;** step (new step or extend commands); include env exports + artifact upload pattern from a sibling perf job.
+   - **L4 Accuracy**: `test-nightly.yml` **Accuracy Test** step under the same model-type group.
+   - **L3**: `test-merge.yml` + full `source_file_dependencies` + `agents` + `plugins`.
+   - **L2**: `test-ready.yml` + same E2E block requirements.
+   - **Invalid param**: `test-weekly.yml` — note **Invalid parameters Test · H100/L4** group; usually **no YAML change** when only extending existing scripts.
+   If the user asked only for **L4 functional cases**, state explicitly that **Perf / Accuracy / Invalid param** were not included (offer to add if needed). If the plan mixes **success e2e** and **invalid param**, split deliverables across e2e vs `invalid_param_test/` and call out both CI files.
+4. **Run commands (required)** — always include, in fenced `bash` blocks:
+   - **Local — whole file**: `cd tests` then `pytest -s -v <path> …`
+   - **Local — single test** (optional but preferred when the change is one function): `pytest -s -v path::test_func …`
+   - **CI-like** (when not L1 CPU): the same **marker + `--run-level`** pairing the level uses (see [references/test-routing.md](references/test-routing.md) and **Step 6** table above)
+   - **Prerequisites** (one line): e.g. “Requires CUDA L4 + weights cached locally”; “Requires H100 + `HF_TOKEN`”; “Syntax-only validation when vllm is not installed locally”
+5. **Result summary** (pass/fail if you executed tests; if not executed, state that explicitly and what the user should run)
+
+## Additional Resources
+
+- Marker and command routing: [references/test-routing.md](references/test-routing.md)
+- CI pipelines (vllm-omni): [test-ready.yml](../../../.buildkite/test-ready.yml) (L1/L2), [test-merge.yml](../../../.buildkite/test-merge.yml) (L3), [test-nightly.yml](../../../.buildkite/test-nightly.yml) (L4), [test-weekly.yml](../../../.buildkite/test-weekly.yml) (invalid param / reliability)
+- L4 documentation example tests (naming, extraction vs copied scripts, output dirs, skips): [docs/contributing/ci/test_examples/l4_doc_example_tests.inc.md](../../../docs/contributing/ci/test_examples/l4_doc_example_tests.inc.md) — see also [PR #1910](https://github.com/vllm-project/vllm-omni/pull/1910)

.claude/skills/vllm-omni-test/references/test-routing.md

@@ -0,0 +1,372 @@
+# Test Routing Reference
+
+Use this reference to map testing goals to levels, markers, and runnable commands.
+
+**Repo paths** (`.buildkite/`, `docs/contributing/…`): link with repo-relative paths from this file (e.g. `../../../../.buildkite/test-ready.yml`, `../../../../docs/contributing/ci/CI_5levels.md`).
+
+## Model-centric e2e filename convention (L2–L4)
+
+When generating a **new** test module tied to a **specific model** under `tests/e2e/offline_inference/` or `tests/e2e/online_serving/`:
+
+| Level | Pattern | Example (`Qwen/Qwen2.5-Omni-7B`) |
+|-------|---------|----------------------------------|
+| **L2**, **L3** | `test_{lowercase_model_slug}.py` | `test_qwen2_5_omni.py` |
+| **L4** | `test_{lowercase_model_slug}_expansion.py` | `test_qwen2_5_omni_expansion.py` |
+
+**Slug rules for `{lowercase_model_slug}`** (canonical: [SKILL.md](../SKILL.md) § *Naming: generated test module files*):
+
+1. Start from the HuggingFace-style id (e.g. `Qwen/Qwen2.5-Omni-7B`), but **do not** put the org into the filename: use the **repo segment only** (`Qwen2.5-Omni-7B`), not `Qwen_...` / `qwen_qwen2_5_...`.
+2. **Lowercase**; replace `.`, `-`, and whitespace with a single `_` (e.g. `Qwen2.5-Omni-7B` → `qwen2_5_omni`). **Omit** trailing size tokens such as `7b` / `30b` in the basename when a single file covers that model line in the directory (matches `test_qwen2_5_omni.py` in-tree).
+3. If two checkpoints in the same folder need separate modules, add a **minimal** disambiguator (e.g. `_7b` vs `_3b`) **only then**.
+4. **L1** unit tests are **not** bound to this pattern; use `tests/<area>/test_<feature>.py` as today.
+
+## Model type markers (`omni` / `tts` / `diffusion`)
+
+Every model-centric e2e test must declare **exactly one** type marker:
+
+| Marker | Model family | Typical APIs |
+|--------|--------------|--------------|
+| `pytest.mark.omni` | Multimodal LLM (Qwen-Omni, …) | `send_omni_request`, `omni_runner`, stage YAML |
+| `pytest.mark.tts` | TTS / speech synthesis | `send_audio_speech_request`, `/v1/audio/speech` |
+| `pytest.mark.diffusion` | Diffusion generative models | `send_diffusion_request`, `send_video_diffusion_request`, `send_images_generations_http_request`, `OmniDiffusionSamplingParams` |
+
+**Diffusion** uses one marker for all output modalities. Nightly CI splits diffusion by **group**, not by extra markers:
+
+| Nightly group (`test-nightly.yml`) | Diffusion scope | Examples |
+|-----------------------------------|-----------------|----------|
+| **Diffusion X2I(&A&T) Model Test** | x2**i** / x2**a** / x2**t** — image, audio, text (non-video) | Qwen-Image*, BAGEL, FLUX, SD3, Z-Image, LongCat, DreamZero |
+| **Diffusion X2V Model Test** | x2**v** — video only | Wan2.2, HunyuanVideo 1.5, Wan VACE |
+
+PR labels for selective nightly runs: `diffusion-x2iat-test`, `diffusion-x2v-test` (plus `nightly-test` / `NIGHTLY=1`).
+
+## Level and Marker Mapping
+
+| Goal | Suggested Level | Marker baseline | Typical location |
+|------|------------------|-----------------|------------------|
+| Unit logic, regression on pure code path | L1 | `core_model and cpu` | `tests/<component>/test_*.py` |
+
+### L1 unit tests — mocking (`mocker`, not `unittest.mock`)
+
+L1 modules run in **Simple Test** (`test-ready.yml` / `test-merge.yml`: `-m 'core_model and cpu'`). When isolation requires doubles:
+
+- **Use** the **`mocker`** fixture (`pytest-mock`): `mocker.patch`, `mocker.spy`, `mocker.Mock`, `mocker.MagicMock`, `mocker.AsyncMock`.
+- **Use** **`monkeypatch`** for env vars and simple `setattr` / `delattr` without mock objects.
+- **Do not use** `unittest.mock` — no `from unittest.mock import ...`, no `@patch`, no `with patch(...)`.
+
+```python
+# Good (L1)
+def test_handler(mocker):
+    mocker.patch("vllm_omni.pkg.fn", return_value=0)
+
+# Bad (L1)
+from unittest.mock import patch
+
+@patch("vllm_omni.pkg.fn")
+def test_handler(mock_fn): ...
+```
+
+E2E (L2+) should not rely on mocks unless documenting a rare exception; prefer real (or lightweight) execution paths.
+| Basic integration/e2e | L2 | `core_model` + **one of** `omni` / `tts` / `diffusion` | `tests/e2e/...` |
+| Advanced integration | L3 | `advanced_model` + type marker | `tests/e2e/...` |
+| Full function / nightly | L4 | `full_model` (nightly) + type marker | **Function:** `tests/e2e/*_expansion.py`; **Perf:** `tests/dfx/perf/tests/*.json`; **Accuracy:** `tests/e2e/accuracy/` |
+| Invalid HTTP / param validation | Weekly (dfx) | `pytest.mark.slow` + type marker + `H100` or `L4` | `tests/dfx/reliability/invalid_param_test/test_invalid_*.py` |
+
+## Marker Selection Rules
+
+1. **Level** (pick one):
+   - `core_model` — L1/L2 (`test-ready.yml`)
+   - `advanced_model` — L3 (`test-merge.yml`)
+   - `full_model` — L4 nightly (`test-nightly.yml`); some expansion tests still carry both `advanced_model` and `full_model` during migration
+2. **Model type** (pick one): `omni`, `tts`, or `diffusion`
+3. **Cross-cutting** when relevant: `parallel`, `cache`, `example`, `benchmark`
+4. **Hardware**: `cpu`, `cuda`, `rocm`, `npu`, `L4`, `H100`, `distributed_cuda`, …
+5. Multi-card: `@hardware_test(...)` in `tests/helpers/mark.py`
+
+## Command Templates
+
+### Quick local checks
+
+```bash
+cd tests
+pytest -s -v test_xxxx.py
+```
+
+### L1
+
+```bash
+cd tests
+pytest -s -v -m "core_model and cpu"
+```
+
+### L2
+
+```bash
+cd tests
+pytest -s -v -m "core_model and not cpu" --run-level=core_model
+```
+
+### L3 (merge)
+
+```bash
+cd tests
+pytest -s -v -m "advanced_model" --run-level=advanced_model
+```
+
+### L4 (nightly)
+
+**Function** (e2e expansion — default when user asks for “L4 functional cases”):
+
+```bash
+cd tests
+pytest -s -v e2e/online_serving/test_qwen_image_expansion.py -m "full_model and diffusion and H100" --run-level=full_model
+```
+
+**Perf** (dfx benchmark harness — separate pillar; not `omni_server` fixtures):
+
+```bash
+cd tests
+export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
+export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
+pytest -s -v dfx/perf/scripts/run_diffusion_benchmark.py \
+  --test-config-file dfx/perf/tests/test_qwen_image_vllm_omni.json
+```
+
+Broad marker sweep (when no explicit file shard):
+
+```bash
+cd tests
+pytest -s -v -m "full_model" --run-level=full_model
+```
+
+## L4 nightly pillars (Function / Accuracy / Perf)
+
+Nightly **L4** for a model is often **multiple jobs** in `test-nightly.yml`, not a single pytest module:
+
+| Pillar | User says | Deliver | Nightly step |
+|--------|-----------|---------|--------------|
+| **Function** | “L4 functional cases” / “functional” / `*_expansion.py` | `tests/e2e/.../test_<model>_expansion.py` | `· Function Test with H100/L4` |
+| **Perf** | “performance” / “perf” / “benchmark” / “full L4” | `tests/dfx/perf/tests/test_<model>_vllm_omni.json` + runner | `· Perf Test · <Model>` |
+| **Accuracy** | “accuracy” / “similarity” | `tests/e2e/accuracy/test_<model>*.py` | `· Accuracy Test` |
+
+**Rule:** “L4 functional cases” → **Function only** by default. State that Perf/Accuracy are separate pillars; add them only when requested.
+
+**Diffusion perf JSON** (`tests/dfx/perf/tests/test_<slug>_vllm_omni.json`): array of objects with `test_name`, `server_params` (`model`, `serve_args`), `benchmark_params[]`, and per-workload **`baseline`** (`throughput_qps`, `latency_mean`, `peak_memory_mb_mean`). Copy structure from `test_qwen_image_vllm_omni.json` in-tree.
+
+**Diffusion perf nightly step** (under **Diffusion X2I(&A&T)**; copy full `kubernetes` plugins from `· Perf Test · Qwen-Image`):
+
+```yaml
+      - label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · <Model-Display-Name>"
+        key: nightly-diffusion-x2iat-performance-<slug>
+        timeout_in_minutes: 180
+        commands:
+          - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
+          - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
+          - export CACHE_DIT_VERSION=1.3.0
+          - |
+            set +e
+            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_<slug>_vllm_omni.json
+            EXIT=$$?
+            buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
+            buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
+            exit $$EXIT
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          # … full kubernetes block (often multi-GPU) — copy from Perf Test · Qwen-Image
+```
+
+| Model type | Perf runner | Config example |
+|------------|-------------|----------------|
+| Diffusion X2I/X2V | `dfx/perf/scripts/run_diffusion_benchmark.py` | `test_qwen_image_vllm_omni.json` |
+| TTS | `dfx/perf/scripts/run_benchmark.py` | `test_tts.json` (`BENCHMARK_DIR`) |
+| Omni | `run_benchmark.py` / omni configs in-tree | `test_omni*.json` |
+
+Do **not** put throughput/latency baselines inside `test_*_expansion.py` — that belongs in the dfx perf JSON + nightly Perf job.
+
+## Invalid parameter validation (weekly / dfx)
+
+When the user’s test plan includes **invalid parameter validation**, **invalid request bodies**, or **HTTP 4xx** validation against a live server:
+
+1. **Do not** author these in `tests/e2e/online_serving/test_*.py` or `*_expansion.py`. **Move** any drafted `test_*` into `tests/dfx/reliability/invalid_param_test/`.
+2. **Pick script by route** (extend in-tree file):
+
+| Route family | Script |
+|--------------|--------|
+| Omni chat / WS video·realtime | `test_invalid_omni_chat.py` |
+| `/v1/audio/speech` (+ stream / batch / voices) | `test_invalid_audio_speech.py` |
+| Audio diffusion | `test_invalid_audio_diffusion.py` |
+| `/v1/images/generations` | `test_invalid_image_generation.py` |
+| `/v1/images/edits` | `test_invalid_image_editing.py` |
+| `/v1/videos*` | `test_invalid_video_generation.py` |
+| Sleep / wakeup / server control | `test_invalid_server_control.py` |
+
+3. **Style:** `pytestmark = [pytest.mark.slow, pytest.mark.<type>]`; `_PARAMS` + `hardware_marks`; `send_*_http_request` with `err_code` + `err_message`; parametrized `body_spec` rows with `id=`; `_minimal_*_json()` helpers; `_SKIP_ISSUE_3649` when tracked in [#3649](https://github.com/vllm-project/vllm-omni/issues/3649).
+4. **CI:** [`.buildkite/test-weekly.yml`](../../../../.buildkite/test-weekly.yml) group **Reliability Test - Invalid parameters Test** — **not** ready/merge/nightly.
+
+```bash
+# Weekly H100 (diffusion / omni / video invalid-param)
+cd tests
+pytest -s -v dfx/reliability/invalid_param_test/ -m "slow and H100"
+
+# Weekly L4 (e.g. Qwen3-TTS speech invalid-param)
+pytest -s -v dfx/reliability/invalid_param_test/ -m "slow and L4"
+
+# Single script / case
+pytest -s -v dfx/reliability/invalid_param_test/test_invalid_image_generation.py::test_images_generations_invalid_requests -m "slow and H100"
+```
+
+**Trigger:** `WEEKLY=1` or PR label `weekly-test`. Appending cases to existing scripts usually needs **no YAML edit** (directory sweep). No `source_file_dependencies` on weekly steps.
+
+### Platform-targeted examples
+
+```bash
+cd tests
+pytest -s -v -m "core_model and distributed_cuda and L4" --run-level=core_model
+```
+
+### Concrete e2e paths (common in-tree)
+
+Paths are relative to `tests/` after `cd tests`.
+
+| Scenario | Example command |
+|----------|------------------|
+| **Omni** offline L2 | `pytest -s -v e2e/offline_inference/test_qwen2_5_omni.py -m "core_model and omni and not cpu" --run-level=core_model` |
+| **Omni** offline L2 — one test | `pytest -s -v e2e/offline_inference/test_qwen2_5_omni.py::test_text_to_text -m "core_model and omni and not cpu" --run-level=core_model` |
+| **Omni** online L2 | `pytest -s -v e2e/online_serving/test_qwen3_omni.py -m "core_model and omni" --run-level=core_model` |
+| **Omni** offline L2 (Qwen3.5-9B VL) | `pytest -s -v e2e/offline_inference/test_qwen3_5_9b.py -m "core_model and omni and not cpu" --run-level=core_model` |
+| **TTS** online L2 | `pytest -s -v e2e/online_serving/test_qwen3_tts_base.py -m "core_model and tts" --run-level=core_model` |
+| **TTS** online L2 — one test | `pytest -s -v e2e/online_serving/test_qwen3_tts_base.py::test_text_to_audio_001 -m "core_model and tts" --run-level=core_model` |
+| **Diffusion X2I** offline L2 | `pytest -s -v e2e/offline_inference/test_t2i_model.py -m "core_model and diffusion and not cpu" --run-level=core_model` |
+| **Diffusion X2I** L2 online smoke | `pytest -s -v e2e/online_serving/test_qwen_image.py e2e/online_serving/test_bagel.py -m "core_model and diffusion" --run-level=core_model` |
+| **Diffusion X2V** L2 online smoke | `pytest -s -v e2e/online_serving/test_wan22_t2v.py -m "core_model and diffusion" --run-level=core_model` |
+| **Diffusion X2I** L4 online expansion (H100) | `pytest -s -v e2e/online_serving/test_qwen_image_expansion.py -m "full_model and diffusion and H100" --run-level=full_model` |
+| **Diffusion X2I** L4 perf (nightly) | `pytest -s -v dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file dfx/perf/tests/test_qwen_image_vllm_omni.json` |
+| **Diffusion X2I** L4 online expansion (Edit) | `pytest -s -v e2e/online_serving/test_qwen_image_edit_expansion.py -m "full_model and diffusion and H100" --run-level=full_model` |
+| **Diffusion X2V** L4 nightly | `pytest -s -v e2e/online_serving/test_wan22_expansion.py -m "full_model and cuda" --run-level=full_model` |
+| **Invalid param** weekly H100 | `pytest -s -v dfx/reliability/invalid_param_test/ -m "slow and H100"` |
+| **Invalid param** weekly L4 | `pytest -s -v dfx/reliability/invalid_param_test/ -m "slow and L4"` |
+
+### Agent / author completion checklist
+
+When adding or modifying tests, do not stop at “where the file lives” — also deliver:
+
+1. **Local**: `cd tests` + `pytest -s -v <path>` (and `path::test_func` when a single case is enough).
+2. **CI-like**: marker string + `--run-level` matching the pipeline (`core_model` / `advanced_model` / `full_model`).
+3. **L1 mocks**: `mocker` / `monkeypatch` only; no `unittest.mock`.
+4. **API client + assert placement**: **General** validation → implement in `assertions.py`, call **inside** `send_*_request` in `runtime.py`; tests call **`send_*_request` only**. **Special** case validation → `assert_*` in `assertions.py`, called **in the test** after `send_*_request`. Low-level `send_*_http_request` is for negative/dfx tests (`err_code`), not ordinary L2+ success e2e.
+5. **Shared assertions**: logic in **`tests/helpers/assertions.py`**; general checks inside `send_*_request`; special checks only in the test. No `_assert_*` in `test_*.py`.
+6. **One case → one `test_*`**: function name reflects what is validated; no `if case_id == ...` mega-test merging multiple scenarios.
+7. **Fixture scope**: default **`omni_server` + `openai_client`** / **`omni_runner` + `omni_runner_handler`** (module). Use **`omni_server_function` + `openai_client_function`** / **`omni_runner_function` + `omni_runner_handler_function`** only when each `test_*` must spawn a fresh instance. Parametrize name must match fixture (`omni_server` vs `omni_server_function`).
+8. **Type marker**: `omni`, `tts`, or `diffusion` on every model e2e module.
+9. **Diffusion L4 Function**: wire `*_expansion.py` into **X2I(&A&T)** or **X2V** **Function Test** in **`test-nightly.yml` only** — do not add `test-merge.yml` unless the user also requested L3.
+10. **Diffusion L4 Perf** (only when requested): add `tests/dfx/perf/tests/test_<slug>_vllm_omni.json` + **Perf Test · &lt;Model&gt;** step (artifact upload); not part of “L4 functional cases” by default.
+11. **E2E Buildkite (L2/L3 only)**: `test-ready.yml` / `test-merge.yml` steps need **`source_file_dependencies`** + full **`agents` + `plugins`**. **L4 nightly** uses explicit file lists or perf scripts in `test-nightly.yml` (no merge job).
+12. **Invalid param**: cases in **`tests/dfx/reliability/invalid_param_test/`** (route-matching script), `send_*_http_request` + `err_code`, `pytest.mark.slow` + `H100`/`L4`; CI = **`test-weekly.yml`** only — do not put in e2e or nightly.
+13. **Prerequisites**: GPU tier, HF cache/token, and any module `skipif` / platform-only YAML.
+
+## Buildkite pipeline mapping
+
+| Level | Repo file | Model-type grouping |
+|-------|-----------|---------------------|
+| L1, L2 | [`.buildkite/test-ready.yml`](../../../../.buildkite/test-ready.yml) | Steps prefixed **Omni ·**, **TTS ·**, **Diffusion ·** under **E2E Test** — **`source_file_dependencies` required** |
+| L3 | [`.buildkite/test-merge.yml`](../../../../.buildkite/test-merge.yml) | Per-model E2E steps; `-m "advanced_model and …"` — **`source_file_dependencies` required** |
+| L4 | [`.buildkite/test-nightly.yml`](../../../../.buildkite/test-nightly.yml) | **Omni / TTS / Diffusion X2I(&A&T) / X2V** — each group may have **Function**, **Accuracy**, **Perf**, **Doc** steps |
+| Invalid param (weekly) | [`.buildkite/test-weekly.yml`](../../../../.buildkite/test-weekly.yml) | **Invalid parameters Test · H100** / **· L4** — sweeps `tests/dfx/reliability/invalid_param_test/` |
+
+### `source_file_dependencies` (E2E Test only — ready & merge)
+
+Required on each step inside the **E2E Test** group in `test-ready.yml` and `test-merge.yml`. Typical entries:
+
+| Category | Path pattern |
+|----------|----------------|
+| Test module | `tests/e2e/online_serving/test_<slug>.py`, `tests/e2e/offline_inference/test_<slug>.py` |
+| API client helpers | `tests/helpers/runtime.py` (when the e2e job uses newly added `send_*` methods) |
+| Shared assert helpers | `tests/helpers/assertions.py` (when the e2e job uses newly added `assert_*` helpers) |
+| AR / omni model | `vllm_omni/model_executor/models/<family>/` |
+| Diffusion model | `vllm_omni/diffusion/models/<family>/` |
+| Stage processor | `vllm_omni/model_executor/stage_input_processors/<family>.py` |
+| Deploy config | `vllm_omni/deploy/<name>.yaml` or `vllm_omni/deploy/ci/<name>.yaml` |
+
+Copy the dependency block from the closest in-tree E2E step for the same model family. Update the list whenever the pytest command or server YAML changes.
+
+### E2E `agents` + `plugins` blocks (required — do not truncate)
+
+When generating or documenting E2E steps, **always include the complete `plugins` section**. Set **`timeout_in_minutes`** on the step for the job deadline; run **`pytest` directly** in `commands` (no `timeout 40m bash -c` wrapper around pytest).
+
+Common patterns in `test-ready.yml` / `test-merge.yml`:
+
+| Queue | Plugin | Required fields |
+|-------|--------|-----------------|
+| `mithril-h100-pool` | `kubernetes` | `resources.limits.nvidia.com/gpu`, `volumeMounts` (`devshm`, `hf-cache`), `env` (`HF_HOME`, `HF_TOKEN` via `secretKeyRef`), `nodeSelector: gpu-h100-sxm`, `volumes` (`devshm` emptyDir, `hf-cache` hostPath `/mnt/hf-cache`) |
+| `gpu_1_queue` / `gpu_4_queue` | `docker#v5.2.0` | `image`, `always-pull`, `propagate-environment`, `shm-size: "8gb"`, `environment: HF_HOME`, `volumes: /fsx/hf_cache` (add `HF_TOKEN` when the step needs hub access) |
+
+**H100 kubernetes template** (copy from **Diffusion · Bagel Test** / **Diffusion · Qwen Image Test** in-tree):
+
+```yaml
+  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
+```
+
+**L4 docker template** (copy from **TTS · Qwen3-TTS Base Test** in `test-merge.yml`):
+
+```yaml
+  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"
+        volumes:
+          - "/fsx/hf_cache:/fsx/hf_cache"
+```
+
+When extending an existing step (e.g. add `tests/e2e/offline_inference/test_qwen_image.py` to **Diffusion · Qwen Image Test**), update `source_file_dependencies` and `commands` only; **keep the existing `agents` + `plugins` block unchanged** unless the hardware tier changes.
+
+The root [`.buildkite/pipeline.yml`](../../../../.buildkite/pipeline.yml) decides **which** child file is uploaded. To run **L3** on a feature branch, comment out `if: build.branch == "main"` on the merge upload step. To run **L4** without `NIGHTLY=1`, comment out the nightly upload step’s `if: build.env("NIGHTLY") == "1"` and the same `if` on relevant steps inside `test-nightly.yml`. Revert such edits before merging.
+
+Platform-specific pipelines (e.g. AMD) follow the same level → file pairing under `.buildkite/`.
+
+## Diffusion RFC (#1832) Alignment Tips
+
+For **X2I(&A&T)** coverage planning:
+
+- Prioritize high-value feature combinations with minimal case count.
+- Split into lightweight L1/L2 validation plus L4 `*_expansion.py` under the **X2I** nightly shards.
+- **X2V** models (Wan, HunyuanVideo) stay in the **X2V** group — do not merge into X2I sweeps.
+
+If hardware is insufficient, provide an executable reduced case plus a deferred full CI/nightly plan.

.gitattributes

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

.github/CODEOWNERS

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

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

@@ -74,28 +74,21 @@ 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 import OmniLLM, create_ar_stage_config, create_dit_stage_config
+      from vllm_omni.entrypoints.omni import Omni
+      from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+      from vllm import SamplingParams
 
-      # Create stage configurations
-      ar_config = create_ar_stage_config(
-          stage_id=0,
-          model_path="Qwen/Qwen3-0.6B",
-          input_modalities=["text"],
-          output_modalities=["text"]
+      omni = Omni(
+          model="Qwen/Qwen-Image",
+          stage_configs_path="/path/to/stage_configs.yaml",
       )
 
-      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")
+      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)
       ```
 
       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=label%3ARFC+sort%3Aupdated-desc) for reference.
+      #### Please take a look at previous [RFCs](https://github.com/vllm-project/vllm-omni/issues?q=in%3Atitle%20RFC%20sort%3Aupdated-desc) for reference.
 - type: textarea
   attributes:
     label: Motivation.
@@ -21,7 +21,7 @@ body:
     description: >
       The proposed change of the RFC.
     value: |
-      Please provide the detailed design document of the RFC using the [template](https://docs.google.com/document/d/12YxSsVeD1jvL-InClkeAEnZyWFDndz_65JmXvsamuV4/edit?tab=t.0#heading=h.4ef4szrsgspp).
+      Please provide the detailed design document of the RFC using the [template](https://docs.google.com/document/d/1jcgR3cDaUQH3VczD4ZcKaJAoYWHjCmnYzHkCNyz-9fk/edit?usp=sharing).
   validations:
     required: true
 - type: textarea

.github/PULL_REQUEST_TEMPLATE.md

@@ -1,21 +1,15 @@
 <!-- markdownlint-disable -->
-PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTTOM) HAVE BEEN CONSIDERED.
+PLEASE FILL IN THE PR DESCRIPTION HERE.
 
 ## Purpose
 
 ## Test Plan
 
+**vLLM Version:**
+
+**vLLM-Omni Commit:**
+
 ## Test Result
 
----
-<details>
-<summary> Essential Elements of an Effective PR Description Checklist </summary>
-
-- [ ] The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
-- [ ] The test plan, such as providing test command.
-- [ ] The test results, such as pasting the results comparison before and after, or e2e results
-- [ ] (Optional) The necessary documentation update, such as updating `supported_models.md` and `examples` for a new model.
-- [ ] (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:** read [CONTRIBUTING.md](https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md) and run the [precheck-pr skill](https://github.com/vllm-project/vllm-omni/blob/main/.claude/skills/precheck-pr/SKILL.md) with the code agent for a self-check against project conventions.
+(anything written below this line will be removed by GitHub Actions)

.github/mergify.yml

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

.github/workflows/build_wheel.yml

@@ -7,6 +7,7 @@ on:
     paths-ignore:
       - 'docs/**'
       - 'examples/**'
+      - 'apps/**'
       - '.github/**'
       - '.buildkite/**'
       - '**.md'
@@ -20,6 +21,7 @@ on:
     paths-ignore:
       - 'docs/**'
       - 'examples/**'
+      - 'apps/**'
       - '.github/**'
       - '.buildkite/**'
       - '*.yml'
@@ -37,10 +39,10 @@ jobs:
         python-version: ["3.11", "3.12"]
 
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v6
 
       - name: Set up Python ${{ matrix.python-version }}
-        uses: actions/setup-python@v5
+        uses: actions/setup-python@v6
         with:
           python-version: ${{ matrix.python-version }}
 
@@ -55,7 +57,7 @@ jobs:
           bash scripts/build_wheel.sh --python python
 
       - name: Upload Wheel Artifact
-        uses: actions/upload-artifact@v4
+        uses: actions/upload-artifact@v7
         with:
           name: vllm-omni-wheel-py${{ matrix.python-version }}
           path: dist/*.whl

.github/workflows/pre-commit.yml

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

.gitignore

@@ -83,6 +83,9 @@ target/
 profile_default/
 ipython_config.py
 
+# uv
+uv.lock
+
 # pyenv
 #   For a library or package, you might want to ignore these files since the code is
 #   intended to run in multiple environments; otherwise, check them in:
@@ -153,6 +156,27 @@ dmypy.json
 # Cython debug symbols
 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
+
+# Codex
+AGENTS.md
+.codex
+.codex/
+
 # cursor
 .cursor/
 
@@ -180,6 +204,7 @@ checkpoints/
 # Cache directories
 cache/
 !vllm_omni/diffusion/cache/
+!tests/diffusion/cache/
 .cache/
 diffusion_cache/
 kv_cache/
@@ -214,6 +239,8 @@ data/
 datasets/
 *.csv
 *.json
+!tests/dfx/perf/tests/*.json
+!apps/ComfyUI-vLLM-Omni/example_workflows/*.json
 *.jsonl
 *.parquet
 
@@ -233,3 +260,16 @@ configs/development.yaml
 Dockerfile.dev
 discussion
 tmp_test
+
+# Auto-generated version file (created by setuptools_scm during build)
+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/

.pre-commit-config.yaml

@@ -8,6 +8,7 @@ repos:
       # list of supported hooks: https://pre-commit.com/hooks.html
       - id: check-yaml
         args: ["--unsafe"]
+        exclude: "examples/online_serving/chart-helm/templates/"
       - id: debug-statements
       - id: end-of-file-fixer
       - id: mixed-line-ending

.readthedocs.yml

@@ -10,6 +10,11 @@ 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

CONTRIBUTING.md

@@ -1,3 +1,5 @@
 # Contributing to vLLM-Omni
 
 You may find information about contributing to vLLM-Omni on [Contributing](https://vllm-omni.readthedocs.io/en/latest/contributing/)
+
+Before submitting a PR, run the [precheck-pr skill](https://github.com/vllm-project/vllm-omni/blob/main/.claude/skills/precheck-pr/SKILL.md) with the code agent for a self-check against project conventions.

README.md

@@ -1,6 +1,6 @@
 <p align="center">
   <picture>
-    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/logos/vllm-omni-logo.png">
+    <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/logos/vllm-omni-logo.png">
     <img alt="vllm-omni" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/logos/vllm-omni-logo.png" width=55%>
   </picture>
 </p>
@@ -9,13 +9,20 @@ Easy, fast, and cheap omni-modality model serving for everyone
 </h3>
 
 <p align="center">
-| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> | <a href="docs/assets/WeChat.jpg"><b>WeChat</b></a> |
+| <a href="https://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> |
 </p>
 
+
 ---
 
 *Latest News* 🔥
-- [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), please check our latest [design](https://docs.google.com/presentation/d/1qv4qMW1rKAqDREMXiUDLIgqqHQe7TDPj/edit?usp=sharing&ouid=110473603432222024453&rtpof=true&sd=true).
+- [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.
 - [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.
 
 ---
@@ -50,8 +57,9 @@ vLLM-Omni is flexible and easy to use with:
 
 vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, including:
 
-- Omni-modality models (e.g. Qwen-Omni)
-- Multi-modality generation models (e.g. Qwen-Image)
+- **Omni-modality models** (e.g. Qwen3-Omni, Cosmos, HunyuanImage, BAGEL)
+- **TTS models** (e.g. Qwen3-TTS, VoxCPM2, Ming-Omni-TTS, CosyVoice3)
+- **Diffusion models** — image, video, and audio generation (e.g. Qwen-Image, Wan2.2, FLUX)
 
 ## Getting Started
 
@@ -60,12 +68,26 @@ Visit our [documentation](https://vllm-omni.readthedocs.io/en/latest/) to learn
 - [Installation](https://vllm-omni.readthedocs.io/en/latest/getting_started/installation/)
 - [Quickstart](https://vllm-omni.readthedocs.io/en/latest/getting_started/quickstart/)
 - [List of Supported Models](https://vllm-omni.readthedocs.io/en/latest/models/supported_models/)
+- [Deployment Recipes](https://recipes.vllm.ai) for vLLM-Omni model serving
 
 ## Contributing
 
 We welcome and value any contributions and collaborations.
 Please check out [Contributing to vLLM-Omni](https://vllm-omni.readthedocs.io/en/latest/contributing/) for how to get involved.
 
+## Citation
+
+If you use vLLM-Omni for your research, please cite our [paper](https://arxiv.org/abs/2602.02204):
+
+```bibtex
+@article{yin2026vllmomni,
+  title={vLLM-Omni: Fully Disaggregated Serving for Any-to-Any Multimodal Models},
+  author={Peiqi Yin, Jiangyun Zhu, Han Gao, Chenguang Zheng, Yongxiang Huang, Taichang Zhou, Ruirui Yang, Weizhi Liu, Weiqing Chen, Canlin Guo, Didan Deng, Zifeng Mo, Cong Wang, James Cheng, Roger Wang, Hongsheng Liu},
+  journal={arXiv preprint arXiv:2602.02204},
+  year={2026}
+}
+```
+
 ## Join the Community
 Feel free to ask questions, provide feedbacks and discuss with fellow users of vLLM-Omni in `#sig-omni` slack channel at [slack.vllm.ai](https://slack.vllm.ai) or vLLM user forum at [discuss.vllm.ai](https://discuss.vllm.ai).
 

apps/ComfyUI-vLLM-Omni/.gitignore

@@ -0,0 +1,115 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# OSX useful to ignore
+*.DS_Store
+.AppleDouble
+.LSOverride
+
+# Thumbnails
+._*
+
+# Files that might appear in the root of a volume
+.DocumentRevisions-V100
+.fseventsd
+.Spotlight-V100
+.TemporaryItems
+.Trashes
+.VolumeIcon.icns
+.com.apple.timemachine.donotpresent
+
+# Directories potentially created on remote AFP share
+.AppleDB
+.AppleDesktop
+Network Trash Folder
+Temporary Items
+.apdisk
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+env/
+venv/
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*,cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+
+# Sphinx documentation
+docs/_build/
+
+# IntelliJ Idea
+.idea
+*.iml
+*.ipr
+*.iws
+
+# PyBuilder
+target/
+
+# Cookiecutter
+output/
+python_boilerplate/
+cookiecutter-pypackage-env/
+
+# vscode settings
+.history/
+*.code-workspace
+
+# Frontend extension
+node_modules/
+.env
+.env.local
+.env.development.local
+.env.test.local
+.env.production.local
+npm-debug.log*
+yarn-debug.log*
+yarn-error.log*
+node.zip
+.vscode/
+.claude/
+.codemate/

apps/ComfyUI-vLLM-Omni/LICENSE

@@ -0,0 +1,15 @@
+Apache Software License 2.0
+
+Copyright (c) 2026, Zeyu Huang
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.

apps/ComfyUI-vLLM-Omni/README.md

@@ -0,0 +1,184 @@
+# vLLM-Omni
+
+vLLM-Omni offers a ComfyUI integration on top of its online serving API.
+It can send model inference requests to either a locally running vLLM-Omni service or a remote one.
+
+## Requirement
+
+- Python 3.12 or above
+- [ComfyUI installed](https://docs.comfy.org/installation/system_requirements)
+- [vLLM-Omni installed](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/) on either the same device or another device discoverable via the internet.
+- No need to install additional packages apart from those already required by ComfyUI.
+
+> [!TIP]
+> If you run both ComfyUI and vLLM-Omni on the same device, you can create separate virtual environments and use different Python versions for them.
+
+
+## Installation
+
+Copy this folder to the `custom_nodes` subfolder of your ComfyUI installation. Your directory should look like `ComfyUI/custom_nodes/ComfyUI-vLLM-Omni`.
+
+If you are running ComfyUI during copying, you should restart ComfyUI to load this extension.
+
+> [!TIP]
+> You can use utility websites such as https://download-directory.github.io/ to download a subdirectory of a repo. Also checkout community discussions (e.g., https://stackoverflow.com/questions/7106012/download-a-single-folder-or-directory-from-a-github-repository) for more info.
+
+On the device and virtual environment you run ComfyUI, launch ComfyUI with
+```bash
+cd ComfyUI
+
+# The regular way
+python main.py
+
+# If you are mainly using this node, launch it faster with
+python main.py --cpu
+```
+
+On the device and virtual environment you run vLLM-Omni, start a model service with
+```bash
+vllm serve The_Model_ID_to_Serve --omni --port 8000
+```
+
+Check **ComfyUI's sidebar -> Node Library**. There should be a new folder named **vLLM-Omni**.
+If no, check your shell running the ComfyUI process. There may be some error messages before the line `Import times for custom nodes:` and the line `To see the GUI go to: http://127.0.0.1:8188`.
+
+## Quickstart
+
+This extension offers the following nodes based on the output modalities (at **ComfyUI sidebar -> Node Library**):
+
+- **Generate Image** for text-to-image and image-to-image tasks
+- **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
+
+This extension also offers example workflows (at **ComfyUI sidebar -> Templates -> vLLM-Omni**)
+
+> [!NOTE]
+> The node UI and feature designs are intended to match vLLM-Omni online serving interfaces. It cannot offer more than what the interfaces support.
+
+To build a simple workflow yourself,
+
+- Drag a generation node onto the canvas.
+- Depending on your need, grab built-in multimedia file loader nodes, such as **image->Load Image**, **image->video->Load Video**, **audio->Load Audio**
+- Depending on your need, grab built-in multimedia file preview nodes, such as **image->Preview Image**, **image->video->Save Video**, **audio->Preview Audio**, **utils->Preview as Text**.
+- If you want to tune sampling parameters, grab corresponding nodes from **vLLM-Omni-> Sampling Params**.
+    - For multi-stage models, you can connect multiple **AR Sampling Params** and **Diffusion Sampling Params** nodes to a **Multi-Stage Sampling Params List** node, and connect this node to the generation node.
+    - For some multi-stage models like BAGEL, [only one stage's sampling parameters are exposed and tunable via vLLM-Omni's online serving API](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/examples/online_serving/bagel/). Thus, these models are treated as single-stage ones. Please check the vLLM-Omni documentation on how to correctly set each model's sampling parameters.
+    - For multi-stage models where all stages are either autoregression or diffusion, you can also connect only a single Sampling Params node, indicating that this set of sampling parameters will be used for all stages.
+
+## Screenshots and Examples
+
+### Multimodal understanding (e.g., Qwen Omni series, BAGEL)
+
+(Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni Multimodal Understanding**)
+
+<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-understanding.jpg">
+    <img alt="vLLM-Omni multimodal understanding" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-understanding.jpg" width=55%>
+  </picture>
+</p>
+
+> [!TIP]
+> Although this node enables all-modality input, you should check whether the specific model you host and request for supports the modalities you connect to the node.
+
+You can configure per-stage sampling parameters for multi-stage models.
+
+<p align="center">
+  <picture>
+    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-multi-stage.jpg">
+    <img alt="vLLM-Omni multiple stages" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-multi-stage.jpg" width=55%>
+  </picture>
+</p>
+
+### Text-to-image and image-to-image generation (e.g., Z-Image-Turbo, Qwen-Image-Edit, BAGEL)
+
+(Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni Image Generation**)
+
+<p align="center">
+  <picture>
+    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-image-generation.jpg">
+    <img alt="vLLM-Omni image generation" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-image-generation.jpg" width=55%>
+  </picture>
+</p>
+
+> [!TIP]
+> The node automatically choose text-to-image or image-to-image API endpoints depending on whether you connect an image input or not.
+
+### 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**)
+
+<p align="center">
+  <picture>
+    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-tts.jpg">
+    <img alt="vLLM-Omni TTS" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-tts.jpg" width=55%>
+  </picture>
+</p>
+
+> [!TIP]
+> There is a dedicated node for VoiceClone tasks with reference audio input. Other simple text-to-speech tasks should use the regular TTS node.
+
+### Chaining multiple model services
+
+(Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni Chaining Services**)
+
+<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-chaining-services.jpg">
+    <img alt="vLLM-Omni TTS" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-chaining-services.jpg" width=55%>
+  </picture>
+</p>
+
+## Develop
+
+Follow the [development convention and rules of vLLM-Omni](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/).
+
+## Limitation and Non-Goals
+
+- Single server mode only. No automatic load balancing or failover.
+- Features set is bounded to vLLM-Omni's online service capability, including
+    - The types of models supported in online mode,
+    - The types of sampling parameters supported in the online mode,
+    - The ways to send files (primarily through full-length base64 in JSON payload),
+    - Figuring out errors in the payload (such as unsupported fields by a specific model) if the endpoint does not explicitly return an error,
+    - (The lack of) Authentication
+    - (The lack of) Progress indicator
+
+## Support
+
+If you are new to ComfyUI, please check out [its documentation](https://docs.comfy.org/) for usage instructions.
+
+If you are new to vLLM-Omni, please also check out [its documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/) for usage instructions.
+
+Whenever you find an issue or problem, please
+
+- First find out if this is an upstream limitation of vLLM-Omni's online serving mode, by [checking their documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/examples/).
+- [Open an issue](https://github.com/vllm-project/vllm-omni/issues) that clearly describes this ComfyUI or online service problem.
+
+## Acknowledgements
+
+Features
+
+- https://github.com/dougbtv/comfyui-vllm-omni/ The official reference implementation for ComfyUI integration with vLLM-Omni's DALL-E compatible image generation API.
+- https://github.com/Comfy-Org/ComfyUI/tree/master/comfy_extras ComfyUI's built-in node implementations.
+
+UI/UX design references
+
+- https://github.com/sgl-project/sglang/pull/15271 SGLang Diffusion's official ComfyUI integration for image and video generation.
+- https://github.com/SXQBW/ComfyUI-Qwen-Omni A third party ComfyUI integration for Qwen Omni series.
+- https://github.com/flybirdxx/ComfyUI-Qwen-TTS https://github.com/DarioFT/ComfyUI-Qwen3-TTS Third  party ComfyUI integrations for Qwen TTS series.

apps/ComfyUI-vLLM-Omni/__init__.py

@@ -0,0 +1,61 @@
+"""Top-level package for comfyui_vllm_omni."""  # noqa: N999  # This is not a python library intended to be imported
+
+__all__ = [
+    "NODE_CLASS_MAPPINGS",
+    "NODE_DISPLAY_NAME_MAPPINGS",
+    "WEB_DIRECTORY",
+]
+
+__author__ = """vLLM-Omni Team"""
+__email__ = "vllm-omni@vllm.ai"
+__version__ = "0.0.1"
+
+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,
+    # === Params ===
+    "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",
+    # === Params ===
+    "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/__init__.py

@@ -0,0 +1 @@
+# noqa: N999   # This is not a python library intended to be imported

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

@@ -0,0 +1,736 @@
+from typing import Literal
+
+import torch
+from comfy_api.input import AudioInput, VideoInput
+
+from .utils.api_client import VLLMOmniClient
+from .utils.logger import get_logger
+from .utils.models import lookup_model_spec
+from .utils.types import (
+    AudioFormat,
+    AutoregressionSamplingParams,
+    DiffusionSamplingParams,
+    QwenTTSModelSpecificParams,
+    WanModelSpecificParams,
+)
+from .utils.validators import (
+    add_sampling_parameters_to_stage,
+    validate_model_and_sampling_params_types,
+)
+
+logger = get_logger(__name__)
+
+
+class _VLLMOmniGenerateBase:
+    """Base class for vLLM-Omni generation nodes with shared functionality."""
+
+    CATEGORY = "vLLM-Omni"
+
+    @classmethod
+    def VALIDATE_INPUTS(cls, url, model) -> str | Literal[True]:
+        """
+        Can only validate this model's own input. Cannot check inputs from other nodes.
+        See: https://docs.comfy.org/custom-nodes/backend/server_overview#validate_inputs
+        """
+        if not url:
+            return "URL must not be empty"
+        if not model:
+            return "Model must not be empty"
+        return True
+
+
+class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "url": ("STRING", {"default": "http://localhost:8000/v1"}),
+                "model": ("STRING", {"default": "Tongyi-MAI/Z-Image-Turbo"}),
+                "prompt": ("STRING", {"multiline": True}),
+                "negative_prompt": ("STRING", {"multiline": True, "default": ""}),
+                "width": ("INT", {"default": 512, "min": 64, "max": 2048}),
+                "height": ("INT", {"default": 512, "min": 64, "max": 2048}),
+            },
+            "optional": {
+                "image": ("IMAGE",),
+                "mask": ("MASK",),
+                # "video": ("VIDEO",),
+                # "audio": ("AUDIO",),
+                "sampling_params": ("SAMPLING_PARAMS",),
+                "lora": ("REMOTE_LORA",),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("image",)
+    FUNCTION = "generate"
+
+    async def generate(
+        self,
+        url: str,
+        model: str,
+        prompt: str,
+        width: int,
+        height: int,
+        negative_prompt: str | None = None,
+        image: torch.Tensor | None = None,
+        mask: torch.Tensor | None = None,
+        audio: AudioInput | None = None,  # Hidden & unused
+        video: VideoInput | None = None,  # Hidden & unused
+        sampling_params: dict | list[dict] | None = None,
+        lora: dict | None = None,
+        **kwargs,
+    ):
+        if kwargs:
+            logger.info("Uncaught kwargs: %s", kwargs)
+        logger.debug("Got sampling params: %s", sampling_params)
+        validate_model_and_sampling_params_types(model, sampling_params)
+        if image is None and mask is not None:
+            raise ValueError("Mask input provided without an image input.")
+
+        client = VLLMOmniClient(url)
+
+        spec, pattern = lookup_model_spec(model)
+        is_bagel = pattern is not None and "bagel" in pattern.lower()
+
+        # Prefer DALL-E compatible API for simple (one-stage) diffusion models
+        if (spec is None or spec["stages"] == ["diffusion"]) and not is_bagel:
+            # The number of sampling parameter groups should have been validated.
+            # Now, simply convert single-item list to dict.
+            if isinstance(sampling_params, list):
+                sampling_params = sampling_params[0]
+            if audio is None and image is None and video is None:
+                # No multimodal input --- use DALL-E image generation
+                logger.info("Using DALL-E image generation endpoint")
+                output = await client.generate_image(
+                    model=model,
+                    prompt=prompt,
+                    width=width,
+                    height=height,
+                    negative_prompt=negative_prompt,
+                    sampling_params=sampling_params,
+                    lora=lora,
+                )
+                return (output,)
+            elif image is not None and audio is None and video is None:
+                # Image and text input --- use DALL-E image edit
+                logger.info("Using DALL-E image edit endpoint")
+                output = await client.edit_image(
+                    model=model,
+                    prompt=prompt,
+                    image=image,
+                    width=width,
+                    height=height,
+                    negative_prompt=negative_prompt,
+                    mask=mask,
+                    sampling_params=sampling_params,
+                    lora=lora,
+                )
+                return (output,)
+
+        logger.info("Using chat completion endpoint")
+        sampling_params = add_sampling_parameters_to_stage(
+            model, sampling_params, "diffusion", width=width, height=height
+        )
+        logger.debug("Edited sampling params: %s", sampling_params)
+
+        output = await client.generate_image_chat_completion(
+            model=model,
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            image=image,
+            audio=audio,
+            video=video,
+            sampling_params=sampling_params,
+            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):
+        return {
+            "required": {
+                "url": ("STRING", {"default": "http://localhost:8000/v1"}),
+                "model": ("STRING", {"default": "Qwen/Qwen2.5-Omni-7B"}),
+                "prompt": ("STRING", {"multiline": True}),
+                "output_text": ("BOOLEAN", {"default": True}),
+                "output_audio": ("BOOLEAN", {"default": True}),
+                "use_audio_in_video": ("BOOLEAN", {"default": True}),
+            },
+            "optional": {
+                "image": ("IMAGE",),
+                "video": ("VIDEO",),
+                "audio": ("AUDIO",),
+                "sampling_params": ("SAMPLING_PARAMS",),
+            },
+        }
+
+    RETURN_TYPES = ("STRING", "AUDIO")
+    RETURN_NAMES = ("text_response", "audio_response")
+    FUNCTION = "generate"
+
+    @classmethod
+    def VALIDATE_INPUTS(cls, url, model, output_text, output_audio) -> str | Literal[True]:  # type: ignore[reportIncompatibleMethodOverride]
+        super_validation = super().VALIDATE_INPUTS(url, model)
+        if isinstance(super_validation, str):
+            return super_validation
+        if not output_text and not output_audio:
+            return "At least one of output_text or output_audio must be True."
+        return True
+
+    async def generate(
+        self,
+        url: str,
+        model: str,
+        prompt: str,
+        image: torch.Tensor | None = None,
+        audio: AudioInput | None = None,
+        video: VideoInput | None = None,
+        sampling_params: dict | list[dict] | None = None,
+        output_text: bool = True,
+        output_audio: bool = True,
+        use_audio_in_video: bool = True,
+        **kwargs,
+    ) -> tuple[str, AudioInput]:
+        if kwargs:
+            logger.info("Uncaught kwargs: %s", kwargs)
+        logger.debug("Got sampling params: %s", sampling_params)
+        validate_model_and_sampling_params_types(model, sampling_params)
+
+        client = VLLMOmniClient(url)
+        spec, pattern = lookup_model_spec(model)
+        is_bagel = pattern is not None and "bagel" in pattern.lower()
+
+        if is_bagel:
+            # A lot of special handlings here...
+            if output_audio:
+                raise ValueError("BAGEL models do not support audio output.")
+            if audio is not None or video is not None:
+                raise ValueError("BAGEL models do not support audio or video input.")
+            (
+                text_response,
+                _,
+            ) = await client.generate_understanding_chat_completion(
+                model=model,
+                prompt=prompt,
+                image=image,
+                audio=None,
+                video=None,
+                sampling_params=sampling_params,
+                modalities=["text"],
+            )
+        else:
+            modalities = []
+            if output_text:
+                modalities.append("text")
+            if output_audio:
+                modalities.append("audio")
+
+            if use_audio_in_video and video is not None:
+                use_audio_in_video = True
+            else:
+                use_audio_in_video = False
+
+            (
+                text_response,
+                audio,
+            ) = await client.generate_understanding_chat_completion(
+                model=model,
+                prompt=prompt,
+                image=image,
+                audio=audio,
+                video=video,
+                sampling_params=sampling_params,
+                modalities=modalities,
+                # == extra kwargs ==
+                mm_processor_kwargs={"use_audio_in_video": use_audio_in_video},
+            )
+
+        if text_response is None:
+            text_response = ""
+        if audio is None:
+            channels = 1
+            duration = 1
+            sample_rate = 44100
+            num_samples = int(round(duration * sample_rate))
+            waveform = torch.zeros((1, channels, num_samples), dtype=torch.float32)
+            audio = {"waveform": waveform, "sample_rate": sample_rate}
+
+        return (text_response, audio)
+
+
+class VLLMOmniTTS(_VLLMOmniGenerateBase):
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "url": ("STRING", {"default": "http://localhost:8000/v1"}),
+                "model": (
+                    "STRING",
+                    {"default": "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice"},
+                ),
+                "input": ("STRING", {"multiline": True}),
+                "voice": ("STRING", {"default": "Vivian"}),
+                "response_format": (["mp3", "opus", "aac", "flac", "wav", "pcm"],),
+                "speed": (
+                    "FLOAT",
+                    {"default": 1.0, "min": 0.25, "max": 4.0, "step": 0.01},
+                ),
+            },
+            "optional": {
+                "model_specific_params": ("TTS_PARAMS",),
+            },
+        }
+
+    RETURN_TYPES = ("AUDIO",)
+    RETURN_NAMES = ("audio",)
+    FUNCTION = "generate"
+
+    async def generate(
+        self,
+        url: str,
+        model: str,
+        input: str,
+        voice: str,
+        response_format: AudioFormat,
+        speed: float,
+        model_specific_params: dict | None,
+        **kwargs,
+    ) -> tuple[AudioInput]:
+        logger.info("Got extra kwargs in TTS: %s", kwargs)
+
+        is_qwen_tts = "qwen3-tts" in model.lower()
+        if not is_qwen_tts and isinstance(model_specific_params, QwenTTSModelSpecificParams):
+            raise ValueError(
+                "You have provided Qwen-specific TTS params."
+                "However, the model appears to not be a Qwen TTS model (no 'Qwen3-TTS' in model name)."
+            )
+
+        combined_params = {**kwargs, **(model_specific_params or {})}
+
+        client = VLLMOmniClient(url)
+
+        audio = await client.generate_speech(
+            model=model,
+            input=input,
+            voice=voice,
+            response_format=response_format,
+            speed=speed,
+            **combined_params,
+        )
+        return (audio,)
+
+
+class VLLMOmniVoiceClone(_VLLMOmniGenerateBase):
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "url": ("STRING", {"default": "http://localhost:8000/v1"}),
+                "model": ("STRING", {"default": "Qwen/Qwen3-TTS-12Hz-1.7B-Base"}),
+                "input": ("STRING", {"multiline": True}),
+                "voice": ("STRING", {"default": "Vivian"}),
+                "response_format": (["mp3", "opus", "aac", "flac", "wav", "pcm"],),
+                "speed": (
+                    "FLOAT",
+                    {"default": 1.0, "min": 0.25, "max": 4.0, "step": 0.01},
+                ),
+                "ref_audio": ("AUDIO",),
+                "ref_text": ("STRING", {"multiline": True}),
+                "x_vector_only_mode": ("BOOLEAN", {"default": False}),
+            },
+            "optional": {
+                "model_specific_params": ("TTS_PARAMS",),
+            },
+        }
+
+    RETURN_TYPES = ("AUDIO",)
+    RETURN_NAMES = ("audio",)
+    FUNCTION = "generate"
+
+    async def generate(
+        self,
+        url: str,
+        model: str,
+        input: str,
+        voice: str,
+        response_format: AudioFormat,
+        speed: float,
+        ref_audio: AudioInput,
+        ref_text: str,
+        x_vector_only_mode: bool,
+        model_specific_params: dict | None,
+        **kwargs,
+    ):
+        is_qwen_tts = "qwen3-tts" in model.lower()
+        if not is_qwen_tts and isinstance(model_specific_params, QwenTTSModelSpecificParams):
+            raise ValueError(
+                "You have provided Qwen-specific TTS params."
+                "However, the model appears to not be a Qwen TTS model (no 'Qwen3-TTS' in model name)."
+            )
+
+        combined_params = {
+            "ref_audio": ref_audio,
+            "ref_text": ref_text,
+            "x_vector_only_mode": x_vector_only_mode,
+            **kwargs,
+            **(model_specific_params or {}),
+        }
+
+        client = VLLMOmniClient(url)
+
+        audio = await client.generate_speech(
+            model=model,
+            input=input,
+            voice=voice,
+            response_format=response_format,
+            speed=speed,
+            **combined_params,
+        )
+        return (audio,)
+
+
+class VLLMOmniARSampling:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "max_tokens": ("INT", {"default": 100, "min": 1, "max": 10000}),
+                "temperature": (
+                    "FLOAT",
+                    {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.01},
+                ),
+                "top_p": (
+                    "FLOAT",
+                    {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01},
+                ),
+                "repetition_penalty": (
+                    "FLOAT",
+                    {"default": 1.0, "min": 0.0, "max": 5.0, "step": 0.01},
+                ),
+                # === Put seed at last. ===
+                # Whenever a field named "seed" is present, ComfyUI adds another field called "control after generate"
+                "seed": (
+                    "INT",
+                    {
+                        "default": -1,
+                        "min": -1,
+                        "step": 1,
+                        "tooltip": "-1 means to not provide a seed.",
+                    },
+                ),
+            }
+        }
+
+    RETURN_TYPES = ("SAMPLING_PARAMS",)
+    RETURN_NAMES = ("AR sampling params",)
+    FUNCTION = "get_params"
+    CATEGORY = "vLLM-Omni/Sampling Params"
+
+    def get_params(self, seed, **kwargs):
+        params = AutoregressionSamplingParams(kwargs)
+        if seed >= 0:
+            params["seed"] = seed
+        return (params,)
+
+
+class VLLMOmniDiffusionSampling:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "n": (
+                    "INT",
+                    {
+                        "default": 1,
+                        "min": 0,
+                        "max": 10,
+                        "step": 1,
+                        "tooltip": "Number of images to generate",
+                    },
+                ),
+                "num_inference_steps": (
+                    "INT",
+                    {
+                        "default": 50,
+                        "min": 1,
+                        "max": 1000,
+                        "tooltip": "Number of denoising steps (higher = better quality, slower).",
+                    },
+                ),
+                "guidance_scale": (
+                    "FLOAT",
+                    {
+                        "default": 7.5,
+                        "min": 0.0,
+                        "max": 20.0,
+                        "step": 0.1,
+                        "tooltip": "Classifier-free guidance scale (higher = more prompt adherence).",
+                    },
+                ),
+                "true_cfg_scale": (
+                    "FLOAT",
+                    {
+                        "default": 1.0,
+                        "min": 0.0,
+                        "max": 20.0,
+                        "step": 0.5,
+                        "tooltip": "True CFG scale for advanced control (model-specific).",
+                    },
+                ),
+                "vae_use_slicing": (
+                    "BOOLEAN",
+                    {
+                        "default": False,
+                        "tooltip": "Enable VAE slicing for reduced memory usage (slight quality trade-off)",
+                    },
+                ),
+                "vae_use_tiling": (
+                    "BOOLEAN",
+                    {
+                        "default": False,
+                        "tooltip": "Enable VAE tiling for reduced memory usage (slight quality trade-off)",
+                    },
+                ),
+                # === Put seed at last. ===
+                # Whenever a field named "seed" is present, ComfyUI adds another field called "control after generate"
+                "seed": (
+                    "INT",
+                    {
+                        "default": -1,
+                        "min": -1,
+                        "step": 1,
+                        "tooltip": "-1 means to not provide a seed.",
+                    },
+                ),
+            }
+        }
+
+    RETURN_TYPES = ("SAMPLING_PARAMS",)
+    RETURN_NAMES = ("diffusion sampling params",)
+    FUNCTION = "get_params"
+    CATEGORY = "vLLM-Omni/Sampling Params"
+
+    def get_params(self, seed, **kwargs):
+        params = DiffusionSamplingParams(kwargs)
+        if seed >= 0:
+            params["seed"] = seed
+        return (params,)
+
+
+class VLLMOmniSamplingParamsList:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "param1": ("SAMPLING_PARAMS",),
+            },
+            "optional": {
+                "param2": ("SAMPLING_PARAMS",),
+                "param3": ("SAMPLING_PARAMS",),
+            },
+        }
+
+    RETURN_TYPES = ("SAMPLING_PARAMS",)
+    RETURN_NAMES = ("param list",)
+    FUNCTION = "aggregate"
+    CATEGORY = "vLLM-Omni/Sampling Params"
+
+    def aggregate(self, param1: dict, param2: dict | None = None, param3: dict | None = None):
+        for i, p in enumerate((param1, param2, param3)):
+            if isinstance(p, list):
+                raise ValueError(
+                    f"Input {i} is a Multi-Stage Sampling Params List. "
+                    f"Expected a single sampling parameters node (either AR or Diffusion)."
+                )
+
+        params = [param1]
+        if param2 is not None:
+            params.append(param2)
+        if param3 is not None:
+            params.append(param3)
+        return (params,)
+
+
+class 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):
+        return {
+            "required": {
+                "task_type": (
+                    ["CustomVoice", "VoiceDesign", "Base"],
+                    {"default": "CustomVoice"},
+                ),
+                "language": (
+                    ["Auto", "Chinese", "English", "Japanese", "Korean"],
+                    {"default": "Auto"},
+                ),
+                "instructions": ("STRING", {"multiline": True}),
+                "max_new_tokens": ("INT", {"default": 2048, "min": 1}),
+            }
+        }
+
+    RETURN_TYPES = ("TTS_PARAMS",)
+    RETURN_NAMES = ("Qwen TTS params",)
+    FUNCTION = "get_params"
+    CATEGORY = "vLLM-Omni/TTS Params"
+
+    def get_params(self, **kwargs):
+        return (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

@@ -0,0 +1,585 @@
+"""
+An high-level API client adapter that forwards ComfyUI inputs to vLLM-Omni's REST API,
+and transforms the API responses back to ComfyUI formats.
+
+The image generation part is derived from dougbtv/comfyui-vllm-omni by Doug (@dougbtv).
+Original source at https://github.com/dougbtv/comfyui-vllm-omni, distributed under the MIT License.
+"""
+
+import asyncio
+import json
+from typing import Any
+
+import aiohttp
+import av.error
+import torch
+from comfy_api.input import AudioInput, VideoInput
+
+from .format import (
+    audio_to_base64,
+    base64_to_audio,
+    base64_to_image_tensor,
+    bytes_to_audio,
+    bytes_to_video,
+    image_tensor_to_base64,
+    image_tensor_to_png_bytes,
+    video_to_base64,
+)
+from .logger import get_logger, pretty_printer
+from .models import lookup_model_spec
+from .types import AudioFormat
+
+logger = get_logger(__name__)
+
+
+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
+    ):
+        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,
+        *,
+        model: str,
+        prompt: str,
+        width: int,
+        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] = {
+            "model": model,
+            "prompt": prompt,
+            "size": size,
+            "response_format": "b64_json",
+        }
+        if negative_prompt:
+            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"
+        async with aiohttp.ClientSession(timeout=self.timeout) as session:
+            try:
+                async with session.post(
+                    url,
+                    json=payload,
+                    headers={"Content-Type": "application/json"},
+                ) as response:
+                    if not response.ok:
+                        error_text = await response.text()
+                        raise (ValueError if response.status < 500 else RuntimeError)(
+                            f"vLLM-Omni API returned status {response.status}: {error_text}"
+                        )
+
+                    try:
+                        data = await response.json()
+                    except aiohttp.ContentTypeError as e:
+                        raise RuntimeError(f"Invalid JSON response from vLLM-Omni: {e}")
+                    if "data" not in data:
+                        raise RuntimeError("API response missing 'data' field - expected OpenAI DALL-E format")
+                    if not data["data"]:
+                        raise RuntimeError("API returned empty data array")
+
+                    image_tensors = []
+                    for idx, img in enumerate(data["data"]):
+                        if "b64_json" not in img:
+                            raise RuntimeError(f"API returned image #{idx} without 'b64_json' field")
+                        base64_str = img["b64_json"]
+                        tensor = base64_to_image_tensor(base64_str)
+                        image_tensors.append(tensor)
+                        logger.debug("Image #%d has shape %s", idx, tensor.shape)
+
+                    batch_tensor = torch.stack(image_tensors, dim=0)
+                    logger.debug("batch_tensor output has shape: %s", batch_tensor.shape)
+                    return batch_tensor
+
+            except aiohttp.ClientError as e:
+                raise RuntimeError(f"Network error connecting to vLLM-Omni at {url}: {e}")
+
+    async def edit_image(
+        self,
+        *,
+        model: str,
+        prompt: str,
+        image: torch.Tensor,
+        width: int,
+        height: int,
+        negative_prompt: str | None = None,
+        mask: torch.Tensor | None = None,
+        sampling_params: dict | None = None,
+        lora: dict | None = None,
+    ) -> torch.Tensor:
+        """Run image editing via DALLE API"""
+        await self._check_model_exist(model)
+
+        size = f"{width}x{height}"
+        image_filename = "image.png"  # Required for multipart form
+        form = aiohttp.FormData()
+        form.add_field("model", model)
+        form.add_field(
+            "image",
+            image_tensor_to_png_bytes(image, image_filename),
+            filename=image_filename,
+            content_type="image/png",
+        )
+        form.add_field("prompt", prompt)
+        form.add_field("size", size)
+        if negative_prompt:
+            form.add_field("negative_prompt", negative_prompt)
+        if sampling_params is not None:
+            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(
+                "mask",
+                image_tensor_to_png_bytes(mask, mask_filename),
+                filename=mask_filename,
+                content_type="image/png",
+            )
+
+        url = self.base_url + "/images/edits"
+        async with aiohttp.ClientSession(timeout=self.timeout) as session:
+            try:
+                async with session.post(url, data=form) as response:
+                    if not response.ok:
+                        error_text = await response.text()
+                        raise (ValueError if response.status < 500 else RuntimeError)(
+                            f"vLLM-Omni API returned status {response.status}: {error_text}"
+                        )
+
+                    try:
+                        data = await response.json()
+                    except aiohttp.ContentTypeError as e:
+                        raise RuntimeError(f"Invalid JSON response from vLLM-Omni: {e}")
+
+                    if "data" not in data:
+                        raise RuntimeError("API response missing 'data' field - expected OpenAI DALL-E format")
+                    if not data["data"]:
+                        raise RuntimeError("API returned empty data array")
+
+                    image_tensors = []
+                    for idx, img in enumerate(data["data"]):
+                        if "b64_json" not in img:
+                            raise RuntimeError(f"API returned image #{idx} without 'b64_json' field")
+                        base64_str = img["b64_json"]
+                        tensor = base64_to_image_tensor(base64_str)
+                        image_tensors.append(tensor)
+
+                    return torch.stack(image_tensors, dim=0)
+
+            except aiohttp.ClientError as e:
+                raise RuntimeError(f"Network error connecting to vLLM-Omni at {url}: {e}")
+
+    async def generate_image_chat_completion(
+        self,
+        *,
+        model: str,
+        prompt: str,
+        negative_prompt: str | None = None,
+        image: torch.Tensor | None = None,
+        audio: AudioInput | None = None,
+        video: VideoInput | None = None,
+        sampling_params: dict | list[dict] | None = None,
+        lora: dict | None = None,
+    ) -> torch.Tensor:
+        payload = VLLMOmniClient._prepare_chat_completion_messages(
+            model=model,
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            image=image,
+            audio=audio,
+            video=video,
+            sampling_params=sampling_params,
+            modalities=["image"],
+            # === below are additional `extra_body` fields, handled by **kwargs ===
+            lora=lora,
+        )
+        choices = await self._generate_base_chat_completion(model, payload)
+
+        image_tensors = []
+        for idx, img_content in enumerate(choices[0]["message"]["content"]):
+            base64_str = img_content.get("image_url", {}).get("url", "")
+            if not base64_str:
+                raise RuntimeError(f"API returned image #{idx} without image url")
+            tensor = base64_to_image_tensor(base64_str)
+            image_tensors.append(tensor)
+
+        return torch.stack(image_tensors, dim=0)
+
+    async def generate_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,
+        *,
+        model: str,
+        prompt: str,
+        image: torch.Tensor | None = None,
+        audio: AudioInput | None = None,
+        video: VideoInput | None = None,
+        sampling_params: dict | list[dict] | None = None,
+        modalities: list[str] = ["text", "audio"],
+        **extra_body,
+    ) -> tuple[str | None, AudioInput | None]:
+        # Response may contain two choices: one with text, one with audio
+        payload = VLLMOmniClient._prepare_chat_completion_messages(
+            model=model,
+            prompt=prompt,
+            negative_prompt=None,
+            image=image,
+            audio=audio,
+            video=video,
+            sampling_params=sampling_params,
+            modalities=modalities,
+            **extra_body,
+        )
+
+        choices = await self._generate_base_chat_completion(model, payload)
+        text_response = None
+        audio_base64 = None
+        for choice in choices:
+            try:
+                text_response = choice["message"]["content"]
+            except (KeyError, TypeError):
+                # Either this case (text response) or the audio response case will be hit. Checking None's later.
+                pass
+            try:
+                audio_base64 = choice["message"]["audio"]["data"]
+            except (KeyError, TypeError):
+                # Either this case (text response) or the audio response case will be hit. Checking None's later.
+                pass
+        if audio_base64 is None and text_response is None:
+            raise RuntimeError(
+                "API response missing both '.message.audio' and 'message.content' fields."
+                f"The choices object is {choices}"
+            )
+        if audio_base64 is not None:
+            audio = base64_to_audio(audio_base64)
+            logger.debug(
+                "audio sample rate %d, audio shape %s, duration in second %f",
+                audio["sample_rate"],
+                audio["waveform"].shape,
+                audio["waveform"].shape[2] / audio["sample_rate"],
+            )
+        else:
+            audio = None
+        return text_response, audio
+
+    async def generate_speech(
+        self,
+        *,
+        model: str,
+        input: str,
+        voice: str,
+        response_format: AudioFormat,
+        speed: float,
+        **extra_params,
+    ) -> AudioInput:
+        await self._check_model_exist(model)
+
+        ref_audio: AudioInput | None = extra_params.pop("ref_audio", None)
+
+        payload = {
+            "model": model,
+            "input": input,
+            "voice": voice,
+            "response_format": response_format,
+            "speed": speed,
+            **extra_params,
+        }
+
+        if ref_audio is not None:
+            audio_base64 = audio_to_base64(ref_audio)
+            payload["ref_audio"] = audio_base64
+
+        logger.debug("Omni TTS payload: %s", pretty_printer.pformat(payload))
+
+        url = self.base_url + "/audio/speech"
+        async with aiohttp.ClientSession(timeout=self.timeout) as session:
+            try:
+                async with session.post(
+                    url,
+                    json=payload,
+                    headers={"Content-Type": "application/json"},
+                ) as response:
+                    if not response.ok:
+                        error_text = await response.text()
+                        raise (ValueError if response.status < 500 else RuntimeError)(
+                            f"vLLM-Omni API returned status {response.status}: {error_text}"
+                        )
+
+                    try:
+                        audio_bytes = await response.read()
+                    except aiohttp.ContentTypeError as e:
+                        raise RuntimeError(f"Invalid JSON response from vLLM-Omni: {e}")
+
+                    try:
+                        audio = bytes_to_audio(audio_bytes)
+                    except av.error.InvalidDataError as e:
+                        raise ValueError(
+                            f"Invalid audio data received from vLLM-Omni: {e}"
+                            "Check if you have input unsupported arguments (such as 'voice')"
+                        )
+                    return audio
+
+            except aiohttp.ClientError as e:
+                raise RuntimeError(f"Network error connecting to vLLM-Omni at {url}: {e}")
+
+    async def _generate_base_chat_completion(self, model: str, payload: dict[str, Any]) -> list[dict[str, Any]]:
+        logger.debug("Omni payload: %s", pretty_printer.pformat(payload))
+        await self._check_model_exist(model)
+
+        url = self.base_url + "/chat/completions"
+        async with aiohttp.ClientSession(timeout=self.timeout) as session:
+            try:
+                async with session.post(
+                    url,
+                    json=payload,
+                    headers={"Content-Type": "application/json"},
+                ) as response:
+                    if not response.ok:
+                        error_text = await response.text()
+                        raise (ValueError if response.status < 500 else RuntimeError)(
+                            f"vLLM-Omni API returned status {response.status}: {error_text}"
+                        )
+
+                    try:
+                        data = await response.json()
+                    except aiohttp.ContentTypeError as e:
+                        raise RuntimeError(f"Invalid JSON response from vLLM-Omni: {e}")
+
+                    logger.debug(
+                        "chat completion response: %s",
+                        pretty_printer.pformat(data),
+                    )
+
+                    try:
+                        return data["choices"]
+                    except (KeyError, TypeError):
+                        raise RuntimeError("Invalid JSON response from vLLM-Omni: missing 'choices' field")
+
+            except aiohttp.ClientError as e:
+                raise RuntimeError(f"Network error connecting to vLLM-Omni at {self.base_url}: {e}")
+
+    async def _check_model_exist(self, model: str):
+        url = self.base_url + "/models"
+        async with aiohttp.ClientSession(timeout=self.timeout) as session:
+            try:
+                async with session.get(
+                    url,
+                    headers={"Content-Type": "application/json"},
+                ) as response:
+                    if not response.ok:
+                        error_text = await response.text()
+                        raise (ValueError if response.status < 500 else RuntimeError)(
+                            f"vLLM-Omni API returned status {response.status} "
+                            f"when getting hosted model list: {error_text}"
+                        )
+
+                    try:
+                        data = await response.json()
+                    except aiohttp.ContentTypeError as e:
+                        raise RuntimeError(f"Invalid JSON response when getting hosted model list from vLLM-Omni: {e}")
+
+            except aiohttp.ClientError as e:
+                raise RuntimeError(f"Network error connecting to vLLM-Omni at {self.base_url}: {e}")
+        try:
+            model_list = data["data"]
+            model_found = next((True for m in model_list if m["id"] == model), False)
+        except (KeyError, TypeError):
+            raise RuntimeError(f"Invalid JSON response of the hosted model list: {data}")
+
+        if not model_found:
+            raise ValueError(f"Model {model} not served at {self.base_url}.")
+
+    @staticmethod
+    def _prepare_chat_completion_messages(
+        *,
+        model: str,
+        prompt: str,
+        negative_prompt: str | None,
+        image: torch.Tensor | None = None,
+        audio: AudioInput | None = None,
+        video: VideoInput | None = None,
+        sampling_params: dict | list[dict] | None = None,
+        modalities: list[str] | None = None,  # diffusion don't have this field
+        **extra_body,
+    ):
+        message_content: list[dict] = [{"type": "text", "text": prompt}]
+        if image is not None:
+            message_content.append(
+                {
+                    "type": "image_url",
+                    "image_url": {"url": image_tensor_to_base64(image)},
+                }
+            )
+        if audio is not None:
+            message_content.append({"type": "audio_url", "audio_url": {"url": audio_to_base64(audio)}})
+        if video is not None:
+            message_content.append({"type": "video_url", "video_url": {"url": video_to_base64(video)}})
+        messages = [{"role": "user", "content": message_content}]
+
+        payload: dict[str, Any] = {"messages": messages, "model": model}
+        if modalities:
+            payload["modalities"] = modalities
+
+        combined_extra_body: dict[str, Any] = {}
+        if sampling_params is not None:
+            spec, _ = lookup_model_spec(model)
+            is_single_sampling_param = isinstance(sampling_params, dict) or len(sampling_params) == 1
+
+            if (spec is None and is_single_sampling_param) or (spec is not None and spec["stages"] == ["diffusion"]):
+                # Diffusion format: extra_body directly contains sampling params.
+                # Validation has already taken care of matching sampling params' types and length. Safe to take [0].
+                # * Use this mode if the model is a simple one-stage diffusion model.
+                # * Fallback to this mode if model is not registered and a single sampling param is provided.
+                sampling_params = sampling_params if isinstance(sampling_params, dict) else sampling_params[0]
+                combined_extra_body: dict[str, Any] = sampling_params.copy()
+                if "n" in combined_extra_body:
+                    combined_extra_body["num_outputs_per_prompt"] = combined_extra_body.pop("n")
+            else:
+                # AR format: the payload has a sampling_params_list field, containing a list.
+                sampling_params_list = sampling_params if isinstance(sampling_params, list) else [sampling_params]
+                payload["sampling_params_list"] = sampling_params_list
+
+        if negative_prompt:
+            combined_extra_body["negative_prompt"] = negative_prompt
+
+        if extra_body:
+            combined_extra_body.update(extra_body)
+
+        # Add extra_body only if it has any content.
+        if combined_extra_body:
+            payload["extra_body"] = combined_extra_body
+
+        # Place to inject any model-specific payload adjustment
+        spec, _ = lookup_model_spec(model)
+        if spec:
+            preprocessor = spec.get("payload_preprocessor", None)
+            if preprocessor is not None:
+                payload = preprocessor(payload)
+
+        return payload

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

@@ -0,0 +1,304 @@
+"""Image/tensor format helpers.
+
+The image generation part is derived from dougbtv/comfyui-vllm-omni by Doug (@dougbtv).
+Original source at https://github.com/dougbtv/comfyui-vllm-omni, distributed under the MIT License.
+"""
+
+import base64
+import mimetypes
+from 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
+
+from .logger import get_logger
+
+logger = get_logger(__name__)
+
+
+def base64_to_image_tensor(base64_str: str, mode: str = "RGB") -> torch.Tensor:
+    """
+    Convert base64-encoded image to ComfyUI image tensor.
+
+    Args:
+        base64_str: Base64-encoded image string
+        mode: PIL image mode (default RGB for transparency support)
+
+    Returns:
+        torch.Tensor with shape (1, H, W, C) in float32 [0, 1] range
+
+    Raises:
+        ValueError: If base64 string is invalid or image cannot be decoded
+    """
+    if base64_str.startswith("data:image"):
+        _, base64_str = base64_str.split(",", 1)
+
+    try:
+        # Decode base64 to bytes
+        image_bytes = base64.b64decode(base64_str)
+    except Exception as e:
+        raise ValueError(f"Invalid base64 string: {e}")
+
+    # Create BytesIO object for PIL
+    image_bytesio = BytesIO(image_bytes)
+
+    # Open with PIL and convert to desired mode
+    try:
+        pil_image = Image.open(image_bytesio)
+        pil_image = pil_image.convert(mode)
+    except Exception as e:
+        raise RuntimeError(f"Failed to open image: {e}")
+
+    image_array = np.asarray(pil_image).astype(np.float32) / 255.0
+    image_tensor = torch.from_numpy(image_array)
+    return image_tensor
+
+
+def image_tensor_to_png_bytes(tensor: torch.Tensor, filename: str = "image.png") -> BytesIO:
+    """
+    Convert ComfyUI image tensor to PNG BytesIO for multipart upload.
+
+    This function converts a ComfyUI IMAGE tensor to a PNG-encoded BytesIO object
+    suitable for multipart/form-data upload. The BytesIO object has its .name
+    attribute set, which is required by aiohttp for file uploads.
+
+    Args:
+        tensor: ComfyUI IMAGE tensor with shape (B, H, W, C), dtype float32, range [0, 1]
+        filename: Name attribute to set on BytesIO (default: "image.png")
+
+    Returns:
+        BytesIO object containing PNG-encoded image with .name attribute set
+
+    Raises:
+        ValueError: If tensor format is invalid (not 4D, wrong dtype, etc.)
+    """
+    if tensor.ndim != 4:
+        raise ValueError(f"Expected 4D tensor with shape (B, H, W, C), got {tensor.ndim}D tensor")
+
+    image_tensor = tensor[0]  # Shape: (H, W, C)
+    image_np = (image_tensor.cpu().numpy() * 255.0).astype(np.uint8)
+    pil_image = Image.fromarray(image_np)
+
+    # Save to BytesIO as image file
+    img_bytes = BytesIO()
+    # Set name attribute (required for multipart upload and mimetype detection)
+    img_bytes.name = filename
+    try:
+        pil_image.save(img_bytes)
+    except Exception as e:
+        raise RuntimeError(f"Failed to save image as file: {e}")
+
+    # Reset position to beginning
+    img_bytes.seek(0)
+
+    return img_bytes
+
+
+def image_tensor_to_base64(tensor: torch.Tensor, filename: str = "image.png") -> str:
+    """
+    Convert ComfyUI image tensor to base64-encoded image string.
+
+    Args:
+        tensor: ComfyUI IMAGE tensor with shape (B, H, W, C), dtype float32, range [0, 1]
+        filename: Name attribute to set on BytesIO (default: "image.png")
+        format: File format of the output image file buffer (default: "PNG")
+
+    Returns:
+        Base64-encoded image string
+
+    Raises:
+        ValueError: If tensor format is invalid (not 4D, wrong dtype, etc.)
+    """
+    img_bytes = image_tensor_to_png_bytes(tensor, filename)
+    img_bytes.seek(0)
+    byte_data = img_bytes.read()
+    base64_str = base64.b64encode(byte_data).decode("utf-8")
+    mime_type = mimetypes.guess_type(filename)[0] or "application/octet-stream"
+    return f"data:{mime_type};base64,{base64_str}"
+
+
+def video_to_bytes(video: VideoInput, filename: str = "video.mp4") -> BytesIO:
+    output_buffer = BytesIO()
+    output_buffer.name = filename
+    video.save_to(output_buffer)
+    output_buffer.seek(0)
+    return output_buffer
+
+
+def video_to_base64(video: VideoInput, filename: str = "video.mp4") -> str:
+    video_buffer = video_to_bytes(video, filename)
+    video_buffer.seek(0)
+    byte_data = video_buffer.read()
+    base64_str = base64.b64encode(byte_data).decode("utf-8")
+    mime_type = mimetypes.guess_type(filename)[0] or "application/octet-stream"
+    return f"data:{mime_type};base64,{base64_str}"
+
+
+def 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"]
+    format = filename.rsplit(".", maxsplit=1)[1]
+    layout = "mono" if waveform.shape[0] == 1 else "stereo"
+
+    output_buffer = BytesIO()
+    output_buffer.name = filename
+    output_container = av.open(output_buffer, mode="w", format=format)
+    if format == "opus":
+        out_stream = output_container.add_stream("libopus", rate=sample_rate, layout=layout)
+        if quality == "64k":
+            out_stream.bit_rate = 64000  # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py
+        elif quality == "96k":
+            out_stream.bit_rate = 96000  # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py
+        elif quality == "128k":
+            out_stream.bit_rate = 128000  # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py
+        elif quality == "192k":
+            out_stream.bit_rate = 192000  # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py
+        elif quality == "320k":
+            out_stream.bit_rate = 320000  # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py
+    elif format == "mp3":
+        out_stream = output_container.add_stream("libmp3lame", rate=sample_rate, layout=layout)
+        if quality == "V0":
+            out_stream.codec_context.qscale = 1  # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py
+        elif quality == "128k":
+            out_stream.bit_rate = 128000  # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py
+        elif quality == "320k":
+            out_stream.bit_rate = 320000  # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py
+    else:  # format == "flac":
+        out_stream = output_container.add_stream("flac", rate=sample_rate, layout=layout)
+
+    frame = av.AudioFrame.from_ndarray(
+        waveform.movedim(0, 1).reshape(1, -1).float().numpy(),
+        format="flt",
+        layout=layout,
+    )
+    frame.sample_rate = sample_rate
+    frame.pts = 0
+    output_container.mux(out_stream.encode(frame))  # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py
+    # Flush encoder
+    output_container.mux(out_stream.encode(None))  # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py
+    output_container.close()
+    output_buffer.seek(0)
+
+    return output_buffer
+
+
+def audio_to_base64(audio: AudioInput, filename: str = "audio.mp3", quality: str = "128k") -> str:
+    audio_buffer = audio_to_bytes(audio, filename, quality)
+    audio_buffer.seek(0)
+    byte_data = audio_buffer.read()
+    base64_str = base64.b64encode(byte_data).decode("utf-8")
+    mime_type = mimetypes.guess_type(filename)[0] or "application/octet-stream"
+    return f"data:{mime_type};base64,{base64_str}"
+
+
+def bytes_to_audio(audio_bytes: bytes) -> AudioInput:
+    """
+    Convert audio bytes to ComfyUI audio tensor.
+
+    Args:
+        audio_bytes: Audio file bytes
+    Returns:
+        torch.Tensor with shape (B, C, T) in float32 range [-1, 1]
+    """
+    audio_buffer = BytesIO(audio_bytes)
+    waveform, sample_rate = nodes_audio.load(audio_buffer)  # type: ignore # Although expect string argument, it calls av.open underneath, which supports BytesIO (file-like)
+    return {"waveform": waveform.unsqueeze(0), "sample_rate": sample_rate}
+
+
+def base64_to_audio(base64_str: str) -> AudioInput:
+    """
+    Convert base64-encoded audio to ComfyUI audio tensor.
+
+    Args:
+        base64_str: Base64-encoded audio string
+    Returns:
+        torch.Tensor with shape (B, C, T) in float32 range [-1, 1]
+    """
+    if base64_str.startswith("data:audio"):
+        _, base64_str = base64_str.split(",", 1)
+
+    try:
+        # Decode base64 to bytes
+        audio_bytes = base64.b64decode(base64_str)
+    except Exception as e:
+        raise ValueError(f"Invalid base64 string: {e}")
+
+    return bytes_to_audio(audio_bytes)

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

@@ -0,0 +1,125 @@
+"""Centralized logger configuration for vLLM-Omni ComfyUI."""
+
+import logging
+import pprint
+import sys
+from typing import Any
+
+
+def get_logger(name: str) -> logging.Logger:
+    """
+    Get or create a logger with proper formatting.
+
+    Args:
+        name: Logger name (typically __name__ of the calling module)
+
+    Returns:
+        Configured logger instance
+    """
+    logger = logging.getLogger(name)
+
+    # Only configure if not already configured
+    if not logger.handlers:
+        logger.setLevel(logging.DEBUG)
+
+        # Create console handler
+        handler = logging.StreamHandler(sys.stdout)
+        handler.setLevel(logging.INFO)
+
+        # Create formatter
+        formatter = logging.Formatter(
+            fmt="(ComfyUI-vLLM-Omni) [%(levelname)s] %(asctime)s [%(filename)s:%(lineno)s] %(message)s",
+            datefmt="%Y-%m-%d %H:%M:%S",
+        )
+        handler.setFormatter(formatter)
+
+        # Add handler to logger
+        logger.addHandler(handler)
+
+        # Prevent propagation to root logger
+        logger.propagate = False
+
+    return logger
+
+
+class OmitBase64PrettyPrinter(pprint.PrettyPrinter):
+    """
+    A PrettyPrinter that redacts specific field names with '...'
+    wherever they appear in nested structures.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def _format(self, obj: Any, stream, indent: int, allowance: int, context, level: int) -> None:
+        # Check if this is a dict with redacted keys
+        if isinstance(obj, dict):
+            # Create a copy with redacted values
+            display_obj = {}
+            for key, value in obj.items():
+                if key == "data" or key == "url":
+                    if value.startswith("data:"):
+                        base64_header = value.split(",", 1)[0]
+                        display_obj[key] = f"{base64_header},***"
+                    elif value.startswith("http://") or value.startswith("https://"):
+                        display_obj[key] = value
+                    elif len(value) > 10:
+                        display_obj[key] = f"{value[:10]}***"
+                    else:
+                        display_obj[key] = value
+                else:
+                    display_obj[key] = value
+            obj = display_obj
+
+        # Handle list/tuple/set containing dicts that might have redacted keys
+        # (pprint will recursively call _format on nested items, so this
+        # handles arbitrary nesting automatically)
+
+        super()._format(obj, stream, indent, allowance, context, level)
+
+
+pretty_printer = OmitBase64PrettyPrinter()
+
+
+# ========== PPrint EXAMPLES ==========
+
+if __name__ == "__main__":
+    data = {
+        "messages": [
+            {
+                "role": "system",
+                "content": (
+                    "You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group,"
+                    "capable of perceiving auditory and visual inputs, as well as generating text and speech."
+                ),
+            },
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "text",
+                        "text": "What sound is it, and what is the drawing about?",
+                    },
+                    {
+                        "type": "text",
+                        "text": "What sound is it, and what is the drawing about?",
+                    },
+                    {
+                        "type": "image_url",
+                        "image_url": {"url": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAwAAAA"},
+                    },
+                    {
+                        "type": "audio_url",
+                        "audio_url": {"url": "data:audio/mpeg;base64,SUQzBAAAAAAAIlRTU0UAAAAOAAADT="},
+                    },
+                ],
+            },
+        ],
+        "extra_body": {"mm_processor_kwargs": {"use_audio_in_video": False}},
+        "modalities": ["text"],
+    }
+
+    # Create printer that redacts 'password' and 'token' fields
+
+    print("\nRedactingPrettyPrinter (hides secrets):")
+    pretty_printer.pprint(data)

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

@@ -0,0 +1,118 @@
+import re
+
+from .types import Modality, ModelMode, Spec
+
+
+def _bagel_payload_preprocessor(payload: dict) -> dict:
+    try:
+        for message in payload["messages"]:
+            for content in message["content"]:
+                if content["type"] == "text":
+                    content["text"] = "<|im_start|>" + content["text"] + "<|im_end|>"
+    except (KeyError, TypeError):
+        raise RuntimeError("Internal Error: malformatted BAGEL payload")
+    return payload
+
+
+def _qwen25_payload_preprocessor(payload: dict) -> dict:
+    if payload["messages"][0]["role"] != "system":
+        payload["messages"] = [
+            {
+                "role": "system",
+                "content": (
+                    "You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group,"
+                    "capable of perceiving auditory and visual inputs, as well as generating text and speech."
+                ),
+            },
+            *payload["messages"],
+        ]
+    return payload
+
+
+_MODEL_PIPELINE_SPECS: dict[str, Spec] = {
+    r"BAGEL-7B-MoT": {
+        "stages": [
+            "diffusion"  # The vLLM-Omni interface treats it as a single-stage diffusion model
+        ],
+        "modes": [
+            {
+                "mode": ModelMode.UNDERSTANDING,
+                "input_modalities": [Modality.TEXT, Modality.IMAGE],
+            }
+        ],
+        "payload_preprocessor": _bagel_payload_preprocessor,
+    },
+    r"Qwen2.5-Omni*": {
+        "stages": ["autoregression", "autoregression", "autoregression"],
+        "payload_preprocessor": _qwen25_payload_preprocessor,
+        "modes": [
+            {
+                "mode": ModelMode.UNDERSTANDING,
+                "input_modalities": [
+                    Modality.TEXT,
+                    Modality.IMAGE,
+                    Modality.VIDEO,
+                    Modality.AUDIO,
+                ],
+            }
+        ],
+    },
+    r"Qwen3-Omni*": {
+        "stages": ["autoregression", "autoregression", "autoregression"],
+        "modes": [
+            {
+                "mode": ModelMode.UNDERSTANDING,
+                "input_modalities": [
+                    Modality.TEXT,
+                    Modality.IMAGE,
+                    Modality.VIDEO,
+                    Modality.AUDIO,
+                ],
+            }
+        ],
+    },
+}
+# Convert dict keys to regex patterns
+MODEL_PIPELINE_SPECS: dict[re.Pattern, Spec] = {}
+for k, v in _MODEL_PIPELINE_SPECS.items():
+    MODEL_PIPELINE_SPECS[re.compile(k)] = v
+del _MODEL_PIPELINE_SPECS
+
+
+def lookup_model_spec(model: str) -> tuple[Spec | None, str | None]:
+    try:
+        last_component = model.rstrip("/").rsplit("/", 1)[-1]
+    except IndexError:
+        last_component = model
+    for pattern, spec in MODEL_PIPELINE_SPECS.items():
+        if pattern.search(last_component):
+            return spec, pattern.pattern
+    return None, None
+
+
+# ============== DEMONSTRATION ==============
+
+if __name__ == "__main__":
+    test_paths = [
+        "Qwen/Qwen2.5-Omni-7B",
+        "MyModels/Qwen2.5-Omni-3B",
+        "/root/home/Qwen2.5-Omni-7B",
+        "Qwen/Qwen3-Omni",
+        "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+        "Custom/Path/UnknownModel-Instruct",
+        "Not/Matching/Anything",
+    ]
+
+    test_payload = {"messages": [{"role": "user", "content": "prompt"}]}
+
+    print("Testing registry lookups:\n")
+    for path in test_paths:
+        spec, _ = lookup_model_spec(path)
+        if spec:
+            if preprocessor := spec.get("payload_preprocessor"):
+                result = preprocessor(test_payload)
+                print(f"✓ {path:<40} → {result}")
+            else:
+                print(f"✓ {path:<40} → No preprocessor")
+        else:
+            print(f"✗ {path:<40} → No match")

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

@@ -0,0 +1,55 @@
+from collections.abc import Callable
+from enum import Enum, auto
+from typing import (
+    Any,
+    Literal,
+    NotRequired,
+    TypeAlias,
+    TypedDict,
+)
+
+AudioFormat: TypeAlias = Literal["mp3", "opus", "aac", "flac", "wav", "pcm"]
+
+
+class AutoregressionSamplingParams(dict):
+    pass
+
+
+class DiffusionSamplingParams(dict):
+    pass
+
+
+class QwenTTSModelSpecificParams(dict):
+    pass
+
+
+class WanModelSpecificParams(dict):
+    pass
+
+
+class ModelMode(Enum):
+    IMAGE_GENERATION = auto()
+    VIDEO_GENERATION = auto()
+    AUDIO_GENERATION = auto()
+    UNDERSTANDING = auto()
+
+
+class Modality(Enum):
+    TEXT = auto()  # maybe not useful. Prompt is always required
+    IMAGE = auto()
+    VIDEO = auto()
+    AUDIO = auto()
+
+
+class ModelModeSpec(TypedDict):
+    mode: ModelMode
+    input_modalities: list[Modality]
+
+
+PayloadPreprocessor: TypeAlias = Callable[[dict[str, Any]], dict[str, Any]]
+
+
+class Spec(TypedDict):
+    stages: list[Literal["diffusion", "autoregression"]]
+    modes: list[ModelModeSpec]
+    payload_preprocessor: NotRequired[PayloadPreprocessor]

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

@@ -0,0 +1,93 @@
+from .logger import get_logger
+from .models import lookup_model_spec
+from .types import AutoregressionSamplingParams, DiffusionSamplingParams
+
+logger = get_logger(__name__)
+
+
+def validate_model_and_sampling_params_types(
+    model_name: str,
+    sampling_param_list: dict | list[dict] | None = None,
+):
+    # Check if model name exists
+    if not model_name:
+        raise ValueError("Model name must not be empty.")
+
+    # Skip if no spec or no sampling params
+    pipeline_spec, _ = lookup_model_spec(model_name)
+    if pipeline_spec is None:
+        logger.info(f"skipping sampling params check because spec for {model_name} is not found")
+        return
+    if sampling_param_list is None:
+        return
+
+    # Check the number of stages and their data types
+    stages = pipeline_spec["stages"]
+    if isinstance(sampling_param_list, list):
+        # Check that the lengths match
+        if len(stages) != len(sampling_param_list):
+            raise ValueError(
+                f"Sampling parameter list length {len(sampling_param_list)} does not match "
+                f"number of stages {len(stages)} for model {model_name}."
+            )
+        # Check that each stage's type match
+        for i, sp in enumerate(sampling_param_list):
+            if not _check_sampling_param_matches_stage(sp, stages[i]):
+                raise ValueError(
+                    f"Sampling parameter type ({sp.__class__.__name__}) does not match "
+                    f"stage type ({stages[i]}) at index {i} for model {model_name}."
+                )
+    elif isinstance(sampling_param_list, dict):
+        # Check that the provided single sampling param matches all stages
+        for i, stage in enumerate(stages):
+            if not _check_sampling_param_matches_stage(sampling_param_list, stage):
+                raise ValueError(
+                    f"Provided single sampling parameter type ({sampling_param_list.__class__.__name__}) must match "
+                    f"the types of all stages of the model. "
+                    f"However, stage {i} of model {model_name} is of type {stage}."
+                )
+
+
+def add_sampling_parameters_to_stage(
+    model_name: str,
+    sampling_param_list: dict | list[dict] | None,
+    stage_type: str,
+    /,
+    **params_to_add,
+) -> dict | list[dict]:
+    """
+    Given a model's name and the sampling parameter list to query this model,
+    add arbitrary additional parameters to the sampling parameters of all stages of the given type.
+    """
+    pipeline_spec, _ = lookup_model_spec(model_name)
+    if not pipeline_spec:
+        logger.warning(
+            f"Since the model {model_name} is not in our list, we cannot ensure if "
+            f"the fields ({tuple(params_to_add.keys())}) are added to the correct stage's sampling params. "
+            f"We will do it heuristically."
+        )
+        pipeline_spec = {"stages": ["diffusion"]}
+
+    stages = pipeline_spec["stages"]
+    if isinstance(sampling_param_list, dict):
+        sampling_param_list = 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
+
+    return sampling_param_list
+
+
+def _check_sampling_param_matches_stage(sampling_param: dict, stage_type: str) -> bool:
+    if stage_type == "autoregression":
+        return isinstance(sampling_param, AutoregressionSamplingParams)
+    if stage_type == "diffusion":
+        return isinstance(sampling_param, DiffusionSamplingParams)
+    raise RuntimeError(f"Internal error: unknown stage type {stage_type}.")

apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni Chaining Services.json

@@ -0,0 +1,552 @@
+{
+    "id": "6643e5fd-fa2a-4f25-935a-483173a8097c",
+    "revision": 0,
+    "last_node_id": 8,
+    "last_link_id": 6,
+    "nodes": [
+        {
+            "id": 1,
+            "type": "PreviewImage",
+            "pos": [
+                1446.2005205859512,
+                -316.7359686049902
+            ],
+            "size": [
+                305.6628685610174,
+                307.0336884172169
+            ],
+            "flags": {},
+            "order": 5,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "images",
+                    "name": "images",
+                    "type": "IMAGE",
+                    "link": 1
+                }
+            ],
+            "outputs": [],
+            "properties": {
+                "Node name for S&R": "PreviewImage"
+            },
+            "widgets_values": []
+        },
+        {
+            "id": 8,
+            "type": "PreviewImage",
+            "pos": [
+                1010.1848219273943,
+                -311.54216706114363
+            ],
+            "size": [
+                305.6628685610174,
+                307.0336884172169
+            ],
+            "flags": {},
+            "order": 4,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "images",
+                    "name": "images",
+                    "type": "IMAGE",
+                    "link": 6
+                }
+            ],
+            "outputs": [],
+            "properties": {
+                "Node name for S&R": "PreviewImage"
+            },
+            "widgets_values": []
+        },
+        {
+            "id": 7,
+            "type": "VLLMOmniDiffusionSampling",
+            "pos": [
+                234.61958293819296,
+                9.405681270043619
+            ],
+            "size": [
+                284.205078125,
+                226
+            ],
+            "flags": {},
+            "order": 0,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "n",
+                    "name": "n",
+                    "type": "INT",
+                    "widget": {
+                        "name": "n"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "num_inference_steps",
+                    "name": "num_inference_steps",
+                    "type": "INT",
+                    "widget": {
+                        "name": "num_inference_steps"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "guidance_scale",
+                    "name": "guidance_scale",
+                    "type": "FLOAT",
+                    "widget": {
+                        "name": "guidance_scale"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "true_cfg_scale",
+                    "name": "true_cfg_scale",
+                    "type": "FLOAT",
+                    "widget": {
+                        "name": "true_cfg_scale"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "vae_use_slicing",
+                    "name": "vae_use_slicing",
+                    "type": "BOOLEAN",
+                    "widget": {
+                        "name": "vae_use_slicing"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "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": [
+                        5
+                    ]
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "VLLMOmniDiffusionSampling"
+            },
+            "widgets_values": [
+                1,
+                50,
+                1,
+                1,
+                false,
+                false,
+                1525,
+                "randomize"
+            ]
+        },
+        {
+            "id": 4,
+            "type": "VLLMOmniDiffusionSampling",
+            "pos": [
+                666.8380026154548,
+                268.86271068330126
+            ],
+            "size": [
+                284.205078125,
+                226
+            ],
+            "flags": {},
+            "order": 1,
+            "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": [
+                        2
+                    ]
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "VLLMOmniDiffusionSampling"
+            },
+            "widgets_values": [
+                4,
+                50,
+                7,
+                1,
+                false,
+                false,
+                42,
+                "fixed"
+            ]
+        },
+        {
+            "id": 5,
+            "type": "VLLMOmniGenerateImage",
+            "pos": [
+                984.723613585788,
+                63.376900027553276
+            ],
+            "size": [
+                416.56628685610167,
+                372.1662621294205
+            ],
+            "flags": {},
+            "order": 3,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "image",
+                    "name": "image",
+                    "shape": 7,
+                    "type": "IMAGE",
+                    "link": 4
+                },
+                {
+                    "localized_name": "mask",
+                    "name": "mask",
+                    "shape": 7,
+                    "type": "MASK",
+                    "link": null
+                },
+                {
+                    "localized_name": "sampling_params",
+                    "name": "sampling_params",
+                    "shape": 7,
+                    "type": "SAMPLING_PARAMS",
+                    "link": 2
+                },
+                {
+                    "localized_name": "url",
+                    "name": "url",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "url"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "model",
+                    "name": "model",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "model"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "prompt",
+                    "name": "prompt",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "prompt"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "negative_prompt",
+                    "name": "negative_prompt",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "negative_prompt"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "width",
+                    "name": "width",
+                    "type": "INT",
+                    "widget": {
+                        "name": "width"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "height",
+                    "name": "height",
+                    "type": "INT",
+                    "widget": {
+                        "name": "height"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "image",
+                    "name": "image",
+                    "type": "IMAGE",
+                    "links": [
+                        1
+                    ]
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "VLLMOmniGenerateImage"
+            },
+            "widgets_values": [
+                "http://localhost:8001/v1",
+                "/home/models/Qwen/Qwen-Image-Edit",
+                "A high-quality, high contrast, stylized portrait of the object in the uploaded reference image. Pop art and doodle style, with abundant scribbling patterns, such as a teal crown, orange lightning bolts, colorful handwritten scripts.",
+                "Realistic",
+                800,
+                800
+            ]
+        },
+        {
+            "id": 6,
+            "type": "VLLMOmniGenerateImage",
+            "pos": [
+                541.2702717429818,
+                -160.70392744708548
+            ],
+            "size": [
+                416.56628685610167,
+                372.1662621294205
+            ],
+            "flags": {},
+            "order": 2,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "image",
+                    "name": "image",
+                    "shape": 7,
+                    "type": "IMAGE",
+                    "link": null
+                },
+                {
+                    "localized_name": "mask",
+                    "name": "mask",
+                    "shape": 7,
+                    "type": "MASK",
+                    "link": null
+                },
+                {
+                    "localized_name": "sampling_params",
+                    "name": "sampling_params",
+                    "shape": 7,
+                    "type": "SAMPLING_PARAMS",
+                    "link": 5
+                },
+                {
+                    "localized_name": "url",
+                    "name": "url",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "url"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "model",
+                    "name": "model",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "model"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "prompt",
+                    "name": "prompt",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "prompt"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "negative_prompt",
+                    "name": "negative_prompt",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "negative_prompt"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "width",
+                    "name": "width",
+                    "type": "INT",
+                    "widget": {
+                        "name": "width"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "height",
+                    "name": "height",
+                    "type": "INT",
+                    "widget": {
+                        "name": "height"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "image",
+                    "name": "image",
+                    "type": "IMAGE",
+                    "links": [
+                        4,
+                        6
+                    ]
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "VLLMOmniGenerateImage"
+            },
+            "widgets_values": [
+                "http://localhost:8000/v1",
+                "/home/models/Tongyi-MAI/Z-Image-Turbo",
+                "A headshot of a cute Siamese kitty. Blurred background due to wide aperture. Close-up look. Realistic.",
+                "Cartoonish.",
+                800,
+                800
+            ]
+        }
+    ],
+    "links": [
+        [
+            1,
+            5,
+            0,
+            1,
+            0,
+            "IMAGE"
+        ],
+        [
+            2,
+            4,
+            0,
+            5,
+            2,
+            "SAMPLING_PARAMS"
+        ],
+        [
+            4,
+            6,
+            0,
+            5,
+            0,
+            "IMAGE"
+        ],
+        [
+            5,
+            7,
+            0,
+            6,
+            2,
+            "SAMPLING_PARAMS"
+        ],
+        [
+            6,
+            6,
+            0,
+            8,
+            0,
+            "IMAGE"
+        ]
+    ],
+    "groups": [],
+    "config": {},
+    "extra": {
+        "workflowRendererVersion": "LG",
+        "ds": {
+            "scale": 1.0426432563169903,
+            "offset": [
+                -90.4071374799346,
+                575.2948920069742
+            ]
+        }
+    },
+    "version": 0.4
+}

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

@@ -0,0 +1 @@
+{"id":"91f75acc-8040-40f6-865a-2e8a7cfd6672","revision":0,"last_node_id":11,"last_link_id":23,"nodes":[{"id":3,"type":"PreviewImage","pos":[1281.8455167767304,-69.02638461454333],"size":[305.6628685610174,307.0336884172169],"flags":{},"order":4,"mode":0,"inputs":[{"localized_name":"图像","name":"images","type":"IMAGE","link":22}],"outputs":[],"properties":{"Node name for S&R":"PreviewImage"},"widgets_values":[]},{"id":11,"type":"VLLMOmniGenerateImage","pos":[816.0962515873642,-137.51112584854445],"size":[400,278],"flags":{},"order":3,"mode":0,"inputs":[{"localized_name":"image","name":"image","shape":7,"type":"IMAGE","link":21},{"localized_name":"mask","name":"mask","shape":7,"type":"MASK","link":null},{"localized_name":"sampling_params","name":"sampling_params","shape":7,"type":"SAMPLING_PARAMS","link":23},{"localized_name":"url","name":"url","type":"STRING","widget":{"name":"url"},"link":null},{"localized_name":"model","name":"model","type":"STRING","widget":{"name":"model"},"link":null},{"localized_name":"prompt","name":"prompt","type":"STRING","widget":{"name":"prompt"},"link":null},{"localized_name":"negative_prompt","name":"negative_prompt","type":"STRING","widget":{"name":"negative_prompt"},"link":null},{"localized_name":"width","name":"width","type":"INT","widget":{"name":"width"},"link":null},{"localized_name":"height","name":"height","type":"INT","widget":{"name":"height"},"link":null}],"outputs":[{"localized_name":"image","name":"image","type":"IMAGE","links":[22]}],"properties":{"Node name for S&R":"VLLMOmniGenerateImage"},"widgets_values":["http://localhost:8000/v1","Qwen/Qwen-Image-Edit","Put this figure in a realistic mountain view","",512,512]},{"id":10,"type":"MarkdownNote","pos":[227.99306819575278,-231.24143306069843],"size":[240.3326478922474,136.3505820791881],"flags":{},"order":1,"mode":0,"inputs":[],"outputs":[],"title":"Note: Task and Input","properties":{},"widgets_values":["vLLM-Omni nodes are categorized based on the output modality. The \"Generate Image\" node supports both text-to-image generation or image-to-image generation (a.k.a. image editing). The node will route to the correct endpoint depending on whether an input image is present or not."],"color":"#432","bgcolor":"#000"},{"id":4,"type":"LoadImage","pos":[496.31859627609606,-229.71277089860084],"size":[270,314],"flags":{},"order":2,"mode":0,"inputs":[{"localized_name":"图像","name":"image","type":"COMBO","widget":{"name":"image"},"link":null},{"localized_name":"选择文件上传","name":"upload","type":"IMAGEUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"图像","name":"IMAGE","type":"IMAGE","links":[21]},{"localized_name":"遮罩","name":"MASK","type":"MASK","links":null}],"properties":{"Node name for S&R":"LoadImage"},"widgets_values":["example.png","image"]},{"id":8,"type":"VLLMOmniDiffusionSampling","pos":[478.59266934006774,183.67711984955648],"size":[284.205078125,226],"flags":{},"order":0,"mode":0,"inputs":[{"localized_name":"n","name":"n","type":"INT","widget":{"name":"n"},"link":null},{"localized_name":"num_inference_steps","name":"num_inference_steps","type":"INT","widget":{"name":"num_inference_steps"},"link":null},{"localized_name":"guidance_scale","name":"guidance_scale","type":"FLOAT","widget":{"name":"guidance_scale"},"link":null},{"localized_name":"true_cfg_scale","name":"true_cfg_scale","type":"FLOAT","widget":{"name":"true_cfg_scale"},"link":null},{"localized_name":"vae_use_slicing","name":"vae_use_slicing","type":"BOOLEAN","widget":{"name":"vae_use_slicing"},"link":null},{"localized_name":"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":[23]}],"properties":{"Node name for S&R":"VLLMOmniDiffusionSampling"},"widgets_values":[4,50,1,1,false,false,42,"randomize"]}],"links":[[21,4,0,11,0,"IMAGE"],[22,11,0,3,0,"IMAGE"],[23,8,0,11,2,"SAMPLING_PARAMS"]],"groups":[{"id":1,"title":"Input","bounding":[213.1706010087147,-313.20095750667554,560.8246471437027,407.4246532472182],"color":"#3f789e","font_size":24,"flags":{}}],"config":{},"extra":{"workflowRendererVersion":"LG","ds":{"scale":1.1469075819486894,"offset":[-55.18047513099182,220.2553505195962]}},"version":0.4}

apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni Multimodal Understanding.json

@@ -0,0 +1,761 @@
+{
+    "id": "1c99f525-0a37-45ba-a28a-7df7c3af66b4",
+    "revision": 0,
+    "last_node_id": 12,
+    "last_link_id": 14,
+    "nodes": [
+        {
+            "id": 1,
+            "type": "VLLMOmniUnderstanding",
+            "pos": [
+                1191.2177053682556,
+                144.66829928181377
+            ],
+            "size": [
+                400,
+                268
+            ],
+            "flags": {},
+            "order": 8,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "image",
+                    "name": "image",
+                    "shape": 7,
+                    "type": "IMAGE",
+                    "link": 1
+                },
+                {
+                    "localized_name": "video",
+                    "name": "video",
+                    "shape": 7,
+                    "type": "VIDEO",
+                    "link": 2
+                },
+                {
+                    "localized_name": "audio",
+                    "name": "audio",
+                    "shape": 7,
+                    "type": "AUDIO",
+                    "link": 3
+                },
+                {
+                    "localized_name": "sampling_params",
+                    "name": "sampling_params",
+                    "shape": 7,
+                    "type": "SAMPLING_PARAMS",
+                    "link": 14
+                },
+                {
+                    "localized_name": "url",
+                    "name": "url",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "url"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "model",
+                    "name": "model",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "model"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "prompt",
+                    "name": "prompt",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "prompt"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "output_text",
+                    "name": "output_text",
+                    "type": "BOOLEAN",
+                    "widget": {
+                        "name": "output_text"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "output_audio",
+                    "name": "output_audio",
+                    "type": "BOOLEAN",
+                    "widget": {
+                        "name": "output_audio"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "use_audio_in_video",
+                    "name": "use_audio_in_video",
+                    "type": "BOOLEAN",
+                    "widget": {
+                        "name": "use_audio_in_video"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "text_response",
+                    "name": "text_response",
+                    "type": "STRING",
+                    "links": [
+                        8
+                    ]
+                },
+                {
+                    "localized_name": "audio_response",
+                    "name": "audio_response",
+                    "type": "AUDIO",
+                    "links": [
+                        9
+                    ]
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "VLLMOmniUnderstanding"
+            },
+            "widgets_values": [
+                "http://localhost:8000/v1",
+                "Qwen/Qwen2.5-Omni-7B",
+                "",
+                true,
+                true,
+                true
+            ]
+        },
+        {
+            "id": 3,
+            "type": "LoadVideo",
+            "pos": [
+                729.5984141255855,
+                -198.631920454299
+            ],
+            "size": [
+                282.798828125,
+                233.0743408203125
+            ],
+            "flags": {},
+            "order": 0,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "file",
+                    "name": "file",
+                    "type": "COMBO",
+                    "widget": {
+                        "name": "file"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "choose file to upload",
+                    "name": "upload",
+                    "type": "IMAGEUPLOAD",
+                    "widget": {
+                        "name": "upload"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "VIDEO",
+                    "name": "VIDEO",
+                    "type": "VIDEO",
+                    "links": [
+                        2
+                    ]
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "LoadVideo"
+            },
+            "widgets_values": [
+                "draw.mp4",
+                "image"
+            ]
+        },
+        {
+            "id": 4,
+            "type": "LoadAudio",
+            "pos": [
+                729.8037086965753,
+                99.86963519703949
+            ],
+            "size": [
+                282.798828125,
+                136
+            ],
+            "flags": {},
+            "order": 1,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "audio",
+                    "name": "audio",
+                    "type": "COMBO",
+                    "widget": {
+                        "name": "audio"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "audioUI",
+                    "name": "audioUI",
+                    "type": "AUDIO_UI",
+                    "widget": {
+                        "name": "audioUI"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "choose file to upload",
+                    "name": "upload",
+                    "type": "AUDIOUPLOAD",
+                    "widget": {
+                        "name": "upload"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "AUDIO",
+                    "name": "AUDIO",
+                    "type": "AUDIO",
+                    "links": [
+                        3
+                    ]
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "LoadAudio"
+            },
+            "widgets_values": [
+                "Megan-Fox.mp3",
+                null,
+                null
+            ]
+        },
+        {
+            "id": 5,
+            "type": "VLLMOmniARSampling",
+            "pos": [
+                510.3517536642828,
+                658.073751009259
+            ],
+            "size": [
+                270,
+                178
+            ],
+            "flags": {},
+            "order": 2,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "max_tokens",
+                    "name": "max_tokens",
+                    "type": "INT",
+                    "widget": {
+                        "name": "max_tokens"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "temperature",
+                    "name": "temperature",
+                    "type": "FLOAT",
+                    "widget": {
+                        "name": "temperature"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "top_p",
+                    "name": "top_p",
+                    "type": "FLOAT",
+                    "widget": {
+                        "name": "top_p"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "repetition_penalty",
+                    "name": "repetition_penalty",
+                    "type": "FLOAT",
+                    "widget": {
+                        "name": "repetition_penalty"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "seed",
+                    "name": "seed",
+                    "type": "INT",
+                    "widget": {
+                        "name": "seed"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "AR sampling params",
+                    "name": "AR sampling params",
+                    "type": "SAMPLING_PARAMS",
+                    "links": [
+                        12
+                    ]
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "VLLMOmniARSampling"
+            },
+            "widgets_values": [
+                100,
+                1,
+                1,
+                1,
+                -1,
+                "randomize"
+            ]
+        },
+        {
+            "id": 7,
+            "type": "VLLMOmniARSampling",
+            "pos": [
+                503.33235181647115,
+                419.34158016181806
+            ],
+            "size": [
+                270,
+                178
+            ],
+            "flags": {},
+            "order": 3,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "max_tokens",
+                    "name": "max_tokens",
+                    "type": "INT",
+                    "widget": {
+                        "name": "max_tokens"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "temperature",
+                    "name": "temperature",
+                    "type": "FLOAT",
+                    "widget": {
+                        "name": "temperature"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "top_p",
+                    "name": "top_p",
+                    "type": "FLOAT",
+                    "widget": {
+                        "name": "top_p"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "repetition_penalty",
+                    "name": "repetition_penalty",
+                    "type": "FLOAT",
+                    "widget": {
+                        "name": "repetition_penalty"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "seed",
+                    "name": "seed",
+                    "type": "INT",
+                    "widget": {
+                        "name": "seed"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "AR sampling params",
+                    "name": "AR sampling params",
+                    "type": "SAMPLING_PARAMS",
+                    "links": [
+                        5,
+                        13
+                    ]
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "VLLMOmniARSampling"
+            },
+            "widgets_values": [
+                100,
+                1,
+                1,
+                1,
+                -1,
+                "randomize"
+            ]
+        },
+        {
+            "id": 8,
+            "type": "VLLMOmniSamplingParamsList",
+            "pos": [
+                820.6056617389042,
+                426.38372037182273
+            ],
+            "size": [
+                263.066015625,
+                66
+            ],
+            "flags": {},
+            "order": 7,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "param1",
+                    "name": "param1",
+                    "type": "SAMPLING_PARAMS",
+                    "link": 5
+                },
+                {
+                    "localized_name": "param2",
+                    "name": "param2",
+                    "shape": 7,
+                    "type": "SAMPLING_PARAMS",
+                    "link": 12
+                },
+                {
+                    "localized_name": "param3",
+                    "name": "param3",
+                    "shape": 7,
+                    "type": "SAMPLING_PARAMS",
+                    "link": 13
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "param list",
+                    "name": "param list",
+                    "type": "SAMPLING_PARAMS",
+                    "links": [
+                        14
+                    ]
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "VLLMOmniSamplingParamsList"
+            },
+            "widgets_values": []
+        },
+        {
+            "id": 11,
+            "type": "MarkdownNote",
+            "pos": [
+                826.2328280438272,
+                569.1890318701705
+            ],
+            "size": [
+                333.8220435590464,
+                261.63596728060656
+            ],
+            "flags": {},
+            "order": 4,
+            "mode": 0,
+            "inputs": [],
+            "outputs": [],
+            "title": "Note: Sampling Parameters",
+            "properties": {},
+            "widgets_values": [
+                "## Sampling Parameter Types\n\nThere are two types of sampling parameters: one for autoregression and one for diffusion.\nYou should ensure that you have chosen the correct type of sampling parameters for the model you request.\n\n## Stages & Shorthand\n\nFor multi-stage models such as Qwen Omni, you can either\n- connect one sampling parameter node, which is applied to all stages.\n- connect exactly the same number of sampling parameter nodes to a \"Multi-Stage Sampling Parameter List\" node, then connect this node to the primary request node.\n\nNote that this shorthand is intended to stay consistent with the [online serving API](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/examples/online_serving/qwen2_5_omni/)"
+            ],
+            "color": "#432",
+            "bgcolor": "#000"
+        },
+        {
+            "id": 12,
+            "type": "MarkdownNote",
+            "pos": [
+                378.9866207003777,
+                152.59550252215752
+            ],
+            "size": [
+                319.7287574247016,
+                107.15904081906785
+            ],
+            "flags": {},
+            "order": 6,
+            "mode": 0,
+            "inputs": [],
+            "outputs": [],
+            "title": "Note: Input",
+            "properties": {},
+            "widgets_values": [
+                "Note that not all models support every modality as input. For example, `ByteDance-Seed/BAGEL-7B-MoT` in Multimodality Understanding mode only support text and image input.\n\nYou should ensure that the input are supported by the model. You can check the corresponding [online serving documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/examples/online_serving/bagel/) for confirmation."
+            ],
+            "color": "#432",
+            "bgcolor": "#000"
+        },
+        {
+            "id": 2,
+            "type": "LoadImage",
+            "pos": [
+                394.4674804308822,
+                -207.6987397548834
+            ],
+            "size": [
+                282.798828125,
+                314
+            ],
+            "flags": {},
+            "order": 5,
+            "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": [
+                        1
+                    ]
+                },
+                {
+                    "localized_name": "MASK",
+                    "name": "MASK",
+                    "type": "MASK",
+                    "links": null
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "LoadImage"
+            },
+            "widgets_values": [
+                "example.png",
+                "image"
+            ]
+        },
+        {
+            "id": 10,
+            "type": "PreviewAudio",
+            "pos": [
+                1664.548345556043,
+                297.5921292054054
+            ],
+            "size": [
+                270,
+                88
+            ],
+            "flags": {},
+            "order": 10,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "audio",
+                    "name": "audio",
+                    "type": "AUDIO",
+                    "link": 9
+                },
+                {
+                    "localized_name": "audioUI",
+                    "name": "audioUI",
+                    "type": "AUDIO_UI",
+                    "widget": {
+                        "name": "audioUI"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [],
+            "properties": {
+                "Node name for S&R": "PreviewAudio"
+            },
+            "widgets_values": []
+        },
+        {
+            "id": 9,
+            "type": "ShowText|pysssss",
+            "pos": [
+                1649.2506875091847,
+                66.22823888292349
+            ],
+            "size": [
+                318.7188464232943,
+                173.38502269972975
+            ],
+            "flags": {},
+            "order": 9,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "text",
+                    "name": "text",
+                    "type": "STRING",
+                    "link": 8
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "STRING",
+                    "name": "STRING",
+                    "shape": 6,
+                    "type": "STRING",
+                    "links": null
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "ShowText|pysssss"
+            },
+            "widgets_values": []
+        }
+    ],
+    "links": [
+        [
+            1,
+            2,
+            0,
+            1,
+            0,
+            "IMAGE"
+        ],
+        [
+            2,
+            3,
+            0,
+            1,
+            1,
+            "VIDEO"
+        ],
+        [
+            3,
+            4,
+            0,
+            1,
+            2,
+            "AUDIO"
+        ],
+        [
+            5,
+            7,
+            0,
+            8,
+            0,
+            "SAMPLING_PARAMS"
+        ],
+        [
+            8,
+            1,
+            0,
+            9,
+            0,
+            "STRING"
+        ],
+        [
+            9,
+            1,
+            1,
+            10,
+            0,
+            "AUDIO"
+        ],
+        [
+            12,
+            5,
+            0,
+            8,
+            1,
+            "SAMPLING_PARAMS"
+        ],
+        [
+            13,
+            7,
+            0,
+            8,
+            2,
+            "SAMPLING_PARAMS"
+        ],
+        [
+            14,
+            8,
+            0,
+            1,
+            3,
+            "SAMPLING_PARAMS"
+        ]
+    ],
+    "groups": [
+        {
+            "id": 1,
+            "title": "Sampling Parameters",
+            "bounding": [
+                480.1649301181556,
+                341.08402513937995,
+                692.4113568972277,
+                510.48648853403665
+            ],
+            "color": "#3f789e",
+            "font_size": 24,
+            "flags": {}
+        },
+        {
+            "id": 2,
+            "title": "Input",
+            "bounding": [
+                344.2364817813528,
+                -287.5850484183313,
+                704.8832238768634,
+                559.5124009832894
+            ],
+            "color": "#3f789e",
+            "font_size": 24,
+            "flags": {}
+        }
+    ],
+    "config": {},
+    "extra": {
+        "workflowRendererVersion": "LG",
+        "ds": {
+            "scale": 0.9478575057427204,
+            "offset": [
+                33.81037136029557,
+                307.1974296197726
+            ]
+        }
+    },
+    "version": 0.4
+}

apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni TTS.json

@@ -0,0 +1 @@
+{"id":"c0a11241-f0e6-4f5c-8488-b3d344a9df3d","revision":0,"last_node_id":7,"last_link_id":3,"nodes":[{"id":4,"type":"PreviewAudio","pos":[690.1759063763212,226.76444957981545],"size":[270,88],"flags":{},"order":5,"mode":0,"inputs":[{"localized_name":"audio","name":"audio","type":"AUDIO","link":2},{"localized_name":"audioUI","name":"audioUI","type":"AUDIO_UI","widget":{"name":"audioUI"},"link":null}],"outputs":[],"properties":{"Node name for S&R":"PreviewAudio"},"widgets_values":[]},{"id":5,"type":"LoadAudio","pos":[-83.61093919682321,103.80510898465587],"size":[282.798828125,136],"flags":{},"order":0,"mode":0,"inputs":[{"localized_name":"audio","name":"audio","type":"COMBO","widget":{"name":"audio"},"link":null},{"localized_name":"audioUI","name":"audioUI","type":"AUDIO_UI","widget":{"name":"audioUI"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"AUDIOUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"AUDIO","name":"AUDIO","type":"AUDIO","links":[3]}],"properties":{"Node name for S&R":"LoadAudio"},"widgets_values":["Megan-Fox.mp3",null,null]},{"id":3,"type":"VLLMOmniVoiceClone","pos":[248.63466280064262,168.13112893706375],"size":[400,306],"flags":{},"order":4,"mode":0,"inputs":[{"localized_name":"ref_audio","name":"ref_audio","type":"AUDIO","link":3},{"localized_name":"model_specific_params","name":"model_specific_params","shape":7,"type":"TTS_PARAMS","link":1},{"localized_name":"url","name":"url","type":"STRING","widget":{"name":"url"},"link":null},{"localized_name":"model","name":"model","type":"STRING","widget":{"name":"model"},"link":null},{"localized_name":"input","name":"input","type":"STRING","widget":{"name":"input"},"link":null},{"localized_name":"voice","name":"voice","type":"STRING","widget":{"name":"voice"},"link":null},{"localized_name":"response_format","name":"response_format","type":"COMBO","widget":{"name":"response_format"},"link":null},{"localized_name":"speed","name":"speed","type":"FLOAT","widget":{"name":"speed"},"link":null},{"localized_name":"ref_text","name":"ref_text","type":"STRING","widget":{"name":"ref_text"},"link":null},{"localized_name":"x_vector_only_mode","name":"x_vector_only_mode","type":"BOOLEAN","widget":{"name":"x_vector_only_mode"},"link":null}],"outputs":[{"localized_name":"audio","name":"audio","type":"AUDIO","links":[2]}],"properties":{"Node name for S&R":"VLLMOmniVoiceClone"},"widgets_values":["http://localhost:8000/v1","Qwen/Qwen3-TTS-12Hz-1.7B-Base","Someone just spilled a cup of coffee on my jacket this morning in the subway. Now I have to wear this stained jacket for a whole day!","","mp3",1,"",true]},{"id":6,"type":"MarkdownNote","pos":[255.52667408769662,27.26842775454861],"size":[381.58204108637847,88],"flags":{},"order":2,"mode":0,"inputs":[],"outputs":[],"title":"Note: TTS Nodes","properties":{},"widgets_values":["Apart from the Voice Cloning node, there is also another TTS node for simple text-to-speech tasks (without extra inputs for reference audio)."],"color":"#432","bgcolor":"#000"},{"id":7,"type":"MarkdownNote","pos":[-196.99453003423977,598.3693693727904],"size":[382.4539509209102,88],"flags":{},"order":3,"mode":0,"inputs":[],"outputs":[],"title":"Note: Model-Specific Parameters","properties":{},"widgets_values":["TTS models often require some tailor-made parameters. If you need to customize these parameters, grab one that matches the model you are requesting from \"TTS Params\" subfolder."],"color":"#432","bgcolor":"#000"},{"id":2,"type":"VLLMOmniQwenTTSParams","pos":[-207.8925310560497,342.7942867650348],"size":[400,200],"flags":{},"order":1,"mode":0,"inputs":[{"localized_name":"task_type","name":"task_type","type":"COMBO","widget":{"name":"task_type"},"link":null},{"localized_name":"language","name":"language","type":"COMBO","widget":{"name":"language"},"link":null},{"localized_name":"instructions","name":"instructions","type":"STRING","widget":{"name":"instructions"},"link":null},{"localized_name":"max_new_tokens","name":"max_new_tokens","type":"INT","widget":{"name":"max_new_tokens"},"link":null}],"outputs":[{"localized_name":"Qwen TTS params","name":"Qwen TTS params","type":"TTS_PARAMS","links":[1]}],"properties":{"Node name for S&R":"VLLMOmniQwenTTSParams"},"widgets_values":["Base","Auto","Super angry",2048]}],"links":[[1,2,0,3,1,"TTS_PARAMS"],[2,3,0,4,0,"AUDIO"],[3,5,0,3,0,"AUDIO"]],"groups":[{"id":1,"title":"Model-Specific Parameters","bounding":[-228.38328407738007,264.42790274716276,443.9887293164586,439.6853171248897],"color":"#3f789e","font_size":24,"flags":{}}],"config":{},"extra":{"workflowRendererVersion":"LG","ds":{"scale":1.1469075819486911,"offset":[898.8819468322333,258.7179979718389]}},"version":0.4}

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

@@ -0,0 +1,513 @@
+{
+    "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
+}

apps/ComfyUI-vLLM-Omni/web/main.js

@@ -0,0 +1,21 @@
+/**
+ * @file This file is intended to add dynamic fields to vLLM-Omni nodes
+ * based on widget (in-node form fields) and input (connection link) values and changes.
+ * However, this functionality is currently disabled/commented out.
+ * Because it introduces too much complexity,
+ * and it may even conflict with the current backend (Python) validation for unknown reasons (pending ComfyUI upstream fixes).
+ */
+
+import { app } from "../../scripts/app.js";
+app.registerExtension({
+    name: "vllm.vllm_omni",
+    async beforeRegisterNodeDef(nodeType, nodeData, app) {
+        if (!nodeData.name.startsWith("VLLMOmni")) {
+            return
+        }
+        // Stub frontend plugin for now
+    },
+    async setup() {
+        console.info("vLLM-Omni Setup complete!")
+    },
+})

benchmarks/README.md

@@ -1,42 +1,60 @@
-# Benchmarks Overview and Architecture
+# Benchmarks
 
-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`).
+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.
 
-## 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.
+## Benchmark families
 
-## 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.
+### [TTS](tts/README.md) — Text-to-Speech
 
-## 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.
+Model-agnostic serving benchmarks for TTS models, including Qwen3-TTS and VoxCPM2.
 
-## 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).
+- **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)
 
-- **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).
+### [Diffusion](diffusion/README.md) — Image and Video Generation
 
-- **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).
+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).
 
-## 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.
+- **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
 
-## 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.
+### [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.

benchmarks/__init__.py

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

benchmarks/accuracy/README.md

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

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

benchmarks/accuracy/common.py

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