update release

Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>

.buildkite/nightly-release-pipeline.yaml

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

.buildkite/pipeline-intel.yaml

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

.buildkite/pipeline.yml

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

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

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

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

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

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

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

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

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

.buildkite/scripts/upload_pipeline_with_skip_ci.sh

@@ -1,137 +0,0 @@
-#!/usr/bin/env bash
-# Evaluate docs-only skip-ci and upload continuation steps from the same `.buildkite/pipeline.yml`
-# (YAML document after the first `---`). Buildkite `if` is evaluated at upload time.
-set -euo pipefail
-
-ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
-PIPELINE_YML="${ROOT}/.buildkite/pipeline.yml"
-
-# Prints a single digit to stdout: 1 = skip image CI, 0 = run. Logs go to stderr.
-is_docs_only_change() {
-  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.yaml" ]]; then
-      continue
-    fi
-    return 1
-  done
-
-  [[ "${has_any}" -eq 1 ]]
-}
-
-resolve_skip_ci() {
-  local is_pr_build=0
-  local files
-  local base_branch base_ref
-
-  if [[ "${BUILDKITE_PULL_REQUEST:-false}" != "false" && -n "${BUILDKITE_PULL_REQUEST:-}" ]]; then
-    is_pr_build=1
-  fi
-
-  if [[ "${is_pr_build}" -eq 1 ]]; then
-    base_branch="${BUILDKITE_PULL_REQUEST_BASE_BRANCH:-main}"
-    if ! git rev-parse --verify "origin/${base_branch}" >/dev/null 2>&1; then
-      echo "resolve_skip_ci: origin/${base_branch} not found locally; trying fetch" >&2
-      git fetch --depth=200 origin "${base_branch}" >/dev/null 2>&1 || true
-    fi
-
-    base_ref=""
-    if git rev-parse --verify "origin/${base_branch}" >/dev/null 2>&1; then
-      base_ref="origin/${base_branch}"
-    elif git rev-parse --verify "${base_branch}" >/dev/null 2>&1; then
-      base_ref="${base_branch}"
-    else
-      echo "resolve_skip_ci: cannot resolve PR base ${base_branch}; skip-ci=0" >&2
-      echo -n 0
-      return 0
-    fi
-
-    if ! files="$(git diff --name-only "${base_ref}...${BUILDKITE_COMMIT}" 2>/dev/null)"; then
-      echo "resolve_skip_ci: failed to compute PR changed files; skip-ci=0" >&2
-      echo -n 0
-      return 0
-    fi
-  elif [[ "${BUILDKITE_BRANCH:-}" == "main" ]]; then
-    if ! git rev-parse --verify "${BUILDKITE_COMMIT}^" >/dev/null 2>&1; then
-      echo "resolve_skip_ci: commit has no parent on main; skip-ci=0" >&2
-      echo -n 0
-      return 0
-    fi
-    if ! files="$(git diff --name-only "${BUILDKITE_COMMIT}^..${BUILDKITE_COMMIT}" 2>/dev/null)"; then
-      echo "resolve_skip_ci: failed to compute main changed files; skip-ci=0" >&2
-      echo -n 0
-      return 0
-    fi
-  else
-    echo "resolve_skip_ci: not PR/main build; skip-ci=0" >&2
-    echo -n 0
-    return 0
-  fi
-
-  if is_docs_only_change <<< "${files}"; then
-    echo "resolve_skip_ci: docs-only change detected; skip-ci=1" >&2
-    echo -n 1
-    return 0
-  fi
-
-  echo "resolve_skip_ci: non-doc changes detected; skip-ci=0" >&2
-  echo -n 0
-}
-
-SKIP_CI="$(resolve_skip_ci)"
-
-if [[ ! -f "${PIPELINE_YML}" ]]; then
-  echo "upload_pipeline_with_skip_ci: missing ${PIPELINE_YML}" >&2
-  exit 1
-fi
-
-export ROOT SKIP_CI PIPELINE_YML
-python3 <<'PY' | buildkite-agent pipeline upload
-import os
-import pathlib
-
-path = pathlib.Path(os.environ["PIPELINE_YML"])
-text = path.read_text(encoding="utf-8")
-sep = "\n---\n"
-if sep not in text:
-    raise SystemExit(
-        "upload_pipeline_with_skip_ci: .buildkite/pipeline.yml must contain a '\\n---\\n' separator "
-        "(document 1 = bootstrap, document 2 = uploaded steps)"
-    )
-_, continuation = text.split(sep, 1)
-
-skip = os.environ.get("SKIP_CI") == "1"
-# When docs-only skip-ci: skip default CI image, but still build for L4 nightly (PR label nightly-test or
-# main NIGHTLY=1), otherwise upload-nightly (depends_on image-build) would be skipped too.
-nightly_only = (
-    '(build.pull_request.labels includes "nightly-test") '
-    '|| (build.branch == "main" && build.env("NIGHTLY") == "1")'
-)
-# Placeholder in pipeline.yml is `if: __IMAGE_BUILD_IF__` (valid YAML); replace value only.
-if skip:
-    rep = f"'{nightly_only}'"
-    ready_rep = "'false'"
-    merge_rep = "'false'"
-else:
-    rep = "'true'"
-    ready_rep = "'build.branch != \"main\" && build.pull_request.labels includes \"ready\"'"
-    merge_rep = "'(build.branch == \"main\" && build.env(\"NIGHTLY\") != \"1\") || (build.branch != \"main\" && build.pull_request.labels includes \"merge-test\")'"
-rendered = (
-    continuation
-    .replace("__IMAGE_BUILD_IF__", rep)
-    .replace("__UPLOAD_READY_IF__", ready_rep)
-    .replace("__UPLOAD_MERGE_IF__", merge_rep)
-)
-print(rendered, end="")
-PY

.buildkite/test-amd-merge.yml

@@ -32,6 +32,7 @@ steps:
   mirror_hardwares: [amdproduction]
   grade: Blocking
   commands:
+    - export GPU_ARCHS=gfx942
     - export VLLM_LOGGING_LEVEL=DEBUG
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
     - |
@@ -54,7 +55,7 @@ steps:
 #     - export GPU_ARCHS=gfx942
 #     - export VLLM_LOGGING_LEVEL=DEBUG
 #     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-#     - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_expansion.py -m "advanced_model and diffusion and L4" --run-level advanced_model
+#     - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_model.py
 
 - label: "Diffusion Cache Backend Test"
   agent_pool: mi325_1
@@ -62,12 +63,13 @@ steps:
   mirror_hardwares: [amdproduction]
   grade: Blocking
   commands:
+    - export GPU_ARCHS=gfx942
     - export VLLM_LOGGING_LEVEL=DEBUG
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
     - timeout 15m pytest -s -v -m "core_model and cache and diffusion and not distributed_cuda and L4"
 
-- label: "Diffusion Sequence Parallelism Test (Need 4 GPUs)"
-  agent_pool: mi325_4
+- label: "Diffusion Sequence Parallelism Test"
+  agent_pool: mi325_2
   depends_on: amd-build
   mirror_hardwares: [amdproduction]
   grade: Blocking
@@ -75,7 +77,6 @@ steps:
     - export VLLM_LOGGING_LEVEL=DEBUG
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
     - timeout 20m pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py
-    - timeout 20m pytest -s -v tests/diffusion/distributed/test_ulysses_uaa_perf.py
 
 # merge-only tests
 - label: "Diffusion Tensor Parallelism Test"
@@ -94,14 +95,22 @@ steps:
   commands:
     - timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py
 
-- label: "Engine Test"
-  agent_pool: mi325_1
+- label: "Benchmark & Engine Test"
+  agent_pool: mi325_2
   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
+    - |
+      timeout 20m bash -c '
+        set +e
+        pytest -s -v tests/benchmarks/test_serve_cli.py
+        EXIT1=\$?
+        pytest -s -v tests/engine/test_async_omni_engine_abort.py
+        EXIT2=\$?
+        exit \$((EXIT1 | EXIT2))
+      '
 
 - label: "Omni Model Test Qwen2-5-Omni"
   agent_pool: mi325_2
@@ -112,7 +121,6 @@ steps:
     - export VLLM_LOGGING_LEVEL=DEBUG
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
     - timeout 20m pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
-    - timeout 20m pytest -s -v tests/e2e/online_serving/test_qwen2_5_omni.py -m "advanced_model" --run-level "advanced_model"
 
 - label: "Omni Model Test Qwen3-Omni"
   agent_pool: mi325_2
@@ -123,10 +131,11 @@ steps:
     - export VLLM_LOGGING_LEVEL=DEBUG
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
     - export VLLM_TEST_CLEAN_GPU_MEMORY=1
-    - timeout 30m pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py tests/e2e/online_serving/test_mimo_audio.py -m "advanced_model" --run-level "advanced_model"
+    - timeout 10m pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
+    - timeout 20m pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "advanced_model" --run-level "advanced_model"
 
 - label: "Qwen3-TTS CustomVoice E2E Test"
-  agent_pool: mi325_1
+  agent_pool: mi325_2
   depends_on: amd-build
   mirror_hardwares: [amdproduction]
   grade: Blocking
@@ -136,21 +145,21 @@ steps:
         export VLLM_LOGGING_LEVEL=DEBUG
         export VLLM_WORKER_MULTIPROC_METHOD=spawn
         export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
-        pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py tests/e2e/offline_inference/test_qwen3_tts_customvoice.py -m "advanced_model" --run-level "advanced_model"
+        pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m "advanced_model" --run-level "advanced_model" && pytest -s -v tests/e2e/offline_inference/test_qwen3_tts_customvoice.py
       '
 
 - label: "Qwen3-TTS Base E2E Test"
-  agent_pool: mi325_1
+  agent_pool: mi325_2
   depends_on: amd-build
   mirror_hardwares: [amdproduction]
   grade: Blocking
   commands:
     - |
-      timeout 30m bash -c '
+      timeout 20m bash -c '
         export VLLM_LOGGING_LEVEL=DEBUG
         export VLLM_WORKER_MULTIPROC_METHOD=spawn
         export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
-        pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py tests/e2e/offline_inference/test_qwen3_tts_base.py -m "advanced_model" --run-level "advanced_model"
+        pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py -m "advanced_model" --run-level "advanced_model" && pytest -s -v tests/e2e/offline_inference/test_qwen3_tts_base.py
       '
 
 - label: "Diffusion Image Edit Test"
@@ -164,58 +173,43 @@ steps:
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
     - timeout 20m pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py
 
-# TODO: Bagel test on ROCm is very unstable. @tjtanaa
-# Need to debug before reneable numerical changes across large PRs
-# # split Bagel Model Test with H100 (Real Weights) into three tests
-# - 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_TEST_CLEAN_GPU_MEMORY=1
-#     - export VLLM_LOGGING_LEVEL=DEBUG
-#     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-#     - export VLLM_ROCM_USE_AITER_RMSNORM=0
-#     - timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py -m "advanced_model" --run-level "advanced_model" -k "shared_memory" -k "rocm"
-
-# - label: "Bagel Img2Img Model Test (2/3)"
-#   agent_pool: mi325_1
-#   depends_on: amd-build
-#   mirror_hardwares: [amdproduction]
-#   grade: Blocking
-#   commands:
-#     - export GPU_ARCHS=gfx942
-#     - export VLLM_TEST_CLEAN_GPU_MEMORY=1
-#     - export VLLM_LOGGING_LEVEL=DEBUG
-#     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-#     - export VLLM_ROCM_USE_AITER_RMSNORM=0
-#     - timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py -m "advanced_model" --run-level "advanced_model" -k "rocm"
-
-# - label: "Bagel Online Serving Test (3/3)"
-#   agent_pool: mi325_1
-#   depends_on: amd-build
-#   mirror_hardwares: [amdproduction]
-#   grade: Blocking
-#   commands:
-#     - export GPU_ARCHS=gfx942
-#     - export VLLM_TEST_CLEAN_GPU_MEMORY=1
-#     - export VLLM_IMAGE_FETCH_TIMEOUT=60
-#     - export VLLM_LOGGING_LEVEL=DEBUG
-#     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-#     - export VLLM_ROCM_USE_AITER_RMSNORM=0
-#     - timeout 40m pytest -s -v tests/e2e/online_serving/test_bagel_online.py -m "advanced_model" --run-level "advanced_model" -k "rocm"
-
-- label: "Voxtral-TTS E2E Test"
+# split Bagel Model Test with H100 (Real Weights) into three tests
+- label: "Bagel Text2Img Model Test"
   agent_pool: mi325_1
   depends_on: amd-build
   mirror_hardwares: [amdproduction]
   grade: Blocking
   commands:
-    - |
-      timeout 20m bash -c '
-        export VLLM_LOGGING_LEVEL=DEBUG
-        export VLLM_WORKER_MULTIPROC_METHOD=spawn
-        pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py tests/e2e/offline_inference/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
-      '
+    - export GPU_ARCHS=gfx942
+    - export VLLM_TEST_CLEAN_GPU_MEMORY=1
+    - export VLLM_LOGGING_LEVEL=DEBUG
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - export VLLM_ROCM_USE_AITER_RMSNORM=0
+    - timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py -m "advanced_model" --run-level "advanced_model" -k "shared_memory" -k "rocm"
+
+- label: "Bagel Img2Img Model Test"
+  agent_pool: mi325_1
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export GPU_ARCHS=gfx942
+    - export VLLM_TEST_CLEAN_GPU_MEMORY=1
+    - export VLLM_LOGGING_LEVEL=DEBUG
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - export VLLM_ROCM_USE_AITER_RMSNORM=0
+    - timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py -m "advanced_model" --run-level "advanced_model" -k "rocm"
+
+- label: "Bagel Online Serving Test"
+  agent_pool: mi325_1
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export GPU_ARCHS=gfx942
+    - export VLLM_TEST_CLEAN_GPU_MEMORY=1
+    - export VLLM_IMAGE_FETCH_TIMEOUT=60
+    - export VLLM_LOGGING_LEVEL=DEBUG
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - export VLLM_ROCM_USE_AITER_RMSNORM=0
+    - timeout 40m pytest -s -v tests/e2e/online_serving/test_bagel_online.py -m "advanced_model" --run-level "advanced_model" -k "rocm"

.buildkite/test-amd-ready.yaml

@@ -9,37 +9,13 @@ steps:
     - export VLLM_ROCM_USE_AITER=0
     - "timeout 20m pytest -v -s -m 'core_model and cpu' --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
 
-- label: "Voxtral TTS CUDA Unit Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - timeout 10m pytest -s -v tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py
-
 - label: "Diffusion Model Test"
-  agent_pool: mi325_1
+  agent_pool: mi325_2
   depends_on: amd-build
   mirror_hardwares: [amdproduction]
   grade: Blocking
   commands:
-    - timeout 30m pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "core_model and diffusion" --run-level "core_model"
-
-- label: "Diffusion Batching Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - timeout 20m pytest -s -v tests/e2e/offline_inference/test_qwen_image_diffusion_batching.py -m "core_model and diffusion" --run-level "core_model"
-
-- label: "Custom Pipeline Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - timeout 20m pytest -s -v tests/e2e/offline_inference/custom_pipeline/ -m "core_model"
+    - timeout 20m pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "core_model and diffusion" --run-level "core_model"
 
 - label: "Diffusion Model CPU offloading Test"
   agent_pool: mi325_1
@@ -47,6 +23,7 @@ steps:
   mirror_hardwares: [amdproduction]
   grade: Blocking
   commands:
+    - export GPU_ARCHS=gfx942
     - export VLLM_LOGGING_LEVEL=DEBUG
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
     - |
@@ -69,7 +46,7 @@ steps:
 #     - export GPU_ARCHS=gfx942
 #     - export VLLM_LOGGING_LEVEL=DEBUG
 #     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-#     - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_expansion.py -m "advanced_model and diffusion and L4" --run-level advanced_model
+#     - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_model.py
 
 - label: "Diffusion Cache Backend Test"
   agent_pool: mi325_1
@@ -100,58 +77,47 @@ steps:
   commands:
     - timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py
 
-- label: "Engine Test"
-  agent_pool: mi325_1
+- label: "Benchmark & Engine Test"
+  agent_pool: mi325_2
   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
+      timeout 30m bash -c '
+        set +e
+        pytest -s -v tests/benchmarks/test_serve_cli.py
+        EXIT1=\$?
+        pytest -s -v tests/engine/test_async_omni_engine_abort.py
+        EXIT2=\$?
+        exit \$((EXIT1 | EXIT2))
       '
 
-
-# NOTE: This test is not running any thing. It is skipped and deselected.
-# Currently it is = 1 skipped, 1 deselected, 17 warnings in 0.03s ======
-# - label: "Omni Model Test Qwen2-5-Omni"
-#   agent_pool: mi325_2
-#   depends_on: amd-build
-#   mirror_hardwares: [amdproduction]
-#   grade: Blocking
-#   commands:
-#     - export VLLM_LOGGING_LEVEL=DEBUG
-#     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-#     - timeout 20m pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py -m "core_model" --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_LOGGING_LEVEL=DEBUG
-#     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-#     - export VLLM_TEST_CLEAN_GPU_MEMORY=1
-#     - timeout 10m pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
-#     - timeout 20m pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "core_model" --run-level "core_model"
-
-- label: "MiMo-Audio E2E Test with H100"
-  agent_pool: mi325_1
+- label: "Omni Model Test Qwen2-5-Omni"
+  agent_pool: mi325_2
   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"
-      '
+    - export VLLM_LOGGING_LEVEL=DEBUG
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - timeout 17m pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
+
+- label: "Omni Model Test Qwen3-Omni"
+  agent_pool: mi325_2
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export VLLM_LOGGING_LEVEL=DEBUG
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - export VLLM_TEST_CLEAN_GPU_MEMORY=1
+    - timeout 10m pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
+    - timeout 10m pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "core_model" --run-level "core_model"
 
 - label: "Qwen3-TTS E2E Test"
-  agent_pool: mi325_1
+  agent_pool: mi325_2
   depends_on: amd-build
   mirror_hardwares: [amdproduction]
   grade: Blocking
@@ -159,21 +125,7 @@ steps:
     - export VLLM_LOGGING_LEVEL=DEBUG
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
     - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
-    - timeout 30m pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m "core_model" --run-level "core_model"
-
-- label: "Voxtral-TTS E2E Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - |
-      timeout 20m bash -c '
-        export VLLM_LOGGING_LEVEL=DEBUG
-        export VLLM_WORKER_MULTIPROC_METHOD=spawn
-        pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
-        pytest -s -v tests/e2e/offline_inference/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
-      '
+    - timeout 20m pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m "core_model" --run-level "core_model"
 
 - label: "Diffusion Image Edit Test"
   agent_pool: mi325_1
@@ -186,55 +138,42 @@ steps:
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
     - timeout 20m pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py
 
-# TODO: Bagel test on ROCm is very unstable. @tjtanaa
-# Need to debug before reneable numerical changes across large PRs
-# - label: "Bagel Text2Img Model Test"
-#   agent_pool: mi325_1
-#   depends_on: amd-build
-#   mirror_hardwares: [amdproduction]
-#   grade: Blocking
-#   commands:
-#     - export GPU_ARCHS=gfx942
-#     - export VLLM_TEST_CLEAN_GPU_MEMORY=1
-#     - export VLLM_LOGGING_LEVEL=DEBUG
-#     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-#     - export VLLM_ROCM_USE_AITER_RMSNORM=0
-#     - timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py -m "core_model" --run-level "core_model" -k "rocm"
-
-# - label: "Bagel Img2Img Model Test"
-#   agent_pool: mi325_1
-#   depends_on: amd-build
-#   mirror_hardwares: [amdproduction]
-#   grade: Blocking
-#   commands:
-#     - export GPU_ARCHS=gfx942
-#     - export VLLM_TEST_CLEAN_GPU_MEMORY=1
-#     - export VLLM_LOGGING_LEVEL=DEBUG
-#     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-#     - export VLLM_ROCM_USE_AITER_RMSNORM=0
-#     - timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py -m "core_model" --run-level "core_model" -k "rocm"
-
-# - label: "Bagel Online Serving Test"
-#   agent_pool: mi325_1
-#   depends_on: amd-build
-#   mirror_hardwares: [amdproduction]
-#   grade: Blocking
-#   commands:
-#     - export GPU_ARCHS=gfx942
-#     - export VLLM_TEST_CLEAN_GPU_MEMORY=1
-#     - export VLLM_IMAGE_FETCH_TIMEOUT=60
-#     - export VLLM_LOGGING_LEVEL=DEBUG
-#     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-#     - export VLLM_ROCM_USE_AITER_RMSNORM=0
-#     - timeout 40m pytest -s -v tests/e2e/online_serving/test_bagel_online.py -m "core_model" --run-level "core_model" -k "rocm"
-
-- label: "CosyVoice3-TTS E2E Test"
+- label: "Bagel Text2Img Model 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.py -m "core_model" --run-level "core_model"
-      '
+    - export GPU_ARCHS=gfx942
+    - export VLLM_TEST_CLEAN_GPU_MEMORY=1
+    - export VLLM_LOGGING_LEVEL=DEBUG
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - export VLLM_ROCM_USE_AITER_RMSNORM=0
+    - timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py -m "core_model" --run-level "core_model" -k "rocm"
+
+- label: "Bagel Img2Img Model Test"
+  agent_pool: mi325_1
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export GPU_ARCHS=gfx942
+    - export VLLM_TEST_CLEAN_GPU_MEMORY=1
+    - export VLLM_LOGGING_LEVEL=DEBUG
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - export VLLM_ROCM_USE_AITER_RMSNORM=0
+    - timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py -m "core_model" --run-level "core_model" -k "rocm"
+
+- label: "Bagel Online Serving Test"
+  agent_pool: mi325_1
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export GPU_ARCHS=gfx942
+    - export VLLM_TEST_CLEAN_GPU_MEMORY=1
+    - export VLLM_IMAGE_FETCH_TIMEOUT=60
+    - export VLLM_LOGGING_LEVEL=DEBUG
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - export VLLM_ROCM_USE_AITER_RMSNORM=0
+    - timeout 40m pytest -s -v tests/e2e/online_serving/test_bagel_online.py -m "core_model" --run-level "core_model" -k "rocm"

.buildkite/test-amd.yaml

@@ -117,17 +117,3 @@ steps:
     - export VLLM_LOGGING_LEVEL=DEBUG
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
     - pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py
-
-
-- label: "Omni Sleep Mode Test"
-  timeout_in_minutes: 40
-  agent_pool: mi325_2
-  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_TEST_CLEAN_GPU_MEMORY="1"
-    - pytest -s -v tests/e2e/offline_inference/test_omni_sleep_mode.py -m "advanced_model and omni and MI325" --run-level "advanced_model"

.buildkite/test-merge.yml

@@ -1,8 +1,3 @@
-env:
-  VLLM_WORKER_MULTIPROC_METHOD: spawn
-  HF_HUB_DOWNLOAD_TIMEOUT: 300
-  HF_HUB_ETAG_TIMEOUT: 60
-
 steps:
   - label: "Simple Unit Test"
     depends_on: upload-merge-pipeline
@@ -61,8 +56,26 @@ steps:
     timeout_in_minutes: 20
     depends_on: upload-merge-pipeline
     commands:
-      # Single pytest session for one combined summary at end of log.
-      - pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
+      - pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
+      - pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
+    agents:
+      queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Audio Generation Model Test"
+    timeout_in_minutes: 20
+    depends_on: upload-merge-pipeline
+    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:
@@ -95,10 +108,11 @@ steps:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
   - label: "Diffusion Sequence Parallelism Test"
-    timeout_in_minutes: 25
+    timeout_in_minutes: 20
     depends_on: upload-merge-pipeline
     commands:
-      - pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py tests/diffusion/distributed/test_ulysses_uaa_perf.py
+      - pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py
+      - pytest -s -v tests/diffusion/distributed/test_ulysses_uaa_perf.py
     agents:
       queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
     plugins:
@@ -156,10 +170,11 @@ steps:
     commands:
       - |
         timeout 15m bash -c '
+                export VLLM_WORKER_MULTIPROC_METHOD=spawn
                 pytest -s -v tests/engine/test_async_omni_engine_abort.py
         '
     agents:
-      queue: "gpu_1_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
+      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
@@ -177,7 +192,9 @@ steps:
     depends_on: upload-merge-pipeline
     commands:
       - export VLLM_LOGGING_LEVEL=DEBUG
-      - pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py tests/e2e/online_serving/test_qwen2_5_omni.py -m "advanced_model" --run-level "advanced_model"
+      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
+      - pytest -s -v tests/e2e/online_serving/test_qwen2_5_omni.py -m "advanced_model" --run-level "advanced_model"
     agents:
       queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
     plugins:
@@ -197,8 +214,9 @@ steps:
       - |
         timeout 20m bash -c '
           export VLLM_LOGGING_LEVEL=DEBUG
+          export VLLM_WORKER_MULTIPROC_METHOD=spawn
           export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
-          pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py tests/e2e/offline_inference/test_qwen3_tts_customvoice.py -m "advanced_model" --run-level "advanced_model"
+          pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m "advanced_model" --run-level "advanced_model" && pytest -s -v tests/e2e/offline_inference/test_qwen3_tts_customvoice.py
         '
     agents:
       queue: "gpu_1_queue"
@@ -219,8 +237,9 @@ steps:
       - |
         timeout 20m bash -c '
           export VLLM_LOGGING_LEVEL=DEBUG
+          export VLLM_WORKER_MULTIPROC_METHOD=spawn
           export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
-          pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py tests/e2e/offline_inference/test_qwen3_tts_base.py -m "advanced_model" --run-level "advanced_model"
+          pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py -m "advanced_model" --run-level "advanced_model" && pytest -s -v tests/e2e/offline_inference/test_qwen3_tts_base.py
         '
     agents:
       queue: "gpu_1_queue"
@@ -239,48 +258,11 @@ steps:
     timeout_in_minutes: 30
     depends_on: upload-merge-pipeline
     commands:
+      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
       - export VLLM_TEST_CLEAN_GPU_MEMORY="1"
-      - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py tests/e2e/online_serving/test_mimo_audio.py -m "advanced_model" --run-level "advanced_model"
-    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: "Audio Streaming Input Test with H100"
-    timeout_in_minutes: 30
-    depends_on: upload-merge-pipeline
-    commands:
-      - export VLLM_TEST_CLEAN_GPU_MEMORY="1"
-      - pytest -s -v tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py -m "advanced_model" --run-level "advanced_model"
+      - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
+      - pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "advanced_model" --run-level "advanced_model"
+      - pytest -s -v tests/e2e/online_serving/test_mimo_audio.py -m "advanced_model" --run-level "advanced_model"
     agents:
       queue: "mithril-h100-pool"
     plugins:
@@ -319,6 +301,7 @@ steps:
     timeout_in_minutes: 20
     depends_on: upload-merge-pipeline
     commands:
+      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
       - pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py
     agents:
       queue: "mithril-h100-pool"
@@ -361,10 +344,12 @@ steps:
       - |
         timeout 55m bash -c '
           set -e
+          export VLLM_WORKER_MULTIPROC_METHOD=spawn
           export VLLM_TEST_CLEAN_GPU_MEMORY=1
           export VLLM_IMAGE_FETCH_TIMEOUT=60
           pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py -m "advanced_model" --run-level "advanced_model" -k "shared_memory"
-          pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py tests/e2e/online_serving/test_bagel_online.py -m "advanced_model" --run-level "advanced_model"
+          pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py -m "advanced_model" --run-level "advanced_model"
+          pytest -s -v tests/e2e/online_serving/test_bagel_online.py -m "advanced_model" --run-level "advanced_model"
         '
     agents:
       queue: "mithril-h100-pool"
@@ -400,46 +385,6 @@ steps:
                   path: /mnt/hf-cache
                   type: DirectoryOrCreate
 
-  - label: "Omni Sleep Mode Test with H100"
-    timeout_in_minutes: 30
-    depends_on: upload-merge-pipeline
-    commands:
-      - export VLLM_TEST_CLEAN_GPU_MEMORY="1"
-      - pytest -s -v tests/e2e/offline_inference/test_omni_sleep_mode.py -m "advanced_model and H100 and omni" --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: "Voxtral-TTS E2E Test"
     timeout_in_minutes: 20
     depends_on: upload-merge-pipeline
@@ -447,7 +392,9 @@ steps:
       - |
         timeout 20m bash -c '
           export VLLM_LOGGING_LEVEL=DEBUG
-          pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py tests/e2e/offline_inference/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
+          export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
+          pytest -s -v tests/e2e/offline_inference/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
         '
     agents:
       queue: "mithril-h100-pool"
@@ -467,11 +414,6 @@ steps:
                 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:

.buildkite/test-ready.yml

@@ -1,8 +1,3 @@
-env:
-  VLLM_WORKER_MULTIPROC_METHOD: spawn
-  HF_HUB_DOWNLOAD_TIMEOUT: 300
-  HF_HUB_ETAG_TIMEOUT: 60
-
 steps:
   - label: "Simple Unit Test"
     depends_on: upload-ready-pipeline
@@ -21,10 +16,11 @@ steps:
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
-  - label: "CUDA Unit Test with single card"
+  - label: "Voxtral TTS CUDA Unit Test"
+    timeout_in_minutes: 10
     depends_on: upload-ready-pipeline
     commands:
-      - timeout 10m pytest -v -s -m 'core_model and cuda and L4 and not distributed_cuda' --ignore=tests/e2e --ignore=tests/engine/test_async_omni_engine_abort.py --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml
+      - "timeout 10m pytest -s -v tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py"
     agents:
       queue: "gpu_1_queue"
     plugins:
@@ -37,22 +33,6 @@ steps:
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
-  - label: "CUDA Unit Test with multi cards"
-    depends_on: upload-ready-pipeline
-    commands:
-      - timeout 10m pytest -v -s -m 'core_model and cuda and L4 and distributed_cuda' --ignore=tests/e2e --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml
-    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
-          environment:
-            - "HF_HOME=/fsx/hf_cache"
-          volumes:
-            - "/fsx/hf_cache:/fsx/hf_cache"
-
   - label: "Diffusion Model Test"
     depends_on: upload-ready-pipeline
     commands:
@@ -120,6 +100,23 @@ steps:
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
+  - label: "Audio Generation Model Test"
+    depends_on: upload-ready-pipeline
+    commands:
+      - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_model.py
+    agents:
+      queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
   - label: "Diffusion Cache Backend Test"
     depends_on: upload-ready-pipeline
     commands:
@@ -155,16 +152,12 @@ steps:
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
-
-  - label: "Engine Test"
+  - label: "Diffusion GPU Worker Test"
     depends_on: upload-ready-pipeline
     commands:
-      - |
-        timeout 15m bash -c '
-                pytest -s -v tests/engine/test_async_omni_engine_abort.py
-        '
+      - timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py
     agents:
-      queue: "gpu_1_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
+      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
@@ -177,11 +170,58 @@ steps:
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
+
+  - label: "Engine Test"
+    depends_on: upload-ready-pipeline
+    commands:
+      - |
+        timeout 15m bash -c '
+                export VLLM_WORKER_MULTIPROC_METHOD=spawn
+                pytest -s -v tests/engine/test_async_omni_engine_abort.py
+        '
+    agents:
+      queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          shm-size: "8gb"
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+
+  - label: "Omni Model Test"
+    depends_on: upload-ready-pipeline
+    commands:
+      - |
+        timeout 17m bash -c '
+          export VLLM_LOGGING_LEVEL=DEBUG
+          export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          pytest -s -v tests/e2e/online_serving/test_qwen2_5_omni.py -m "core_model" --run-level "core_model"
+        '
+    agents:
+      queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
   - label: "Omni Model Test with H100"
     depends_on: upload-ready-pipeline
     commands:
       - |
         timeout 20m bash -c '
+          export VLLM_WORKER_MULTIPROC_METHOD=spawn
           pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "core_model" --run-level "core_model"
         '
     agents:
@@ -219,6 +259,7 @@ steps:
       - |
         timeout 30m bash -c '
           export VLLM_LOGGING_LEVEL=DEBUG
+          export VLLM_WORKER_MULTIPROC_METHOD=spawn
           pytest -s -v tests/e2e/online_serving/test_mimo_audio.py -m "core_model" --run-level "core_model"
         '
     agents:
@@ -230,7 +271,7 @@ steps:
               - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
                 resources:
                   limits:
-                    nvidia.com/gpu: 1
+                    nvidia.com/gpu: 2
                 volumeMounts:
                   - name: devshm
                     mountPath: /dev/shm
@@ -261,6 +302,7 @@ steps:
       - |
         timeout 20m bash -c '
           export VLLM_LOGGING_LEVEL=DEBUG
+          export VLLM_WORKER_MULTIPROC_METHOD=spawn
           export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
           pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m "core_model" --run-level "core_model"
         '
@@ -278,105 +320,6 @@ steps:
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
-  - label: "VoxCPM E2E Test"
-    timeout_in_minutes: 20
-    depends_on: upload-ready-pipeline
-    commands:
-      - |
-        timeout 20m bash -c '
-          pip install voxcpm
-          export VLLM_LOGGING_LEVEL=DEBUG
-          export VLLM_WORKER_MULTIPROC_METHOD=spawn
-          pytest -s -v tests/e2e/offline_inference/test_voxcpm.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: "VoxCPM2 Native AR E2E Test"
-    timeout_in_minutes: 20
-    depends_on: upload-ready-pipeline
-    commands:
-      - |
-        timeout 20m bash -c '
-          pip install voxcpm
-          export VLLM_LOGGING_LEVEL=DEBUG
-          export VLLM_WORKER_MULTIPROC_METHOD=spawn
-          pytest -s -v tests/e2e/offline_inference/test_voxcpm2.py -m "core_model" --run-level "core_model"
-        '
-    agents:
-      queue: "gpu_1_queue"
-    plugins:
-      - docker#v5.2.0:
-          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-          always-pull: true
-          propagate-environment: true
-          shm-size: "8gb"
-          environment:
-            - "HF_HOME=/fsx/hf_cache"
-            - "HF_TOKEN"
-          volumes:
-            - "/fsx/hf_cache:/fsx/hf_cache"
-
-  - label: "OmniVoice E2E Test"
-    timeout_in_minutes: 20
-    depends_on: upload-ready-pipeline
-    commands:
-      - |
-        timeout 20m bash -c '
-          export VLLM_LOGGING_LEVEL=DEBUG
-          pytest -s -v tests/e2e/online_serving/test_omnivoice.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"
-          volumes:
-            - "/fsx/hf_cache:/fsx/hf_cache"
-
-  - label: "Qwen3-TTS Base E2E Test (ModelRunner V2)"
-    depends_on: upload-ready-pipeline
-    soft_fail:
-      - exit_status: 1
-    commands:
-      - |
-        timeout 20m bash -c '
-          export VLLM_LOGGING_LEVEL=DEBUG
-          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: "Voxtral-TTS E2E Test"
     timeout_in_minutes: 20
     depends_on: upload-ready-pipeline
@@ -384,7 +327,9 @@ steps:
       - |
         timeout 20m bash -c '
           export VLLM_LOGGING_LEVEL=DEBUG
-          pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "core_model" --run-level "core_model"
+          export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
+          pytest -s -v tests/e2e/offline_inference/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
         '
     agents:
       queue: "mithril-h100-pool"
@@ -420,6 +365,7 @@ steps:
   #   commands:
   #     - |
   #       timeout 20m bash -c '
+  #         export VLLM_WORKER_MULTIPROC_METHOD=spawn
   #         pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py
   #       '
   #   agents:
@@ -456,6 +402,7 @@ steps:
     commands:
       - |
         timeout 30m bash -c '
+          export VLLM_WORKER_MULTIPROC_METHOD=spawn
           export VLLM_TEST_CLEAN_GPU_MEMORY=1
           pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py -m "core_model" --run-level "core_model"
         '
@@ -498,6 +445,7 @@ steps:
     commands:
       - |
         timeout 30m bash -c '
+          export VLLM_WORKER_MULTIPROC_METHOD=spawn
           export VLLM_TEST_CLEAN_GPU_MEMORY=1
           pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py -m "core_model" --run-level "core_model"
         '
@@ -540,6 +488,7 @@ steps:
     commands:
       - |
         timeout 40m bash -c '
+          export VLLM_WORKER_MULTIPROC_METHOD=spawn
           export VLLM_TEST_CLEAN_GPU_MEMORY=1
           export VLLM_IMAGE_FETCH_TIMEOUT=60
           pytest -s -v tests/e2e/online_serving/test_bagel_online.py -m "core_model" --run-level "core_model"
@@ -577,45 +526,3 @@ steps:
                 hostPath:
                   path: /mnt/hf-cache
                   type: DirectoryOrCreate
-
-  - label: "CosyVoice3-TTS E2E Test"
-    timeout_in_minutes: 20
-    depends_on: upload-ready-pipeline
-    commands:
-      - |
-        timeout 20m bash -c '
-          pytest -s -v tests/e2e/online_serving/test_cosyvoice3_tts.py -m "core_model" --run-level "core_model"
-        '
-    agents:
-      queue: "mithril-h100-pool"
-    plugins:
-      - kubernetes:
-          podSpec:
-            containers:
-              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-                resources:
-                  limits:
-                    nvidia.com/gpu: 1
-                volumeMounts:
-                  - name: devshm
-                    mountPath: /dev/shm
-                  - name: hf-cache
-                    mountPath: /root/.cache/huggingface
-                env:
-                  - name: HF_HOME
-                    value: /root/.cache/huggingface
-                  - name: HF_TOKEN
-                    valueFrom:
-                      secretKeyRef:
-                        name: hf-token-secret
-                        key: token
-            nodeSelector:
-              node.kubernetes.io/instance-type: gpu-h100-sxm
-            volumes:
-              - name: devshm
-                emptyDir:
-                  medium: Memory
-              - name: hf-cache
-                hostPath:
-                  path: /mnt/hf-cache
-                  type: DirectoryOrCreate

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

.claude/skills/readme.md

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

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

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

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

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

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

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

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

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

.github/ISSUE_TEMPLATE/750-RFC.yml

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

.github/scripts/pr_reviewer.py

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

.github/workflows/pr-reviewer.yml

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

.gitignore

@@ -83,9 +83,6 @@ target/
 profile_default/
 ipython_config.py
 
-# uv
-uv.lock
-
 # pyenv
 #   For a library or package, you might want to ignore these files since the code is
 #   intended to run in multiple environments; otherwise, check them in:
@@ -158,23 +155,10 @@ cython_debug/
 
 # Claude
 CLAUDE.md
-/.claude/*
-!.claude/skills/
-!.claude/skills/readme.md
-!.claude/skills/add-diffusion-model/
-!.claude/skills/add-diffusion-model/SKILL.md
-!.claude/skills/add-diffusion-model/references/
-!.claude/skills/add-diffusion-model/references/*.md
-!.claude/skills/add-tts-model/
-!.claude/skills/add-tts-model/SKILL.md
-!.claude/skills/review-pr/
-!.claude/skills/review-pr/SKILL.md
-!.claude/skills/review-pr/references/
-!.claude/skills/review-pr/references/*.md
+.claude/
 
 # Codex
 AGENTS.md
-.codex
 .codex/
 
 # cursor
@@ -204,7 +188,6 @@ checkpoints/
 # Cache directories
 cache/
 !vllm_omni/diffusion/cache/
-!tests/diffusion/cache/
 .cache/
 diffusion_cache/
 kv_cache/
@@ -264,6 +247,3 @@ tmp_test
 vllm_omni/_version.py
 # output files
 *.wav
-.worktrees/
-# CI overlay yamls materialized from tests/utils.py:_CI_OVERLAYS at test time
-tests/.ci_generated/

benchmarks/accuracy/README.md

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

benchmarks/accuracy/image_to_image/README.md

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

benchmarks/build_dataset/download_process_data_seedtts.md

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

benchmarks/build_dataset/seed_tts_design/en/meta.lst

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

benchmarks/build_dataset/seed_tts_smoke/en/meta.lst

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

benchmarks/diffusion/backends.py

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

benchmarks/diffusion/diffusion_benchmark_serving.py

@@ -12,15 +12,15 @@ Supports multiple backends:
     - v1/videos: Use /v1/videos endpoint
 
 Usage:
-    # Video (v1/videos backend)
+    # Video (vllm-omni backend)
     t2v:
     python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
-        --backend v1/videos --dataset vbench --task t2v --num-prompts 10 \
+        --backend vllm-omni --dataset vbench --task t2v --num-prompts 10 \
         --height 480 --width 640 --fps 16 --num-frames 80
 
     i2v:
     python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
-        --backend v1/videos --dataset vbench --task i2v --num-prompts 10
+        --backend vllm-omni --dataset vbench --task i2v --num-prompts 10
 
 
     # Image (vllm-omni backend)
@@ -49,7 +49,7 @@ Usage:
         --backend openai --dataset vbench --task t2i --num-prompts 10 \
         --height 1024 --width 1024 --port 3000
 
-    # Video (v1/videos)
+    # Video (v1/vedeos)
     t2v:
     python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
         --backend v1/videos --dataset random --task t2v --num-prompts 1 \
@@ -558,7 +558,6 @@ class RandomDataset(BaseDataset):
         super().__init__(args, api_url, model)
         self.num_prompts = args.num_prompts
         self.enable_negative_prompt = enable_negative_prompt
-        self.num_input_images = max(1, args.num_input_images)
         self.random_request_config = getattr(args, "random_request_config", None)
         if self.random_request_config:
             self.random_request_config = json.loads(self.random_request_config)
@@ -581,7 +580,11 @@ class RandomDataset(BaseDataset):
 
         # Random image generate
         if self.args.task in ["i2v", "ti2v", "ti2i", "i2i"]:
-            self._random_image_path = self._generate_random_image_paths()
+            img = Image.new("RGB", (512, 512), (255, 255, 255))
+
+            image_path = os.path.join(tempfile.gettempdir(), "diffusion_benchmark_random_image.png")
+            self._random_image_path = [image_path]
+            img.save(image_path)
         else:
             self._random_image_path = None
 
@@ -616,18 +619,6 @@ class RandomDataset(BaseDataset):
     def get_requests(self) -> list[RequestFuncInput]:
         return [self[i] for i in range(len(self))]
 
-    def _generate_random_image_paths(self) -> list[str]:
-        image_paths: list[str] = []
-        for image_idx in range(self.num_input_images):
-            img = Image.new("RGB", (512, 512), (255, 255, 255))
-            image_path = os.path.join(
-                tempfile.gettempdir(),
-                f"diffusion_benchmark_random_image_{image_idx}.png",
-            )
-            img.save(image_path)
-            image_paths.append(image_path)
-        return image_paths
-
 
 def _compute_expected_latency_ms_from_base(req: RequestFuncInput, args, base_time_ms: float | None) -> float | None:
     """Compute expected execution time (ms) based on a base per-step-per-frame unit time.
@@ -1124,15 +1115,6 @@ if __name__ == "__main__":
             '{"width":768,"height":768,"num_inference_steps":20,"weight":0.85}]'
         ),
     )
-    parser.add_argument(
-        "--num-input-images",
-        type=int,
-        default=1,
-        help=(
-            "Number of synthetic input images to attach for image-conditioned tasks "
-            "(i2v, ti2v, ti2i, i2i) when using random dataset."
-        ),
-    )
 
     args = parser.parse_args()
 

benchmarks/fish-speech/bench_voice_cache.py

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

benchmarks/fish-speech/fish_bench_utils.py

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

benchmarks/qwen3-omni/README.md

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

benchmarks/qwen3-omni/transformers/eval_qwen3_moe_omni_transformers.sh

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

benchmarks/qwen3-omni/transformers/qwen3_omni_moe_model.py

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

benchmarks/qwen3-omni/transformers/qwen3_omni_moe_transformers.py

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

benchmarks/qwen3-omni/vllm_omni/eval_qwen3_moe_omni.sh

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

benchmarks/qwen3-tts/README.md

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

benchmarks/qwen3-tts/plot_results.py

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

benchmarks/qwen3-tts/results/.gitignore

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

benchmarks/qwen3-tts/run_benchmark.sh

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

benchmarks/qwen3-tts/transformers/bench_tts_hf.py

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

benchmarks/qwen3-tts/vllm_omni/bench_async_chunk.py

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

benchmarks/qwen3-tts/vllm_omni/bench_tts_serve.py

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

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

@@ -1,3 +1,5 @@
+# Qwen3-TTS batch_size=1 config (streaming with async_chunk)
+# 2-stage pipeline: Talker -> Code2Wav
 async_chunk: true
 stage_args:
   - stage_id: 0
@@ -5,85 +7,87 @@ stage_args:
     is_comprehension: true
     runtime:
       devices: "0"
-      max_batch_size: 1
     engine_args:
-      dtype: bfloat16
-      model_stage: latent_generator
-      model_arch: VoxCPMForConditionalGeneration
-      # Optional persistent HF-compatible config dir for native VoxCPM models.
-      hf_config_path: ${oc.env:VLLM_OMNI_VOXCPM_HF_CONFIG_PATH,}
+      max_num_seqs: 1
+      model_stage: qwen3_tts
+      model_arch: Qwen3TTSTalkerForConditionalGeneration
       worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      enforce_eager: true
+      enforce_eager: false
       trust_remote_code: true
-      async_scheduling: false
+      async_scheduling: true
       enable_prefix_caching: false
       engine_output_type: latent
-      gpu_memory_utilization: 0.75
+      gpu_memory_utilization: 0.3
       distributed_executor_backend: "mp"
-      max_num_batched_tokens: 4096
+      max_num_batched_tokens: 512
       max_model_len: 4096
-      custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.voxcpm.latent2vae_async_chunk
+      custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_tts.talker2code2wav_async_chunk
     output_connectors:
       to_stage_1: connector_of_shared_memory
     default_sampling_params:
-      temperature: 0.0
-      top_p: 1.0
-      top_k: -1
+      temperature: 0.9
+      top_k: 50
       max_tokens: 4096
       seed: 42
       detokenize: false
-      repetition_penalty: 1.0
-    final_output: false
+      repetition_penalty: 1.05
+      stop_token_ids: [2150]
 
   - stage_id: 1
     stage_type: llm
     runtime:
       devices: "0"
-      max_batch_size: 1
     engine_args:
-      dtype: float32
-      model_stage: vae
-      model_arch: VoxCPMForConditionalGeneration
-      hf_config_path: ${oc.env:VLLM_OMNI_VOXCPM_HF_CONFIG_PATH,}
+      max_num_seqs: 1
+      model_stage: code2wav
+      model_arch: Qwen3TTSCode2Wav
       worker_type: generation
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       enforce_eager: true
       trust_remote_code: true
-      async_scheduling: false
+      async_scheduling: true
       enable_prefix_caching: false
       engine_output_type: audio
-      gpu_memory_utilization: 0.1
+      gpu_memory_utilization: 0.3
       distributed_executor_backend: "mp"
       max_num_batched_tokens: 8192
-      max_model_len: 4096
+      max_model_len: 32768
     engine_input_source: [0]
-    input_connectors:
-      from_stage_0: connector_of_shared_memory
     final_output: true
     final_output_type: audio
+    input_connectors:
+      from_stage_0: connector_of_shared_memory
+    tts_args:
+      max_instructions_length: 500
     default_sampling_params:
       temperature: 0.0
       top_p: 1.0
       top_k: -1
-      max_tokens: 1
+      max_tokens: 65536
       seed: 42
       detokenize: true
       repetition_penalty: 1.0
 
 runtime:
   enabled: true
+  defaults:
+    window_size: -1
+    max_inflight: 1
 
   connectors:
     connector_of_shared_memory:
       name: SharedMemoryConnector
       extra:
         shm_threshold_bytes: 65536
-        codec_streaming: false
+        codec_streaming: true
         connector_get_sleep_s: 0.01
         connector_get_max_wait_first_chunk: 3000
         connector_get_max_wait: 300
+        codec_chunk_frames: 25
+        codec_left_context_frames: 25
 
   edges:
     - from: 0
       to: 1
+      window_size: -1

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

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

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

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

benchmarks/qwen3-tts/vllm_omni/plot_async_chunk.py

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

benchmarks/qwen3-tts/vllm_omni/run_async_chunk_benchmark.sh

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

benchmarks/tts/README.md

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

benchmarks/tts/bench_tts.py

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

benchmarks/tts/bench_voxcpm_offline.py

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

benchmarks/tts/model_configs.yaml

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

benchmarks/tts/plot_results.py

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

docker/Dockerfile.ci

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

docker/Dockerfile.cuda

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

docker/Dockerfile.npu

@@ -1,20 +1,20 @@
 ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
-ARG VLLM_ASCEND_TAG=v0.18.0rc1
+ARG VLLM_ASCEND_TAG=v0.17.0rc1
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 
-# WORKDIR /vllm-workspace/vllm
-# RUN git fetch origin --tags && git checkout v0.18.0
+WORKDIR /vllm-workspace/vllm
+RUN git fetch origin --tags && git checkout v0.18.0
 
-# WORKDIR /vllm-workspace/vllm-ascend
-# RUN git fetch origin releases/v0.18.0 && git checkout d781902ce9dbda8ab1e11bb0f2f0c1bc508fee7a
-# # Install vllm-ascend
-# # Append `libascend_hal.so` path (devlib) to LD_LIBRARY_PATH
-# RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
-#     source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
-#     source /usr/local/Ascend/nnal/atb/set_env.sh && \
-#     export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
-#     python3 -m pip install -v -e /vllm-workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/ && \
-#     python3 -m pip cache purge
+WORKDIR /vllm-workspace/vllm-ascend
+RUN git fetch origin releases/v0.18.0 && git checkout d781902ce9dbda8ab1e11bb0f2f0c1bc508fee7a
+# Install vllm-ascend
+# Append `libascend_hal.so` path (devlib) to LD_LIBRARY_PATH
+RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
+    source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
+    source /usr/local/Ascend/nnal/atb/set_env.sh && \
+    export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
+    python3 -m pip install -v -e /vllm-workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/ && \
+    python3 -m pip cache purge
 
 ARG APP_DIR=/vllm-workspace/vllm-omni
 WORKDIR ${APP_DIR}

docker/Dockerfile.npu.a3

@@ -1,20 +1,20 @@
 ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
-ARG VLLM_ASCEND_TAG=v0.18.0rc1-a3
+ARG VLLM_ASCEND_TAG=v0.17.0rc1-a3
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 
-# WORKDIR /vllm-workspace/vllm
-# RUN git fetch origin --tags && git checkout v0.18.0
+WORKDIR /vllm-workspace/vllm
+RUN git fetch origin --tags && git checkout v0.18.0
 
-# WORKDIR /vllm-workspace/vllm-ascend
-# RUN git fetch origin releases/v0.18.0 && git checkout d781902ce9dbda8ab1e11bb0f2f0c1bc508fee7a
-# # Install vllm-ascend
-# # Append `libascend_hal.so` path (devlib) to LD_LIBRARY_PATH
-# RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
-#     source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
-#     source /usr/local/Ascend/nnal/atb/set_env.sh && \
-#     export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
-#     python3 -m pip install -v -e /vllm-workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/ && \
-#     python3 -m pip cache purge
+WORKDIR /vllm-workspace/vllm-ascend
+RUN git fetch origin releases/v0.18.0 && git checkout d781902ce9dbda8ab1e11bb0f2f0c1bc508fee7a
+# Install vllm-ascend
+# Append `libascend_hal.so` path (devlib) to LD_LIBRARY_PATH
+RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
+    source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
+    source /usr/local/Ascend/nnal/atb/set_env.sh && \
+    export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
+    python3 -m pip install -v -e /vllm-workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/ && \
+    python3 -m pip cache purge
 
 ARG APP_DIR=/vllm-workspace/vllm-omni
 WORKDIR ${APP_DIR}

docker/Dockerfile.rocm

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

docker/Dockerfile.xpu

@@ -15,7 +15,9 @@ RUN apt clean && apt-get update -y && \
     apt-get install -y --no-install-recommends --fix-missing \
     curl \
     espeak-ng \
+    ffmpeg \
     git \
+    libsndfile1 \
     libsm6 \
     libxext6 \
     libgl1 \
@@ -74,7 +76,7 @@ ENV UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
 ENV UV_INDEX_STRATEGY="unsafe-best-match"
 ENV UV_LINK_MODE="copy"
 
-ARG VLLM_VERSION=v0.19.0
+ARG VLLM_VERSION=v0.18.0
 RUN git clone -b ${VLLM_VERSION} https://github.com/vllm-project/vllm
 WORKDIR /workspace/vllm
 

docs/.nav.yml

@@ -10,7 +10,6 @@ nav:
       - Image Generation: serving/image_generation_api.md
       - Image Edit: serving/image_edit_api.md
       - Text to Speech: serving/speech_api.md
-      - Streaming Video Input: serving/video_stream_api.md
   - Examples:
     - examples/README.md
     - Offline Inference:
@@ -35,7 +34,6 @@ nav:
     - Online Serving:
       - BAGEL-7B-MoT: user_guide/examples/online_serving/bagel.md
       - vLLM-Omni Helm Chart: user_guide/examples/online_serving/chart-helm.md
-      - Diffusers Backend Adapter Example: user_guide/examples/online_serving/diffusers_pipeline_adapter.md
       - Fish Speech S2 Pro: user_guide/examples/online_serving/fish_speech.md
       - GLM-Image Online Serving: user_guide/examples/online_serving/glm_image.md
       - Image-To-Image: user_guide/examples/online_serving/image_to_image.md
@@ -66,8 +64,6 @@ nav:
         - FP8: user_guide/diffusion/quantization/fp8.md
         - Int8: user_guide/diffusion/quantization/int8.md
         - GGUF: user_guide/diffusion/quantization/gguf.md
-      - Attention Backends: user_guide/diffusion/attention_backends.md
-      - Frame Interpolation: user_guide/diffusion/frame_interpolation.md
       - Parallelism:
         - Overview: user_guide/diffusion/parallelism/overview.md
         - CFG Parallel: user_guide/diffusion/parallelism/cfg_parallel.md
@@ -84,6 +80,7 @@ nav:
 - Developer Guide:
   - General:
     - contributing/README.md
+    - pr_reviewer.md
     - glob: contributing/*
       flatten_single_child_sections: true
   - Model Implementation:
@@ -100,16 +97,14 @@ nav:
       - design/feature/disaggregated_inference.md
       - design/feature/ray_based_execution.md
       - design/feature/omni_connectors/
-      - design/feature/prefix_caching.md
       - design/feature/cfg_parallel.md
-      - design/feature/expert_parallel.md
       - design/feature/sequence_parallel.md
       - design/feature/tensor_parallel.md
       - design/feature/vae_parallel.md
       - design/feature/hsdp.md
       - design/feature/cache_dit.md
       - design/feature/teacache.md
-      - design/feature/async_chunk.md
+      - design/feature/async_chunk_design.md
       - design/feature/vae_parallel.md
       - design/feature/diffusion_step_execution.md
     - Module Design:

docs/api/README.md

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

docs/cli/serve.md

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

docs/community/governance.md

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

docs/configuration/README.md

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

docs/configuration/pd_disaggregation.md

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

docs/configuration/stage_configs.md

@@ -3,210 +3,7 @@
 In vLLM-Omni, the target model is separated into multiple stages, which are processed by different LLMEngines, DiffusionEngines or other types of engines. Depending on different types of stages, such as Autoregressive (AR) stage or Diffusion transformer (DiT) stage, each can choose corresponding schedulers, model workers to load with the Engines in a plug-in fashion.
 
 !!! note
-    Default deploy config YAMLs (for example, `vllm_omni/deploy/qwen2_5_omni.yaml`, `vllm_omni/deploy/qwen3_omni_moe.yaml`, and `vllm_omni/deploy/qwen3_tts.yaml`) are bundled and loaded automatically when neither `--stage-configs-path` nor `--deploy-config` is provided — the model registry resolves the right pipeline + deploy YAML by `model_type`. The bundled defaults have been verified on 1xH100 for Qwen2.5-Omni and 2xH100 for Qwen3-Omni. Models that have not yet migrated to the new schema continue to use the legacy `vllm_omni/model_executor/stage_configs/<model>.yaml` files via `--stage-configs-path`.
-
-## New deploy schema reference
-
-The new deploy schema lives under `vllm_omni/deploy/` and is paired with a frozen `PipelineConfig` registered by the model's `pipeline.py`. Each deploy YAML has these top-level fields:
-
-| Field | Type | Required | Default | Description |
-|-------|------|----------|---------|-------------|
-| `base_config` | str (path) | optional | — | Overlay parent (relative or absolute). `stages:` / `platforms:` deep-merged by stage_id; other scalars overlay-wins. Intended for user-authored overlays; prod yamls stay flat. |
-| `async_chunk` | bool | optional | `true` | Enable chunked streaming between stages. Pin to `false` if the pipeline runs end-to-end. |
-| `connectors` | dict | optional | `null` | Named connector specs (`{name, extra}`). Referenced by each stage's `input_connectors` / `output_connectors`. See [Connector schema](#connector-schema). |
-| `edges` | list | optional | `null` | Explicit edge list for the KV transfer graph. Auto-derived from stage inputs if omitted. |
-| `stages` | list | required | — | Per-stage engine args + wiring (see [Stage fields](#stage-fields)). |
-| `platforms` | dict | optional | `null` | Keyed by `npu` / `rocm` / `xpu`, each contains a `stages:` list with per-platform overrides applied on top of the CUDA defaults. |
-| `pipeline` | str | optional | `null` | Override the auto-detected pipeline registry key (used for structural variants like `qwen2_5_omni_thinker_only`). |
-| `trust_remote_code` | bool | optional | `true` | **Pipeline-wide.** Trust HF remote code on model load; applies to every stage. |
-| `distributed_executor_backend` | str \| null | optional | `null` | **Pipeline-wide.** Distributed executor backend forwarded to vLLM (`"mp"`, `"ray"`, `"external_launcher"`). If omitted, vLLM auto-selects backend from runtime topology. |
-| `dtype` | str \| null | optional | `null` | **Pipeline-wide.** Model dtype for every stage. |
-| `quantization` | str \| null | optional | `null` | **Pipeline-wide.** Quantization method for every stage. |
-| `enable_prefix_caching` | bool | optional | `false` | **Pipeline-wide.** Prefix cache toggle applied to every stage. |
-| `enable_chunked_prefill` | bool \| null | optional | `null` | **Pipeline-wide.** Chunked prefill toggle applied to every stage. |
-| `data_parallel_size` | int | optional | `1` | **Pipeline-wide.** DP degree for every stage. |
-| `pipeline_parallel_size` | int | optional | `1` | **Pipeline-wide.** PP degree for every stage. |
-
-Note: for diffusion path, `distributed_executor_backend` currently defaults to
-`mp`, and `ray` / `external_launcher` are not fully supported yet.
-
-### Stage fields
-
-Each entry under `stages:` accepts any `StageDeployConfig` field directly (no nested `engine_args:`). Only fields whose value legitimately varies across stages live here; pipeline-wide settings (trust_remote_code, distributed_executor_backend, dtype, quantization, prefix/chunked prefill, DP/PP sizes) are declared at the top level and applied to every stage. Unknown keys fall through to `engine_extras:` and are forwarded to the engine.
-
-| Field | Type | Required | Default | Description |
-|-------|------|----------|---------|-------------|
-| `stage_id` | int | required | — | Stage identity; matched against `PipelineConfig.stages[*].stage_id`. |
-| `max_num_seqs` | int | optional | `64` | Max concurrent sequences per stage. |
-| `gpu_memory_utilization` | float | optional | `0.9` | Per-stage memory budget. |
-| `tensor_parallel_size` | int | optional | `1` | TP degree for this stage. |
-| `enforce_eager` | bool | optional | `false` | Disable CUDA graphs. |
-| `max_num_batched_tokens` | int | optional | `32768` | Prefill budget. |
-| `max_model_len` | int \| null | optional | `null` | Per-stage context length (auto-sets `VLLM_ALLOW_LONG_MAX_MODEL_LEN=1` when larger than HF default). |
-| `async_scheduling` | bool \| null | optional | `null` | Per-stage async scheduling toggle. |
-| `devices` | str | optional | `"0"` | `CUDA_VISIBLE_DEVICES`-style device list. |
-| `output_connectors` | dict \| null | optional | `null` | Keyed by `to_stage_<n>`; values are names registered under top-level `connectors:`. |
-| `input_connectors` | dict \| null | optional | `null` | Keyed by `from_stage_<n>`; values are names registered under top-level `connectors:`. |
-| `default_sampling_params` | dict \| null | optional | `null` | Baseline sampling params. Deep-merged with pipeline `sampling_constraints` (pipeline wins). |
-| `engine_extras` | dict | optional | `{}` | Catch-all for keys not listed above; deep-merged across overlays. Also carries per-stage overrides of pipeline-wide settings (e.g. stage-specific `dtype`). |
-
-### Connector schema
-
-Each entry under top-level `connectors:` follows this shape:
-
-```yaml
-connectors:
-  <connector_name>:
-    name: <ConnectorClassName>     # required — class registered in vllm_omni.distributed
-    extra:                         # optional — forwarded to the connector's __init__
-      <key>: <value>
-      ...
-```
-
-| Connector class | Use case | `extra` keys |
-|-----------------|----------|--------------|
-| `SharedMemoryConnector` | Same-host KV transfer between stages (default for bundled YAMLs). | `shm_threshold_bytes` (int, default `65536`). |
-| `MooncakeStoreConnector` | Cross-host KV transfer over TCP. Required for multi-node deployments. | `host`, `metadata_server`, `master`, `segment` (int bytes), `localbuf` (int bytes), `proto` (`"tcp"` / `"rdma"`). |
-
-A stage references a connector by name in its `input_connectors` / `output_connectors`:
-
-```yaml
-connectors:
-  shm:
-    name: SharedMemoryConnector
-
-stages:
-  - stage_id: 0
-    output_connectors: {to_stage_1: shm}
-  - stage_id: 1
-    input_connectors:  {from_stage_0: shm}
-```
-
-### CLI flags introduced in this refactor
-
-| Flag | Description |
-|------|-------------|
-| `--deploy-config PATH` | Load a new-schema deploy YAML. Takes precedence over `--stage-configs-path`. **Optional** — when omitted, the bundled `vllm_omni/deploy/<model_type>.yaml` is auto-loaded by the model registry. |
-| `--stage-overrides JSON` | Per-stage JSON overrides, e.g. `'{"0":{"gpu_memory_utilization":0.5}}'`. Per-stage values always win over global flags. |
-| `--async-chunk` / `--no-async-chunk` | Flip the deploy YAML's `async_chunk:` bool. Unset (default) leaves the YAML value in force. |
-| `--stage-configs-path` | **Deprecated.** Accepts legacy `stage_args` yamls and (auto-detected) new deploy yamls; emits a deprecation warning. Migrate to `--deploy-config`. To be removed in a follow-up PR. |
-
-### Stage-Based CLI Paradigm
-
-The stage-based CLI paradigm facilitates the execution of discrete pipeline stages within isolated processes:
-
-- **Stage 0** typically encapsulates the orchestrator and the primary API server. Invocation requires `--stage-id 0`,
-  `--omni-master-address`, `--omni-master-port`, and standard port declarations (e.g., `--port`).
-- **Worker Stages** operate without a distinct API server (i.e., using `--headless`), are assigned sequential `--stage-id` identifiers, and must reference the corresponding
-  `--omni-master-address` and `--omni-master-port` parameters to successfully register with Stage 0.
-
-For migrated architectures, the system automatically resolves and loads the bundled deployment YAML. Consequently, the primary execution path
-does **not** necessitate the explicit definition of `--deploy-config`:
-
-```bash
-vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
-    --port 8091 \
-    --stage-id 0 \
-    --omni-master-address 127.0.0.1 \
-    --omni-master-port 26000
-
-vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
-    --stage-id 1 \
-    --headless \
-    --omni-master-address 127.0.0.1 \
-    --omni-master-port 26000
-```
-
-When instantiating a custom deployment YAML conforming to the updated schema, append the `--deploy-config /path/to/override.yaml` directive
-to all node invocations. For legacy architectures (e.g., BAGEL) configured via deprecated `stage_args:` schemas, continue to specify the relevant configuration via `--stage-configs-path /path/to/config.yaml`.
-
-In the context of standard initialization architectures, utilizing the `--stage-overrides` parameter operates as the optimal methodology
-for delineating stage-specific tuning from the CLI interface:
-
-```bash
-vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
-    --stage-overrides '{"1": {"gpu_memory_utilization": 0.5}}'
-```
-
-Conversely, in the context of the **stage-based CLI** paradigm, given that each execution process exclusively instantiates a single pipeline stage, configuration override attributes
-can be defined uniformly via explicit CLI flags on the corresponding instantiation command, rendering composite `--stage-overrides` JSON strings unnecessary:
-
-```bash
-vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
-    --stage-id 1 \
-    --headless \
-    --gpu-memory-utilization 0.5 \
-    --omni-master-address 127.0.0.1 \
-    --omni-master-port 26000
-```
-
-### Precedence
-
-From highest to lowest:
-
-1. Per-stage flags (`--stage-overrides` JSON, `--stage-<id>-<key>` if registered)
-2. Explicit global CLI flags (`--gpu-memory-utilization 0.85`, etc.)
-3. Platform section (`platforms.npu.stages`, etc.) on top of the base `stages:`
-4. Overlay YAML (via `base_config:`) on top of the base YAML
-5. Parser defaults
-
-### Worked override example
-
-Starting from the bundled `vllm_omni/deploy/qwen3_omni_moe.yaml`:
-
-```yaml
-# vllm_omni/deploy/qwen3_omni_moe.yaml (excerpt)
-async_chunk: true
-stages:
-  - stage_id: 0
-    gpu_memory_utilization: 0.9
-    max_num_seqs: 32
-  - stage_id: 1
-    gpu_memory_utilization: 0.7
-    max_num_seqs: 16
-```
-
-A user-authored overlay that inherits the base and overrides only stage 1:
-
-```yaml
-# my_overrides.yaml
-base_config: /path/to/vllm_omni/deploy/qwen3_omni_moe.yaml
-stages:
-  - stage_id: 1
-    gpu_memory_utilization: 0.5     # smaller GPU
-```
-
-Launched with both an explicit global flag and a per-stage override:
-
-```bash
-vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
-    --deploy-config my_overrides.yaml \
-    --max-model-len 16384 \
-    --stage-overrides '{"0": {"max_num_seqs": 8}}'
-```
-
-Within the stage-based CLI paradigm, equivalent configuration parameters can inherently be passed directly
-as command-line arguments to the designated single-stage process instantiation:
-
-```bash
-vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
-    --stage-id 0 \
-    --max-num-seqs 8 \
-    --omni-master-address 127.0.0.1 \
-    --omni-master-port 26000
-```
-
-Effective config per stage after the merge:
-
-| Stage | Field | Final value | Source |
-|-------|-------|-------------|--------|
-| 0 | `gpu_memory_utilization` | `0.9` | base YAML (overlay didn't touch stage 0) |
-| 0 | `max_num_seqs` | `8` | per-stage CLI (`--stage-overrides`) — wins over base `32` |
-| 0 | `max_model_len` | `16384` | global CLI |
-| 1 | `gpu_memory_utilization` | `0.5` | overlay YAML — wins over base `0.7` |
-| 1 | `max_num_seqs` | `16` | base YAML (overlay didn't touch this field) |
-| 1 | `max_model_len` | `16384` | global CLI |
-| 2 | (all defaults) | — | base YAML (no overrides apply) |
+    Default stage config YAMLs (for example, `vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml` and `vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml`) are bundled and loaded automatically when `stage_configs_path` is not provided. They have been verified to work on 1xH100 for Qwen2.5-Omni and 2xH100 for Qwen3-Omni.
 
 Therefore, as a core part of vLLM-Omni, the stage configs for a model have several main functions:
 
@@ -216,14 +13,9 @@ Therefore, as a core part of vLLM-Omni, the stage configs for a model have sever
 - Input and output dependencies for each stage.
 - Default input parameters.
 
-To override specific parameters, explicitly inject the customized configuration schema
-in both online and offline instantiation flows. Prioritize the `--deploy-config` flag
-when loading the new-schema deploy YAML schemas, reserving the `--stage-configs-path` parameter
-exclusively to maintain compatibility with legacy `stage_args` YAML constructs.
+If users want to modify some part of it. The custom stage_configs file can be input as input argument in both online and offline. Just like examples below:
 
-Examples:
-
-For offline (Assume necessary dependencies have been imported):
+For offline (Assume necessary dependencies have ben imported):
 ```python
 model_name = "Qwen/Qwen2.5-Omni-7B"
 omni = Omni(model=model_name, stage_configs_path="/path/to/custom_stage_configs.yaml")
@@ -231,13 +23,7 @@ omni = Omni(model=model_name, stage_configs_path="/path/to/custom_stage_configs.
 
 For online serving:
 ```bash
-vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --deploy-config /path/to/deploy_config.yaml
-```
-
-Legacy online serving:
-
-```bash
-vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
+vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
 ```
 !!! important
     We are actively iterating on the definition of stage configs, and we welcome all feedbacks from both community users and developers to help us shape the development!
@@ -249,7 +35,7 @@ stage_args:
   - stage_id: 0 # mark the unique id for each stage
     runtime: # The disaggregated configuration
       process: true  # Run this stage in a separate process
-      devices: "0" # Logical device index for this stage (mapped through CUDA_VISIBLE_DEVICES / ASCEND_RT_VISIBLE_DEVICES if set)
+      devices: "0" # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device)
     engine_args: # Engine arguments for a certain engine
       model_stage: thinker
       max_num_seqs: 1
@@ -328,12 +114,16 @@ stage_args:
 # Top-level runtime config (concise): default windows and stage edges
 runtime:
   enabled: true
-
+  defaults:
+    window_size: -1             # Simplified: trigger downstream only after full upstream completion
+    max_inflight: 1             # Simplified: process serially within each stage
   edges:
     - from: 0                   # thinker → talker: trigger only after receiving full input (-1)
       to: 1
+      window_size: -1
     - from: 1                   # talker → code2wav: trigger only after receiving full input (-1)
       to: 2
+      window_size: -1
 
 ```
 
@@ -365,9 +155,7 @@ Default: `true`
 
 #### `runtime.devices`
 
-Logical device indices for this stage, specified as a string. Values are **logical indices** (`0`, `1`, `2`, ...) — not physical GPU IDs — and are mapped through the platform's visibility env var (`CUDA_VISIBLE_DEVICES` on CUDA, `ASCEND_RT_VISIBLE_DEVICES` on NPU) before being applied via `torch.cuda.set_device()` (or the equivalent).
-
-Example: if `CUDA_VISIBLE_DEVICES=0,2,4` is set in the environment, then `devices: "0"` selects physical GPU 0 (the first visible), `devices: "1"` selects physical GPU 2, and `devices: "0,1"` makes physical GPUs 0 and 2 available to the stage. If no visibility env var is set, logical and physical IDs coincide.
+Visible devices for this stage, specified as a string. This controls which GPU devices are available to the stage process, similar to setting `CUDA_VISIBLE_DEVICES` or using `torch.cuda.set_device()`. For example, `"0"` uses GPU 0, `"1"` uses GPU 1, and `"0,1"` makes both GPUs 0 and 1 visible.
 
 Default: `"0"`
 

docs/configuration/stage_configs/qwen2_5_omni.yaml

@@ -0,0 +1,94 @@
+# stage config for running qwen2.5-omni with AsyncOmniEngine + Orchestrator runtime.
+stage_args:
+  - stage_id: 0
+    runtime:
+      process: true            # Run this stage in a separate process
+      devices: "0"            # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device)
+    engine_args:
+      model_stage: thinker
+      max_num_seqs: 1
+      model_arch: Qwen2_5OmniForConditionalGeneration
+      worker_type: ar
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      gpu_memory_utilization: 0.8
+      enforce_eager: true  # Now we only support eager mode
+      trust_remote_code: true
+      engine_output_type: latent
+      enable_prefix_caching: false
+    is_comprehension: true
+    final_output: true
+    final_output_type: text
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 2048
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.1
+  - stage_id: 1
+    runtime:
+      process: true
+      devices: "1"
+    engine_args:
+      model_stage: talker
+      max_num_seqs: 1
+      model_arch: Qwen2_5OmniForConditionalGeneration
+      worker_type: ar
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      gpu_memory_utilization: 0.8
+      enforce_eager: true
+      trust_remote_code: true
+      enable_prefix_caching: false
+      engine_output_type: latent
+    engine_input_source: [0]
+    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen2_5_omni.thinker2talker
+    default_sampling_params:
+      temperature: 0.9
+      top_p: 0.8
+      top_k: 40
+      max_tokens: 2048
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.05
+      stop_token_ids: [8294]
+  - stage_id: 2
+    runtime:
+      process: true
+      devices: "0"            # Example: use a different GPU than the previous stage; use "0" if single GPU
+    engine_args:
+      model_stage: code2wav
+      max_num_seqs: 1
+      model_arch: Qwen2_5OmniForConditionalGeneration
+      worker_type: generation
+      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
+      gpu_memory_utilization: 0.15
+      enforce_eager: true
+      trust_remote_code: true
+      enable_prefix_caching: false
+      engine_output_type: audio
+    engine_input_source: [1]
+    final_output: true
+    final_output_type: audio
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 2048
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.1
+
+# Top-level runtime config (concise): default windows and stage edges
+runtime:
+  enabled: true
+  defaults:
+    window_size: -1             # Simplified: trigger downstream only after full upstream completion
+    max_inflight: 1             # Simplified: process serially within each stage
+  edges:
+    - from: 0                   # thinker → talker: trigger only after receiving full input (-1)
+      to: 1
+      window_size: -1
+    - from: 1                   # talker → code2wav: trigger only after receiving full input (-1)
+      to: 2
+      window_size: -1

docs/contributing/ci/CI_5levels.md

@@ -86,8 +86,7 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
         /tests/e2e/online_serving/test_{model_name}_expansion.py<br>
         /tests/e2e/offline_inference/test_{model_name}_expansion.py<br>
         <strong>Performance:</strong><br>
-        /tests/dfx/perf/tests/test_qwen_omni.json (Omni), test_tts.json (TTS),<br>
-        and /tests/dfx/perf/tests/test_{diffusion_model}_vllm_omni.json (Diffusion)<br>
+        /tests/dfx/perf/tests/test.json<br>
         <strong>Doc Test:</strong><br>
         tests/example/online_serving/test_{model_name}.py<br>
         tests/example/offline_inference/test_{model_name}.py
@@ -105,8 +104,7 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
       <td> Depends on reality </td>
       <td>
         <strong>Stability:</strong><br>
-        /tests/dfx/stability/tests/test_qwen3_omni.json<br>
-        /tests/dfx/stability/tests/test_wan22.json<br>
+        /tests/dfx/stability/tests/test.json<br>
         <strong>Reliability:</strong><br>
         tests/e2e/reliability/test_{model_name}.py
       </td>
@@ -170,6 +168,7 @@ vllm_omni/                                    tests/
 │   └── arg_utils.py                            │   └── test_arg_utils.py               ⬜
 │
 ├── entrypoints/                        →     ├── entrypoints/
+│   ├── async_omni_diffusion.py                 │   ├── test_async_omni_diffusion_config.py  ✅
 │   ├── stage_utils.py                          │   ├── test_stage_utils.py            ✅
 │   ├── cli/                                     │   ├── cli/                           (benchmarks/test_serve_cli.py covers CLI serve)
 │   │   └── ...                                  │   │   └── test_*.py                  ⬜
@@ -232,7 +231,8 @@ vllm_omni/                                    tests/
                                                │   ├── test_qwen3_omni_expansion.py
                                                │   ├── test_mimo_audio.py
                                                │   ├── test_image_gen_edit.py
-                                               │   └── test_images_generations_lora.py
+                                               │   ├── test_images_generations_lora.py
+                                               │   └── stage_configs/
                                                └── offline_inference/                  ✅
                                                    ├── test_qwen2_5_omni.py
                                                    ├── test_qwen3_omni.py
@@ -243,17 +243,16 @@ vllm_omni/                                    tests/
                                                    ├── test_zimage_tensor_parallel.py
                                                    ├── test_cache_dit.py
                                                    ├── test_teacache.py
-                                                   ├── test_stable_audio_expansion.py
+                                                   ├── test_stable_audio_model.py
                                                    ├── test_diffusion_cpu_offload.py
                                                    ├── test_diffusion_layerwise_offload.py
                                                    ├── test_diffusion_lora.py
                                                    ├── test_sequence_parallel.py
-                                                   └── stage_configs/                  (legacy schema, still
-                                                       ├── bagel_*.yaml                 present for unmigrated
-                                                       └── npu/, rocm/, etc.            models)
-
-# Migrated models (qwen3_omni_moe, qwen2_5_omni, qwen3_tts) live under
-# vllm_omni/deploy/ instead — see docs/configuration/stage_configs.md.
+                                                   └── stage_configs/
+                                                       ├── qwen2_5_omni_ci.yaml
+                                                       ├── qwen3_omni_ci.yaml
+                                                       ├── bagel_*.yaml
+                                                       └── npu/, rocm/, etc.
 ```
 
 
@@ -273,7 +272,7 @@ Before entering specific testing levels, the project establishes two common spec
 
 L1 and L2 level testing form the foundation of the quality assurance system. L1 level testing focuses on verifying the internal logic correctness of code units (e.g., functions, classes), ensuring each independent component behaves as designed.
 
-L2 level testing builds upon L1 by introducing GPU resources and verifying that the end-to-end (E2E) process of the model in basic deployment scenarios is smooth. For example, it uses dummy models to confirm that core interfaces like the inference pipeline, output format, and streaming response work properly. The common goal of these two levels is to provide developers with rapid feedback, discovering and fixing issues early in the development cycle.
+L2 level testing builds upon L1 by introducing GPU resources and verifying that the end-to-end (E2E) process of the model in basic deployment scenarios is smooth. For example, it uses dummy models to confirm that core interfaces like the inference pipeline, output format, and streaming response work properly. The common goal of these two levels is to provide developers with rapid feedback, discovering and fixing issues early in the development cycle  .
 
 
 
@@ -419,13 +418,13 @@ L3 level testing executes after code is merged into the main branch. Its core pu
 
     **Explanation**:
 
-    @pytest.mark.advanced_model: Marks the test as L3 merge level, indicating deep validation with real models. @pytest.mark.full_model: Marks L4 nightly-only suites (e.g. `test_*_expansion.py`, doc examples).
+    @pytest.mark.advanced_model: Marks the test as L3 or L4 level, indicating that this test case performs deep validation, using real models for performance, integration, and accuracy testing. This forms a "basic-advanced" correspondence with the core_model mark at the L2 level.
 
     @pytest.mark.core_model: Marks the test as L1 or L2 level, indicating that this test case validates the basic functionality of the core model. It uses mock weights and only checks if the relevant interface functions correctly.
 
     @pytest.mark.parametrize: A parameterization decorator that allows abstracting test data into parameters, enabling reuse of the same test logic across different data configurations. indirect=True indicates that parameters will be passed to the fixture for processing.
 
-    **Notes**: If you believe the test case only needs to execute basic run logic at the PR-level CI, you can mark it only with @pytest.mark.core_model. If you believe it only needs to execute deep validation at merge (L3), use @pytest.mark.advanced_model. For L4 nightly-only expansion and doc-example tests, use @pytest.mark.full_model with `--run-level full_model`. If the test case needs both basic run and deep validation, mark with @pytest.mark.core_model and the appropriate L3/L4 marker (`advanced_model` and/or `full_model`).
+    **Notes**: If you believe the test case only needs to execute basic run logic at the PR-level CI, you can mark it only with @pytest.mark.core_model. If you believe it only needs to execute deep validation run logic at the merge or nightly level, you can mark it only with @pytest.mark.advanced_model. If you believe the test case needs to accommodate both basic run and deep validation test logic, you should mark it with both @pytest.mark.core_model and @pytest.mark.advanced_model.
 
     **2.4.2 Test Function Definition and Documentation**
 
@@ -517,11 +516,9 @@ L3 level testing executes after code is merged into the main branch. Its core pu
 
     **Single Request**: The comment clearly states this is a single-request completion test. For concurrent testing, it can be extended to multiple requests using request_num = n.
 
-    **Implicit Validation**: The `send_omni_request` and `send_diffusion_request` methods internally includes validation logic dynamically selected based on the --run-level parameter: core_model performs basic validation, while advanced_model and full_model perform deep validation.
+    **Implicit Validation**: The `send_omni_request` and `send_diffusion_request` methods internally includes validation logic dynamically selected based on the --run-level parameter: core_model performs basic validation, while advanced_model performs deep validation.
 
--   ***Run Command (L3 merge)***: `pytest -s -v /tests/e2e/online_serving/test_{model_name}.py -m advanced_model --run-level=advanced_model`
-
--   ***Run Command (L4 nightly expansion)***: `pytest -s -v /tests/e2e/online_serving/test_{model_name}_expansion.py -m full_model --run-level=full_model`
+-   ***Run Command***: `pytest -s -v /tests/e2e/online_serving/test_{model_name}.py -m advanced_model --run-level=advanced_model`
 
 ## Chapter 3: L4 Level Testing - Full Functionality, Performance, and Documentation Testing
 
@@ -534,13 +531,13 @@ L4 level testing is a comprehensive quality audit before a version release. It e
 ### 3.2 Testing Content and Scope
 
 -   ***Full Functionality Testing***: Executes all test cases defined in `test_{model_name}_expansion.py`, covering all implemented features, positive flows, boundary conditions, and exception handling.
--   ***Performance Testing***: Uses `tests/dfx/perf/tests/test_qwen_omni.json`, `tests/dfx/perf/tests/test_tts.json`, and diffusion configs in the form `tests/dfx/perf/tests/test_*_vllm_omni.json` (passed to `run_benchmark.py` via `--test-config-file`) to drive performance testing tools for stress, load, and endurance tests, collecting metrics like throughput, response time, and resource utilization.
+-   ***Performance Testing***: Uses the `tests/dfx/perf/tests/test.json` configuration file to drive performance testing tools for stress, load, and endurance tests, collecting metrics like throughput, response time, and resource utilization.
 -   ***Documentation Testing***: Verifies whether the example code provided to users is runnable and its results match the description.
 
 ### 3.3 Test Directory and Execution Files
 
 -   ***Functional Testing***: Same directories as L3.
--   ***Performance Test Configuration***: `tests/dfx/perf/tests/test_qwen_omni.json`, `tests/dfx/perf/tests/test_tts.json`, and diffusion configs `tests/dfx/perf/tests/test_*_vllm_omni.json` (e.g. `test_qwen_image_vllm_omni.json`)
+-   ***Performance Test Configuration***: `tests/dfx/perf/tests/test.json`
 -   ***Documentation Example Tests***:
 -   -   `tests/example/online_serving/test_{model_name}.py`
     -   `tests/example/offline_inference/test_{model_name}.py`
@@ -575,12 +572,12 @@ L5 level testing focuses on the performance of model services under ***long-runn
 
 ### 4.2 Testing Content and Scope
 
--   ***Long-term Stability (Stability) Testing***: Uses JSON under `tests/dfx/stability/tests/` (for example `test_qwen3_omni.json` and `test_wan22.json`) to run the service under moderate load for an extended period (e.g., over 12 hours), monitoring whether metrics like memory/VRAM usage, response time, and throughput degrade over time, and whether the service process remains stable.
+-   ***Long-term Stability (Stability) Testing***: Uses the `tests/dfx/stability/tests/test.json` configuration to run the service under moderate load for an extended period (e.g., over 12 hours), monitoring whether metrics like memory/VRAM usage, response time, and throughput degrade over time, and whether the service process remains stable.
 -   ***Reliability Testing***: Uses `tests/e2e/reliability/test_{model_name}.py` to actively simulate various fault and abnormal scenarios, such as: dependent service interruption, abnormal input data, network flicker, hardware resource preemption, etc., to verify the system's fault tolerance, self-healing, and graceful degradation capabilities.
 
 ### 4.3 Test Directory and Execution Files
 
--   ***Stability Test Configuration***: `tests/dfx/stability/tests/test_qwen3_omni.json`, `tests/dfx/stability/tests/test_wan22.json` (one JSON per model / runner family)
+-   ***Stability Test Configuration***: `tests/dfx/stability/tests/test.json`
 -   ***Reliability Test Suite***: `tests/e2e/reliability/test_{model_name}.py`
 
 ### 4.4 Execution Method and Example
@@ -591,7 +588,7 @@ L5 level testing focuses on the performance of model services under ***long-runn
 <details>
 <summary> Test Examples</summary>
 
-When you want to add L5-level stability test cases, add or extend the appropriate JSON file under `tests/dfx/stability/tests/` (for example `test_qwen3_omni.json` for Omni bench traffic, or `test_wan22.json` for diffusion `/v1/videos` workloads). The following illustrates the Qwen3-Omni shape:
+When you want to add L5-level stability test cases, you can refer to the following format for case addition in `tests/dfx/stability/tests/test.json`:
 
 ```json
 {
@@ -662,7 +659,7 @@ All other optional parameters follow the same rules as the in Chapter 3.4.
 
 </details>
 
--   -   ***Stability***: `pytest -s -v tests/dfx/stability/scripts/test_stability_qwen3_omni.py` or `pytest -s -v tests/dfx/stability/scripts/test_stability_wan22.py` (or add `test_stability_<model>.py` alongside a matching JSON config)
+-   -   ***Stability***: `pytest -s -v tests/dfx/stability/scripts/test_{model_name}.py`
     -   ***Reliability***: `pytest -s -v tests/e2e/reliability/test_{model_name}.py`
 
 ## Summary

docs/contributing/ci/test_examples/l4_functionality_tests.inc.md

@@ -37,10 +37,10 @@ Currently all the features are available in online serving mode. Hence, only nee
 **Code Style**
 
 - Validation: test that the multimodal output files of your model have the correct shapes. `OpenAIClientHandler.send_diffusion_request` should have taken care of this.
-- Test marks: always add `full_model` and `diffusion` for L4 nightly `test_*_expansion.py` cases. Add GPU-related marks if needed. Ref: [Markers for Tests](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/tests_markers/).
+- Test marks: always add `advanced_model` and `diffusion`. Add GPU-related marks if needed. Ref: [Markers for Tests](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/tests_markers/).
 - To maximize code reuse, you may refer to
     - `tests/conftest.py` for `omni_server` (running server in subprocess) and `openai_client` fixtures (sending requests and validating output), `generate_synthetic_image` and `assert_XXX_valid` helper.
-    - `tests/helpers/mark.py` for `@hardware_test(...)` and `hardware_marks`.
+    - `tests/utils.py` for `@hardware_test(...)` and `hardware_marks`.
     - [Parametrizing tests (pytest doc)](https://docs.pytest.org/en/stable/example/parametrize.html) to reuse test function implementation for different cases.
 - Doc: add a concise docstring for each test function.
 - Reference L4 test implementation: [tests/e2e/online_serving/test_qwen_image_edit_expansion.py](https://github.com/vllm-project/vllm-omni/blob/main/tests/e2e/online_serving/test_qwen_image_edit_expansion.py).

docs/contributing/ci/test_examples/l4_performance_tests.inc.md

@@ -1,4 +1,4 @@
-When you want to add L4-level ***performance test*** cases, you can refer to the following format for case addition in `tests/dfx/perf/tests/test_qwen_omni.json`, `tests/dfx/perf/tests/test_tts.json`, or diffusion configs such as `tests/dfx/perf/tests/test_*_vllm_omni.json` (selected via `pytest ... run_benchmark.py --test-config-file <path>`):
+When you want to add L4-level ***performance test*** cases, you can refer to the following format for case addition in tests/dfx/perf/tests/test.json:
 
 ```JSON
 {

docs/contributing/ci/test_guide.md

@@ -42,63 +42,32 @@ Our test scripts use the pytest framework. First, please use `git clone https://
     ```
     The latest test commands for various test suites can be found in the [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-ready.yml).
 
-=== "L3 level"
+=== "L3 level & L4 level"
 
     ```bash
+    cd tests
     pytest -s -v -m "advanced_model" --run-level=advanced_model
     ```
+    If you only want to run L3 test case, you can use:
+    ```bash
+    pytest -s -v e2e/ --ignore-glob='*expansion.py' -m "advanced_model" --run-level=advanced_model
+    ```
     If you only want to run a specific test case, you can use:
     ```bash
     pytest -s -v test_xxxx.py --run-level=advanced_model
     ```
     If you only want to run specific test cases on a particular platform, you can use:
     ```bash
-    pytest -s -v -m "advanced_model and distributed_cuda and L4"  --run-level=advanced_model
+    pytest -s -v -m "core_model and distributed_cuda and L4"  --run-level=core_model
     ```
+    Note: To run performance tests, use:
+    ```bash
+    pytest -s -v perf/scripts/run_benchmark.py
+    ```
+
     The latest L3 test commands for various test suites can be found in the [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-merge.yml).
 
-
-=== "L4 level"
-
-    ```bash
-    cd tests
-    pytest -s -v -m "full_model" --run-level=full_model
-    ```
-    If you only want to run a specific test case, you can use:
-    ```bash
-    pytest -s -v test_xxxx.py --run-level=full_model
-    ```
-    If you only want to run specific test cases on a particular platform, you can use:
-    ```bash
-    pytest -s -v -m "full_model and distributed_cuda and L4"  --run-level=full_model
-    ```
-    Note: To run performance tests (defaults to ``test_qwen_omni.json``; use ``--test-config-file tests/dfx/perf/tests/test_tts.json`` for TTS):
-    ```bash
-    pytest -s -v tests/dfx/perf/scripts/run_benchmark.py
-    ```
-    The latest L4 (nightly) test commands use the `full_model` marker and `--run-level full_model` (see [test-nightly.yml](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-nightly.yml) and [test-nightly-diffusion.yml](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-nightly-diffusion.yml)). Example:
-
-    ```bash
-    cd tests
-    pytest -s -v -m "full_model and omni and H100" --run-level=full_model
-    ```
-
-=== "L5 level"
-
-    L5 includes stability and reliability testing. Typical commands:
-
-    ```bash
-    cd tests
-
-    # Stability: Qwen3-Omni
-    pytest -s -v dfx/stability/scripts/test_stability_qwen3_omni.py
-
-    # Stability: Wan2.2 (v1/videos diffusion benchmark loop)
-    pytest -s -v dfx/stability/scripts/test_stability_wan22.py
-
-    ```
-
-    The latest L5 commands for CI can be found in the [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-ready.yml).
+    The latest L4 test commands for various test suites can be found in the [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-nightly.yml).
 
 You can find more information about markers in the documentation: [marker doc](./tests_markers.md)
 

docs/contributing/ci/tests_markers.md

@@ -8,8 +8,7 @@ Defined in `pyproject.toml`:
 | Marker             | Description                                               |
 | ------------------ | --------------------------------------------------------- |
 | `core_model`       | L1&L2 tests (run in each PR)                              |
-| `advanced_model`   | L3 tests (run on each merge to main)                 |
-| `full_model`       | L4 tests (run nightly) |
+| `advanced_model`   | L3&L4 level tests (run in each merge or nightly)          |
 | `diffusion`        | Diffusion model tests                                     |
 | `omni`             | Omni model tests                                          |
 | `cache`            | Cache backend tests                                       |
@@ -39,7 +38,7 @@ Defined in `pyproject.toml`:
 ### Example usage for markers
 
 ```python
-from tests.helpers.mark import hardware_test
+from tests.utils import hardware_test
 
 @pytest.mark.core_model
 @pytest.mark.omni
@@ -54,7 +53,7 @@ def test_video_to_audio()
 
 ### Decorator: `@hardware_test`
 
-This decorator is intended to make hardware-aware, cross-platform test authoring easier and more robust for CI/CD environments. The `hardware_test` decorator in `vllm-omni/tests/helpers/mark.py` performs the following actions:
+This decorator is intended to make hardware-aware, cross-platform test authoring easier and more robust for CI/CD environments. The `hardware_test` decorator in `vllm-omni/tests/utils.py` performs the following actions:
 
 1. **Applies platform and resource markers**  
    Adds the appropriate pytest markers for each specified hardware platform (e.g., `cuda`, `rocm`, `xpu`, `npu`) and resource type (e.g., `L4`, `H100`, `MI325`, `B60`, `A2`, `A3`).
@@ -106,7 +105,7 @@ This decorator is intended to make hardware-aware, cross-platform test authoring
 `hardware_marks` returns a list of pytest mark objects with the same signature as `@hardware_test`. Use it when you need more flexibility, such as attaching hardware marks to individual `pytest.param` entries rather than an entire test function.
 
 ```python
-from tests.helpers.mark import hardware_marks
+from tests.utils import hardware_marks
 
 MULTI_CARD_MARKS = hardware_marks(
     res={"cuda": "H100", "rocm": "MI325", "npu": "A2"}, num_cards=2
@@ -134,9 +133,9 @@ If you want to add support for a new platform (e.g., "tpu" for a new accelerator
        "distributed_tpu: Tests that require multiple TPU devices",
    ]
    ```
-2. **Implement a marker construction function for your platform** in `vllm-omni/tests/helpers/mark.py`:
+2. **Implement a marker construction function for your platform** in `vllm-omni/tests/utils.py`:
    ```python
-   # In vllm-omni/tests/helpers/mark.py
+   # In vllm-omni/tests/utils.py
 
    def tpu_marks(*, res: str, num_cards: int):
        test_platform = pytest.mark.tpu
@@ -176,4 +175,4 @@ If you want to add support for a new platform (e.g., "tpu" for a new accelerator
 - Plug into `hardware_marks`  
 - You're done: tests using `@hardware_test` or `hardware_marks` with your platform now automatically get the correct markers, distribution, and isolation!
 
-See code in `vllm-omni/tests/helpers/mark.py` for existing examples (`cuda_marks`, `rocm_marks`, `npu_marks`).
+See code in `vllm-omni/tests/utils.py` for existing examples (`cuda_marks`, `rocm_marks`, `npu_marks`).

docs/contributing/ci/tests_style.md

@@ -73,6 +73,7 @@ vllm_omni/                                    tests/
 │   └── arg_utils.py                            │   └── test_arg_utils.py               ⬜
 │
 ├── entrypoints/                        →     ├── entrypoints/
+│   ├── async_omni_diffusion.py                 │   ├── test_async_omni_diffusion_config.py  ✅
 │   ├── stage_utils.py                          │   ├── test_stage_utils.py            ✅
 │   ├── cli/                                     │   ├── cli/                           (benchmarks/test_serve_cli.py covers CLI serve)
 │   │   └── ...                                  │   │   └── test_*.py                  ⬜
@@ -135,7 +136,8 @@ vllm_omni/                                    tests/
                                                │   ├── test_qwen3_omni_expansion.py
                                                │   ├── test_mimo_audio.py
                                                │   ├── test_image_gen_edit.py
-                                               │   └── test_images_generations_lora.py
+                                               │   ├── test_images_generations_lora.py
+                                               │   └── stage_configs/
                                                └── offline_inference/                  ✅
                                                    ├── test_qwen2_5_omni.py
                                                    ├── test_qwen3_omni.py
@@ -146,18 +148,17 @@ vllm_omni/                                    tests/
                                                    ├── test_zimage_tensor_parallel.py
                                                    ├── test_cache_dit.py
                                                    ├── test_teacache.py
-                                                   ├── test_stable_audio_expansion.py
+                                                   ├── test_stable_audio_model.py
                                                    ├── test_diffusion_cpu_offload.py
                                                    ├── test_diffusion_layerwise_offload.py
                                                    ├── test_diffusion_lora.py
                                                    ├── test_sequence_parallel.py
                                                    ├── test_qwen_image_edit_expansion.py
-                                                   └── stage_configs/                  (legacy schema, still present
-                                                       ├── bagel_*.yaml                 for unmigrated models)
+                                                   └── stage_configs/
+                                                       ├── qwen2_5_omni_ci.yaml
+                                                       ├── qwen3_omni_ci.yaml
+                                                       ├── bagel_*.yaml
                                                        └── npu/, rocm/, etc.
-
-# Migrated models (qwen3_omni_moe, qwen2_5_omni, qwen3_tts) live under
-# vllm_omni/deploy/ instead — see docs/configuration/stage_configs.md.
 examples/                                     tests
 │                                             └── examples
 ├── online_serving/                     →         ├── online_serving/
@@ -221,13 +222,14 @@ from pathlib import Path
 import openai
 import pytest
 
-from tests.helpers.media import (
-    convert_audio_bytes_to_text,
+from tests.conftest import (
+    OmniServer,
+    convert_audio_to_text,
     cosine_similarity_text,
+    dummy_messages_from_mix_data,
     generate_synthetic_video,
+    merge_base64_and_convert_to_text,
 )
-from tests.helpers.runtime import OmniServer, dummy_messages_from_mix_data
-from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
 from vllm_omni.platforms import current_omni_platform
 
 # Edit: model name and stage config path
@@ -235,7 +237,7 @@ models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]
 
 #If you use the default configuration file, you can directly use the following address.
 def get_default_config():
-    return get_deploy_config_path("ci/qwen3_omni_moe.yaml")
+    return str(Path(__file__).parent.parent / "stage_configs" / "qwen3_omni_ci.yaml")
 
 #If you need to modify the configuration file, you can use modify_stage_config.
 def get_chunk_config():
@@ -404,7 +406,7 @@ def test_mix_to_text_audio_001(client: openai.OpenAI, omni_server, request) -> N
     # PURPOSE: Verify text and audio outputs convey the same information
     # CUSTOMIZATION: Adjust similarity threshold (0.9) based on accuracy requirements
     assert audio_data is not None, "No audio output is generated"
-    audio_content = convert_audio_bytes_to_text(audio_data)
+    audio_content = merge_base64_and_convert_to_text(audio_data)
     print(f"text content is: {text_content}")
     print(f"audio content is: {audio_content}")
     similarity = cosine_similarity_text(audio_content.lower(), text_content.lower())
@@ -427,7 +429,7 @@ from pathlib import Path
 import pytest
 from vllm.assets.video import VideoAsset
 
-from tests.helpers.mark import hardware_test
+from tests.utils import hardware_test
 from ..multi_stages.conftest import OmniRunner
 
 # Optional: set process start method for workers

docs/contributing/model/adding_diffusion_model.md

@@ -802,7 +802,7 @@ omni = Omni(model="your-model", enable_layerwise_offload=True)
 
 ```python
 class WanTransformer3DModel(nn.Module):
-    _layerwise_offload_blocks_attrs = ["blocks"]  # Attribute name containing transformer blocks
+    _layerwise_offload_blocks_attr = "blocks"  # Attribute name containing transformer blocks
 
     def __init__(self):
         self.blocks = nn.ModuleList([...])  # Transformer blocks
@@ -813,16 +813,16 @@ class WanTransformer3DModel(nn.Module):
 
 ---
 
-### Diffusion Pipeline Profiler (Performance Profiling)
+### Diffusion Timing (Performance Profiling)
 When adapting a new diffusion model, it is often useful to analyze the latency of key components such as text encoding, diffusion denoising, and VAE decoding.
 vLLM-Omni provides a timing utility via `DiffusionPipelineProfilerMixin` to help developers quickly identify performance bottlenecks.
 
 !!! info
-      `DiffusionPipelineProfilerMixin` is different from using `torch.profiler` for diffusion models, as introduced in this [tutorial](https://github.com/vllm-project/vllm-omni/blob/main/docs/contributing/profiling.md). `DiffusionPipelineProfilerMixin` only prints the timing information of multiple functions (such as `vae.decode`), while `torch.profiler` saves detailed GPU/CPU computation time, call/execution steps.
+      `DiffusionPipelineProfilerMixin` is different from using `torch.profiler` for diffusion models, as introduced in this [tutorial](https://github.com/vllm-project/vllm-omni/blob/main/docs/contributing/profiling.md#3-profiling-diffusion-models). `DiffusionPipelineProfilerMixin` only prints the timing information of multiple functions (such as `vae.decode`), while `torch.profiler` saves detailed GPU/CPU computation time, call/execution steps.
 
 This tool automatically measures the execution time of selected pipeline modules and prints the results in the logs.
 
-**Enabling Diffusion Pipeline Profiler**
+**Enabling Diffusion Timing**
 
 
 Enable timing by setting:
@@ -843,7 +843,7 @@ If not specified, the default targets are used:
 **Adding DiffusionPipelineProfilerMixin to a Pipeline**
 To enable timing support in your pipeline, inherit from DiffusionPipelineProfilerMixin.
 ```python
-from vllm_omni.diffusion.profiler import DiffusionPipelineProfilerMixin
+from vllm_omni.diffusion.utils.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 
 class YourModelPipeline(nn.Module, DiffusionPipelineProfilerMixin):
     # Optional: Specify custom timing targets
@@ -862,9 +862,7 @@ class YourModelPipeline(nn.Module, DiffusionPipelineProfilerMixin):
         ...
 
         # initialize timing profiler
-        self.setup_diffusion_pipeline_profiler(
-            enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler
-        )
+        self.setup_diffusion_pipeline_profiler(enable_diffusion_pipeline_profiler)
 ```
 The mixin dynamically wraps selected methods and records their execution time during inference.
 
@@ -908,9 +906,9 @@ tokenizer.forward
 
 When enabled, timing logs appear like this:
 ```
-[DiffusionPipelineProfiler] text_encoder.forward took 0.018s
-[DiffusionPipelineProfiler] diffuse took 2.412s
-[DiffusionPipelineProfiler] vae.decode took 0.063s
+[DiffusionTiming] text_encoder.forward took 0.018s
+[DiffusionTiming] diffuse took 2.412s
+[DiffusionTiming] vae.decode took 0.063s
 ```
 These measurements help identify bottlenecks during model adaptation and optimization
 

docs/contributing/model/adding_omni_model.md

@@ -313,7 +313,7 @@ The registry uses lazy loading, so the model class is imported only when needed.
 
 ## Stage Configuration
 
-Create a YAML configuration file in `vllm_omni/deploy/`. For a complete example, see the [Qwen3-Omni configuration file](gh-file:vllm_omni/deploy/qwen3_omni_moe.yaml).
+Create a YAML configuration file in `vllm_omni/model_executor/stage_configs/`. For a complete example, see the [Qwen3-Omni configuration file](gh-file:vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml).
 
 ### Key Configuration Fields
 
@@ -408,17 +408,18 @@ Understanding the data structures is crucial for implementing stage transitions:
 
 **Input to your function:**
 - `stage_list[source_stage_id].engine_outputs`: List of `EngineCoreOutput` objects
--   -  Each contains `outputs`: List of `RequestOutput` objects
-    - Each `RequestOutput` has:
--   -  - `token_ids`: Generated token IDs
-       - `multimodal_output`: Dict with keys like `"code_predictor_codes"`, etc.These are the hidden states or intermediate outputs from the model's forward pass
-       - `prompt_token_ids`: Original prompt token IDs
+  - Each contains `outputs`: List of `RequestOutput` objects
+  - Each `RequestOutput` has:
+    - `token_ids`: Generated token IDs
+    - `multimodal_output`: Dict with keys like `"code_predictor_codes"`, etc.
+      - These are the hidden states or intermediate outputs from the model's forward pass
+    - `prompt_token_ids`: Original prompt token IDs
 
 **Output from your function:**
 - Must return `list[OmniTokensPrompt]` where each `OmniTokensPrompt` contains:
--   - `prompt_token_ids`: List[int] - Token IDs for the next stage
-    - `additional_information`: Dict[str, Any] - Optional metadata (e.g., embeddings, hidden states)
-    - `multi_modal_data`: Optional multimodal data if needed
+  - `prompt_token_ids`: List[int] - Token IDs for the next stage
+  - `additional_information`: Dict[str, Any] - Optional metadata (e.g., embeddings, hidden states)
+  - `multi_modal_data`: Optional multimodal data if needed
 
 ### How Model Outputs Are Stored
 
@@ -613,7 +614,7 @@ For a complete reference implementation, see:
 - **Thinker**: `vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py`
 - **Talker**: `vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py`
 - **Code2Wav**: `vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py`
-- **Stage config**: `vllm_omni/deploy/qwen3_omni_moe.yaml`
+- **Stage config**: `vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml`
 - **Input processors**: `vllm_omni/model_executor/stage_input_processors/qwen3_omni.py`
 - **Registry**: `vllm_omni/model_executor/models/registry.py`
 - **Testing**: `vllm_omni/tests/e2e/offline_inference/test_qwen3_omni.py`

docs/contributing/model/adding_tts_model.md

@@ -1,93 +1,20 @@
 # Adding a TTS Model
 
-This guide walks through adding a new TTS model to vLLM-Omni. Two patterns are
-supported:
-
-- **Two-stage pipeline** (e.g. Qwen3-TTS, Fish Speech): an AR code-predictor stage
-  feeds an audio decoder stage via the `async_chunk` framework. This is the standard
-  pattern for maximum streaming performance.
-- **Single-stage AR model** (e.g. MOSS-TTS-Nano): the model runs entirely inside one
-  AR worker and streams audio chunks directly from its own `inference_stream()` generator.
-
-Qwen3-TTS is used as the reference for the two-stage pattern. For the single-stage
-pattern, refer to MOSS-TTS-Nano.
+This guide walks through adding a new TTS model to vLLM-Omni, using **Qwen3-TTS**
+as a reference. Qwen3-TTS demonstrates the standard two-stage TTS pipeline and the
+key optimizations all TTS models in this repo should follow.
 
 ## Table of Contents
 
 1. [Overview](#overview)
-2. [Cross-Cutting Invariants](#cross-cutting-invariants)
-3. [Directory Structure](#directory-structure)
-4. [Step-by-Step Implementation](#step-by-step-implementation)
-5. [Key Components](#key-components)
-6. [Model Registration](#model-registration)
-7. [Stage Configuration](#stage-configuration)
-8. [Stage Input Processors](#stage-input-processors)
-9. [Online Serving Integration](#online-serving-integration)
-10. [Single-Stage Models](#single-stage-models)
-11. [Testing](#testing)
-12. [Pre-commit and DCO](#pre-commit-and-dco)
-13. [Summary](#summary)
-
-## 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). Each has surfaced as a silent
-bug in a shipped PR — check them at the end of every phase, not just at the start.
-
-**I1. Streaming output contract.** Pick one per-step semantics for `forward()` and
-document it in the docstring:
-
-- *Delta*: yield only new audio samples produced this step. Preferred — linear cost.
-- *Cumulative*: re-decode from step 0 every call. O(N²); only acceptable when the
-  codec exposes no streaming decode.
-
-If you choose delta, audit the full chain: `forward()` returns the new chunk →
-`_consolidate_multimodal_tensors()` in `vllm_omni/engine/output_processor.py`
-concatenates the audio key into a single tensor at finish → streaming consumers
-receive per-step chunks, offline consumers receive the concatenated tensor. A
-mismatch (consolidator skips the key with `continue`, or consumers expect a list
-but receive a tensor) is invisible in offline RTF benchmarks — users hear replays
-or truncation only under live playback.
-
-**I2. Multimodal output consumer hygiene.** `outputs[0].outputs[0].multimodal_output[key]`
-can be `Tensor`, `list[Tensor]` (pre-consolidation snapshot), `np.ndarray`, or
-scalar. In every test, example, and benchmark:
-
-- Never write `dict.get("a") or dict.get("b")` on tensor values — Python evaluates
-  the tensor's truthiness and raises `Boolean value of Tensor with more than one
-  value is ambiguous`. Use explicit `if x is None` chains.
-- 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 step, per-frame vocoder):
-
-- No `tensor.item()`, `.cpu()`, or `.tolist()` — each triggers a GPU→CPU sync; a
-  10-step × 60-frame × 4-op loop creates 2400 syncs per request.
-- Prefer `dst.copy_(src)` over `dst.fill_(src.item())` for scalar-into-buffer writes.
-- Whole-model `torch.compile(Model.forward, fullgraph=False)` usually outperforms
-  per-submodule compile — fewer dispatch boundaries, larger fusion regions. Measure
-  before choosing granularity.
-- No Python control flow that depends on tensor values; use `torch.where` or masking.
-
-Profile before optimizing.
-
-**I4. Validation pyramid.** Offline RTF alone is necessary but not sufficient. A
-new TTS model must pass all three levels:
-
-| Layer | Catches | Tool |
-|-------|---------|------|
-| Offline RTF / duration | Throughput regressions, missing audio, wrong sample rate | `end2end.py`, pytest e2e |
-| Browser streaming playback | Delta-vs-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 with 4+ parallel prompts |
-
-**I5. Per-request state belongs to the request.** If the model caches anything
-across `forward()` calls (streaming generators, codec buffers, sliding-window pads,
-CUDA graph state), key it by `info.get("_omni_req_id")` and free the entry on
-request finish. A shared buffer silently corrupts audio across concurrent requests —
-the symptom is crosstalk or truncation under load, nothing in single-request tests.
+2. [Directory Structure](#directory-structure)
+3. [Step-by-Step Implementation](#step-by-step-implementation)
+4. [Key Components](#key-components)
+5. [Model Registration](#model-registration)
+6. [Stage Configuration](#stage-configuration)
+7. [Stage Input Processors](#stage-input-processors)
+8. [Testing](#testing)
+9. [Summary](#summary)
 
 ## Overview
 
@@ -101,7 +28,7 @@ and can be placed on different devices. Qwen3-TTS has two stages:
 
 Each stage is a separate model class configured independently via YAML. The two stages
 are connected by the `async_chunk` framework, which enables inter-stage streaming for
-low first-packet latency (see [Async Chunk Design](../../design/feature/async_chunk.md)).
+low first-packet latency (see [Async Chunk Design](../../design/feature/async_chunk_design.md)).
 
 ### Without async_chunk (batch mode)
 
@@ -193,18 +120,8 @@ vllm_omni/model_executor/stage_configs/
 | `models/qwen3_tts/qwen3_tts.py` | Unified model class |
 | `models/qwen3_tts/qwen3_tts_code_predictor_vllm.py` | Stage 0 - optimized AR |
 | `models/qwen3_tts/qwen3_tts_code2wav.py` | Stage 1 - decoder |
-| `deploy/qwen3_tts.yaml` (new schema) | Deploy config (async_chunk enabled) — paired with `models/qwen3_tts/pipeline.py` for the frozen topology |
-
-> **Chunked vs end-to-end modes**: `qwen3_tts` registers a single
-> pipeline whose stage 1 declares alternate processor functions — an
-> `async_chunk_process_next_stage_input_func` (per-chunk streaming, used
-> when `deploy.async_chunk=True`) and a `sync_process_input_func`
-> (batch-end, used when `deploy.async_chunk=False`). The loader selects
-> one at merge time based on the bool, so `--no-async-chunk` alone
-> switches modes — no variant yaml or variant pipeline registration is
-> needed. Pipelines that only make sense in one mode (e.g.
-> `qwen3_omni_moe` is always chunked) can keep using the unconditional
-> `custom_process_*` fields.
+| `stage_configs/qwen3_tts.yaml` | Stage config (async_chunk enabled) |
+| `stage_configs/qwen3_tts_batch.yaml` | Batch mode config |
 | `stage_input_processors/qwen3_tts.py` | Stage transition processors |
 
 ## Step-by-Step Implementation
@@ -629,302 +546,6 @@ Recommended test cases for a new TTS model:
 
 Reference test: `tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py`
 
-### E2E Online Serving Tests (`tests/e2e/online_serving/test_<your_model>.py`)
-
-The `omni_server` fixture in `tests/conftest.py` is **module-scoped**. Each distinct
-`OmniServerParams` id in the same test file forces the fixture to tear the server
-down and spawn a new one mid-module. A few rules that save real CI debugging time:
-
-- **Prefer a single `OmniServerParams` set per file.** If you need to exercise two
-  deploy variants (e.g. `model.yaml` and `model_async_chunk.yaml`), either use one
-  variant and exercise streaming via request args, or split into two test files so
-  each file does exactly one server lifecycle. Mid-module teardown/restart is the
-  fragile path and surfaces startup races first.
-- **Never depend on server-side fetching of external URLs** for reference audio or
-  other fixture data. CI runners (and China-hosted dev boxes) routinely fail on
-  SSL/DNS for public URLs. Inline the payload as a `data:audio/wav;base64,...`
-  ref_audio value — the serving layer accepts both forms.
-- **Don't roll your own readiness probe.** The harness already waits for HTTP 200
-  on `/health` before releasing the server to the test. If your model needs extra
-  warmup signals, expose them through `/health` rather than adding `time.sleep(...)`
-  inside the test. (Bare TCP `connect_ex` probes were insufficient; see
-  `tests/conftest.py` `OmniServer.wait_for_ready`.)
-- **Use `core_model` marker + H100 hardware_test** to match the `test-ready.yml`
-  pipeline so your test is picked up by the `ready` label, not only nightly.
-
-## Online Serving Integration
-
-To expose your model through the `/v1/audio/speech` OpenAI-compatible endpoint, add
-**all five** of the following integration points to
-`vllm_omni/entrypoints/openai/serving_speech.py` in a **single commit**. Adding them
-piecemeal causes partial-integration failures that are hard to debug.
-
-### 1. Stage constant
-
-Near the top of the file, alongside the other `_*_TTS_MODEL_STAGES` constants:
-
-```python
-_YOUR_MODEL_TTS_MODEL_STAGES = {"your_model_stage_key"}
-```
-
-### 2. Union into `_TTS_MODEL_STAGES`
-
-Add to the `_TTS_MODEL_STAGES` set union:
-
-```python
-_TTS_MODEL_STAGES: set[str] = (
-    ...
-    | _YOUR_MODEL_TTS_MODEL_STAGES
-)
-```
-
-### 3. Model type detection
-
-In `_detect_tts_model_type()`, add before the final `return None`:
-
-```python
-if model_stage in _YOUR_MODEL_TTS_MODEL_STAGES:
-    return "your_model"
-```
-
-### 4. Request validation dispatch
-
-In `_validate_tts_request()`, add before the fallback `return`:
-
-```python
-if self._tts_model_type == "your_model":
-    return self._validate_your_model_request(request)
-```
-
-### 5. Validation and parameter-builder methods
-
-Add two new methods:
-
-```python
-def _validate_your_model_request(
-    self, request: OpenAICreateSpeechRequest
-) -> str | None:
-    """Validate YourModel request. Returns an error string or 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: OpenAICreateSpeechRequest
-) -> dict[str, Any]:
-    """Build additional_information dict for YourModel."""
-    params: dict[str, Any] = {"text": [request.input]}
-    if request.voice is not None:
-        params["voice"] = [request.voice]
-    # Add any other model-specific fields here
-    return params
-```
-
-Then wire `_build_your_model_params` into the request-dispatch block in
-`_create_tts_request()` (search for the equivalent `_build_*_params` call for an
-existing model to find the right location). If the model supports voice cloning
-(`ref_audio` → `prompt_audio_path`, `ref_text` → `prompt_text`), add those mappings
-here too — follow any existing `_build_<model>_params` in `serving_speech.py` (e.g.
-`_build_moss_tts_params` for the voice-cloning variant) for the pattern.
-
-> **Two dispatch patterns coexist:** Fish Speech uses a `self._is_fish_speech` boolean
-> checked *before* `elif self._is_tts`. All newer models use the `_tts_model_type`
-> string pattern shown above. For new models, always use the string pattern — do not
-> add new `_is_*` boolean flags.
-
-> **Note on unused variables:** Only extract parameters in `_build_your_model_params`
-> that you actually pass to the model's generate / `inference_stream` call. Extracting
-> a variable without forwarding it will trigger a `ruff F841` pre-commit failure.
-
-### Merge conflicts
-
-`serving_speech.py` is modified by every new model PR and is the most common source of
-rebase conflicts. When rebasing onto `main` and a conflict appears here, the resolution
-is always to **keep both** the upstream model's additions and your own — never discard
-either side. After resolving:
-
-```bash
-git add vllm_omni/entrypoints/openai/serving_speech.py
-git rebase --continue
-```
-
-## Single-Stage Models
-
-Some TTS models (e.g. MOSS-TTS-Nano) do not use a two-stage pipeline. Instead the
-entire AR LM and audio decoder run inside a single AR worker, streaming audio chunks
-directly from the model's own generator.
-
-### Directory structure
-
-```
-vllm_omni/model_executor/models/your_model_name/
-    __init__.py
-    modeling_your_model_name.py    # unified class: load_weights + forward + streaming
-
-vllm_omni/model_executor/stage_configs/your_model_name.yaml
-```
-
-No stage input processor is needed.
-
-### Stage config
-
-Use a single stage with `worker_type: ar`. The `is_comprehension: true` field and the
-top-level `async_chunk: false` are required — omitting them causes silent
-misclassification in the serving layer. Set `max_num_seqs` to at least 4 for
-concurrent production use.
-
-```yaml
-# stage_configs/your_model_name.yaml
-async_chunk: false
-
-stage_args:
-  - stage_id: 0
-    stage_type: llm
-    is_comprehension: true          # required for serving_speech.py dispatch
-    runtime:
-      devices: "0"
-    engine_args:
-      model_stage: your_model_stage_key
-      model_arch: YourModelForCausalLM
-      worker_type: ar
-      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      engine_output_type: audio
-      max_num_seqs: 4               # min 4 for concurrent requests; default 1 causes gaps
-    final_output: true
-    final_output_type: audio
-```
-
-### Generator-based streaming pattern
-
-This is the MOSS-TTS-Nano pattern, distinct from VoxCPM2's vLLM-native AR pattern
-(see `plan/voxcpm2_native_ar_design.md` for that variant). Load model weights in
-`load_weights()` (not `__init__`) so vLLM finishes distributed initialisation before
-any CUDA allocations. Stream via a per-request generator stored in an instance dict:
-
-```python
-class YourModelForCausalLM(nn.Module):
-
-    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
-        super().__init__()
-        self._lm: nn.Module | None = None          # populated in load_weights()
-        self._stream_gens: dict[str, Any] = {}     # 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 [{}]
-        # Return empty output during 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(...)
-
-        outputs, last_flags = [], []
-        for info in infos:
-            request_key = str(info.get("_omni_req_id", "0"))  # set by vLLM, not user code
-            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, "is_last": last_flags})
-
-    def _create_stream_gen(self, info: dict):
-        """Yield (waveform_tensor, is_last) from the model's inference_stream().
-
-        Handle both incremental ("audio" events) and batch ("result" event) models:
-        some upstream implementations emit one "result" event with the full waveform
-        instead of incremental "audio" events. Both paths must be covered.
-        """
-        for event in self._lm.inference_stream(...):
-            if event["type"] == "audio":
-                yield event["waveform"], False
-            elif event["type"] == "result":
-                # Fallback for models that don't emit incremental audio events
-                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 when the last chunk has been yielded so the AR
-        # scheduler ends the request at the right time.
-        ...
-```
-
-For an in-tree reference, look for any single-stage AR model under
-`vllm_omni/model_executor/models/` (for example
-`moss_tts_nano/modeling_moss_tts_nano.py` once its integration has landed).
-
-## Pre-commit and DCO
-
-All contributions must pass the pre-commit checks and the Developer Certificate of
-Origin (DCO) sign-off before merging.
-
-### Running pre-commit
-
-Install the hooks once with `pre-commit install`. Then run before committing:
-
-```bash
-pre-commit run --files \
-  vllm_omni/model_executor/models/your_model_name/*.py \
-  vllm_omni/entrypoints/openai/serving_speech.py \
-  vllm_omni/model_executor/models/registry.py \
-  tests/e2e/offline_inference/test_your_model_name.py \
-  tests/e2e/online_serving/test_your_model_name.py
-```
-
-When pre-commit **modifies files**, it exits with a non-zero code but the reformatting
-is correct. Stage the modified files and commit again — do not revert the changes.
-
-Common failures and fixes:
-
-| Check | Cause | Fix |
-|-------|-------|-----|
-| `ruff F841` | Local variable assigned but never used | Remove the extraction or forward it to the model call |
-| `ruff E402` | Module-level import not at top of file | Move import to the top-level import block |
-| `ruff format` | Line length, spacing, or quote style | Accept the auto-fix, stage, and re-commit |
-
-### DCO sign-off
-
-Every commit must carry a `Signed-off-by` trailer. Use the `-s` flag when committing:
-
-```bash
-git commit -s -m "feat(your-model): add YourModel TTS support"
-```
-
-Or configure git to add it automatically:
-
-```bash
-git config format.signOff true
-```
-
-To fix a missing sign-off on the most recent commit:
-
-```bash
-git commit --amend -s --no-edit
-git push origin your-branch --force-with-lease
-```
-
-> The DCO check verifies that the commit author email matches the `Signed-off-by` email.
-> Make sure `git config user.email` is set to the address associated with your GitHub
-> account before committing.
-
 ## Adding a Model Recipe
 
 After implementing and testing your model, add a model recipe to the
@@ -936,19 +557,15 @@ for the expected format.
 
 Adding a TTS model to vLLM-Omni involves:
 
-1. **Create model directory** with AR stage, decoder stage, and unified class (two-stage)
-   or a single unified class with generator-based streaming (single-stage)
+1. **Create model directory** with AR stage, decoder stage, and unified class
 2. **AR stage** - use vLLM's native decoder layers with fused QKV; do not wrap HF directly
 3. **Decoder stage** - thin wrapper around your audio decoder; implement `chunked_decode_streaming()`
 4. **Unified class** - dispatches on `model_stage`; same structure as `Qwen3TTSModelForGeneration`
 5. **Register** all stage classes in `registry.py`
-6. **YAML configs** - provide both batch and `async_chunk` variants (two-stage), or a single-stage AR config
-7. **Stage input processor** - buffer Stage 0 outputs and forward in chunks of 25 (two-stage only)
-8. **Online serving** - add all 5 integration points to `serving_speech.py` in one commit
-9. **Tests** - cover single request, batching, and streaming
-10. **Pre-commit + DCO** - run `pre-commit` before pushing; sign every commit with `git commit -s`
-11. **Model recipe** - add to [vllm-project/recipes](https://github.com/vllm-project/recipes)
-12. **Invariants** - re-check I1–I5 (streaming contract, consumer hygiene, hot-loop discipline, validation pyramid, per-request state) at the end of every phase
+6. **YAML configs** - provide both batch and `async_chunk` variants
+7. **Stage input processor** - buffer Stage 0 outputs and forward in chunks of 25
+8. **Tests** - cover single request, batching, and async_chunk streaming
+9. **Model recipe** - add to [vllm-project/recipes](https://github.com/vllm-project/recipes)
 
 ### Qwen3-TTS Reference Files
 
@@ -957,12 +574,11 @@ Adding a TTS model to vLLM-Omni involves:
 | `models/qwen3_tts/qwen3_tts.py` | Unified model class |
 | `models/qwen3_tts/qwen3_tts_code_predictor_vllm.py` | AR stage with vLLM fused ops |
 | `models/qwen3_tts/qwen3_tts_code2wav.py` | Decoder stage with `chunked_decode_streaming()` |
-| `models/qwen3_tts/pipeline.py` | Frozen pipeline topology (registered at import time) |
-| `deploy/qwen3_tts.yaml` | Deploy config (user-editable, async_chunk + SharedMemoryConnector) |
+| `stage_configs/qwen3_tts.yaml` | Stage configuration |
 | `stage_input_processors/qwen3_tts.py` | Stage transition processors |
 
 For more information, see:
 
 - [Architecture Overview](../../design/architecture_overview.md)
-- [Async Chunk Design](../../design/feature/async_chunk.md)
+- [Async Chunk Design](../../design/feature/async_chunk_design.md)
 - [Stage Configuration Guide](../../configuration/stage_configs.md)

docs/contributing/profiling.md

@@ -1,286 +1,216 @@
-# Profiling Diffusion Models
+# Profiling vLLM-Omni
 
-> **Warning:** Profiling is for development and debugging only. It adds significant overhead and should not be enabled in production.
+> **Warning:** Profiling incurs significant overhead. Use only for development and debugging, never in production.
 
-Diffusion profiling supports two backends through `profiler_config`:
+vLLM-Omni uses the PyTorch Profiler to analyze performance across both **multi-stage omni-modality models** and **diffusion models**.
 
-- `torch`: detailed CPU/CUDA traces, operator tables, and optional memory snapshots
-- `cuda`: low-overhead CUDA range control for NVIDIA Nsight Systems (`nsys`)
+### 1. Configure Profiling in the Stage YAML
 
-## 1. Configure `profiler_config`
-
-Use `profiler_config` to enable profiling for a diffusion model. For diffusion usage, pass it directly to `Omni(...)` or `vllm serve`.
-
-Minimal torch-profiler config:
+Enable profiling by adding `profiler_config` under `engine_args` for the stage(s) you want to profile in your stage config YAML:
 
 ```yaml
-profiler_config:
-  profiler: torch
-  torch_profiler_dir: ./perf
+stage_args:
+  - stage_id: 0
+    stage_type: llm
+    engine_args:
+      # ... other engine args ...
+      profiler_config:
+        profiler: torch
+        torch_profiler_dir: ./perf
 ```
 
-Supported fields:
-
 | Field | Description |
 |---|---|
-| `profiler` | Profiler backend. Supported values: `torch`, `cuda`. Use `torch` for `trace.json`, Excel operator tables, and optional memory snapshots. Use `cuda` for Nsight Systems only. |
-| `torch_profiler_dir` | Output directory for torch-profiler artifacts. Required when `profiler: torch`. |
-| `torch_profiler_use_gzip` | Compress `trace_rank*.json` into `trace_rank*.json.gz`. |
-| `torch_profiler_record_shapes` | Record input shapes and add a `by_shape` sheet to `ops_rank*.xlsx`. |
-| `torch_profiler_with_stack` | Record call stacks, add a `by_stack` sheet to `ops_rank*.xlsx`, and export `stacks_cpu_rank*.txt` and `stacks_cuda_rank*.txt`. |
-| `torch_profiler_with_memory` | Enable memory profiling and attempt to dump `memory_snapshot_rank*.pickle`. The pickle is only generated when the current backend supports memory history and snapshot APIs. |
-| `torch_profiler_with_flops` | Enable FLOPs collection in `torch.profiler`. This does not add a separate output file. |
-| `torch_profiler_dump_cuda_time_total` | Export an additional text summary `profiler_out_<rank>.txt` sorted by `self_cuda_time_total`. |
-| `delay_iterations` | Number of worker iterations to skip before profiling starts. |
-| `max_iterations` | Maximum number of worker iterations to capture before auto-stop. |
-| `wait_iterations` | Torch-profiler wait iterations before warmup. |
-| `warmup_iterations` | Torch-profiler warmup iterations. |
-| `active_iterations` | Torch-profiler active iterations. |
+| `profiler` | Profiler backend to use. Currently supports `torch`. |
+| `torch_profiler_dir` | Directory where trace files are saved. Created automatically if it doesn't exist. |
 
-### Minimal configurations by output
+> **Tip:** Only enable `profiler_config` on stages you actually need to profile. Stages without it will not start a profiler, keeping overhead minimal.
 
-Only collect trace output:
+### 2. Profiling Omni-Modality Models
+
+**Selective Stage Profiling**
+
+It is highly recommended to profile specific stages to prevent producing overly large trace files:
 
 ```python
-profiler_config = {
-    "profiler": "torch",
-    "torch_profiler_dir": "./perf",
-}
+# Profile all stages
+omni_llm.start_profile()
+
+# Only profile Stage 1
+omni_llm.start_profile(stages=[1])
+
+# Stage 0 (Thinker) and Stage 2 (Audio Decoder) for qwen omni
+omni_llm.start_profile(stages=[0, 2])
 ```
 
-Outputs:
+> **Important:** Always pass the same `stages` list to both `start_profile()` and `stop_profile()`. If you omit `stages` from `stop_profile()`, it defaults to stopping all stages — including ones that were never started — which will produce errors.
 
-- `trace_rank*.json`
-- `ops_rank*.xlsx` with a `summary` sheet
-
-Collect compressed trace output:
+**Python Usage**: Wrap your generation logic with `start_profile()` and `stop_profile()`.
 
 ```python
-profiler_config = {
-    "profiler": "torch",
-    "torch_profiler_dir": "./perf",
-    "torch_profiler_use_gzip": True,
-}
+profiler_stages = [0]  # Only profile the stages you need
+
+# 1. Start profiling
+omni.start_profile(stages=profiler_stages)
+
+# Initialize generator
+omni_generator = omni.generate(prompts, sampling_params_list, py_generator=args.py_generator)
+
+total_requests = len(prompts)
+processed_count = 0
+
+# Main Processing Loop
+for stage_outputs in omni_generator:
+
+    # ... [Output processing logic for text/audio would go here] ...
+
+    # Update count to track when to stop profiling
+    processed_count += len(stage_outputs.request_output)
+
+    # 2. Check if all requests are done to stop the profiler safely
+    if profiler_enabled and processed_count >= total_requests:
+        print(f"[Info] Processed {processed_count}/{total_requests}. Stopping profiler inside active loop...")
+
+        # Stop the profiler while workers are still active
+        # Pass the same stages list used in start_profile()
+        omni_llm.stop_profile(stages=profiler_stages)
+
+        # Wait for traces to flush to disk
+        print("[Info] Waiting 30s for workers to write trace files to disk...")
+        time.sleep(30)
+        print("[Info] Trace export wait time finished.")
+
+omni_llm.close()
 ```
 
-Outputs:
-
-- `trace_rank*.json.gz`
-- `ops_rank*.xlsx` with a `summary` sheet
-
-Collect trace and full operator tables:
-
-```python
-profiler_config = {
-    "profiler": "torch",
-    "torch_profiler_dir": "./perf",
-    "torch_profiler_record_shapes": True,
-    "torch_profiler_with_stack": True,
-}
-```
-
-Outputs:
-
-- `trace_rank*.json`
-- `ops_rank*.xlsx` with `summary`, `by_shape`, and `by_stack`
-- `stacks_cpu_rank*.txt`
-- `stacks_cuda_rank*.txt`
-
-Collect trace, operator tables, and memory snapshots:
-
-```python
-profiler_config = {
-    "profiler": "torch",
-    "torch_profiler_dir": "./perf",
-    "torch_profiler_record_shapes": True,
-    "torch_profiler_with_stack": True,
-    "torch_profiler_with_memory": True,
-}
-```
-
-Outputs:
-
-- `trace_rank*.json`
-- `ops_rank*.xlsx` with `summary`, `by_shape`, and `by_stack`
-- `stacks_cpu_rank*.txt`
-- `stacks_cuda_rank*.txt`
-- `memory_snapshot_rank*.pickle` when supported by the current backend
-
-### Full torch-profiler configuration
-
-If you want to enable the commonly used torch-profiler options together:
-
-```python
-profiler_config = {
-    "profiler": "torch",
-    "torch_profiler_dir": "./perf",
-    "torch_profiler_use_gzip": False,
-    "torch_profiler_record_shapes": True,
-    "torch_profiler_with_stack": True,
-    "torch_profiler_with_memory": True,
-    "torch_profiler_with_flops": False,
-    "torch_profiler_dump_cuda_time_total": False,
-    "delay_iterations": 0,
-    "max_iterations": 0,
-    "wait_iterations": 0,
-    "warmup_iterations": 0,
-    "active_iterations": 0,
-}
-```
-
-## 2. Profiling Diffusion with PyTorch Profiler
-
-Single-stage diffusion models use `start_profile()` / `stop_profile()` controls. The profiler only writes artifacts after profiling has been started and then stopped.
-
-```python
-from vllm_omni import Omni
-
-omni = Omni(
-    model="Wan-AI/Wan2.2-I2V-A14B-Diffusers",
-    profiler_config={
-        "profiler": "torch",
-        "torch_profiler_dir": "./perf",
-    },
-)
-
-omni.start_profile()
-...
-omni.stop_profile()
-```
-
-For diffusion offline example scripts under `examples/offline_inference/`, pass `--profiler-config` as a JSON object. The script enables profiling when this argument is set and wraps generation with `start_profile()` / `stop_profile()`.
-
-Example:
 
+**CLI Usage** (using `end2end.py`):
 ```bash
-python examples/offline_inference/image_to_video/image_to_video.py \
-  --model Wan-AI/Wan2.2-I2V-A14B-Diffusers \
-  --image input.jpg \
-  --prompt "A cat playing with yarn" \
-  --profiler-config '{
-    "profiler": "torch",
-    "torch_profiler_dir": "./perf",
-    "torch_profiler_record_shapes": true,
-    "torch_profiler_with_stack": true
-  }'
+# Profile only Stage 0 (Thinker)
+python end2end.py --output-wav output_audio \
+    --query-type text --enable-profiler --profiler-stages 0
+
+# Profile Stage 0 and Stage 2
+python end2end.py --output-wav output_audio \
+    --query-type text --enable-profiler --profiler-stages 0 2
+
+# Profile all stages (omit --profiler-stages)
+python end2end.py --output-wav output_audio \
+    --query-type text --enable-profiler
 ```
 
-Examples:
+**Examples**:
 
-1. [Image edit example](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/image_to_image/image_edit.py)
-2. [Image to video example](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/image_to_video)
+1. **Qwen2.5-Omni**:  [https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen2_5_omni/end2end.py](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen2_5_omni/end2end.py)
 
-## 3. Profiling Diffusion with Nsight Systems (`nsys`)
+2. **Qwen3-Omni**:   [https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen3_omni/end2end.py](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen3_omni/end2end.py)
 
-For Nsight Systems, use `profiler: cuda` and wrap the process with `nsys profile`.
+### 3. Profiling diffusion models
 
+Diffusion profiling is End-to-End, capturing encoding, denoising loops, and decoding. Standalone diffusion scripts use `--profiler-dir` to enable profiling.
+
+**CLI Usage:**
 ```bash
-nsys profile \
-  --trace-fork-before-exec=true \
-  --cuda-graph-trace=node \
-  --capture-range=cudaProfilerApi \
-  --capture-range-end=repeat \
-  -o diffusion_trace \
-  python image_to_video.py ...
+python image_to_video.py \
+    --model Wan-AI/Wan2.2-I2V-A14B-Diffusers \
+    --image qwen-bear.png \
+    --prompt "A cat playing with yarn, smooth motion" \
+    --profiler-dir \
+    \
+    # Minimize Spatial Dimensions (Optional but helpful):
+    #    Drastically reduces memory usage so the profiler doesn't
+    #    crash due to overhead, though for accurate performance
+    #    tuning you often want target resolutions.
+    --height 48 \
+    --width 64 \
+    \
+    # Minimize Temporal Dimension (Frames):
+    #    Video models process 3D tensors (Time, Height, Width).
+    #    Reducing frames to the absolute minimum (2) keeps the
+    #    tensor size small, ensuring the trace file doesn't become
+    #    multi-gigabytes in size.
+    --num-frames 2 \
+    \
+    # Minimize Iteration Loop (Steps):
+    #    This is the most critical setting for profiling.
+    #    Diffusion models run the same loop X times.
+    #    Profiling 2 steps gives you the exact same performance
+    #    data as 50 steps, but saves minutes of runtime and
+    #    prevents the trace viewer from freezing.
+    --num-inference-steps 2 \
+    \
+    --guidance-scale 5.0 \
+    --guidance-scale-high 6.0 \
+    --boundary-ratio 0.875 \
+    --flow-shift 12.0 \
+    --fps 16 \
+    --output i2v_output.mp4
 ```
 
-The Python process being profiled must create the diffusion engine with:
+> **Note:** For diffusion stages within a multi-stage omni pipeline, use `profiler_config` in the stage YAML instead (see Section 1).
 
-```python
-profiler_config = {"profiler": "cuda"}
-```
+**Examples**:
 
-Then call `start_profile()` before the requests you want to capture and `stop_profile()` after them. The diffusion worker processes open and close the CUDA capture range themselves, so `nsys` sees the actual GPU work instead of only the parent process.
+1. **Qwen image edit**:  [https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/image_to_image/image_edit.py](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/image_to_image/image_edit.py)
 
-## 4. Profiling Online Serving
+2. **Wan-AI/Wan2.2-I2V-A14B-Diffusers**:   [https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/image_to_video](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/image_to_video)
 
-When `profiler_config.profiler` is set for a diffusion model, the server exposes:
+### 4. Profiling Online Serving
 
-- `POST /start_profile`
-- `POST /stop_profile`
-
-### Start the server
-
-Single-stage diffusion serving with torch profiler:
+When `profiler_config` is set in the stage YAML, the server automatically exposes `/start_profile` and `/stop_profile` HTTP endpoints.
 
+**1. Start the server** with a stage YAML that has `profiler_config` enabled:
 ```bash
-vllm serve Wan-AI/Wan2.2-I2V-A14B-Diffusers \
-  --omni \
-  --port 8091 \
-  --profiler-config '{
-    "profiler": "torch",
-    "torch_profiler_dir": "/tmp/vllm_profile_wan22_i2v",
-    "torch_profiler_with_stack": true,
-    "torch_profiler_with_flops": false,
-    "torch_profiler_use_gzip": true,
-    "torch_profiler_dump_cuda_time_total": false,
-    "torch_profiler_record_shapes": true,
-    "torch_profiler_with_memory": true,
-    "delay_iterations": 0,
-    "max_iterations": 0,
-    "wait_iterations": 0,
-    "warmup_iterations": 0,
-    "active_iterations": 0
-  }'
-```
-
-Single-stage diffusion serving with Nsight Systems:
-
-```bash
-nsys profile \
-  --trace-fork-before-exec=true \
-  --cuda-graph-trace=node \
-  --capture-range=cudaProfilerApi \
-  --capture-range-end=repeat \
-  -o serving_trace \
-  vllm serve Wan-AI/Wan2.2-I2V-A14B-Diffusers \
+vllm serve Qwen/Qwen2.5-Omni-7B \
     --omni \
-    --port 8091 \
-    --profiler-config '{"profiler": "cuda"}'
+    --stage-configs-path qwen2_5_omni.yaml \
+    --port 8091
 ```
 
-### Control capture
-
-Example profiling flow for an online Qwen-Image request:
+Or for one stage diffusion models:
 
 ```bash
-# Start profiling.
+vllm serve Wan-AI/Wan2.2-I2V-A14B-Diffusers --omni --port 8091 --profiler-config '{"profiler": "torch", "torch_profiler_dir": "./vllm_profile"}'
+```
+
+**2. Start profiling** by sending a POST request:
+```bash
+# Profile all stages that have profiler_config set
 curl -X POST http://localhost:8091/start_profile
 
-# Send a Qwen-Image generation request while profiling is active.
-curl http://localhost:8091/v1/images/generations \
-  -H "Content-Type: application/json" \
-  -d '{
-    "model": "Qwen/Qwen-Image",
-    "prompt": "A red vintage bicycle parked beside a quiet canal at sunset"
-  }'
-
-# Stop profiling and flush profiler artifacts.
-curl -X POST http://localhost:8091/stop_profile
+# Profile specific stages only
+curl -X POST http://localhost:8091/start_profile \
+    -H "Content-Type: application/json" \
+    -d '{"stages": [0]}'
 ```
 
-## 5. Diffusion Pipeline Profiler
+**3. Send your inference requests** as normal while the profiler is running.
 
-For lightweight per-stage pipeline timing such as `vae.decode` or `diffuse`, see [Diffusion Pipeline Profiler](model/adding_diffusion_model.md#diffusion-pipeline-profiler-performance-profiling). That utility logs stage durations only and does not generate torch-profiler artifacts such as `trace.json`, Excel tables, or memory snapshots.
+**4. Stop profiling** and collect traces:
+```bash
+# Stop all stages
+curl -X POST http://localhost:8091/stop_profile
 
-## 6. Analyze Results
+# Stop specific stages (must match the stages you started)
+curl -X POST http://localhost:8091/stop_profile \
+    -H "Content-Type: application/json" \
+    -d '{"stages": [0]}'
+```
 
-Torch-profiler output:
+Trace files are written to the `torch_profiler_dir` specified in your stage YAML.
 
-- Chrome/Perfetto trace: `trace_rank*.json` or `trace_rank*.json.gz`
-- Excel workbook: `ops_rank*.xlsx` with `summary`, and optional `by_shape` / `by_stack` sheets
-- Stack exports: `stacks_cpu_rank*.txt` and `stacks_cuda_rank*.txt` when stack capture is enabled
-- Memory snapshot: `memory_snapshot_rank*.pickle` when memory capture is enabled and supported by the backend
-- Optional CUDA-time text summary: `profiler_out_<rank>.txt` when `torch_profiler_dump_cuda_time_total` is enabled
+> **Important:** Always stop the same stages you started. Stopping a stage that was never started will produce errors.
 
-CUDA profiler / Nsight Systems output:
+### 5. Analyzing Traces
 
-- `.nsys-rep` report files written by `nsys -o ...`
+Output files are saved to the `torch_profiler_dir` specified in your stage YAML config.
 
-Recommended viewers:
+**Output**
+**Chrome Trace** (`.json.gz`): Visual timeline of kernels and stages. Open in Perfetto UI.
 
-- [Perfetto](https://ui.perfetto.dev/) for torch traces
-- `nsys stats <report>.nsys-rep` for CLI summaries
-- Nsight Systems GUI for CUDA kernel timelines
+**Viewing Tools:**
 
-For upstream background on the underlying vLLM profiling infrastructure, see the [vLLM profiling guide](https://docs.vllm.ai/en/stable/contributing/profiling/).
+- [Perfetto](https://ui.perfetto.dev/) (recommended)
+- `chrome://tracing` (Chrome only)
+
+**Note**: vLLM-Omni reuses the PyTorch Profiler infrastructure from vLLM. See the official vLLM profiler documentation:  [vLLM Profiling Guide](https://docs.vllm.ai/en/stable/contributing/profiling/)

docs/design/feature/async_chunk_design.md

@@ -1,4 +1,4 @@
-# Async Chunk
+# Async Chunk Design
 
 ## Table of Contents
 
@@ -19,7 +19,7 @@ The `async_chunk` feature enables asynchronous, chunked processing of data acros
 
 For qwen3-omni:
 - **Thinker → Talker**: Per decode step (typically chunk_size=1)
-- **Talker → Code2Wav**: Accumulated to `codec_chunk_frames` (default=25) before sending. During the initial phase, a dynamic initial chunk size (IC) is automatically selected based on server load to reduce TTFP. Use the per-request `initial_codec_chunk_frames` API field to override.
+- **Talker → Code2Wav**: Accumulated to `codec_chunk_frames` (default=25) before sending. During the initial phase, a dynamic initial chunk size (IC) is automatically selected based on server load to reduce TTFA. Use the per-request `initial_codec_chunk_frames` API field to override.
 - **Code2Wav**: Streaming decode with code2wav chunk_size
 
 With `async_chunk`:
@@ -75,85 +75,26 @@ Enabling **async_chunk** (False→True) sharply reduces time-to-first-audio (TTF
 </p>
 
 ## Architecture
+### Data Flow
 
-### Async Chunk Pipeline Overview
+#### Sequential Flow
+<p align="center">
+  <picture>
+    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/qwen3-omni-non-async-chunk.png">
+    <img alt="Data Flow between stages" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/qwen3-omni-non-async-chunk.png" width=100%>
+  </picture>
+</p>
 
-The following diagram illustrates the **Async Chunk Architecture** for multi-stage models (e.g., Qwen3-Omni with Thinker → Talker → Code2Wav), showing how data flows through the 4-stage pipeline with parallel processing and dual-stream output:
+#### Async Chunk Flow
 
 <p align="center">
   <picture>
     <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/qwen3-omni-async-chunk.png">
-    <img alt="Async Chunk Pipeline Architecture" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/qwen3-omni-async-chunk.png" width=100%>
+    <img alt="Data Flow between stages" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/qwen3-omni-async-chunk.png" width=100%>
   </picture>
 </p>
 
-**Diagram Legend:**
-
-| Step | Stage Type | Description |
-|------|-----------|------------|
-| `prefill` | Initialization | Context processing, KV cache initialization |
-| `decode` | Autoregressive | Token-by-token generation in AR stages |
-| `codes` | Audio Encoding | RVQ codec codes from Talker stage |
-| `output` | Final Output | Text chunks or audio waveforms |
-
-### Data Flow
-
-#### Stage 0: Thinker (Multimodal Understanding + Text Generation)
-- **Prefill**: Processes multimodal input (text/image/audio/video), initializes KV cache
-- **Decode Loop**: Generates text tokens autoregressively
-- **Chunk Triggers**: Each decode step (typically `chunk_size=1`) can trigger downstream processing
-- **Dual Output**:
-  - **Text Stream**: `text_0`, `text_1`, `text_2`... `text_n` streamed to output
-  - **Hidden States**: Passed to Talker stage for audio synthesis
-
-#### Stage 1: Talker (Text → RVQ Audio Codes)
-- **Prefill**: Receives hidden states from Thinker as semantic condition
-- **Decode Loop**: Generates RVQ codec codes autoregressively
-- **Accumulation**: Codes accumulate to `codec_chunk_frames` (default=25) before forwarding
-- **Dynamic IC**: Initial chunk size auto-selected based on server load to optimize TTFP
-- **Output**: `codes` blocks (chunk 0, 1, ... n) sent to Code2Wav
-
-#### Stage 2: Code2Wav (Vocoder Decoder)
-- **Non-Autoregressive**: Processes RVQ codes in parallel batches
-- **Streaming Decode**: Converts codes to audio waveforms chunk-by-chunk
-- **Batching**: Supports batched inference for multiple concurrent requests
-- **Output**: Audio segments `audio_0`, `audio_1`, ... `audio_n`
-
-#### Stage 3: Output (Dual Stream)
-- **Text Streaming**: `text_0` → `text_1` → `text_2` → ... (user sees response in real-time)
-- **Audio Streaming**: `audio_0` → `audio_1` → ... (user hears audio progressively)
-
-### Execution Timeline
-
-```
-Timeline: Parallel vs Sequential
-
-Sequential (async_chunk=false):
-[Thinker: ████████████████████]  (2.0s)
-                            [Talker: ████████████████████]  (3.0s)
-                                                        [Code2Wav: ████]  (1.0s)
-Total: 6.0s, TTFP: 6.0s
-
-Async Chunk (async_chunk=true):
-[Thinker: ████░░░░████░░░░████]  (2.0s, streaming)
-     [Talker: ░░████░░░░████░░]  (3.0s, parallel)
-         [Code2Wav: ░░░░████░░]  (1.0s, batched)
-Total: ~3.5s, TTFP: ~0.5s
-
-█ = Active computation  ░ = Waiting/idle
-```
-
-#### Sequential Flow (for comparison)
-<p align="center">
-  <picture>
-    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/qwen3-omni-non-async-chunk.png">
-    <img alt="Sequential Data Flow" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/qwen3-omni-non-async-chunk.png" width=100%>
-  </picture>
-</p>
-
-In sequential mode, each stage must wait for the previous stage to complete entirely before starting.
-
-### Async Chunk System Architecture
+### Async Chunk architecture
 <p align="center">
   <picture>
     <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/async-chunk-architecture.png">

docs/design/feature/cfg_parallel.md

@@ -25,9 +25,7 @@ In standard Classifier-Free Guidance, each diffusion step requires two forward p
 1. **Positive/Conditional**: Guided by the text prompt
 2. **Negative/Unconditional**: Typically using empty or negative prompt
 
-Some models require 3 or more CFG branches (see [N-Branch CFG](#n-branch-cfg-3-branches)).
-
-CFG-Parallel eliminates this bottleneck by distributing the forward passes across different GPU ranks, allowing them to execute simultaneously rather than sequentially.
+CFG-Parallel eliminates this bottleneck by distributing the two forward passes across different GPU ranks, allowing them to execute simultaneously rather than sequentially.
 
 ### Architecture
 
@@ -35,11 +33,9 @@ vLLM-omni provides `CFGParallelMixin` that encapsulates all CFG parallel logic.
 
 | Method | Purpose | Automatic Behavior |
 |--------|---------|-------------------|
-| [`predict_noise_maybe_with_cfg()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Predict noise with 2-branch CFG | Detects parallel mode, distributes computation, gathers results |
-| [`predict_noise_with_multi_branch_cfg()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Predict noise with N-branch CFG | Round-robin dispatches N branches across M GPUs |
+| [`predict_noise_maybe_with_cfg()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Predict noise with CFG | Detects parallel mode, distributes computation, gathers results |
 | [`scheduler_step_maybe_with_cfg()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Step scheduler | All ranks step locally (no broadcast needed) |
-| [`combine_cfg_noise()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Combine 2-branch predictions | Applies CFG formula with optional normalization |
-| [`combine_multi_branch_cfg_noise()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Combine N-branch predictions | Override for custom multi-branch combine logic |
+| [`combine_cfg_noise()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Combine positive/negative | Applies CFG formula with optional normalization |
 | [`predict_noise()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Forward pass wrapper | Override for custom transformer calls |
 | [`cfg_normalize_function()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Normalize CFG output | Override for custom normalization |
 
@@ -61,22 +57,6 @@ vLLM-omni provides `CFGParallelMixin` that encapsulates all CFG parallel logic.
 
 - All ranks compute the scheduler step locally — no broadcast needed because `predict_noise_maybe_with_cfg` already ensures all ranks have identical noise predictions after `all_gather` + local combine.
 
-### N-Branch CFG (3+ branches)
-
-Some models require more than 2 CFG branches. For example, Bagel and OmniGen2 use 3 branches, DreamID Omni uses 4 branches.
-
-`predict_noise_with_multi_branch_cfg()` handles these by automatically dispatching N branches across M GPUs using round-robin (rule: branch `i` → rank `i % M`):
-
-| Branches (N) | GPUs (M) | Dispatch |
-|:---:|:---:|:---|
-| 3 | 2 | `[[0, 2], [1]]` |
-| 3 | 3 | `[[0], [1], [2]]` |
-| 4 | 2 | `[[0, 2], [1, 3]]` |
-| 4 | 3 | `[[0, 3], [1], [2]]` |
-| 4 | 4 | `[[0], [1], [2], [3]]` |
-
-When a rank handles multiple branches, it runs them sequentially. After `all_gather`, all ranks execute `combine_multi_branch_cfg_noise()` locally, producing identical results.
-
 ---
 
 ## Step-by-Step Implementation
@@ -118,7 +98,6 @@ class YourModelPipeline(nn.Module, CFGParallelMixin):
 - `positive_kwargs`: transformer arguments for conditional (text-guided) prediction
 - `negative_kwargs`: transformer arguments for unconditional prediction (set to `None` if CFG disabled)
 - For image editing pipelines, add `output_slice=image_seq_len` to extract the generative image portion
-- For models with 3+ CFG branches, see [Multi-Branch CFG](#multi-branch-cfg-3-branches) in the Customization section
 
 ### Step 2: Call `diffuse`
 
@@ -192,42 +171,20 @@ class LongCatImagePipeline(nn.Module, CFGParallelMixin):
 ```
 
 
-### Multi-Branch CFG (3+ branches)
+### Override `combine_cfg_noise()` for Multi-Output Models
 
-For models with 3 or more CFG branches, use `predict_noise_with_multi_branch_cfg()` instead of `predict_noise_maybe_with_cfg()`, and override `combine_multi_branch_cfg_noise()` for custom combine logic. This interface also works for standard 2-branch CFG — just pass 2 branches in `branches_kwargs`.
-
-**Example (3-branch with dual guidance scale):**
+When `predict_noise()` returns a tuple (e.g., video + audio), the default `combine_cfg_noise()` applies CFG to every element. Override it to apply different logic per element — for example, CFG on video but positive-only on audio:
 
 ```python
-class YourMultiBranchPipeline(nn.Module, CFGParallelMixin):
-    def combine_multi_branch_cfg_noise(self, predictions, true_cfg_scale, cfg_normalize=False):
-        text_scale = true_cfg_scale["text"]
-        image_scale = true_cfg_scale["image"]
-        pos, ref, uncond = predictions
-        return uncond + image_scale * (ref - uncond) + text_scale * (pos - ref)
-
-    def diffuse(self, ...):
-        for i, t in enumerate(timesteps):
-            positive_kwargs = {...}   # conditional prompt
-            ref_neg_kwargs = {...}    # negative prompt + reference
-            uncond_kwargs = {...}     # unconditional
-
-            noise_pred = self.predict_noise_with_multi_branch_cfg(
-                do_true_cfg=do_true_cfg,
-                true_cfg_scale={"text": text_guidance_scale, "image": image_guidance_scale},
-                branches_kwargs=[positive_kwargs, ref_neg_kwargs, uncond_kwargs],
-            )
-            latents = self.scheduler_step_maybe_with_cfg(noise_pred, t, latents, do_true_cfg)
-
-        return latents
+class MyVideoAudioPipeline(nn.Module, CFGParallelMixin):
+    def combine_cfg_noise(self, positive_noise_pred, negative_noise_pred, scale, normalize):
+        (video_pos, audio_pos) = positive_noise_pred
+        (video_neg, audio_neg) = negative_noise_pred
+        video_combined = super().combine_cfg_noise(video_pos, video_neg, scale, normalize)
+        return (video_combined, audio_pos)  # audio: positive only, no CFG
 ```
 
-### Override Combine Functions
-
-There are two combine functions for different scenarios:
-
-- **`combine_cfg_noise()`** — Used by `predict_noise_maybe_with_cfg()`. Override when `predict_noise()` returns a tuple (e.g., video + audio) and you need per-element CFG logic.
-- **`combine_multi_branch_cfg_noise()`** — Used by `predict_noise_with_multi_branch_cfg()`. Override to implement custom multi-branch combine formulas (see [Multi-Branch CFG](#multi-branch-cfg-3-branches) above).
+This also requires `predict_noise()` to return a tuple (see [Override predict_noise](#override-predict_noise-for-custom-transformer-calls) above).
 
 ### Implement a Composite Scheduler for Multi-Output Models
 
@@ -346,5 +303,4 @@ Adding CFG-Parallel support:
 
 1. ✅ **Create mixin** - Inherit from `CFGParallelMixin` and implement `diffuse()` method
 2. ✅ **(Optional) Customize** - Override `predict_noise()` or `cfg_normalize_function()` for custom behavior
-3. ✅ **(Optional) Multi-branch** - For 3+ branch models, use `predict_noise_with_multi_branch_cfg()` and override `combine_multi_branch_cfg_noise()`
-4. ✅ **Test** - Verify with `--cfg-parallel-size 2` (or 3/4 for multi-branch) and compare performance
+3. ✅ **Test** - Verify with `--cfg-parallel-size 2` and compare performance

docs/design/feature/expert_parallel.md

@@ -1,221 +0,0 @@
-# Expert Parallel
-
-This section describes how to add Expert Parallel (EP) to a diffusion transformer that uses Mixture-of-Experts (MoE) layers.
-We use **HunyuanImage3.0** as the reference implementation.
-
----
-
-## Table of Contents
-
-- [Overview](#overview)
-- [Step-by-Step Implementation](#step-by-step-implementation)
-- [Testing](#testing)
-- [Reference Implementations](#reference-implementations)
-- [Summary](#summary)
-
----
-
-## Overview
-
-### What is Expert Parallel?
-
-**Expert Parallel** is a parallelism strategy in Mixture-of-Experts (MoE) models that distributes different expert networks across distinct computational devices. Each device holds and computes only a subset of experts (local experts), with tokens dispatched to and gathered from remote devices via collective communication operations (e.g., All-to-All, All-Gather).
-
-| Backend | Description |
-|---------|-------------|
-| `allgather_reducescatter` | Default backend based on allgather/reducescatter primitives, suitable for general EP+DP deployments.|
-
-## Configuration
-
-Enable EP by setting the `--enable-expert-parallel` flag. The EP size is automatically calculated as:
-
-```text
-EP_SIZE = TP_SIZE × SP_SIZE × CFG_SIZE × DP_SIZE
-```
-
-
-Where:
-
-- `TP_SIZE`: Tensor parallel size
-- `SP_SIZE`: Sequence parallel size
-- `CFG_SIZE`: Classifier-free guidance parallel size
-- `DP_SIZE`: Data parallel size
-- `EP_SIZE`: Expert parallel size (computed automatically)
-
-Note:
-- Expert parallelism is only applicable to Mixture-of-Experts (MoE) models.
-- The EP group is created **per pipeline stage**, meaning it includes all ranks that participate in model parallelism except pipeline parallelism.
-- The underlying communication pattern for expert parallelism is **All-to-All** among the ranks in the EP group.
-
-For example, consider a configuration with `TP=2`, `SP=1`, `CFG=2`, and `DP=4` (total 2×1×2×4 = 16 GPUs).
-
-- Expert layers are handled by an EP group of size 16.
-
-- Attention layers use tensor parallelism of size 2 within each of the 8 DP groups (because `DP×CFG×SP = 4×2×1 = 8` groups, each containing the 2 TP ranks). Inside each such group, the attention weights are sharded across the 2 GPUs.
-
-
-## Step-by-Step Implementation
-
-### Step 1: Configure Expert Parallelism Settings
-
-Calculate local experts per rank:
-
-```
-ep_size = 8  # Expert Parallel size (typically equals TP size)
-num_experts = 64
-num_local_experts = num_experts // ep_size  # 8 experts per card
-
-# Check divisibility
-assert num_experts % ep_size == 0, "Experts must be divisible by EP size"
-```
-
-### Step 2: Use Sparse MoE Block to enable EP routing.
-
-Example:
-```
-from vllm.model_executor.layers.linear import ReplicatedLinear
-class HunYuanSparseMoeBlock(nn.Module):
-    def __init__(
-        self,
-        config: PretrainedConfig,
-        layer_id: int = -1,
-        prefix: str = "",
-    ):
-        super().__init__()
-        self.tp_size = get_tensor_model_parallel_world_size()
-        self.n_routed_experts = config.num_experts  # 64
-
-        # Calculate local experts per rank (key for EP)
-        if self.tp_size > self.n_routed_experts:
-            raise ValueError(f"TP size {self.tp_size} > experts {self.n_routed_experts}")
-
-        # Routing gate (replicated on all ranks, computes scores for all tokens to all experts)
-        self.gate = ReplicatedLinear(
-            config.hidden_size,
-            config.num_experts,
-            bias=False,
-            quant_config=None,
-            prefix=f"{prefix}.gate",
-        )
-
-        # EP expert layer (factory loads platform-specific implementation)
-        self.experts = HunyuanFusedMoE(...)
-```
-**Key Points:**
-- gate is **ReplicatedLinear** (replicated on all ranks)
-- experts is created via **HunyuanFusedMoE factory**, which automatically handles EP dispatch
-
-### Step 3: Initialize EP Runtime
-
-Initialize the EP communication context before model loading.
-```
-from vllm.utils.import_utils import resolve_obj_by_qualname
-# Call during __init__ or model loading
-op_name = "hunyuan_fused_moe"
-
-# Prepare EP runtime: establish communication groups, assign local expert indices, init _expert_map
-current_omni_platform.prepare_diffusion_op_runtime(op_name)
-
-# Factory automatically resolves platform implementation (GPU: FusedMoE / NPU: AscendFusedMoE)
-impl = resolve_obj_by_qualname(
-    current_omni_platform.get_diffusion_model_impl_qualname(op_name)
-)
-```
-
-### Step 4: Expert Weight Mapping & Loading
-
-Each rank loads only the expert weights assigned to its local allocation.
-```
-# Get expert parameter mapping (different per rank)
-expert_mapping = HunyuanFusedMoE.make_expert_params_mapping(
-    model=self,
-    ckpt_gate_proj_name="gate_proj",
-    ckpt_down_proj_name="down_proj",
-    ckpt_up_proj_name="up_proj",
-    num_experts=64,
-    num_redundant_experts=0,
-)
-# Returns: [(param_name, weight_name, expert_id, shard_id), ...]
-# Note: Each rank only contains mappings for its local expert_ids
-
-# Filter non-local experts during loading
-for name, loaded_weight in weights:
-    if "mlp.experts" in name:
-        # Parse expert_id from weight name (implementation needed)
-        expert_id = parse_expert_id_from_name(name)
-        local_expert_start = (ep_rank) * num_local_experts
-        local_expert_end = (ep_rank + 1) * num_local_experts
-
-        if not (local_expert_start <= expert_id < local_expert_end):
-            continue  # Skip non-local expert weights
-```
-### Step 5: Forward Pass with EP
-
-Example (MoE Forward):
-```
-def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-    orig_shape = hidden_states.shape
-    hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
-
-    # 1. Global routing computation (all tokens, all expert scores)
-    # hidden_states: [num_tokens, hidden_dim] (full tensor)
-    router_logits, _ = self.gate(hidden_states)  # [num_tokens, num_experts]
-
-    # 2. EP dispatch and compute (HunyuanFusedMoE handles all_to_all internally)
-    # - Dispatch: Send tokens to target ranks based on router_logits
-    # - Local Compute: Each rank processes only its num_local_experts
-    # - Combine: Results returned to original token positions
-    final_hidden_states = self.experts(
-        hidden_states=hidden_states,
-        router_logits=router_logits,
-    )
-
-    # 3. Add shared expert output (not EP, computed on all ranks)
-    if self.shared_mlp is not None:
-        shared_out = self.shared_mlp(hidden_states)
-        final_hidden_states = final_hidden_states + shared_out
-
-    # 4. Tensor Parallel All-Reduce (synchronize across TP group)
-    if self.tp_size > 1:
-        final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel(
-            final_hidden_states
-        )
-
-    return final_hidden_states.view(orig_shape)
-```
-
-## Testing
-After adding Expert Parallel support, test via command line:
-```bash
-cd examples/offline_inference/text_to_image
-python text_to_image.py \
-    --model Your-org/your-model \
-    --prompt "a cup of coffee on the table" \
-    --output "ep_enabled.png" \
-    --num-inference-steps 50 \
-    --guidance-scale 5.0 \
-    --tensor-parallel-size 8 \
-    --seed 1234 \
-    --enable-expert-parallel
-```
-
-vLLM‑Omni currently focuses on core diffusion model inference acceleration, so the Expert Parallel implementation includes only the basic multi‑GPU expert sharding functionality (enabled via --enable-expert-parallel). Advanced features such as communication backend selection (--all2all-backend), load balancing (--enable-eplb and its configuration), and multi‑node deployment belong to the extended capabilities of the main vLLM project and have not yet been integrated into Omni.
-
-## Reference Implementations
-
-Complete examples in the codebase:
-
-| Model | Path | Pattern | Notes |
-|-------|------|---------|-------|
-| **HunyuanImage3.0** | `vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py` | Standard EP | Full implementation with validation |
-| **EP Tests** | `vllm-omni/tests/e2e/offline_inference/test_expert_parallel.py` | E2E testing | EP correctness and performance |
-| **Constraint Tests** | `vllm-omni/tests/diffusion/models/hunyuan_image3/test_hunyuan_fused_moe.py` | Unit testing | Validation logic |
-
----
-## Summary
-
-Adding Expert Parallel support to diffusion model:
-
-1. **Identify MoE layers** - Locate the router and expert networks in each transformer block.
-2. **Validate EP constraints** – Ensure num_experts is divisible by expert_parallel_size.
-3. **Test** - Run with enable-expert-parallel, check memory reduction, speedup, and output quality against single‑GPU baseline.

docs/design/feature/omni_connectors/mooncake_transfer_engine_connector.md

@@ -33,8 +33,8 @@ runtime:
         zmq_port: 50051               # ZMQ base port (see "Port Offset Scheme" below)
         protocol: "rdma"              # "rdma" or "tcp"
         device_name: ""               # RDMA device (e.g., "mlx5_0"), empty for auto-detect
-        memory_pool_size: 4294967296  # 4 GB (CPU); use 2147483648 (2 GB) for GPU
-        memory_pool_device: "cpu"     # "cpu" for pinned memory (recommended), "cuda" for GPUDirect RDMA
+        memory_pool_size: 2147483648  # 2GB memory pool
+        memory_pool_device: "cpu"     # "cpu" for pinned memory, "cuda" for GPUDirect RDMA
 ```
 
 Wire stages to the connector:
@@ -64,8 +64,8 @@ stage_args:
 
 | Parameter | Default | Description |
 |---|---|---|
-| `memory_pool_size` | 4 GB (CPU) / 2 GB (GPU) | Total size of the RDMA-registered memory pool in bytes. Recommended 4 GB for CPU pinned memory; 2 GB for GPU VRAM to conserve device memory. |
-| `memory_pool_device` | `"cpu"` | `"cpu"`: pinned host memory (recommended, works on all topologies). `"cuda"`: GPU VRAM for GPUDirect RDMA (requires NIC-GPU direct PCIe connectivity, PIX topology). |
+| `memory_pool_size` | 1 GB | Total size of the RDMA-registered memory pool in bytes. |
+| `memory_pool_device` | `"cpu"` | `"cpu"`: pinned host memory (recommended). `"cuda"`: GPU VRAM for GPUDirect RDMA (requires NIC-GPU direct PCIe connectivity). |
 
 ### Networking
 
@@ -107,10 +107,10 @@ receiver_connect  = remote_side_channel_port + tp_rank
 
 ## Memory Pool Modes
 
-| Mode | Config | Recommended Pool Size | Data Flow | Best For |
-|---|---|---|---|---|
-| CPU Pinned | `memory_pool_device: "cpu"` | 4 GB | GPU → CPU pool → RDMA → CPU pool → GPU | Most hardware topologies (recommended) |
-| GPUDirect | `memory_pool_device: "cuda"` | 2 GB | GPU → GPU pool → RDMA (NIC reads GPU BAR1) → GPU pool | NIC-GPU direct PCIe (PIX topology) |
+| Mode | Config | Data Flow | Best For |
+|---|---|---|---|
+| CPU Pinned | `memory_pool_device: "cpu"` | GPU → CPU pool → RDMA → CPU pool → GPU | Most hardware topologies (recommended) |
+| GPUDirect | `memory_pool_device: "cuda"` | GPU → GPU pool → RDMA (NIC reads GPU BAR1) → GPU pool | NIC-GPU direct PCIe (PIX topology) |
 
 > **Note**: GPUDirect RDMA requires the NIC and GPU to share a direct PCIe
 > switch (PIX topology). On systems where they are connected via PXB or NODE,

docs/design/feature/prefix_caching.md

@@ -1,164 +0,0 @@
-# Automatic Prefix Caching in Omni Models
-
-
----
-
-## Table of Contents
-
-- [Overview](#overview)
-- [High-Level Approach](#high-level-approach)
-- [Example](#example)
-- [What About Multimodal Inputs?](#what-about-multimodal-inputs)
-
----
-
-### Overview
-
-Prefix caching in the context of kv-cache management is a useful optimization for avoiding redundant computations. The main idea is that we store portions of the kv-cache from processed requests, so that we can reuse them if incoming requests have the same prefix as previous requests.
-
-vLLM manages the kv-cache as blocks, which represent a span of tokens of a fixed length. Blocks are hashable by the content that they contain, which typically means the tokens within the span, but also could be influenced by other factors, e.g., LoRA and multimodal data.
-
-vLLM implements automatic prefix caching for managing its kv-cache, which is best understood by reading the design document [here](https://docs.vllm.ai/en/latest/design/prefix_caching/). vLLM-Omni builds on top of the prefix caching mechanism in a noninvasive way to allow caching between stages in Omni pipelines. This typically means for a given stage we aim to support caching for the following:
-
-- The last hidden states produced by the stage
-- Model / stage specific multimodal data
-
-!!! note "Note 1"
-    This document describes vLLM-Omni's mechanism for caching tensor outputs that are meant to be passed between stages, when requests have common prefixes, similar to the way in which vLLM has prefix caching for the kv-cache. This works in conjunction with vLLM's multimodal encoder caching, but is distinct. See the final section for a concrete example for how they tie together in practice.
-
-### High-Level Approach
-!!! note "Note 2"
-    Prior to reading this section, it's recommended to take a look at the design documents in vLLM for [Automatic Prefix Caching](https://docs.vllm.ai/en/latest/features/automatic_prefix_caching/), which will make some of the concepts more clear.
-
-The main focus of vLLM-Omni's approach to prefix caching stage outputs is to build on vLLM's prefix caching in the least invasive way possible while minimizing impact for cache misses, and consuming a minimal amount of GPU memory. To understand the implementation, there are a few important things to note:
-
-- Between stages, device tensors are generally moved to CPU; this is important since we're just caching the outputs of stages, so it is okay to keep the entire cache on the CPU.
-
-- For a tensor to be considered cacheable, the first dimension (currently) needs to be the same as the token count, as it allows us to reuse block/slot mappings for our externally maintained tensor caches. This allows us to dynamically discover the tensors to be marked as cacheable outputs in each Omni model without having to explicitly specify cacheable output field names in every model.
-
-With this in mind, consider the set of blocks in a 2D layout, where the row represents the index of blocks being considered, and the columns represent the slots corresponding to tokens within each block. Since we know the `num_blocks` and `block_size` from our kv cache config, if we want to cache a tensor with feature size `D`, we can preallocate a CPU tensor of size `(num_blocks, block_size, D)`, and use the same block index and slot mapping to retrieve the corresponding feature vector.
-
-
-### Example
-!!! note "Note 3"
-    Prefix caching in vLLM-Omni currently is only supported on AutoRegressive stages with one kv-cache group. It can be enabled/disabled per-stage via the `enable_prefix_caching` parameter in the model's stage config.
-
-The way in which vLLM-Omni ties into vLLM's prefix caching is best understood by example. Say that we have the following:
-
-- `num_blocks=8`
-- `block_size=4`
-- `hidden_size=2`
-- A stage specific multimodal output tensor named `mm_feature` with feature dimension `16`
-
-The prefix cache flow is then outlined below.
-
-1. When the model is initialized, we can determine the `hidden_size` from the `ModelConfig`, and allocate a cache of size `(num_blocks, block_size, hidden_size)`.
-
-2. Say we process the request `The quick brown fox was tired and slept beneath the shady tree`, which is 12 tokens and evenly divides into 3 blocks as shown below.
-
-```
-         [  The quick brown fox  ] [  was tired and slept ] [beneath the shady tree ]
-Block 1: |<--- block tokens ---->|
-Block 2: |<------- prefix ------>| |<--- block tokens --->|
-Block 3: |<------------------ prefix -------------------->| |<--- block tokens ---->|
-```
-
-When the request processes, we inspect the multimodal outputs and identify the `mm_feature` tensor, which will be of shape `(seq_len, feature_dim)`, i.e., `(12, 16)` in this example. We note that the first axis is dependent on the `seq_len` and add a new cache_tensor of shape `(num_blocks, block_size, feature_dim)` to our multimodal cache for tensors.
-
-
-3. If we lay out the cache as a 2D tensor of shape (`num_blocks`, `block_size`), we'll have something like the following:
-
-```
-0: [  The quick brown fox  ]
-1: [  was tired and slept  ]
-2: [beneath the shady tree ]
-3: [EMPTY]
-...
-7: [EMPTY]
-```
-
-Or, if we flatten it down to 1D,
-```
-0: The
-1: quick
-2: brown
-3: fox
-...
-11: tree
-12: [EMPTY]
-...
-```
-
-which we can think of as row indices into the hidden states tensor if we view it as the 2D shape `(num_blocks x block_size, feature_dim)`. That is, the analogous flattened (from 3D -> 2D) mapping of the cache for hidden states becomes the following.
-```
-0: <hidden states vector of len 2 corresponding to 'The'>
-1: <hidden states vector of len 2 corresponding to 'quick'>
-2: <hidden states vector of len 2 corresponding to 'brown'>
-3: <hidden states vector of len 2 corresponding to 'fox'>
-...
-11: <hidden states vector of len 2 corresponding to 'tree'>
-12: [EMPTY]
-...
-```
-
-Similarly, for the multimodal outputs cache, the flattened coordinates are the same, but the `mm_feature` maps to vectors of length `16` instead of the hidden size of `2`. Note that in practice, we may have multiple  multimodal output tensors per forward pass, which may have different names and different feature dimensions.
-
-
-4. Now, say that we receive a new request `The quick brown fox jumped over the dog`.
-
-```
-         [  The quick brown fox  ] [  jumped over the dog ]
-Block 1: |<--- block tokens ---->|
-Block 2: |<------- prefix ------>| |<--- block tokens --->|
-```
-
-Here, we will have a cache hit for `Block 1` which will be detected by vLLM based on the hash of the first block when it's handling the prefix caching on the kv-cache. As a result, when we get the output from the scheduler, we will see that `num_computed_tokens=4` (corresponding to the cached first block), and we only need to process the remaining 4 new tokens in the new prefill.
-
-Since we have the block indices / slot mappings from the kv cache manager, we can simply mirror the mappings and leverage the same indices for the cached hidden states and multimodal outputs. This allows us to look up the correct tensors from our externally maintained 3D caches.
-
-```
-0: [  The quick brown fox  ] < already in the cache
-1: [  was tired and slept  ]
-2: [beneath the shady tree ]
-3: [ jumped over the dog  ] < added on the second request
-4: [EMPTY]
-...
-7: [EMPTY]
-...
-```
-
-Finally, to pass the full hidden states and multimodal outputs to the next stage, we simply concatenate the cached contents with the corresponding new tensors computed from the current forward call.
-
-
-### What About Multimodal Inputs?
-It's also useful to consider the case about how Omni prefix caching is handled when we have multimodal inputs that don't cleanly end on block boundaries, as well as how this works with multimodal encoder caching in vLLM. For example:
-
-```
-         [   Im0  Im1  Im2  Im3  ] [ Im4  Im5 foo <empty> ]
-Block 1: |<--- block tokens ---->|
-Block 2: |<------- prefix ------>| |<--- block tokens --->|
-```
-
-In this case, only `Block 1` will have outputs stored in the prefix tensor cache, because vLLM does not store partial blocks. This may appear to be a problem at first glance, because the multimodal input is fragmented across a new block that wasn't cached.
-
-In reality, this isn't a big problem for correctness, because vLLM also maintains an encoder cache for multimodal inputs. In other words, after the first pass, we'll have the following:
-
-- The Block 1 hash, which is used for prefix caching
-- The hash describing the image data starting at position 0 and with length 6
-- In vLLM's encoder cache, a mapping from the image hash above to the encoder output
-
-
-To understand what happens, say we get the following input as a second request:
-```
-         [   Im0  Im1  Im2  Im3  ] [  Im4  Im5 bar  baz  ]
-Block 1: |<--- block tokens ---->|
-Block 2: |<------- prefix ------>| |<--- block tokens --->|
-```
-
-First, the scheduler will check for a prefix cache hit, which we will see on `Block 1`. As a result, we will have 4 tokens marked as precomputed, and only see the remaining 4 tokens in the following prefill.
-
-Because we have multimodal data in a scheduled span that isn't fully precomputed, we still need to call the visual encoder. However, since we have the image hash and encoder cache, we will retrieve the encoder outputs for `Im4` and `Im5` as we create the multimodal embeddings.
-
-When we pass our multimodal tensors to the language model component in the same stage, we'll then expect the same outputs, because the prefix caching behaviors in vLLM-Omni / vLLM match, so the LLM will use vLLM's KV cache manager's prefix caching to correctly handle the attention information for `Block 1` while calculating the outputs for `Block 2`, giving us the correct results for processing `Block 2` with the context of `Block 1`.
-
-Finally, we look up the output hidden states/multimodal tensors corresponding to the prefix cache hit `Block 1` and concatenate it with the forward pass result to get the final result, which is expected to be identical to the full hidden states when prefix caching is disabled.

docs/design/feature/teacache.md

@@ -326,41 +326,9 @@ for prompt in tqdm(prompts, desc="Collecting data"):
 
 # Estimate coefficients
 coeffs = estimator.estimate(poly_order=4)
-print(f"Estimated coefficients: {coeffs}")
+print(f"Estimated coefficients: {coeffs.tolist()}")
 ```
 
-Note: some models may require the vLLM context and config to be initialized to initialize vLLM modules. To this end, you may need a workaround like the following to be able to run coefficient estimation.
-```python
-from vllm_omni.diffusion.forward_context import set_forward_context
-from vllm_omni.diffusion.distributed.parallel_state import (
-    init_distributed_environment,
-    initialize_model_parallel,
-)
-from vllm.config import VllmConfig
-...
-
-if __name__ == "__main__":
-    os.environ["MASTER_ADDR"] = "localhost"
-    os.environ["MASTER_PORT"] = "8192"
-    os.environ["LOCAL_RANK"] = "0"
-    os.environ["RANK"] = "0"
-    os.environ["WORLD_SIZE"] = "1"
-
-    vllm_config = VllmConfig()
-    init_distributed_environment()
-    initialize_model_parallel()
-
-    # NOTE: you may have to pass an initialized OmniDiffusionConfig as a kwarg
-    # here to make current sp checks happy; if this is the case, just create one
-    # .from_kwargs() with the model name to get around this check for now,
-    # since your estimator subclass should handle the actual model configuration.
-    #
-    # This will be cleaned up in the future
-    with set_forward_context(vllm_config):
-        <create the estimator + run estimation here>
-```
-
-
 **Data Statistics Guide:**
 
 | Metric | Good Range | Warning Signs |

docs/design/feature/vae_parallel.md

@@ -1,15 +1,14 @@
 # VAE Patch Parallelism
 
 This document describes how to add **VAE Patch Parallelism** support to a diffusion model.
-We use **Qwen-Image** as the reference implementation for decode parallel, and **Wan2.2** for encode parallel.
+We use **Qwen-Image** as the reference implementation.
 
 ---
 
 ## Table of Contents
 
 - [Overview](#overview)
-- [Step-by-Step Implementation (Decode)](#step-by-step-implementation-decode)
-- [Encode Parallel Implementation](#encode-parallel-implementation)
+- [Step-by-Step Implementation](#step-by-step-implementation)
 - [Testing](#testing)
 - [Reference Implementations](#reference-implementations)
 - [Summary](#summary)
@@ -20,13 +19,13 @@ We use **Qwen-Image** as the reference implementation for decode parallel, and *
 
 ### What is Vae Patch parallel?
 
-**VAE Patch Parallelism** is an acceleration technique for both **encoding** and **decoding**. Instead of processing the entire tensor at once, the tensor is:
+**VAE Patch Parallelism** is a decoding acceleration technique. Instead of decoding the entire latent tensor at once, the latent tensor is:
 
 + Split into multiple spatial tiles
 
 + Distributed across multiple ranks
 
-+ Encoded/Decoded in parallel
++ Decoded in parallel
 
 + Merged to reconstruct the final output
 
@@ -36,17 +35,10 @@ This approach:
 
 + Reduces peak memory usage per device
 
-+ Accelerates encoding/decoding latency
-
-### When to Use Encode vs Decode Parallel
-
-| Operation | Use Case | Example |
-|-----------|----------|---------|
-| **Decode Parallel** | Text-to-Image, Text-to-Video | Latent → Image/Video |
-| **Encode Parallel** | Image-to-Video (I2V) | Image → Latent (for conditioning) |
++ Accelerates decoding latency
 
 ### Architecture
-We introduce **DistributedVaeExecutor** as the core component responsible for distributed VAE encoding/decoding.
+We introduce **DistributedVaeExecutor** as the core component responsible for distributed VAE decoding.
 
 The executor is model-agnostic and accepts three function parameters:
 
@@ -92,7 +84,7 @@ Therefore:
 
 + Merge must perform blending to avoid seams
 
-## Step-by-Step Implementation (Decode)
+## Step-by-Step Implementation
 
 ### Step 1: Implement DistributedAutoencoderKLQwenImage
 `QwenImagePipeline` use `AutoencoderKLQwenImage` for vae, so implement a distributed version:
@@ -213,14 +205,14 @@ def tile_merge(self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid
 We need to override tiled_decode, the main logic is:
 + check distributed is enabled
 + select split/exec/merge
-+ Invoke self.distributed_executor.execute to decode
++ Invoke self.distributed_decoder.execute to decode
 ```
 def tiled_decode(self, z: torch.Tensor, return_dict: bool = True):
     if not self.is_distributed_enabled():
         return super().tiled_decode(z, return_dict=return_dict)
 
     logger.info("Decode run with distributed executor")
-    result = self.distributed_executor.execute(
+    result = self.distributed_decoder.execute(
         z,
         DistributedOperator(split=self.tile_split, exec=self.tile_exec, merge=self.tile_merge),
         broadcast_result=True,
@@ -251,166 +243,6 @@ class YourModelPipeline(nn.Module):
 +       ).to(self.device)
 ```
 
-## Encode Parallel Implementation
-
-For models that require VAE encoding (e.g., Image-to-Video), you can also parallelize the encode operation. We use **Wan2.2** as the reference implementation.
-
-### Step 1: Implement encode_tile_split
-
-Similar to decode, split the input tensor into tiles. Key considerations:
-
-+ **Patchify handling**: If the model uses `patch_size`, scale tile parameters accordingly
-+ **Temporal chunking**: Video VAEs may have temporal compression (e.g., 4x)
-
-```python
-def encode_tile_split(self, x: torch.Tensor) -> tuple[list[TileTask], GridSpec]:
-    _, _, num_frames, height, width = x.shape
-    encode_spatial_compression_ratio = self.spatial_compression_ratio
-
-    # Scale tile parameters for patchified coordinate system
-    tile_sample_min_height = self.tile_sample_min_height
-    tile_sample_min_width = self.tile_sample_min_width
-    tile_sample_stride_height = self.tile_sample_stride_height
-    tile_sample_stride_width = self.tile_sample_stride_width
-
-    if self.config.patch_size is not None:
-        # When input is patchified, scale tile parameters accordingly
-        encode_spatial_compression_ratio = self.spatial_compression_ratio // self.config.patch_size
-        tile_sample_min_height = tile_sample_min_height // self.config.patch_size
-        tile_sample_min_width = tile_sample_min_width // self.config.patch_size
-        tile_sample_stride_height = tile_sample_stride_height // self.config.patch_size
-        tile_sample_stride_width = tile_sample_stride_width // self.config.patch_size
-
-    latent_height = height // encode_spatial_compression_ratio
-    latent_width = width // encode_spatial_compression_ratio
-
-    tile_latent_min_height = tile_sample_min_height // encode_spatial_compression_ratio
-    tile_latent_min_width = tile_sample_min_width // encode_spatial_compression_ratio
-    tile_latent_stride_height = tile_sample_stride_height // encode_spatial_compression_ratio
-    tile_latent_stride_width = tile_sample_stride_width // encode_spatial_compression_ratio
-
-    blend_height = tile_latent_min_height - tile_latent_stride_height
-    blend_width = tile_latent_min_width - tile_latent_stride_width
-
-    tiletask_list = []
-    # Use temporal compression ratio from config instead of hardcoding
-    temporal_compression = self.config.scale_factor_temporal
-
-    for i in range(0, height, tile_sample_stride_height):
-        for j in range(0, width, tile_sample_stride_width):
-            time_list = []
-            frame_range = 1 + (num_frames - 1) // temporal_compression
-            for k in range(frame_range):
-                if k == 0:
-                    tile = x[:, :, :1, i : i + tile_sample_min_height, j : j + tile_sample_min_width]
-                else:
-                    tile = x[
-                        :, :,
-                        1 + temporal_compression * (k - 1) : 1 + temporal_compression * k,
-                        i : i + tile_sample_min_height,
-                        j : j + tile_sample_min_width,
-                    ]
-                time_list.append(tile)
-            tiletask_list.append(
-                TileTask(len(tiletask_list), (i // tile_sample_stride_height, j // tile_sample_stride_width),
-                         time_list, workload=time_list[0].shape[3] * time_list[0].shape[4])
-            )
-
-    grid_spec = GridSpec(
-        split_dims=(3, 4),
-        grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1),
-        tile_spec={
-            "latent_height": latent_height, "latent_width": latent_width,
-            "blend_height": blend_height, "blend_width": blend_width,
-            "tile_latent_stride_height": tile_latent_stride_height,
-            "tile_latent_stride_width": tile_latent_stride_width,
-        },
-        output_dtype=self.dtype,
-    )
-    return tiletask_list, grid_spec
-```
-
-### Step 2: Implement encode_tile_exec
-
-```python
-def encode_tile_exec(self, task: TileTask) -> torch.Tensor:
-    """Encode a single sample tile into latent space."""
-    self.clear_cache()
-    time = []
-    for k, tile in enumerate(task.tensor):
-        self._enc_conv_idx = [0]
-        encoded = self.encoder(tile, feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx)
-        encoded = self.quant_conv(encoded)
-        time.append(encoded)
-    result = torch.cat(time, dim=2)
-    self.clear_cache()
-    return result
-```
-
-### Step 3: Implement encode_tile_merge
-
-```python
-def encode_tile_merge(
-    self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid_spec: GridSpec
-) -> torch.Tensor:
-    """Merge encoded tiles into a full latent tensor."""
-    grid_h, grid_w = grid_spec.grid_shape
-    result_rows = []
-    for i in range(grid_h):
-        result_row = []
-        for j in range(grid_w):
-            tile = coord_tensor_map[(i, j)]
-            if i > 0:
-                tile = self.blend_v(coord_tensor_map[(i - 1, j)], tile, grid_spec.tile_spec["blend_height"])
-            if j > 0:
-                tile = self.blend_h(coord_tensor_map[(i, j - 1)], tile, grid_spec.tile_spec["blend_width"])
-            result_row.append(tile[:, :, :,
-                : grid_spec.tile_spec["tile_latent_stride_height"],
-                : grid_spec.tile_spec["tile_latent_stride_width"]])
-        result_rows.append(torch.cat(result_row, dim=-1))
-
-    enc = torch.cat(result_rows, dim=3)[
-        :, :, :, : grid_spec.tile_spec["latent_height"], : grid_spec.tile_spec["latent_width"]
-    ]
-    return enc
-```
-
-### Step 4: Override tiled_encode method
-
-Override `tiled_encode` instead of `encode`. The parent's `_encode()` handles patchify before calling `tiled_encode()`, so input `x` is already patchified.
-
-```python
-def tiled_encode(self, x: torch.Tensor) -> torch.Tensor:
-    """
-    Encode using distributed VAE executor.
-
-    Note: x is already patchified by parent's _encode() before calling this method.
-    """
-    if not self.is_distributed_enabled():
-        return super().tiled_encode(x)
-
-    self.clear_cache()
-    result = self.distributed_executor.execute(
-        x,
-        DistributedOperator(
-            split=self.encode_tile_split,
-            exec=self.encode_tile_exec,
-            merge=self.encode_tile_merge,
-        ),
-        broadcast_result=True,  # Latents needed by all ranks for diffusion
-    )
-    self.clear_cache()
-    return result
-```
-
-**Key differences from decode parallel:**
-
-| Aspect | Decode Parallel | Encode Parallel |
-|--------|-----------------|-----------------|
-| `broadcast_result` | Often `False` (only rank 0 needs output) | `True` (all ranks need latents for diffusion) |
-| Patchify | Applied in merge (unpatchify) | Handled by parent `_encode()` before `tiled_encode()` |
-| Temporal chunking | Frame-by-frame | Chunk-based (e.g., 1 + 4n frames) |
-
 ## Testing
 Verify numerical consistency between:
 + vae_patch_parallel_size = 1
@@ -440,20 +272,18 @@ When vae_patch_parallel_size is larger than the DiT world size, it will automati
 
 Complete examples in the codebase:
 
-| Model | Path | Decode Parallel | Encode Parallel |
-|-------|------|-----------------|-----------------|
-| **Z-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl.py` | ✅ | ❌ |
-| **Wan2.2** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py` | ✅ | ✅ |
-| **Qwen-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py` | ✅ | ❌ |
+| Model | Path | Notes |
+|-------|------|-------|
+| **Z-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl.py` | Distributed AutoencoderKL |
+| **Wan2.2** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py` | Distributed AutoencoderKLWan |
+| **Qwen-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py` | Distributed AutoencoderKLQwenImage |
 
 ---
 
 ## Summary
 
-Adding VAE Patch Parallel support to diffusion model:
+Adding Vae Patch Parallel support to diffusion model:
 
-1. **Implement Distributed VAE** - Inherit from base VAE class and `DistributedVaeMixin`
-2. **Decode Parallel** - Refactor `tiled_decode` into `tile_split`/`tile_exec`/`tile_merge`
-3. **Encode Parallel** (optional) - Implement `encode_tile_split`/`encode_tile_exec`/`encode_tile_merge` for I2V models
-4. **Change VAE model in pipeline** - Use the distributed version
-5. **Test** - Verify numerical consistency with `vae_patch_parallel_size=1` vs `N`
+1. **Implement Distributed Vae** - mainly copy from `diffusers` tiled_decode, and refactor into split/exec/merge
+2. **Change vae model in pipeline to Distributed Vae**
+3. **Test** - Verify with `tensor_parallel_size=N` quality