Clean up PR template

Signed-off-by: yujuancao07 <yujuancao07@gmail.com>

.buildkite/pipeline.yml

@@ -1,14 +1,14 @@
 # Document 1: Buildkite loads only this block on first parse. The next step resolves docs-only skip-ci
 # from git diff, then uploads document 2. When docs-only skip applies, image-build still runs for the nightly
-# exception in upload_pipeline_with_skip_ci.sh (nightly-test label or main NIGHTLY=1); otherwise upload-nightly
+# exception in upload_pipeline.py (nightly-test label or main NIGHTLY=1); otherwise upload-nightly
 # can be skipped together with test-ready/test-merge. Weekly uploads also depend on image-build.
 
-# Document 2: appended after `---`; same file, read by upload_pipeline_with_skip_ci.sh (not evaluated as a second pipeline by Buildkite). Child uploads include test-ready.yml, test-merge.yml, test-nightly.yml, and test-weekly.yml.
+# Document 2: appended after `---`; same file, read by upload_pipeline.py (not evaluated as a second pipeline by Buildkite). Child uploads include test-ready.yml, test-merge.yml, test-nightly.yml, and test-weekly.yml.
 steps:
   - label: ":github: Resolve skip-ci & upload pipeline"
     key: upload-ci-pipeline
     commands:
-      - "bash .buildkite/scripts/upload_pipeline_with_skip_ci.sh"
+      - "python3 .buildkite/scripts/upload_pipeline.py --upload .buildkite/pipeline.yml"
     agents:
       queue: "cpu_queue_premerge"
 
@@ -41,7 +41,7 @@ steps:
     key: upload-merge-pipeline
     if: __UPLOAD_MERGE_IF__
     commands:
-      - buildkite-agent pipeline upload .buildkite/test-merge.yml
+      - python3 .buildkite/scripts/upload_pipeline.py --upload .buildkite/test-merge.yml
     agents:
       queue: "cpu_queue_premerge"
 

.buildkite/rebase-pipeline.yaml

@@ -23,7 +23,7 @@ steps:
     depends_on: image-build
     key: upload-merge-pipeline
     commands:
-      - buildkite-agent pipeline upload .buildkite/test-merge.yml
+      - python3 .buildkite/scripts/upload_pipeline.py --upload .buildkite/test-merge.yml
     agents:
       queue: "cpu_queue_premerge"
 

.buildkite/scripts/upload_pipeline.py

@@ -0,0 +1,308 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+"""Render and optionally upload Buildkite pipeline YAML with diff-aware logic.
+
+Bootstrap mode (pipeline.yml with __IMAGE_BUILD_IF__ placeholders):
+  - Detect docs-only PR/main changes and substitute skip-ci ``if`` expressions.
+
+Test pipeline mode (e.g. test-merge.yml):
+  - Drop steps or groups whose ``source_file_dependencies`` do not match changed files.
+  - ``source_file_dependencies`` may be set on a leaf step or on a group step.
+  - Always strip ``source_file_dependencies`` from uploaded YAML (Buildkite does not
+    recognize this uploader-only key).
+
+Usage:
+  python3 upload_pipeline.py [--upload] <pipeline.yml>
+
+  # Bootstrap (replaces upload_pipeline_with_skip_ci.sh):
+  python3 upload_pipeline.py --upload .buildkite/pipeline.yml
+
+  # Test pipeline (replaces upload_test_pipeline_with_diff_skip.py):
+  python3 upload_pipeline.py --upload .buildkite/test-merge.yml
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import subprocess
+import sys
+from pathlib import Path
+from typing import Any
+
+try:
+    import yaml
+except ImportError:
+    subprocess.check_call(
+        [sys.executable, "-m", "pip", "install", "-q", "pyyaml"],
+    )
+    import yaml
+
+LOG = "upload_pipeline"
+ROOT = Path(__file__).resolve().parent.parent.parent
+DOC_SEP = "\n---\n"
+BOOTSTRAP_MARKER = "__IMAGE_BUILD_IF__"
+
+
+def _log(message: str) -> None:
+    print(f"{LOG}: {message}", file=sys.stderr)
+
+
+def _git(*args: str) -> subprocess.CompletedProcess[str]:
+    return subprocess.run(
+        ["git", *args],
+        check=False,
+        capture_output=True,
+        text=True,
+    )
+
+
+def resolve_changed_files() -> list[str] | None:
+    """Return changed file paths, or None when diff cannot be resolved."""
+    is_pr = os.environ.get("BUILDKITE_PULL_REQUEST", "false") != "false" and os.environ.get(
+        "BUILDKITE_PULL_REQUEST", ""
+    )
+    commit = os.environ.get("BUILDKITE_COMMIT", "")
+
+    if is_pr:
+        base_branch = os.environ.get("BUILDKITE_PULL_REQUEST_BASE_BRANCH", "main")
+        base_ref = f"origin/{base_branch}"
+        if _git("rev-parse", "--verify", base_ref).returncode != 0:
+            _log(f"origin/{base_branch} not found locally; trying fetch")
+            _git("fetch", "--depth=200", "origin", base_branch)
+        if _git("rev-parse", "--verify", base_ref).returncode != 0:
+            if _git("rev-parse", "--verify", base_branch).returncode == 0:
+                base_ref = base_branch
+            else:
+                _log(f"cannot resolve PR base {base_branch}; using safe defaults")
+                return None
+        diff_range = f"{base_ref}...{commit}"
+    elif os.environ.get("BUILDKITE_BRANCH", "") == "main":
+        if _git("rev-parse", "--verify", f"{commit}^").returncode != 0:
+            _log("main commit has no parent; using safe defaults")
+            return None
+        diff_range = f"{commit}^..{commit}"
+    else:
+        _log("not PR/main build; using safe defaults")
+        return None
+
+    result = _git("diff", "--name-only", diff_range)
+    if result.returncode != 0:
+        _log(f"git diff failed ({diff_range}); using safe defaults")
+        return None
+
+    files = [line for line in result.stdout.splitlines() if line.strip()]
+    _log(f"{len(files)} changed file(s)")
+    return files
+
+
+def is_docs_only_change(changed_files: list[str]) -> bool:
+    has_any = False
+    for file_path in changed_files:
+        if not file_path:
+            continue
+        has_any = True
+        if file_path.startswith("docs/"):
+            continue
+        if file_path.endswith(".md"):
+            continue
+        if file_path == "mkdocs.yml":
+            continue
+        return False
+    return has_any
+
+
+def resolve_skip_ci(changed_files: list[str] | None) -> bool:
+    if changed_files is None:
+        _log("skip-ci=0 (could not resolve changed files)")
+        return False
+    if is_docs_only_change(changed_files):
+        _log("docs-only change detected; skip-ci=1")
+        return True
+    _log("non-doc changes detected; skip-ci=0")
+    return False
+
+
+def render_bootstrap_pipeline(text: str, *, skip_ci: bool) -> str:
+    if DOC_SEP in text:
+        _, continuation = text.split(DOC_SEP, 1)
+    else:
+        continuation = text
+
+    nightly_only = (
+        '(build.pull_request.labels includes "nightly-test") || (build.branch == "main" && build.env("NIGHTLY") == "1")'
+    )
+    if skip_ci:
+        image_if = f"'{nightly_only}'"
+        ready_if = "'false'"
+        merge_if = "'false'"
+    else:
+        image_if = "'true'"
+        ready_if = '\'build.branch != "main" && build.pull_request.labels includes "ready"\''
+        merge_if = (
+            '\'(build.branch == "main" && build.env("NIGHTLY") != "1" '
+            '&& build.env("WEEKLY") != "1") || '
+            '(build.branch != "main" && build.pull_request.labels includes "merge-test")\''
+        )
+
+    return (
+        continuation.replace("__IMAGE_BUILD_IF__", image_if)
+        .replace("__UPLOAD_READY_IF__", ready_if)
+        .replace("__UPLOAD_MERGE_IF__", merge_if)
+    )
+
+
+def _matches_dependencies(changed_files: list[str], prefixes: list[str]) -> bool:
+    for path in changed_files:
+        for prefix in prefixes:
+            normalized = prefix.rstrip("/")
+            if path == normalized or path.startswith(f"{normalized}/"):
+                return True
+    return False
+
+
+def _strip_source_file_dependencies(step: dict[str, Any]) -> dict[str, Any]:
+    return {key: value for key, value in step.items() if key != "source_file_dependencies"}
+
+
+def _step_label(step: dict[str, Any]) -> str:
+    return str(step.get("group") or step.get("label") or "<step>")
+
+
+def _filter_steps(steps: list[Any], changed_files: list[str]) -> list[Any]:
+    filtered: list[Any] = []
+    for step in steps:
+        if not isinstance(step, dict):
+            filtered.append(step)
+            continue
+
+        deps = step.get("source_file_dependencies")
+        if deps is not None and not isinstance(deps, list):
+            raise ValueError(
+                f"source_file_dependencies must be a list in step {_step_label(step)!r}",
+            )
+        if deps is not None and not _matches_dependencies(changed_files, deps):
+            _log(
+                f"skip {_step_label(step)!r} (no changes under {deps})",
+            )
+            continue
+
+        nested = step.get("steps")
+        if nested is not None:
+            kept_nested = _filter_steps(nested, changed_files)
+            if not kept_nested:
+                _log(f"omit empty group {_step_label(step)!r}")
+                continue
+            new_step = _strip_source_file_dependencies(step)
+            new_step["steps"] = kept_nested
+            filtered.append(new_step)
+            continue
+
+        if deps is not None:
+            filtered.append(_strip_source_file_dependencies(step))
+        else:
+            filtered.append(step)
+
+    return filtered
+
+
+def _strip_uploader_metadata_from_steps(steps: list[Any]) -> list[Any]:
+    """Remove uploader-only keys while keeping all steps (no diff filtering)."""
+    stripped: list[Any] = []
+    for step in steps:
+        if not isinstance(step, dict):
+            stripped.append(step)
+            continue
+
+        deps = step.get("source_file_dependencies")
+        if deps is not None and not isinstance(deps, list):
+            raise ValueError(
+                f"source_file_dependencies must be a list in step {_step_label(step)!r}",
+            )
+
+        nested = step.get("steps")
+        new_step = _strip_source_file_dependencies(step)
+        if nested is not None:
+            new_step["steps"] = _strip_uploader_metadata_from_steps(nested)
+        stripped.append(new_step)
+
+    return stripped
+
+
+def render_test_pipeline(
+    doc: dict[str, Any],
+    changed_files: list[str] | None,
+) -> dict[str, Any]:
+    steps = doc.get("steps")
+    if not isinstance(steps, list):
+        return doc
+    if changed_files is not None:
+        steps = _filter_steps(steps, changed_files)
+    else:
+        steps = _strip_uploader_metadata_from_steps(steps)
+    return {**doc, "steps": steps}
+
+
+def resolve_pipeline_path(arg: str) -> Path:
+    path = Path(arg)
+    if path.is_absolute():
+        return path
+    return ROOT / path
+
+
+def render_pipeline(path: Path) -> str:
+    text = path.read_text(encoding="utf-8")
+    changed_files = resolve_changed_files()
+
+    if BOOTSTRAP_MARKER in text:
+        rendered = render_bootstrap_pipeline(text, skip_ci=resolve_skip_ci(changed_files))
+        return rendered
+
+    doc = yaml.safe_load(text)
+    if not isinstance(doc, dict):
+        raise ValueError(f"invalid pipeline YAML: {path}")
+
+    doc = render_test_pipeline(doc, changed_files)
+
+    return yaml.safe_dump(doc, sort_keys=False)
+
+
+def upload_to_buildkite(content: str) -> None:
+    subprocess.run(
+        ["buildkite-agent", "pipeline", "upload"],
+        input=content,
+        text=True,
+        check=True,
+    )
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "pipeline",
+        nargs="?",
+        default=".buildkite/pipeline.yml",
+        help="Pipeline YAML path (default: .buildkite/pipeline.yml)",
+    )
+    parser.add_argument(
+        "--upload",
+        action="store_true",
+        help="Pipe rendered YAML to buildkite-agent pipeline upload",
+    )
+    args = parser.parse_args()
+
+    path = resolve_pipeline_path(args.pipeline)
+    if not path.is_file():
+        _log(f"missing pipeline file: {path}")
+        return 1
+
+    rendered = render_pipeline(path)
+    if args.upload:
+        upload_to_buildkite(rendered)
+    else:
+        sys.stdout.write(rendered)
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

.buildkite/scripts/upload_pipeline_with_skip_ci.sh

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

.buildkite/test-amd-merge.yml

@@ -1,54 +1,111 @@
+env:
+  VLLM_WORKER_MULTIPROC_METHOD: spawn
+  HF_HUB_DOWNLOAD_TIMEOUT: 300
+  HF_HUB_ETAG_TIMEOUT: 60
+
 steps:
 
-- label: "Simple · Diffusion & Model Executor Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  timeout_in_minutes: 50
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - export VLLM_ROCM_USE_AITER=0
-    # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
-    # that is not supported on CPU
-    - "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
+- group: ":card_index_dividers: Simple Test"
+  steps:
+    - label: "Simple · Diffusion & Model Executor Test"
+      agent_pool: mi325_1
+      depends_on: amd-build
+      timeout_in_minutes: 50
+      mirror_hardwares: [amdproduction]
+      grade: Blocking
+      commands:
+        - export VLLM_ROCM_USE_AITER=0
+        # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
+        # that is not supported on CPU
+        - "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
 
-- label: "Simple · Other Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  timeout_in_minutes: 50
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - export VLLM_ROCM_USE_AITER=0
-    # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
-    # that is not supported on CPU
-    - "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
+    - label: "Simple · Other Test"
+      agent_pool: mi325_1
+      depends_on: amd-build
+      timeout_in_minutes: 50
+      mirror_hardwares: [amdproduction]
+      grade: Blocking
+      commands:
+        - export VLLM_ROCM_USE_AITER=0
+        # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
+        # that is not supported on CPU
+        - "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
 
-- label: "Diffusion Images API LoRA E2E"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - timeout 20m pytest -s -v tests/diffusion/lora/test_images_generations_lora.py
+- group: ":card_index_dividers: Diffusion Test"
+  steps:
+    - label: "Diffusion · LoRA Test"
+      agent_pool: mi325_1
+      depends_on: amd-build
+      mirror_hardwares: [amdproduction]
+      grade: Blocking
+      commands:
+        - timeout 20m pytest -s -v tests/diffusion/lora/test_images_generations_lora.py
 
-- label: "Diffusion Model CPU offloading Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - export VLLM_LOGGING_LEVEL=DEBUG
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - |
-      timeout 20m bash -c '
-        set +e
-        pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py
-        EXIT1=\$?
-        pytest -s -v tests/diffusion/offloader/test_diffusion_layerwise_offload.py
-        EXIT2=\$?
-        exit \$((EXIT1 | EXIT2))
-      '
+    - label: "Diffusion Model CPU offloading Test"
+      agent_pool: mi325_1
+      depends_on: amd-build
+      mirror_hardwares: [amdproduction]
+      grade: Blocking
+      commands:
+        - export VLLM_LOGGING_LEVEL=DEBUG
+        - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+        - |
+          timeout 20m bash -c '
+            set +e
+            pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py
+            EXIT1=\$?
+            pytest -s -v tests/diffusion/offloader/test_diffusion_layerwise_offload.py
+            EXIT2=\$?
+            exit \$((EXIT1 | EXIT2))
+          '
+
+    - label: "Diffusion Sequence Parallelism Test (Need 4 GPUs)"
+      agent_pool: mi325_4
+      depends_on: amd-build
+      mirror_hardwares: [amdproduction]
+      grade: Blocking
+      commands:
+        - export VLLM_LOGGING_LEVEL=DEBUG
+        - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+        - timeout 20m pytest -s -v tests/diffusion/distributed/test_sequence_parallel.py
+        - timeout 20m pytest -s -v tests/diffusion/distributed/test_ulysses_uaa_perf.py
+
+    # merge-only tests
+    - label: "Diffusion Tensor Parallelism Test"
+      agent_pool: mi325_2
+      depends_on: amd-build
+      mirror_hardwares: [amdproduction]
+      grade: Blocking
+      commands:
+        - pytest -s -v tests/diffusion/distributed/test_zimage_parallelism.py
+
+    - label: "Diffusion Model Test"
+      agent_pool: mi325_1
+      depends_on: amd-build
+      mirror_hardwares: [amdproduction]
+      grade: Blocking
+      commands:
+        - timeout 30m pytest -s -v tests/e2e/offline_inference/test_z_image.py
+
+    - label: "Diffusion GPU Worker Test"
+      agent_pool: mi325_2
+      depends_on: amd-build
+      mirror_hardwares: [amdproduction]
+      grade: Blocking
+      commands:
+        - timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py
+
+    - label: "Diffusion · Wan22 Test"
+      agent_pool: mi325_1
+      depends_on: amd-build
+      mirror_hardwares: [amdproduction]
+      grade: Blocking
+      commands:
+        - |
+          timeout 40m bash -c "
+            export VLLM_IMAGE_FETCH_TIMEOUT=60
+            pytest -s -v tests/e2e/offline_inference/test_wan22_t2v.py tests/e2e/online_serving/test_wan22_t2v.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
+          "
 
 ## ISSUE depends on `diffusers` package: https://github.com/huggingface/diffusers/issues/13274
 # - label: "Audio Generation Model Test"
@@ -62,42 +119,6 @@ steps:
 #     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
 #     - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_expansion.py -m "advanced_model and diffusion and L4" --run-level advanced_model
 
-- label: "Diffusion Sequence Parallelism Test (Need 4 GPUs)"
-  agent_pool: mi325_4
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - export VLLM_LOGGING_LEVEL=DEBUG
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - timeout 20m pytest -s -v tests/diffusion/distributed/test_sequence_parallel.py
-    - timeout 20m pytest -s -v tests/diffusion/distributed/test_ulysses_uaa_perf.py
-
-# merge-only tests
-- label: "Diffusion Tensor Parallelism Test"
-  agent_pool: mi325_2
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - pytest -s -v tests/diffusion/distributed/test_zimage_parallelism.py
-
-- label: "Diffusion Model Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - timeout 30m pytest -s -v tests/e2e/offline_inference/test_z_image.py
-
-- label: "Diffusion GPU Worker Test"
-  agent_pool: mi325_2
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py
-
 - label: "Engine Test"
   agent_pool: mi325_1
   depends_on: amd-build
@@ -111,50 +132,65 @@ steps:
 # This model seems to be important, so I am trying to run it here.
 # Since this is from test-nightly.yml,
 # I keep this non-blocking.
-- label: "Omni Model Test Qwen2-5-Omni"
-  agent_pool: mi325_4
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  commands:
-    - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
-    - pytest -s -v tests/e2e/*/test_qwen2_5_omni_expansion.py --run-level "core_model"
+- group: ":card_index_dividers: E2E Test"
+  steps:
+    - label: "Omni Model Test Qwen2-5-Omni"
+      agent_pool: mi325_4
+      depends_on: amd-build
+      mirror_hardwares: [amdproduction]
+      commands:
+        - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
+        - pytest -s -v tests/e2e/*/test_qwen2_5_omni_expansion.py --run-level "core_model"
 
-- label: "Omni Model Test Qwen3-Omni"
-  agent_pool: mi325_2
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py -m 'advanced_model and cuda' --run-level 'advanced_model'
+    - label: "Omni Model Test Qwen3-Omni"
+      agent_pool: mi325_2
+      depends_on: amd-build
+      mirror_hardwares: [amdproduction]
+      grade: Blocking
+      commands:
+        - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+        - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py -m 'advanced_model and cuda' --run-level 'advanced_model'
 
-- label: "Qwen3-TTS CustomVoice E2E Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - |
-      timeout 20m bash -c '
-        export VLLM_LOGGING_LEVEL=DEBUG
-        export VLLM_WORKER_MULTIPROC_METHOD=spawn
-        export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
-        pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py tests/e2e/offline_inference/test_qwen3_tts_customvoice.py -m "advanced_model" --run-level "advanced_model"
-      '
+    - label: "Qwen3-TTS CustomVoice E2E Test"
+      agent_pool: mi325_1
+      depends_on: amd-build
+      mirror_hardwares: [amdproduction]
+      grade: Blocking
+      commands:
+        - |
+          timeout 20m bash -c '
+            export VLLM_LOGGING_LEVEL=DEBUG
+            export VLLM_WORKER_MULTIPROC_METHOD=spawn
+            export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
+            pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py tests/e2e/offline_inference/test_qwen3_tts_customvoice.py -m "advanced_model" --run-level "advanced_model"
+          '
 
-- label: "Qwen3-TTS Base E2E Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - |
-      timeout 30m bash -c '
-        export VLLM_LOGGING_LEVEL=DEBUG
-        export VLLM_WORKER_MULTIPROC_METHOD=spawn
-        export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
-        pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py tests/e2e/offline_inference/test_qwen3_tts_base.py -m "advanced_model" --run-level "advanced_model"
-      '
+    - label: "Qwen3-TTS Base E2E Test"
+      agent_pool: mi325_1
+      depends_on: amd-build
+      mirror_hardwares: [amdproduction]
+      grade: Blocking
+      commands:
+        - |
+          timeout 30m bash -c '
+            export VLLM_LOGGING_LEVEL=DEBUG
+            export VLLM_WORKER_MULTIPROC_METHOD=spawn
+            export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
+            pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py tests/e2e/offline_inference/test_qwen3_tts_base.py -m "advanced_model" --run-level "advanced_model"
+          '
+
+    - label: "Voxtral-TTS E2E Test"
+      agent_pool: mi325_1
+      depends_on: amd-build
+      mirror_hardwares: [amdproduction]
+      grade: Blocking
+      commands:
+        - |
+          timeout 20m bash -c '
+            export VLLM_LOGGING_LEVEL=DEBUG
+            export VLLM_WORKER_MULTIPROC_METHOD=spawn
+            pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py tests/e2e/offline_inference/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
+          '
 
 # TODO: Bagel test on ROCm is very unstable. @tjtanaa
 # Need to debug before reneable numerical changes across large PRs
@@ -195,28 +231,3 @@ steps:
 #     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
 #     - export VLLM_ROCM_USE_AITER_RMSNORM=0
 #     - timeout 40m pytest -s -v tests/e2e/online_serving/test_bagel.py -m "advanced_model" --run-level "advanced_model" -k "rocm"
-
-- label: "Diffusion · Wan22 Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - |
-      timeout 40m bash -c "
-        export VLLM_IMAGE_FETCH_TIMEOUT=60
-        pytest -s -v tests/e2e/offline_inference/test_wan22_t2v.py tests/e2e/online_serving/test_wan22_t2v.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
-      "
-
-- label: "Voxtral-TTS E2E Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - |
-      timeout 20m bash -c '
-        export VLLM_LOGGING_LEVEL=DEBUG
-        export VLLM_WORKER_MULTIPROC_METHOD=spawn
-        pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py tests/e2e/offline_inference/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
-      '

.buildkite/test-amd-ready.yaml

@@ -1,42 +1,33 @@
+env:
+  VLLM_WORKER_MULTIPROC_METHOD: spawn
+  HF_HUB_DOWNLOAD_TIMEOUT: 300
+  HF_HUB_ETAG_TIMEOUT: 60
+
 steps:
 
-- label: "Simple · Diffusion & Model Executor Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - export VLLM_ROCM_USE_AITER=0
-    # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
-    # that is not supported on CPU
-    - "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
+- group: ":card_index_dividers: Simple Test"
+  steps:
+    - label: "Simple · Diffusion & Model Executor Test"
+      agent_pool: mi325_1
+      depends_on: amd-build
+      mirror_hardwares: [amdproduction]
+      grade: Blocking
+      commands:
+        - export VLLM_ROCM_USE_AITER=0
+        # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
+        # that is not supported on CPU
+        - "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
 
-- label: "Simple · Other Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - export VLLM_ROCM_USE_AITER=0
-    # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
-    # that is not supported on CPU
-    - "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
-
-- label: "Voxtral TTS CUDA Unit Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - timeout 10m pytest -s -v tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py
-
-- label: "Diffusion · Batch Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - timeout 20m pytest -sv tests/diffusion/batching -m 'core_model and cuda' --run-level "core_model"
+    - label: "Simple · Other Test"
+      agent_pool: mi325_1
+      depends_on: amd-build
+      mirror_hardwares: [amdproduction]
+      grade: Blocking
+      commands:
+        - export VLLM_ROCM_USE_AITER=0
+        # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
+        # that is not supported on CPU
+        - "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
 
 - label: "Custom Pipeline Test"
   agent_pool: mi325_1
@@ -46,23 +37,92 @@ steps:
   commands:
     - timeout 20m pytest -s -v tests/e2e/offline_inference/custom_pipeline/ -m "core_model"
 
-- label: "Diffusion Model CPU offloading Test"
+
+- label: "Voxtral TTS CUDA Unit Test"
   agent_pool: mi325_1
   depends_on: amd-build
   mirror_hardwares: [amdproduction]
   grade: Blocking
   commands:
-    - export VLLM_LOGGING_LEVEL=DEBUG
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - |
-      timeout 20m bash -c '
-        set +e
-        pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py
-        EXIT1=\$?
-        pytest -s -v tests/diffusion/offloader/test_diffusion_layerwise_offload.py
-        EXIT2=\$?
-        exit \$((EXIT1 | EXIT2))
-      '
+    - timeout 10m pytest -s -v tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py
+
+- group: ":card_index_dividers: Diffusion Test"
+  steps:
+    - label: "Diffusion · Batch Test"
+      agent_pool: mi325_1
+      depends_on: amd-build
+      mirror_hardwares: [amdproduction]
+      grade: Blocking
+      commands:
+        - timeout 20m pytest -sv tests/diffusion/batching -m 'core_model and cuda' --run-level "core_model"
+
+    - label: "Diffusion · Offloader Test"
+      agent_pool: mi325_1
+      depends_on: amd-build
+      mirror_hardwares: [amdproduction]
+      grade: Blocking
+      commands:
+        - export VLLM_LOGGING_LEVEL=DEBUG
+        - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+        - |
+          timeout 20m bash -c '
+            set +e
+            pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py
+            EXIT1=\$?
+            pytest -s -v tests/diffusion/offloader/test_diffusion_layerwise_offload.py
+            EXIT2=\$?
+            exit \$((EXIT1 | EXIT2))
+          '
+
+    - label: "Diffusion · Cache Test"
+      agent_pool: mi325_1
+      depends_on: amd-build
+      mirror_hardwares: [amdproduction]
+      grade: Blocking
+      commands:
+        - export GPU_ARCHS=gfx942
+        - export VLLM_LOGGING_LEVEL=DEBUG
+        - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+        - timeout 15m pytest -s -v -m "core_model and cache and diffusion and not distributed_cuda and L4"
+
+    - label: "Diffusion Sequence Parallelism Test"
+      agent_pool: mi325_2
+      depends_on: amd-build
+      mirror_hardwares: [amdproduction]
+      grade: Blocking
+      commands:
+        - export VLLM_LOGGING_LEVEL=DEBUG
+        - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+        - timeout 20m pytest -s -v tests/diffusion/distributed/test_sequence_parallel.py -m core_model
+
+    - label: "Diffusion GPU Worker Test"
+      agent_pool: mi325_2
+      depends_on: amd-build
+      mirror_hardwares: [amdproduction]
+      grade: Blocking
+      commands:
+        - timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py
+
+    - label: "Diffusion · Model Test"
+      agent_pool: mi325_1
+      depends_on: amd-build
+      mirror_hardwares: [amdproduction]
+      grade: Blocking
+      commands:
+        - timeout 15m pytest -sv tests/diffusion/models/ -m 'core_model and cuda' --run-level "core_model"
+
+    - label: "Diffusion · Wan22 Test"
+      agent_pool: mi325_1
+      depends_on: amd-build
+      mirror_hardwares: [amdproduction]
+      grade: Blocking
+      commands:
+        - |
+          timeout 40m bash -c "
+            export VLLM_TEST_CLEAN_GPU_MEMORY=1
+            export VLLM_IMAGE_FETCH_TIMEOUT=60
+            pytest -s -v tests/e2e/online_serving/test_wan22_t2v.py -m 'core_model' --run-level 'core_model'
+          "
 
 ## ISSUE depends on `diffusers` package: https://github.com/huggingface/diffusers/issues/13274
 # - label: "Audio Generation Model Test"
@@ -76,56 +136,6 @@ steps:
 #     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
 #     - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_expansion.py -m "advanced_model and diffusion and L4" --run-level advanced_model
 
-- label: "Diffusion Cache Backend Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - export GPU_ARCHS=gfx942
-    - export VLLM_LOGGING_LEVEL=DEBUG
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - timeout 15m pytest -s -v -m "core_model and cache and diffusion and not distributed_cuda and L4"
-
-- label: "Diffusion Sequence Parallelism Test"
-  agent_pool: mi325_2
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - export VLLM_LOGGING_LEVEL=DEBUG
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - timeout 20m pytest -s -v tests/diffusion/distributed/test_sequence_parallel.py -m core_model
-
-- label: "Diffusion GPU Worker Test"
-  agent_pool: mi325_2
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py
-
-- label: "Diffusion · Model Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - timeout 15m pytest -sv tests/diffusion/models/ -m 'core_model and cuda' --run-level "core_model"
-
-- label: "Diffusion · Wan22 Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - |
-      timeout 40m bash -c "
-        export VLLM_TEST_CLEAN_GPU_MEMORY=1
-        export VLLM_IMAGE_FETCH_TIMEOUT=60
-        pytest -s -v tests/e2e/online_serving/test_wan22_t2v.py -m 'core_model' --run-level 'core_model'
-      "
-
 - label: "Engine Test"
   agent_pool: mi325_1
   depends_on: amd-build
@@ -149,43 +159,56 @@ steps:
 #     - timeout 10m pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
 #     - timeout 20m pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "core_model" --run-level "core_model"
 
-- label: "MiMo-Audio E2E Test with H100"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - |
-      timeout 30m bash -c '
-        export VLLM_LOGGING_LEVEL=DEBUG
-        export VLLM_WORKER_MULTIPROC_METHOD=spawn
-        pytest -s -v tests/e2e/online_serving/test_mimo_audio.py -m "core_model" --run-level "core_model"
-      '
+- group: ":card_index_dividers: E2E Test"
+  steps:
+    - label: "MiMo-Audio E2E Test with H100"
+      agent_pool: mi325_1
+      depends_on: amd-build
+      mirror_hardwares: [amdproduction]
+      grade: Blocking
+      commands:
+        - |
+          timeout 30m bash -c '
+            export VLLM_LOGGING_LEVEL=DEBUG
+            export VLLM_WORKER_MULTIPROC_METHOD=spawn
+            pytest -s -v tests/e2e/online_serving/test_mimo_audio.py -m "core_model" --run-level "core_model"
+          '
 
-- label: "Qwen3-TTS E2E Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - export VLLM_LOGGING_LEVEL=DEBUG
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
-    - timeout 30m pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m "core_model" --run-level "core_model"
+    - label: "Qwen3-TTS E2E Test"
+      agent_pool: mi325_1
+      depends_on: amd-build
+      mirror_hardwares: [amdproduction]
+      grade: Blocking
+      commands:
+        - export VLLM_LOGGING_LEVEL=DEBUG
+        - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+        - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
+        - timeout 30m pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m "core_model" --run-level "core_model"
 
-- label: "Voxtral-TTS E2E Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - |
-      timeout 20m bash -c '
-        export VLLM_LOGGING_LEVEL=DEBUG
-        export VLLM_WORKER_MULTIPROC_METHOD=spawn
-        pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
-        pytest -s -v tests/e2e/offline_inference/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
-      '
+    - 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"
+          '
+
+    - label: "CosyVoice3-TTS E2E Test"
+      agent_pool: mi325_1
+      depends_on: amd-build
+      mirror_hardwares: [amdproduction]
+      grade: Blocking
+      commands:
+        - |
+          timeout 20m bash -c '
+            pytest -s -v tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py -m "full_model" --run-level "core_model"
+          '
 
 # TODO: Bagel test on ROCm is very unstable. @tjtanaa
 # Need to debug before reneable numerical changes across large PRs
@@ -225,14 +248,3 @@ steps:
 #     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
 #     - export VLLM_ROCM_USE_AITER_RMSNORM=0
 #     - timeout 40m pytest -s -v tests/e2e/online_serving/test_bagel.py -m "core_model" --run-level "core_model" -k "rocm"
-
-- label: "CosyVoice3-TTS E2E Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - |
-      timeout 20m bash -c '
-        pytest -s -v tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py -m "full_model" --run-level "core_model"
-      '

.buildkite/test-merge.yml

@@ -44,6 +44,9 @@ steps:
     depends_on: upload-merge-pipeline
     steps:
       - label: "Diffusion · LoRA Test"
+        source_file_dependencies:
+          - tests/diffusion/lora/
+          - vllm_omni/diffusion/lora/
         timeout_in_minutes: 30
         commands:
           - pytest -sv tests/diffusion/lora/ -m 'advanced_model and cuda' --run-level "advanced_model"
@@ -61,6 +64,9 @@ steps:
                 - "/fsx/hf_cache:/fsx/hf_cache"
 
   - label: "Distributed Test"
+    source_file_dependencies:
+      - tests/distributed/
+      - vllm_omni/distributed/
     timeout_in_minutes: 60
     commands:
       - |
@@ -105,6 +111,9 @@ steps:
 
   - group: ":card_index_dividers: Entrypoints Test"
     depends_on: upload-merge-pipeline
+    source_file_dependencies:
+      - tests/entrypoints/
+      - vllm_omni/entrypoints/
     steps:
       - label: "Entrypoint Test with H100"
         timeout_in_minutes: 60
@@ -165,6 +174,13 @@ steps:
     depends_on: upload-merge-pipeline
     steps:
       - label: "TTS · Qwen3-TTS CustomVoice Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_qwen3_tts_customvoice.py
+          - tests/e2e/offline_inference/test_qwen3_tts_customvoice.py
+          - vllm_omni/model_executor/models/qwen3_tts/
+          - vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
+          - vllm_omni/deploy/qwen3_tts.yaml
+          - vllm_omni/model_executor/stage_configs/qwen3_tts_uniproc.yaml
         commands:
           - |
             timeout 20m bash -c "
@@ -186,6 +202,12 @@ steps:
                 - "/fsx/hf_cache:/fsx/hf_cache"
 
       - label: "TTS · Qwen3-TTS Base Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_qwen3_tts_base.py
+          - tests/e2e/offline_inference/test_qwen3_tts_base.py
+          - vllm_omni/model_executor/models/qwen3_tts/
+          - vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
+          - vllm_omni/deploy/qwen3_tts.yaml
         commands:
           - |
             timeout 20m bash -c "
@@ -206,7 +228,34 @@ steps:
               volumes:
                 - "/fsx/hf_cache:/fsx/hf_cache"
 
+      - label: "TTS · Ming-Omni-TTS Test"
+        timeout_in_minutes: 20
+        commands:
+          - |
+            timeout 20m bash -c "
+              export VLLM_LOGGING_LEVEL=DEBUG
+              pytest -s -v tests/e2e/online_serving/test_ming_tts.py tests/e2e/offline_inference/test_ming_tts.py -m 'advanced_model and cuda' --run-level 'advanced_model'
+            "
+        agents:
+          queue: "gpu_1_queue"
+        plugins:
+          - docker#v5.2.0:
+              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+              always-pull: true
+              propagate-environment: true
+              shm-size: "8gb"
+              environment:
+                - "HF_HOME=/fsx/hf_cache"
+              volumes:
+                - "/fsx/hf_cache:/fsx/hf_cache"
+
       - label: "TTS · Voxtral-TTS Test"
+        source_file_dependencies:
+          - tests/e2e/offline_inference/test_voxtral_tts.py
+          - tests/e2e/online_serving/test_voxtral_tts.py
+          - vllm_omni/model_executor/models/voxtral_tts/
+          - vllm_omni/model_executor/stage_input_processors/voxtral_tts.py
+          - vllm_omni/deploy/voxtral_tts.yaml
         timeout_in_minutes: 20
         commands:
           - |
@@ -249,6 +298,12 @@ steps:
                       type: DirectoryOrCreate
 
       - label: "TTS · GLM-TTS Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_glm_tts.py
+          - tests/e2e/offline_inference/test_glm_tts.py
+          - vllm_omni/model_executor/models/glm_tts/
+          - vllm_omni/model_executor/stage_input_processors/glm_tts.py
+          - vllm_omni/deploy/glm_tts.yaml
         timeout_in_minutes: 20
         commands:
           - |
@@ -291,6 +346,12 @@ steps:
                       type: DirectoryOrCreate
 
       - label: "Omni · Qwen3-Omni Test"
+        source_file_dependencies:
+          - tests/e2e/offline_inference/test_qwen3_omni.py
+          - tests/e2e/online_serving/test_qwen3_omni.py
+          - vllm_omni/model_executor/models/qwen3_omni/
+          - vllm_omni/model_executor/stage_input_processors/qwen3_omni.py
+          - vllm_omni/deploy/qwen3_omni_moe.yaml
         timeout_in_minutes: 30
         commands:
           - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py -m 'advanced_model and cuda' --run-level 'advanced_model'
@@ -329,6 +390,11 @@ steps:
                       type: DirectoryOrCreate
 
       - label: "Omni · Mimo Audio Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_mimo_audio.py
+          - vllm_omni/model_executor/models/mimo_audio/
+          - vllm_omni/model_executor/stage_input_processors/mimo_audio.py
+          - vllm_omni/deploy/mimo_audio.yaml
         timeout_in_minutes: 30
         commands:
           - pytest -s -v tests/e2e/online_serving/test_mimo_audio.py -m 'advanced_model and cuda' --run-level 'advanced_model'
@@ -366,37 +432,19 @@ steps:
                       path: /mnt/hf-cache
                       type: DirectoryOrCreate
 
-      - label: "Omni · VoxCPM2 Test"
-        timeout_in_minutes: 20
-        commands:
-          - |
-            timeout 20m bash -c "
-              export VLLM_LOGGING_LEVEL=DEBUG
-              export VLLM_WORKER_MULTIPROC_METHOD=spawn
-              pytest -s -v tests/e2e/offline_inference/test_voxcpm2.py -m 'advanced_model and cuda' --run-level 'advanced_model'
-            "
-        agents:
-          queue: "gpu_1_queue"
-        plugins:
-          - docker#v5.2.0:
-              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-              always-pull: true
-              propagate-environment: true
-              shm-size: "8gb"
-              environment:
-                - "HF_HOME=/fsx/hf_cache"
-                - "HF_TOKEN"
-              volumes:
-                - "/fsx/hf_cache:/fsx/hf_cache"
-
-      - label: "TTS · VoxCPM2 Online Test"
+      - label: "TTS · VoxCPM2 Test"
+        source_file_dependencies:
+          - tests/e2e/offline_inference/test_voxcpm2_tts.py
+          - tests/e2e/online_serving/test_voxcpm2_tts.py
+          - vllm_omni/model_executor/models/voxcpm2/
+          - vllm_omni/deploy/voxcpm2.yaml
         timeout_in_minutes: 30
         commands:
           - |
             timeout 30m bash -c "
               export VLLM_LOGGING_LEVEL=DEBUG
               export VLLM_WORKER_MULTIPROC_METHOD=spawn
-              pytest -s -v tests/e2e/online_serving/test_voxcpm2_tts.py -m 'advanced_model and cuda' --run-level 'advanced_model'
+              pytest -s -v tests/e2e/offline_inference/test_voxcpm2_tts.py tests/e2e/online_serving/test_voxcpm2_tts.py -m 'advanced_model and cuda' --run-level 'advanced_model'
             "
         agents:
           queue: "gpu_1_queue"
@@ -413,6 +461,12 @@ steps:
                 - "/fsx/hf_cache:/fsx/hf_cache"
 
       - label: "Omni · Higgs-Audio V2 Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_higgs_audio_v2.py
+          - tests/e2e/offline_inference/test_higgs_audio_v2.py
+          - vllm_omni/model_executor/models/higgs_audio_v2/
+          - vllm_omni/model_executor/stage_input_processors/higgs_audio_v2.py
+          - vllm_omni/deploy/higgs_audio_v2.yaml
         timeout_in_minutes: 30
         commands:
           - |
@@ -438,6 +492,9 @@ steps:
                 - "/fsx/hf_cache:/fsx/hf_cache"
 
       - label: "Diffusion · Qwen Image Edit Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_qwen_image_edit.py
+          - vllm_omni/diffusion/models/qwen_image/
         timeout_in_minutes: 20
         commands:
           - pytest -s -v tests/e2e/online_serving/test_qwen_image_edit.py -m 'advanced_model and cuda' --run-level 'advanced_model'
@@ -476,6 +533,9 @@ steps:
                       type: DirectoryOrCreate
 
       - label: "Diffusion · Longcat Image Edit Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_longcat_image_edit.py
+          - vllm_omni/diffusion/models/longcat_image/
         timeout_in_minutes: 20
         commands:
           - pytest -s -v tests/e2e/online_serving/test_longcat_image_edit.py -m 'advanced_model and cuda' --run-level 'advanced_model'
@@ -514,6 +574,9 @@ steps:
                       type: DirectoryOrCreate
 
       - label: "Diffusion · Longcat Image Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_longcat_image.py
+          - vllm_omni/diffusion/models/longcat_image/
         timeout_in_minutes: 20
         commands:
           - pytest -s -v tests/e2e/online_serving/test_longcat_image.py -m 'advanced_model and cuda' --run-level 'advanced_model'
@@ -552,6 +615,13 @@ steps:
                       type: DirectoryOrCreate
 
       - label: "Diffusion · Bagel Test"
+        source_file_dependencies:
+          - tests/e2e/offline_inference/test_bagel.py
+          - tests/e2e/online_serving/test_bagel.py
+          - vllm_omni/model_executor/models/bagel/
+          - vllm_omni/diffusion/models/bagel/
+          - vllm_omni/model_executor/stage_input_processors/bagel.py
+          - vllm_omni/deploy/bagel.yaml
         timeout_in_minutes: 20
         commands:
           - pytest -s -v tests/e2e/offline_inference/test_bagel.py tests/e2e/online_serving/test_bagel.py -m 'advanced_model and cuda' --run-level 'advanced_model'
@@ -591,6 +661,9 @@ steps:
 
 
       - label: "Diffusion · Qwen Image Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_qwen_image.py
+          - vllm_omni/diffusion/models/qwen_image/
         commands:
           - |
             timeout 40m bash -c '
@@ -632,6 +705,9 @@ steps:
                       type: DirectoryOrCreate
 
       - label: "Diffusion · Qwen Image Layered Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_qwen_image_layered.py
+          - vllm_omni/diffusion/models/qwen_image/
         commands:
           - |
             timeout 40m bash -c "
@@ -673,6 +749,10 @@ steps:
                       type: DirectoryOrCreate
 
       - label: "Diffusion · Wan22 Test"
+        source_file_dependencies:
+          - tests/e2e/offline_inference/test_wan22_t2v.py
+          - tests/e2e/online_serving/test_wan22_t2v.py
+          - vllm_omni/diffusion/models/wan2_2/
         commands:
           - |
             timeout 40m bash -c "
@@ -714,6 +794,9 @@ steps:
                       type: DirectoryOrCreate
 
       - label: "Diffusion · HunyuanVideo 1.5 Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_hunyuan_video_15.py
+          - vllm_omni/diffusion/models/hunyuan_video/
         timeout_in_minutes: 40
         commands:
           - |
@@ -755,6 +838,10 @@ steps:
                       type: DirectoryOrCreate
 
       - label: "Diffusion · Z Image Test"
+        source_file_dependencies:
+          - tests/e2e/offline_inference/test_z_image.py
+          - tests/e2e/online_serving/test_z_image.py
+          - vllm_omni/diffusion/models/z_image/
         timeout_in_minutes: 30
         commands:
           - pytest -s -v tests/e2e/offline_inference/test_z_image.py tests/e2e/online_serving/test_z_image.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
@@ -772,6 +859,10 @@ steps:
                 - "/fsx/hf_cache:/fsx/hf_cache"
 
       - label: "Diffusion · FLUX.2 Klein Test"
+        source_file_dependencies:
+          - tests/e2e/offline_inference/test_flux2_klein.py
+          - tests/e2e/online_serving/test_flux2_klein.py
+          - vllm_omni/diffusion/models/flux2_klein/
         timeout_in_minutes: 30
         commands:
           - pytest -s -v tests/e2e/offline_inference/test_flux2_klein.py tests/e2e/online_serving/test_flux2_klein.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
@@ -789,6 +880,9 @@ steps:
                 - "/fsx/hf_cache:/fsx/hf_cache"
 
       - label: "Diffusion · SD3 Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_sd3.py
+          - vllm_omni/diffusion/models/sd3/
         timeout_in_minutes: 45
         commands:
           - |

.buildkite/test-nightly.yml

@@ -591,6 +591,7 @@ steps:
             tests/e2e/online_serving/test_longcat_image_edit_expansion.py
             tests/e2e/online_serving/test_flux_2_dev_expansion.py
             tests/e2e/online_serving/test_bagel_expansion.py
+            tests/e2e/offline_inference/test_sensenova_u1_expansion.py
             -m "full_model and diffusion and H100 and not distributed_cuda"
             --run-level "full_model"
         agents:

.buildkite/test-ready.yml

@@ -309,37 +309,14 @@ steps:
                       path: /mnt/hf-cache
                       type: DirectoryOrCreate
 
-      - label: "Omni · VoxCPM2 Test"
-        timeout_in_minutes: 20
-        commands:
-          - |
-            timeout 20m bash -c "
-              export VLLM_LOGGING_LEVEL=DEBUG
-              export VLLM_WORKER_MULTIPROC_METHOD=spawn
-              pytest -s -v tests/e2e/offline_inference/test_voxcpm2.py -m 'core_model' --run-level 'core_model'
-            "
-        agents:
-          queue: "gpu_1_queue"
-        plugins:
-          - docker#v5.2.0:
-              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-              always-pull: true
-              propagate-environment: true
-              shm-size: "8gb"
-              environment:
-                - "HF_HOME=/fsx/hf_cache"
-                - "HF_TOKEN"
-              volumes:
-                - "/fsx/hf_cache:/fsx/hf_cache"
-
-      - label: "TTS · VoxCPM2 Online Test"
+      - label: "TTS · VoxCPM2 Test"
         timeout_in_minutes: 30
         commands:
           - |
             timeout 30m bash -c "
               export VLLM_LOGGING_LEVEL=DEBUG
               export VLLM_WORKER_MULTIPROC_METHOD=spawn
-              pytest -s -v tests/e2e/online_serving/test_voxcpm2_tts.py -m 'core_model and tts' --run-level 'core_model'
+              pytest -s -v tests/e2e/offline_inference/test_voxcpm2_tts.py tests/e2e/online_serving/test_voxcpm2_tts.py -m 'core_model' --run-level 'core_model'
             "
         agents:
           queue: "gpu_1_queue"

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

@@ -6,6 +6,13 @@
 {% set docker_image_amd = "rocm/vllm-omni:$BUILDKITE_COMMIT" %}
 {% set default_working_dir = "/app/vllm-omni" %}
 
+{% if env %}
+env:
+{% for name in env %}
+  {{ name }}: "{{ env[name] }}"
+{% endfor %}
+{% endif %}
+steps:
   - group: "AMD Tests"
     depends_on: ~
     steps:
@@ -31,8 +38,15 @@
         agents:
           queue: amd-cpu
 
-    {% for step in steps %}
-    {% if step.mirror_hardwares and mirror_hw in step.mirror_hardwares %}
+{% for step_group in steps %}
+{% if step_group.group %}
+  - group: "{{ step_group.group }}"
+{% if step_group.depends_on %}
+    depends_on: {{ step_group.depends_on }}
+{% endif %}
+    steps:
+{% for step in step_group.steps %}
+{% if step.mirror_hardwares and mirror_hw in step.mirror_hardwares %}
       - label: "{{ step.agent_pool }}: {{ step.label }}"
         depends_on: amd-build
         agents:
@@ -46,8 +60,23 @@
         command: bash .buildkite/scripts/hardware_ci/run-amd-test.sh
         env:
           DOCKER_BUILDKIT: "1"
+{% if step.env %}
+{% for name in step.env %}
+          {{ name }}: "{{ step.env[name] }}"
+{% endfor %}
+{% endif %}
           TEST_COMMAND: |-
             (command rocm-smi || true) && cd {{ (step.working_dir or default_working_dir) | safe }}
+{% if env %}
+{% for name in env %}
+            export {{ name }}="{{ env[name] }}"
+{% endfor %}
+{% endif %}
+{% if step.env %}
+{% for name in step.env %}
+            export {{ name }}="{{ step.env[name] }}"
+{% endfor %}
+{% endif %}
 {% if "mi250" in step.agent_pool %}
             python3 -m pip uninstall -y amd-aiter
 {% endif %}
@@ -61,5 +90,52 @@
         {% else %}
         soft_fail: true
         {% endif%}
+{% endif %}
+{% endfor %}
+{% elif step_group.mirror_hardwares and mirror_hw in step_group.mirror_hardwares %}
+{% set step = step_group %}
+  - label: "{{ step.agent_pool }}: {{ step.label }}"
+    depends_on: amd-build
+    agents:
+      {% if step.agent_pool %}
+      queue: amd_{{ step.agent_pool }}
+      {% else %}
+      queue: amd_mi325_1
+      {% endif %}
+{% set cmd_body = (step.command or (step.commands | join("\n"))) | trim %}
+{% set indented_cmd = cmd_body | replace("\n", "\n        ") %}
+    command: bash .buildkite/scripts/hardware_ci/run-amd-test.sh
+    env:
+      DOCKER_BUILDKIT: "1"
+{% if step.env %}
+{% for name in step.env %}
+      {{ name }}: "{{ step.env[name] }}"
+{% endfor %}
+{% endif %}
+      TEST_COMMAND: |-
+        (command rocm-smi || true) && cd {{ (step.working_dir or default_working_dir) | safe }}
+{% if env %}
+{% for name in env %}
+        export {{ name }}="{{ env[name] }}"
+{% endfor %}
+{% endif %}
+{% if step.env %}
+{% for name in step.env %}
+        export {{ name }}="{{ step.env[name] }}"
+{% endfor %}
+{% endif %}
+{% if "mi250" in step.agent_pool %}
+        python3 -m pip uninstall -y amd-aiter
+{% endif %}
+        {{ indented_cmd | safe }}
+    priority: 100
+    {% if step.timeout_in_minutes %}
+    timeout_in_minutes: {{ step.timeout_in_minutes }}
     {% endif %}
-    {% endfor %}
+    {% if step.grade and step.grade == "Blocking" %}
+    soft_fail: false
+    {% else %}
+    soft_fail: true
+    {% endif%}
+{% endif %}
+{% endfor %}

.github/PULL_REQUEST_TEMPLATE.md

@@ -1,24 +1,15 @@
 <!-- markdownlint-disable -->
-**Before submitting:** run the precheck-pr skill with code agent for a self-check against project conventions — catches dead code, missing benchmarks, and title format issues.
-
-PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTTOM) HAVE BEEN CONSIDERED.
+PLEASE FILL IN THE PR DESCRIPTION HERE.
 
 ## Purpose
 
 ## Test Plan
 
+**vLLM Version:**
+
+**vLLM-Omni Commit:**
+
 ## Test Result
 
----
-<details>
-<summary> Essential Elements of an Effective PR Description Checklist </summary>
-
-- [ ] The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
-- [ ] The test plan. Please provide the test scripts & test commands. Please state the reasons if your codes don't require additional test scripts. For test file guidelines, please check the [test style doc](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/tests_style/)
-- [ ] The test results. Please paste the results comparison before and after, or the e2e results.
-- [ ] (Optional) The necessary documentation update, such as updating `supported_models.md` and `examples` for a new model. **Please run `mkdocs serve` to sync the documentation editions to `./docs`.**
-- [ ] (Optional) Release notes update. If your change is user-facing, please update the release notes draft.
-</details>
-
-**BEFORE SUBMITTING,** please read <https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md>
+**BEFORE SUBMITTING:** read [CONTRIBUTING.md](https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md) and run the [precheck-pr skill](https://github.com/vllm-project/vllm-omni/blob/main/.claude/skills/precheck-pr/SKILL.md) with the code agent for a self-check against project conventions.
 (anything written below this line will be removed by GitHub Actions)

CONTRIBUTING.md

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

README.md

@@ -16,7 +16,7 @@ Easy, fast, and cheap omni-modality model serving for everyone
 ---
 
 *Latest News* 🔥
-- [2026/06] We released [0.22.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.22.0) - an **omnimodal world-model** release aligned with vLLM 0.22, featuring [Nvidia Cosmos3](recipes/cosmos3/Cosmos3-Nano.md)/DreamZero world model support, expanded quantization coverage across Blackwell/NPU/XPU, TTS production improvements, new models including MiniCPM-o 4.5, MOSS-TTS, and Lance, plus RL integration with veRL-Omni.
+- [2026/06] We released [0.22.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.22.0) - an **omnimodal world-model** release aligned with vLLM 0.22, featuring [Nvidia Cosmos3](recipes/cosmos3/Cosmos3-Nano.md)/DreamZero world model support, expanded quantization coverage across Blackwell/NPU/XPU, TTS production improvements, new models including MiniCPM-o 4.5, MOSS-TTS, and Lance, plus RL integration with [VeRL-Omni](https://github.com/verl-project/verl-omni).
 - [2026/05] We released [0.20.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.20.0) - refreshes the serving/runtime stack for large-scale omni workloads, and improves diffusion model performance, quantization, and hardware readiness across CUDA, ROCm, MUSA, NPU, and XPU backends.
 - [2026/03] We released [0.18.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.18.0) - strengthens the core runtime through a large entrypoint refactor and scheduler/runtime cleanups, expands unified quantization and diffusion execution, broadens multimodal model coverage, and improves production readiness across audio, omni, image, video, RL, and multi-platform deployments.
 - [2026/03] Check out our first public [project deepdive](https://youtu.be/sgwNfsNnR9I) at the vLLM Hong Kong Meetup!
@@ -57,8 +57,9 @@ vLLM-Omni is flexible and easy to use with:
 
 vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, including:
 
-- Omni-modality models (e.g. Qwen-Omni)
-- Multi-modality generation models (e.g. Qwen-Image)
+- **Omni-modality models** (e.g. Qwen3-Omni, Cosmos, HunyuanImage, BAGEL)
+- **TTS models** (e.g. Qwen3-TTS, VoxCPM2, Ming-Omni-TTS, CosyVoice3)
+- **Diffusion models** — image, video, and audio generation (e.g. Qwen-Image, Wan2.2, FLUX)
 
 ## Getting Started
 

docs/contributing/README.md

@@ -118,17 +118,13 @@ Only specific types of PRs will be reviewed. The PR title is prefixed appropriat
 
 ### Pre-Check Before Submitting
 
-Before submitting a PR, run the precheck-pr skill with code agent for a self-review against project conventions:
-
-```
-/precheck-pr
-```
+Before submitting a PR, run the [precheck-pr skill](https://github.com/vllm-project/vllm-omni/blob/main/.claude/skills/precheck-pr/SKILL.md) with the code agent for a self-review against project conventions:
 
 The skill offers two modes:
 - **Quick (~3 min):** catches showstoppers — PR title format, missing benchmark claims, rebase status
 - **Full (~10 min):** thorough maintainer-grade review — dead code scan, copy-paste detection, import hygiene
 
-The precheck covers five PR types: Bug Fix, Performance, New Model, Diffusion Model, and General. Each type has a tailored checklist that validates evidence quality (repro steps, A/B benchmarks, registry entries, etc.). See the skill in `.claude/skills/precheck-pr/` for the full checklist.
+The precheck covers five PR types: Bug Fix, Performance, New Model, Diffusion Model, and General. Each type has a tailored checklist that validates evidence quality (repro steps, A/B benchmarks, registry entries, etc.). See the [precheck-pr skill](https://github.com/vllm-project/vllm-omni/blob/main/.claude/skills/precheck-pr/SKILL.md) for the full checklist.
 
 ### Local Test
 Please run the L1 and L2 test cases locally first and attach the results before contacting us to add the "ready" label. Please refer to the [test instructions](./ci/test_guide.md) for running the test cases.

docs/design/feature/disaggregated_inference.md

@@ -8,6 +8,7 @@ Backend-specific setup lives in separate docs:
 - [SharedMemoryConnector](omni_connectors/shared_memory_connector.md)
 - [MooncakeStoreConnector](omni_connectors/mooncake_store_connector.md)
 - [MooncakeTransferEngineConnector](omni_connectors/mooncake_transfer_engine_connector.md)
+- [MoriTransferEngineConnector](omni_connectors/mori_transfer_engine_connector.md)
 - [YuanrongConnector](omni_connectors/yuanrong_connector.md)
 - [YuanrongTransferEngineConnector](omni_connectors/yuanrong_transfer_engine_connector.md)
 
@@ -23,6 +24,7 @@ Current connectors operate in D2H2D (device to host to device) mode.
 | Single node | SharedMemoryConnector | Auto-configured if no connector is specified. |
 | Multi node (Mooncake Store) | MooncakeStoreConnector | TCP-based, requires Mooncake Master + metadata server. |
 | Multi node (Mooncake RDMA) | MooncakeTransferEngineConnector | RDMA/TCP direct transfer with managed memory pool. Fastest. |
+| Multi node (Mori RDMA) | MoriTransferEngineConnector | RDMA direct transfer via Mori IOEngine. |
 | Multi node (Yuanrong) | YuanrongConnector | Requires Yuanrong Datasystem + etcd. |
 | Ascend NPU P2P (Yuanrong TE) | YuanrongTransferEngineConnector | Uses Yuanrong TransferEngine directly. Configure NPU device IPv4 and `memory_pool_device: "npu"`. |
 

docs/design/feature/omni_connectors/mori_transfer_engine_connector.md

@@ -0,0 +1,76 @@
+# MoriTransferEngineConnector
+
+## When to Use
+
+Currently supports intra-node deployment with Mori.
+
+As noted in #1742, inter-node support will be added back in a future
+refactor.
+
+## Mechanism
+
+Uses Mori's `IOEngine` / `MemoryDesc` API for zero-copy RDMA transfers.
+
+- Data Plane: RDMA (InfiniBand/RoCE) with managed memory pool.
+- Control Plane: ZMQ for pull-request handshake and async completion.
+
+## Installation
+
+See the [Mori repository](https://github.com/ROCm/mori) for installation instructions.
+
+## Configuration
+
+Mori is configured through the new deploy-config schema (see
+[`docs/configuration/stage_configs.md`](../../../configuration/stage_configs.md)).
+Define the connector at the top level of the deploy YAML and reference it
+by name from each stage's `input_connectors` / `output_connectors`:
+
+```yaml
+connectors:
+  mori_connector:
+    name: MoriTransferEngineConnector
+    extra:
+      host: "auto"
+      zmq_port: 50051
+      device_name: ""
+      memory_pool_size: 536870912
+      memory_pool_device: "cuda"
+
+stages:
+  - stage_id: 0
+    output_connectors:
+      to_stage_1: mori_connector
+
+  - stage_id: 1
+    input_connectors:
+      from_stage_0: mori_connector
+```
+
+A ready-to-run intra-node example for Qwen3-Omni-MoE on AMD MI300X lives
+at `vllm_omni/deploy/qwen3_omni_moe_mori_intranode.yaml` and can be loaded
+with:
+
+```bash
+vllm-omni serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --log-stats \
+    --deploy-config vllm_omni/deploy/qwen3_omni_moe_mori_intranode.yaml
+```
+
+The yaml wires `MoriTransferEngineConnector` (with `backend_type: xgmi`) to
+the chunk_transfer_adapter path (`async_chunk: true`) so stage-to-stage
+hidden-state and codec-frame streams ship GPU-to-GPU over AMD Infinity
+Fabric instead of SHM.  Qwen2.5-Omni + Mori is not yet functional on
+the chunk path: it needs `thinker2talker_async_chunk` /
+`talker2code2wav_async_chunk` input processors that do not exist yet
+(the orchestrator-level path the upstream PR originally targeted was
+lost during the entrypoints → engine refactor and removed in #1742).
+
+Parameters:
+
+- host: local RDMA IP (`"auto"` for auto-detect).
+- zmq_port: ZMQ base port for control-plane communication.
+- device_name: RDMA device (e.g., `"mlx5_0"`), empty for auto-detect.
+- memory_pool_size: RDMA memory pool size in bytes.
+- memory_pool_device: `"cpu"` (pinned) or `"cuda"` (GPUDirect / XGMI RDMA).
+
+For more details, refer to the
+[Mori repository](https://github.com/ROCm/mori).

docs/models/supported_models.md

@@ -63,6 +63,7 @@ th {
 | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-0.6B-Base` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
+| `MingTTSForConditionalGeneration` | Ming-omni-tts dense 0.5B | `inclusionAI/Ming-omni-tts-0.5B` | ✅︎ | | | |
 | `GLMTTSForConditionalGeneration` | GLM-TTS | `zai-org/GLM-TTS` | ✅︎ | | | |
 | `MossTTSNanoForCausalLM` | MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | ✅︎ | | | |
 | `MossTTSDelayModel` | MOSS-TTS, MOSS-TTSD, MOSS-SoundEffect, MOSS-VoiceGenerator | `OpenMOSS-Team/MOSS-VoiceGenerator` | ✅︎ | | | |

docs/user_guide/diffusion_features.md

@@ -148,7 +148,7 @@ The following tables show which models support each feature:
 | **Wan2.2-S2V**               |     ❌     |     ✅      |           ❌           |       ❌        |         ✅         |         ❌         |   ❌    |             ✅             |  ✅ (encode/decode)   |       ❌        |        ❌         |
 | **Wan2.1-VACE**              |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
 | **LTX-2**                    |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |
-| **LTX-2.3**                  |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
+| **LTX-2.3**                  |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ❌    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
 | **Helios**                   |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **HunyuanVideo-1.5 T2V I2V** |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |  ✅ (encode/decode)   |       ✅        |        ❌         |
 | **DreamID-Omni**             |     ❌     |     ❌      |           ❌           |       ✅        |         ❌         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |

examples/offline_inference/image_to_image/image_edit.py

@@ -143,6 +143,20 @@ def parse_args() -> argparse.Namespace:
         default=None,
         required=False,
     )
+    parser.add_argument(
+        "--width",
+        type=int,
+        default=None,
+        metavar="W",
+        help="Output image width in pixels. Default: None (pipeline's default).",
+    )
+    parser.add_argument(
+        "--height",
+        type=int,
+        default=None,
+        metavar="H",
+        help="Output image height in pixels. Default: None (pipeline's default).",
+    )
     parser.add_argument(
         "--seed",
         type=int,
@@ -236,8 +250,8 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--resolution",
         type=int,
-        default=640,
-        help="Bucket in (640, 1024) to determine the condition and output resolution",
+        default=None,
+        help="Bucket in (640, 1024) to determine the condition and output resolution. If width and height are not provided, this will be set to default 640.",
     )
 
     parser.add_argument(
@@ -384,6 +398,15 @@ def parse_args() -> argparse.Namespace:
 def main():
     args = parse_args()
 
+    if args.resolution and (args.width or args.height):
+        raise ValueError("--resolution and --width/--height cannot be specified together")
+    if args.width is not None and args.width <= 0:
+        raise ValueError("--width must be a positive integer")
+    if args.height is not None and args.height <= 0:
+        raise ValueError("--height must be a positive integer")
+    if not args.width and not args.height and not args.resolution:
+        args.resolution = 640
+
     # Validate input images exist and load them
     input_images = []
     for image_path in args.image:
@@ -464,6 +487,7 @@ def main():
     print(
         f"  Parallel configuration: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}, tensor_parallel_size={args.tensor_parallel_size}, enable_expert_parallel: {args.enable_expert_parallel}"
     )
+    print(f"  Output image size: {args.width or 'default'}x{args.height or 'default'}")
     print(f"{'=' * 60}\n")
 
     generation_start = time.perf_counter()
@@ -473,22 +497,26 @@ def main():
         omni.start_profile()
 
     # Generate edited image
+    sampling_params = OmniDiffusionSamplingParams(
+        width=args.width,
+        height=args.height,
+        generator=generator,
+        true_cfg_scale=args.cfg_scale,
+        guidance_scale=args.guidance_scale,
+        guidance_scale_2=args.guidance_scale_2,
+        num_inference_steps=args.num_inference_steps,
+        num_outputs_per_prompt=args.num_outputs_per_prompt,
+        layers=args.layers,
+    )
+    if args.resolution:
+        sampling_params.resolution = args.resolution
     outputs = omni.generate(
         {
             "prompt": args.prompt,
             "negative_prompt": args.negative_prompt,
             "multi_modal_data": {"image": input_image},
         },
-        OmniDiffusionSamplingParams(
-            generator=generator,
-            true_cfg_scale=args.cfg_scale,
-            guidance_scale=args.guidance_scale,
-            guidance_scale_2=args.guidance_scale_2,
-            num_inference_steps=args.num_inference_steps,
-            num_outputs_per_prompt=args.num_outputs_per_prompt,
-            layers=args.layers,
-            resolution=args.resolution,
-        ),
+        sampling_params,
     )
     generation_end = time.perf_counter()
     generation_time = generation_end - generation_start

examples/offline_inference/text_to_speech/README.md

@@ -17,6 +17,7 @@ list of supported architectures across all modalities, see
 | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | 2 (talker + code2wav) | ✓ | ✓ | — | 24 kHz |
 | Fish Speech S2 Pro | `fishaudio/s2-pro` | dual-AR | ✓ | ✓ | — | 44.1 kHz |
 | GLM-TTS | `zai-org/GLM-TTS` | 2 (AR + DiT) | ✓ (required) | ✓ | — | 24 kHz |
+| Ming-omni-tts | `inclusionAI/Ming-omni-tts-0.5B` | 2 (AR + audio VAE) | ✓ | ✓ | style / IP / dialect / TTA / podcast | 44.1 kHz |
 | Ming-flash-omni-TTS | `Jonathan1909/Ming-flash-omni-2.0` | single (talker only) | — (caption-controlled) | — | style / IP / basic captions | 44.1 kHz |
 | MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | single (AR + codec) | ✓ (required) | ✓ | voice_clone, continuation | 48 kHz |
 | OmniVoice | `k2-fsa/OmniVoice` | 2 (gen + dec) | ✓ | — | voice design, language hint | 24 kHz |
@@ -159,6 +160,46 @@ Streaming requires `async_chunk: true` in the stage config.
 
 ---
 
+## Ming-omni-tts
+
+Dense 0.5B two-stage TTS pipeline (`AR + flow` + audio VAE) at 44.1 kHz. The example covers style, IP voice, music-only generation, text-to-audio events, emotion, dialect, zero-shot cloning, podcast, speech+BGM, and speech+environment-sound cases.
+
+### Quick start
+```bash
+python examples/offline_inference/text_to_speech/ming_tts/end2end.py \
+    --case style \
+    --deploy-config vllm_omni/deploy/ming_tts.yaml \
+    --enforce-eager
+```
+
+### Voice cloning
+```bash
+python examples/offline_inference/text_to_speech/ming_tts/end2end.py \
+    --case zero_shot \
+    --ref-audio /path/to/reference.wav \
+    --ref-text "在此奉劝大家别乱打美白针。" \
+    --deploy-config vllm_omni/deploy/ming_tts.yaml \
+    --enforce-eager
+```
+
+### Streaming
+```bash
+python examples/offline_inference/text_to_speech/ming_tts/end2end.py \
+    --case basic \
+    --ref-audio /path/to/reference.wav \
+    --streaming \
+    --deploy-config vllm_omni/deploy/ming_tts.yaml \
+    --enforce-eager
+```
+
+### Notes
+- `style`, `ip`, `bgm`, and `tta` do not require reference audio.
+- Reference-audio cases use `--ref-audio`; `zero_shot` also requires `--ref-text`.
+- `podcast` uses multiple references via `--ref-audio-paths`.
+- Full case details live in [`ming_tts/README.md`](ming_tts/README.md).
+
+---
+
 ## Ming-flash-omni-TTS
 
 Standalone talker-only deployment of Ming-flash-omni-2.0 at 44.1 kHz. Voice is controlled through caption fields (`风格` / `IP` / `语速`/`基频`/`音量`) rather than reference audio.

examples/offline_inference/text_to_speech/ming_tts/README.md

@@ -0,0 +1,248 @@
+# Ming-omni-tts Offline Inference
+
+`end2end.py` runs Ming dense 0.5B end to end with vLLM-Omni. It uses the in-repo Ming prompt assembly helper directly, so the example request shape matches the real integration instead of a simplified wrapper.
+
+## Files
+
+| File | Purpose |
+|---|---|
+| `end2end.py` | Driver: CLI, case loading, prompt construction, orchestration (~150 lines) |
+| `cases.yaml` | All 11 built-in case definitions (prompt, text, instruction, ref-audio flags, flow controls) |
+| `runner.py` | Engine management and audio output (streaming + blocking paths) |
+
+## Model Overview
+
+Ming dense 0.5B is exposed here as a two-stage offline pipeline:
+
+- **Stage 0**: Qwen2-based AR generation with Ming prompt formatting and inline flow controls
+- **Stage 1**: audio VAE decode to mono 44.1 kHz waveform
+
+`config_ming_tts.py` adapts the checkpoint's HuggingFace config fields
+(LLM, DiT, aggregator, AudioVAE, token ids). `vllm_omni/deploy/ming_tts.yaml`
+selects the vLLM-Omni pipeline and stage runtime topology, including
+connectors, async chunking, memory limits, and sampling defaults.
+
+The example supports both:
+
+- **Blocking eager** via `vllm_omni/deploy/ming_tts.yaml`
+- **Async chunk eager** via `vllm_omni/deploy/ming_tts.yaml` (default `async_chunk: true`)
+
+## Setup
+
+Install vLLM-Omni with the platform requirements for your accelerator:
+
+```bash
+uv pip install -e .
+```
+
+The Ming offline example does not require a separate upstream Ming package.
+Reference-audio cases use the repo dependencies for audio loading,
+resampling, and CampPlus speaker extraction, including `soundfile`,
+`torchaudio`, and `onnxruntime-rocm` in the documented ROCm environment.
+
+The tested ROCm environment is summarized in the
+[Ming recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/inclusionAI/Ming-omni-tts-0.5B.md).
+
+## Supported Cases
+
+These cases cover the upstream dense 0.5B cookbook surface that maps cleanly onto the current vLLM-Omni example:
+
+- `style`: zero-speaker style-conditioned speech
+- `ip`: zero-speaker IP voice generation
+- `bgm`: music-only generation
+- `tta`: text-to-audio event generation with FlowLoss controls
+- `emotion`: reference-audio speech with emotion control
+- `basic`: reference-audio speech with speed / pitch / volume control
+- `dialect`: reference-audio speech with dialect control
+- `zero_shot`: reference-audio cloning with explicit transcript
+- `podcast`: multi-reference dialogue generation with automatic speaker embedding extraction
+- `speech_bgm`: speech with background music conditioning
+- `speech_sound`: speech with environmental sound conditioning
+
+## Quick Start
+
+Run the zero-speaker style example:
+
+```bash
+python examples/offline_inference/text_to_speech/ming_tts/end2end.py \
+    --case style \
+    --deploy-config vllm_omni/deploy/ming_tts.yaml \
+    --enforce-eager
+```
+
+Run zero-shot cloning with a transcript:
+
+```bash
+python examples/offline_inference/text_to_speech/ming_tts/end2end.py \
+    --case zero_shot \
+    --ref-audio /path/to/10002287-00000094.wav \
+    --ref-text "在此奉劝大家别乱打美白针。" \
+    --deploy-config vllm_omni/deploy/ming_tts.yaml \
+    --enforce-eager
+```
+
+Run emotion-controlled speech:
+
+```bash
+python examples/offline_inference/text_to_speech/ming_tts/end2end.py \
+    --case emotion \
+    --ref-audio /path/to/emotion_prompt.wav \
+    --deploy-config vllm_omni/deploy/ming_tts.yaml \
+    --enforce-eager
+```
+
+Run podcast generation with two reference clips:
+
+```bash
+python examples/offline_inference/text_to_speech/ming_tts/end2end.py \
+    --case podcast \
+    --ref-audio-paths /path/to/CTS-CN-F2F-2019-11-11-423-012-A.wav /path/to/CTS-CN-F2F-2019-11-11-423-012-B.wav \
+    --deploy-config vllm_omni/deploy/ming_tts.yaml \
+    --enforce-eager
+```
+
+The script automatically extracts one 192-d speaker embedding per reference WAV using the Ming model's `campplus.onnx`.
+
+If you already have precomputed multi-speaker embeddings, you can override extraction with:
+
+```bash
+--speaker-embedding /path/to/podcast_speaker_embeddings.json
+```
+
+where the JSON is a list of speaker embeddings, one 192-d vector per speaker.
+
+Run text-to-audio event generation:
+
+```bash
+python examples/offline_inference/text_to_speech/ming_tts/end2end.py \
+    --case tta \
+    --deploy-config vllm_omni/deploy/ming_tts.yaml \
+    --enforce-eager
+```
+
+Use async_chunk streaming:
+
+```bash
+python examples/offline_inference/text_to_speech/ming_tts/end2end.py \
+    --case basic \
+    --ref-audio /path/to/10002287-00000095.wav \
+    --streaming \
+    --deploy-config vllm_omni/deploy/ming_tts.yaml \
+    --enforce-eager
+```
+
+`--streaming` uses `AsyncOmni` and the async_chunk stage config. It currently
+supports one prompt per process invocation; use blocking mode for
+`--num-prompts > 1`.
+
+Collect runtime stats and a manifest:
+
+```bash
+python examples/offline_inference/text_to_speech/ming_tts/end2end.py \
+    --case style \
+    --deploy-config vllm_omni/deploy/ming_tts.yaml \
+    --enforce-eager \
+    --enable-stats \
+    --stats-log-file output_audio/ming_style_pipeline.log \
+    --metadata-json output_audio/ming_style_manifest.json
+```
+
+## Reference Fixtures
+
+The upstream Ming cookbook uses these public audio fixtures from `inclusionAI/Ming-omni-tts/data/wavs`:
+
+- `10002287-00000094.wav` for zero-shot cloning
+- `10002287-00000095.wav` for `basic`
+- `emotion_prompt.wav` for `emotion`
+- `yue_prompt.wav` for `dialect`
+- `00000309-00000300.wav` for `speech_bgm` and `speech_sound`
+- `CTS-CN-F2F-2019-11-11-423-012-A.wav` and `CTS-CN-F2F-2019-11-11-423-012-B.wav` for `podcast`
+
+## Validation Matrix
+
+The repo-facing example is intended to cover the same dense TTS workflows used
+by the local Ming validation script:
+
+| Case | Blocking `deploy/ming_tts.yaml` | Async chunk `deploy/ming_tts.yaml` | Extra inputs |
+|---|---:|---:|---|
+| `style` | Yes | Optional smoke test | none |
+| `ip` | Yes | Optional smoke test | none |
+| `bgm` | Yes | Optional smoke test | none |
+| `tta` | Yes | Optional smoke test | none |
+| `emotion` | Yes | Yes | `--ref-audio emotion_prompt.wav` |
+| `basic` | Yes | Yes | `--ref-audio 10002287-00000095.wav` |
+| `dialect` | Yes | Yes | `--ref-audio yue_prompt.wav` |
+| `zero_shot` | Yes | Yes | `--ref-audio 10002287-00000094.wav --ref-text ...` |
+| `podcast` | Yes | Yes | two `--ref-audio-paths` |
+| `speech_bgm` | Yes | Yes | `--ref-audio 00000309-00000300.wav` |
+| `speech_sound` | Yes | Yes | `--ref-audio 00000309-00000300.wav` |
+
+## Validated Outputs
+
+The following measurements are retained from an earlier L4 CUDA validation;
+they are not ROCm benchmark results. Default async_chunk matched blocking
+output frame counts and Stage-1 patch counts for every case:
+
+| Case | Blocking frames / patches / sec | Async chunk frames / patches / sec |
+|---|---:|---:|
+| `style` | 409248 / 29 / 9.28 | 409248 / 29 / 9.28 |
+| `ip` | 183456 / 13 / 4.16 | 183456 / 13 / 4.16 |
+| `bgm` | 1326528 / 94 / 30.08 | 1326528 / 94 / 30.08 |
+| `tta` | 465696 / 33 / 10.56 | 465696 / 33 / 10.56 |
+| `emotion` | 324576 / 23 / 7.36 | 324576 / 23 / 7.36 |
+| `basic` | 211680 / 15 / 4.80 | 211680 / 15 / 4.80 |
+| `dialect` | 239904 / 17 / 5.44 | 239904 / 17 / 5.44 |
+| `zero_shot` | 409248 / 29 / 9.28 | 409248 / 29 / 9.28 |
+| `podcast` | 437472 / 31 / 9.92 | 437472 / 31 / 9.92 |
+| `speech_bgm` | 296352 / 21 / 6.72 | 296352 / 21 / 6.72 |
+| `speech_sound` | 352800 / 25 / 8.00 | 352800 / 25 / 8.00 |
+
+## Key Arguments
+
+| Argument | Description |
+|---|---|
+| `--model` | Hugging Face repo or local Ming checkpoint path |
+| `--deploy-config` | Deploy config YAML. Use `vllm_omni/deploy/ming_tts.yaml` |
+| `--case` | Built-in demo case |
+| `--ref-audio` | Single reference wav path for cloning-style cases |
+| `--ref-audio-paths` | Multiple reference wav paths, used by `podcast` |
+| `--ref-text` | Reference transcript. Required for `zero_shot` |
+| `--instructions` | Free-form Ming instruction string |
+| `--instruction-json` | Structured Ming instruction JSON |
+| `--speaker-embedding` | JSON file containing a 192-d speaker embedding |
+| `--extract-speaker-embeddings` | Force CampPlus speaker extraction from the provided reference audio paths |
+| `--max-decode-steps` | Override `ming_max_decode_steps` |
+| `--num-prompts` | Repeat the same case N times. Outputs are indexed when `N > 1` |
+| `--streaming` | Use `AsyncOmni` and async_chunk transport |
+| `--enforce-eager` | Recommended for Ming dense; non-eager is out of scope |
+| `--enable-stats` / `--log-stats` | Enable vLLM-Omni per-request stats logging |
+| `--stats-log-file` | Optional path for the stats log |
+| `--metadata-json` | Optional path for the run manifest JSON |
+| `--stage-init-timeout` | Per-stage initialization timeout in seconds |
+| `--init-timeout` | Total initialization timeout in seconds |
+| `--batch-timeout` | Batch timeout in seconds |
+| `--worker-backend` | `multi_process` or `ray` |
+| `--ray-address` | Ray cluster address when using `--worker-backend ray` |
+
+## Output
+
+- The script writes one mono 44.1 kHz WAV file per run
+- Default output directory: `output_audio/`
+- Default filename: `ming_<case>.wav`
+- When `--num-prompts > 1`, outputs are indexed as `ming_<case>_00000.wav`, `..._00001.wav`, etc.
+- When stats are enabled, the script can also write:
+  - a stats log file such as `ming_style_pipeline.log`
+  - a manifest JSON with per-output metadata, stage durations, peak memory info,
+    and streaming client latency metrics when `--streaming` is used
+
+## Notes
+
+- `style` and `ip` are zero-speaker paths and do not require a reference clip
+- `emotion`, `basic`, `dialect`, `speech_bgm`, and `speech_sound` require one reference clip
+- `zero_shot` requires both `--ref-audio` and `--ref-text`
+- `podcast` requires at least two reference clips via `--ref-audio-paths`
+- `podcast` automatically extracts one speaker embedding per reference clip
+- `--speaker-embedding` may contain either one 192-d vector or a list of 192-d vectors
+- `--enforce-eager` was used for the validated runs
+- The earlier L4 validation used SDPA for the Ming audio VAE instead of
+  FlashAttention2, which is the preferred default when available.

examples/offline_inference/text_to_speech/ming_tts/cases.yaml

@@ -0,0 +1,110 @@
+style:
+  prompt: "Please generate speech based on the following description.\n"
+  text: "我会一直在这里陪着你，直到你慢慢、慢慢地沉入那个最温柔的梦里……好吗？"
+  instruction:
+    风格: >-
+      这是一种ASMR耳语，属于一种旨在引发特殊感官体验的创意风格。这个女性使用轻柔的普通话进行耳语，声音气音成分重。音量极低，紧贴麦克风，语速极慢，旨在制造触发听者颅内快感的声学刺激。
+  use_zero_spk_emb: true
+  max_decode_steps: 200
+
+ip:
+  prompt: "Please generate speech based on the following description.\n"
+  text: "这款产品的名字，叫变态坑爹牛肉丸。"
+  instruction:
+    IP: "灵小甄"
+  use_zero_spk_emb: true
+  max_decode_steps: 200
+
+bgm:
+  prompt: "Please generate music based on the following description.\n"
+  text: " Genre: 电子舞曲. Mood: 自信 / 坚定. Instrument: 架子鼓. Theme: 节日. Duration: 30s."
+  instruction: null
+  use_zero_spk_emb: false
+  max_decode_steps: 400
+
+tta:
+  prompt: "Please generate audio events based on given text.\n"
+  text: "Thunder and a gentle rain"
+  instruction: null
+  use_zero_spk_emb: false
+  max_decode_steps: 200
+  cfg: 4.5
+  sigma: 0.3
+  temperature: 2.5
+
+emotion:
+  prompt: "Please generate speech based on the following description.\n"
+  text: "我竟然抢到了陈奕迅的演唱会门票！太棒了！终于可以现场听一听他的歌声了！"
+  instruction:
+    情感: "高兴"
+  requires_ref_audio: true
+  auto_extract_speaker_embeddings: true
+  max_decode_steps: 200
+
+basic:
+  prompt: "Please generate speech based on the following description.\n"
+  text: "简单地说，这相当于惠普把消费领域市场拱手相让了。"
+  instruction:
+    语速: "快速"
+    基频: "中"
+    音量: "中"
+  requires_ref_audio: true
+  auto_extract_speaker_embeddings: true
+  max_decode_steps: 200
+
+dialect:
+  prompt: "Please generate speech based on the following description.\n"
+  text: "我觉得社会企业同个人都有责任"
+  instruction:
+    方言: "广粤话"
+  requires_ref_audio: true
+  auto_extract_speaker_embeddings: true
+  max_decode_steps: 200
+
+zero_shot:
+  prompt: "Please generate speech based on the following description.\n"
+  text: "我们的愿景是构建未来服务业的数字化基础设施，为世界带来更多微小而美好的改变。"
+  instruction: null
+  requires_ref_audio: true
+  requires_ref_text: true
+  auto_extract_speaker_embeddings: true
+  max_decode_steps: 200
+
+podcast:
+  prompt: "Please generate speech based on the following description.\n"
+  text: " speaker_1:你可以说一下，就大概说一下，可能虽然我也不知道，我看过那部电影没有。\n speaker_2:就是那个叫什么，变相一节课的嘛。\n speaker_1:嗯。\n speaker_2:一部搞笑的电影。\n speaker_1:一部搞笑的。\n"
+  instruction: null
+  prompt_text: " speaker_1:并且我们还要进行每个月还要考核 笔试的话还要进行笔试，做个，当服务员还要去笔试了\n speaker_2:对啊，这真的很奇怪，就是 单纯的因，单纯自己工资不高，只是因为可能人家那个店比较出名一点，就对你苛刻要求\n"
+  requires_ref_audio_count: 2
+  auto_extract_speaker_embeddings: true
+  max_decode_steps: 200
+
+speech_bgm:
+  prompt: "Please generate speech based on the following description.\n"
+  text: "此次业绩下滑原因，可归结为企业停止服务某些品牌，而带来的负面影响。"
+  instruction:
+    BGM:
+      Genre: "当代古典音乐."
+      Mood: "温暖 / 友善."
+      Instrument: "电吉他"
+      Theme: "节日."
+      SNR: 10.0
+      ENV: null
+  requires_ref_audio: true
+  auto_extract_speaker_embeddings: true
+  max_decode_steps: 200
+
+speech_sound:
+  prompt: "Please generate speech based on the following description.\n"
+  text: "此次业绩下滑原因，可归结为企业停止服务某些品牌，而带来的负面影响。"
+  instruction:
+    BGM:
+      ENV: "Birds chirping"
+      SNR: 10.0
+      Genre: null
+      Mood: null
+      Instrument: null
+      Theme: null
+  requires_ref_audio: true
+  auto_extract_speaker_embeddings: true
+  max_decode_steps: 200

examples/offline_inference/text_to_speech/ming_tts/end2end.py

@@ -0,0 +1,200 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Offline inference demo for Ming-omni-tts via vLLM Omni."""
+
+import json
+import os
+from pathlib import Path
+
+import soundfile as sf
+import torch
+import torchaudio
+import yaml
+from transformers import AutoTokenizer
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+from vllm_omni.model_executor.models.ming_tts.config_ming_tts import (
+    KEY_CFG,
+    KEY_MAX_DECODE_STEPS,
+    KEY_SIGMA,
+    KEY_TEMPERATURE,
+    SAMPLE_RATE,
+)
+from vllm_omni.model_executor.models.ming_tts.prompt_assembly import build_ming_dense_prompt
+from vllm_omni.model_executor.models.ming_tts.speaker_extractor import MingSpeakerEmbeddingExtractor
+
+try:
+    from .runner import (
+        build_manifest,
+        build_sampling_params,
+        resolve_metadata_json,
+        resolve_stats_log_file,
+        run_generation,
+    )
+except ImportError:
+    from runner import (
+        build_manifest,
+        build_sampling_params,
+        resolve_metadata_json,
+        resolve_stats_log_file,
+        run_generation,
+    )
+
+_DEFAULT_MODEL = "inclusionAI/Ming-omni-tts-0.5B"
+_REPO_ROOT = Path(__file__).resolve().parents[4]
+_DEFAULT_DEPLOY_CONFIG = str(_REPO_ROOT / "vllm_omni/deploy/ming_tts.yaml")
+_CASES_FILE = Path(__file__).with_name("cases.yaml")
+
+CASE_DEFAULTS = yaml.safe_load(_CASES_FILE.read_text(encoding="utf-8")) or {}
+if not CASE_DEFAULTS:
+    raise RuntimeError(f"Empty or missing case definitions in {_CASES_FILE}")
+
+
+def _build_parser() -> FlexibleArgumentParser:
+    p = FlexibleArgumentParser(description="Offline Ming-omni-tts example")
+    p.add_argument("--model", default=_DEFAULT_MODEL, help="Model name or local path")
+    p.add_argument("--deploy-config", default=None, help="Deploy config YAML; auto-selected when omitted")
+    p.add_argument("--case", choices=sorted(CASE_DEFAULTS), default="style", help="Built-in demo case")
+    p.add_argument("--text", default=None, help="Override case text")
+    p.add_argument("--prompt", default=None, help="Override the system prompt prefix")
+    p.add_argument("--instructions", default=None, help="Free-form Ming instruction string")
+    p.add_argument(
+        "--instruction-json", default=None, help='Structured Ming instruction JSON, e.g. \'{"方言":"广粤话"}\''
+    )
+    p.add_argument("--ref-audio", default=None, help="Single reference audio path for cloning-style cases")
+    p.add_argument("--ref-audio-paths", nargs="+", default=None, help="Multiple reference audio paths (podcast)")
+    p.add_argument("--ref-text", default=None, help="Reference transcript; required for zero_shot")
+    p.add_argument("--speaker-embedding", default=None, help="Path to a JSON speaker embedding file")
+    p.add_argument(
+        "--extract-speaker-embeddings",
+        action="store_true",
+        help="Extract speaker embeddings from ref audio via campplus.onnx",
+    )
+    p.add_argument("--max-decode-steps", type=int, default=None, help="Override ming_max_decode_steps")
+    p.add_argument("--output-dir", default="output_audio", help="Directory for output wav files")
+    p.add_argument("--output-name", default=None, help="Output wav filename")
+    p.add_argument("--num-prompts", type=int, default=1, help="Repeat the same prompt N times")
+    p.add_argument("--streaming", action="store_true", help="Use AsyncOmni with async_chunk streaming")
+    p.add_argument("--trust-remote-code", action="store_true")
+    p.add_argument("--enforce-eager", action="store_true")
+    p.add_argument(
+        "--log-stats", "--enable-stats", dest="log_stats", action="store_true", help="Enable Omni stats logging"
+    )
+    p.add_argument("--stats-log-file", default=None, help="Path for the Omni stats log file")
+    p.add_argument("--metadata-json", default=None, help="Path for the run manifest JSON")
+    p.add_argument("--stage-init-timeout", type=int, default=300, help="Per-stage init timeout (s)")
+    p.add_argument("--init-timeout", type=int, default=600, help="Total init timeout (s)")
+    p.add_argument("--batch-timeout", type=int, default=5, help="Batch timeout (s)")
+    p.add_argument("--shm-threshold-bytes", type=int, default=65536)
+    p.add_argument("--worker-backend", default="multi_process", choices=["multi_process", "ray"])
+    p.add_argument("--ray-address", default=None, help="Ray cluster address (--worker-backend ray)")
+    return p
+
+
+def _finalize_args(args) -> None:
+    if args.instructions is not None and args.instruction_json is not None:
+        raise RuntimeError("Use either --instructions or --instruction-json, not both")
+    if args.num_prompts < 1:
+        raise RuntimeError("--num-prompts must be at least 1")
+    if args.streaming and args.num_prompts != 1:
+        raise RuntimeError("--streaming currently supports exactly one prompt")
+    if args.deploy_config is None:
+        args.deploy_config = _DEFAULT_DEPLOY_CONFIG
+
+
+def _load_waveform(path: str) -> torch.Tensor:
+    samples, sr = sf.read(path, dtype="float32")
+    wav = torch.as_tensor(samples, dtype=torch.float32)
+    if wav.ndim == 2:
+        wav = wav.mean(dim=1)
+    wav = wav.reshape(1, -1)
+    if int(sr) != SAMPLE_RATE:
+        wav = torchaudio.functional.resample(wav, int(sr), SAMPLE_RATE)
+    return wav
+
+
+def _load_speaker_embedding(path: str) -> torch.Tensor:
+    return torch.as_tensor(json.loads(Path(path).read_text(encoding="utf-8")), dtype=torch.float32)
+
+
+def _ref_audio_paths(args) -> list[str]:
+    if args.ref_audio is not None and args.ref_audio_paths is not None:
+        raise RuntimeError("Use either --ref-audio or --ref-audio-paths, not both")
+    if args.ref_audio_paths is not None:
+        return list(args.ref_audio_paths)
+    return [args.ref_audio] if args.ref_audio else []
+
+
+def _resolve_reference_inputs(args, case: dict, paths: list[str]):
+    required = int(case.get("requires_ref_audio_count", 0))
+    if required > 0 and len(paths) < required:
+        raise RuntimeError(f"Case '{args.case}' needs {required} ref audio paths via --ref-audio-paths")
+    if required <= 0 and case.get("requires_ref_audio") and not paths:
+        raise RuntimeError(f"--ref-audio required for case '{args.case}'")
+    if not paths:
+        return None
+    wavs = [_load_waveform(p) for p in paths]
+    return wavs[0] if len(wavs) == 1 else wavs
+
+
+def _resolve_speaker_embedding(args, case: dict, paths: list[str]):
+    if args.speaker_embedding:
+        return _load_speaker_embedding(args.speaker_embedding)
+    if not (case.get("auto_extract_speaker_embeddings") or args.extract_speaker_embeddings) or not paths:
+        return None
+    embs = MingSpeakerEmbeddingExtractor(args.model).extract_many(paths)
+    if not embs:
+        raise RuntimeError("Speaker extraction produced no embeddings")
+    return embs[0] if len(embs) == 1 else torch.stack(embs, dim=0)
+
+
+def _build_prompt_payload(tokenizer, args):
+    case = CASE_DEFAULTS[args.case]
+    paths = _ref_audio_paths(args)
+    prompt_text = args.ref_text if args.ref_text is not None else case.get("prompt_text")
+    if case.get("requires_ref_text") and not prompt_text:
+        raise RuntimeError(f"--ref-text required for case '{args.case}'")
+    prompt_waveform = _resolve_reference_inputs(args, case, paths) if prompt_text is not None else None
+    speaker_embedding = _resolve_speaker_embedding(args, case, paths)
+    use_zero_spk_emb = bool(case.get("use_zero_spk_emb")) and prompt_waveform is None and speaker_embedding is None
+    runtime_controls = {KEY_MAX_DECODE_STEPS: args.max_decode_steps or case["max_decode_steps"]}
+    for key, field in [(KEY_CFG, "cfg"), (KEY_SIGMA, "sigma"), (KEY_TEMPERATURE, "temperature")]:
+        if field in case:
+            runtime_controls[key] = case[field]
+    instruction = (
+        json.loads(args.instruction_json) if args.instruction_json else (args.instructions or case.get("instruction"))
+    )
+    return build_ming_dense_prompt(
+        tokenizer,
+        prompt=args.prompt or case["prompt"],
+        text=args.text or case["text"],
+        runtime_controls=runtime_controls,
+        instruction=instruction,
+        prompt_text=prompt_text,
+        prompt_waveform=prompt_waveform,
+        speaker_embedding=speaker_embedding,
+        use_zero_spk_emb=use_zero_spk_emb,
+    )
+
+
+def main():
+    args = _build_parser().parse_args()
+    _finalize_args(args)
+    tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=False)
+    prompt_payload = _build_prompt_payload(tokenizer, args)
+    case = CASE_DEFAULTS[args.case]
+    sampling_params_list = build_sampling_params(args.max_decode_steps or case["max_decode_steps"])
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    stats_log_file = resolve_stats_log_file(args)
+    summaries = run_generation(args, prompt_payload, sampling_params_list, output_dir, stats_log_file)
+    metadata_json = resolve_metadata_json(args)
+    manifest = build_manifest(args, prompt_payload, stats_log_file, summaries)
+    if metadata_json is not None:
+        Path(metadata_json).write_text(json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8")
+        print(f"Saved run manifest to {metadata_json}")
+
+
+if __name__ == "__main__":
+    os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
+    main()

examples/offline_inference/text_to_speech/ming_tts/runner.py

@@ -0,0 +1,277 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import asyncio
+import json
+import time
+import uuid
+import wave
+from pathlib import Path
+
+import torch
+from vllm import SamplingParams
+
+from vllm_omni import AsyncOmni, Omni
+from vllm_omni.model_executor.models.ming_tts.config_ming_tts import (
+    KEY_SPEAKER_EMBEDDING,
+    SAMPLE_RATE,
+    TEXT_EOS_TOKEN_ID,
+)
+
+
+def coerce_audio_tensor(audio, *, async_chunk: bool) -> torch.Tensor:
+    if isinstance(audio, list):
+        parts = []
+        for item in audio:
+            tensor = torch.as_tensor(item, dtype=torch.float32).reshape(-1)
+            if tensor.numel() > 0:
+                parts.append(tensor)
+        if not parts:
+            return torch.zeros((0,), dtype=torch.float32)
+        return torch.cat(parts, dim=0)
+
+    return torch.as_tensor(audio, dtype=torch.float32).reshape(-1)
+
+
+def resolve_sr(sr) -> int:
+    if isinstance(sr, list):
+        sr = sr[-1]
+    if hasattr(sr, "item"):
+        return int(sr.item())
+    return int(sr)
+
+
+def extract_sample_rate(multimodal_output: dict) -> int:
+    sr = multimodal_output.get("sr")
+    if sr is None:
+        raise RuntimeError("Expected multimodal_output['sr']")
+    return resolve_sr(sr)
+
+
+def write_wav(path: str, audio: torch.Tensor, sample_rate: int) -> None:
+    audio = audio.clamp(-1.0, 1.0)
+    pcm16 = (audio * 32767.0).round().to(torch.int16).cpu().numpy()
+    with wave.open(path, "wb") as wav_file:
+        wav_file.setnchannels(1)
+        wav_file.setsampwidth(2)
+        wav_file.setframerate(int(sample_rate))
+        wav_file.writeframes(pcm16.tobytes())
+
+
+def request_index(request_id: str | None, fallback: int) -> int:
+    try:
+        return int(request_id)
+    except (TypeError, ValueError):
+        if isinstance(request_id, str):
+            head = request_id.split("_", 1)[0]
+            if head.isdigit():
+                return int(head)
+    return fallback
+
+
+def audio_summary(audio: torch.Tensor, sample_rate: int) -> dict:
+    waveform = audio.detach().cpu().reshape(-1).to(torch.float32)
+    return {
+        "sample_rate": int(sample_rate),
+        "num_samples": int(waveform.numel()),
+        "duration_seconds": float(waveform.numel()) / float(sample_rate),
+        "max_abs_amplitude": float(waveform.abs().max().item()) if waveform.numel() > 0 else 0.0,
+    }
+
+
+def resolve_output_name(output_name: str | None, case: str, index: int, total: int) -> str:
+    if total == 1:
+        return output_name or f"ming_{case}.wav"
+    base = Path(output_name or f"ming_{case}.wav")
+    return f"{base.stem}_{index:05d}{base.suffix or '.wav'}"
+
+
+def resolve_stats_log_file(args) -> str | None:
+    if not args.log_stats:
+        return None
+    if args.stats_log_file:
+        return args.stats_log_file
+    base = Path(args.output_name or f"ming_{args.case}.wav").stem
+    return str(Path(args.output_dir) / f"{base}_pipeline.log")
+
+
+def resolve_metadata_json(args) -> str | None:
+    if args.metadata_json:
+        return args.metadata_json
+    if args.log_stats:
+        base = Path(args.output_name or f"ming_{args.case}.wav").stem
+        return str(Path(args.output_dir) / f"{base}_manifest.json")
+    return None
+
+
+def build_manifest(args, prompt_payload, stats_log_file: str | None, outputs: list[dict]) -> dict:
+    additional_information = {}
+    if isinstance(prompt_payload, dict):
+        additional_information = dict(prompt_payload.get("additional_information", {}))
+    return {
+        "model": args.model,
+        "case": args.case,
+        "streaming": bool(args.streaming),
+        "deploy_config": args.deploy_config,
+        "enforce_eager": bool(args.enforce_eager),
+        "num_prompts": int(args.num_prompts),
+        "log_stats": bool(args.log_stats),
+        "stats_log_file": stats_log_file,
+        "prompt_text": additional_information.get("prompt_text"),
+        "instruction": additional_information.get("instruction"),
+        "speaker_embedding_shape": (
+            list(additional_information[KEY_SPEAKER_EMBEDDING].shape)
+            if KEY_SPEAKER_EMBEDDING in additional_information
+            and hasattr(additional_information[KEY_SPEAKER_EMBEDDING], "shape")
+            else None
+        ),
+        "outputs": outputs,
+        "generated_at_unix": time.time(),
+    }
+
+
+def build_engine_kwargs(args, stats_log_file: str | None) -> dict:
+    kwargs = {
+        "model": args.model,
+        "deploy_config": args.deploy_config,
+        "enforce_eager": args.enforce_eager,
+        "trust_remote_code": args.trust_remote_code,
+        "log_stats": args.log_stats,
+        "stage_init_timeout": args.stage_init_timeout,
+        "init_timeout": args.init_timeout,
+        "batch_timeout": args.batch_timeout,
+        "shm_threshold_bytes": args.shm_threshold_bytes,
+        "worker_backend": args.worker_backend,
+    }
+    if stats_log_file is not None:
+        kwargs["log_file"] = stats_log_file
+    if args.ray_address is not None:
+        kwargs["ray_address"] = args.ray_address
+    return kwargs
+
+
+def build_sampling_params(max_decode_steps: int) -> list[SamplingParams]:
+    return [
+        SamplingParams(
+            temperature=0.0,
+            max_tokens=max_decode_steps + 1,
+            stop_token_ids=[int(TEXT_EOS_TOKEN_ID)],
+        ),
+        SamplingParams(temperature=0.0, max_tokens=1),
+    ]
+
+
+async def run_streaming(args, prompt_payload, sampling_params_list, output_dir: Path, stats_log_file: str | None):
+    engine = AsyncOmni(**build_engine_kwargs(args, stats_log_file))
+    try:
+        all_audio_chunks = []
+        accumulated_samples = 0
+        chunk_idx = 0
+        start_time = time.time()
+        chunk_times = []
+        ttfp_seconds = None
+        final_stage_output = None
+        async for stage_output in engine.generate(
+            prompt=prompt_payload,
+            request_id=str(uuid.uuid4()),
+            sampling_params_list=sampling_params_list,
+        ):
+            final_stage_output = stage_output
+            multimodal_output = stage_output.multimodal_output or {}
+            audio = multimodal_output.get("audio")
+            if audio is None:
+                continue
+
+            finished = stage_output.finished
+            if isinstance(audio, torch.Tensor):
+                if finished:
+                    audio_chunk = audio[accumulated_samples:].float().detach().cpu()
+                else:
+                    audio_chunk = audio.float().detach().cpu()
+            elif isinstance(audio, list):
+                audio_chunk = torch.as_tensor(audio[chunk_idx], dtype=torch.float32).reshape(-1).cpu()
+            else:
+                audio_chunk = torch.as_tensor(audio, dtype=torch.float32).reshape(-1).cpu()
+
+            accumulated_samples += int(audio_chunk.numel())
+            chunk_idx += 1
+            if audio_chunk.numel() > 0:
+                now = time.time()
+                if ttfp_seconds is None:
+                    ttfp_seconds = now - start_time
+                chunk_times.append(now)
+                all_audio_chunks.append(audio_chunk)
+
+        if not all_audio_chunks:
+            raise RuntimeError("Streaming Ming example produced no audio chunks")
+
+        waveform = torch.cat(all_audio_chunks, dim=0)
+        output_name = resolve_output_name(args.output_name, args.case, 0, 1)
+        output_path = str(output_dir / output_name)
+        write_wav(output_path, waveform, SAMPLE_RATE)
+        summary = {
+            "request_id": getattr(final_stage_output, "request_id", None),
+            "stage_id": getattr(final_stage_output, "stage_id", None),
+            "output_path": output_path,
+            "stage_durations": getattr(final_stage_output, "stage_durations", {}),
+            "peak_memory_mb": getattr(final_stage_output, "peak_memory_mb", 0.0),
+            "ttfp_seconds": ttfp_seconds,
+            "mean_inter_chunk_seconds": (
+                sum(t1 - t0 for t0, t1 in zip(chunk_times, chunk_times[1:])) / (len(chunk_times) - 1)
+                if len(chunk_times) > 1
+                else None
+            ),
+        }
+        summary.update(audio_summary(waveform, SAMPLE_RATE))
+        print(f"Saved streaming output to {output_path}")
+        print(json.dumps(summary, ensure_ascii=False, indent=2))
+        return [summary]
+    finally:
+        engine.shutdown()
+
+
+def run_non_streaming(args, prompt_payload, sampling_params_list, output_dir: Path, stats_log_file: str | None):
+    engine = Omni(**build_engine_kwargs(args, stats_log_file))
+    try:
+        outputs = engine.generate(
+            prompts=[prompt_payload for _ in range(args.num_prompts)],
+            sampling_params_list=sampling_params_list,
+            py_generator=False,
+        )
+        summaries = []
+        for fallback_index, output in enumerate(outputs):
+            if output.final_output_type != "audio":
+                continue
+            multimodal_output = output.multimodal_output or {}
+            waveform = coerce_audio_tensor(multimodal_output.get("audio"), async_chunk=False)
+            sample_rate = extract_sample_rate(multimodal_output)
+            output_name = resolve_output_name(
+                args.output_name,
+                args.case,
+                request_index(output.request_id, fallback_index),
+                args.num_prompts,
+            )
+            output_path = str(output_dir / output_name)
+            write_wav(output_path, waveform, sample_rate)
+            summary = {
+                "request_id": output.request_id,
+                "stage_id": output.stage_id,
+                "output_path": output_path,
+                "stage_durations": output.stage_durations,
+                "peak_memory_mb": output.peak_memory_mb,
+            }
+            summary.update(audio_summary(waveform, sample_rate))
+            summaries.append(summary)
+            print(f"Saved output to {output_path}")
+            print(json.dumps(summary, ensure_ascii=False, indent=2))
+        if not summaries:
+            raise RuntimeError("Non-streaming Ming example produced no audio outputs")
+        return summaries
+    finally:
+        engine.close()
+
+
+def run_generation(args, prompt_payload, sampling_params_list, output_dir: Path, stats_log_file: str | None):
+    if args.streaming:
+        return asyncio.run(run_streaming(args, prompt_payload, sampling_params_list, output_dir, stats_log_file))
+    return run_non_streaming(args, prompt_payload, sampling_params_list, output_dir, stats_log_file)

examples/offline_inference/text_to_speech/voxtral_tts/end2end.py

@@ -212,7 +212,7 @@ def run_non_streaming(inputs, sampling_params_list, model_name, args, output_dir
     output_audio_dur = 0.0
 
     for batch_idx, o in enumerate(outputs):
-        audio_tensor = torch.cat(o.multimodal_output["audio"])
+        audio_tensor = torch.cat((o.multimodal_output["audio"],))
         audio_array = audio_tensor.tolist()
         output_audio_dur += float(len(audio_array)) / 24000
         if args.write_audio:

examples/online_serving/text_to_speech/README.md

@@ -16,6 +16,7 @@ For the full list of supported architectures across all modalities, see
 |---|---|---|---|---|---|
 | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | ✓ |
 | GLM-TTS | `zai-org/GLM-TTS` | ✓ (`ref_audio`+`ref_text`, required) | ✓ (PCM stream) | — | ✓ |
+| Ming-omni-tts | `inclusionAI/Ming-omni-tts-0.5B` | ✓ (`ref_audio` / `speaker_embedding`) | ✓ (PCM stream) | IP labels + structured `instructions` | — |
 | Ming-flash-omni-TTS | `Jonathan1909/Ming-flash-omni-2.0` | — (caption-controlled) | — | caption fields (`instructions`) | — |
 | MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | ✓ (`ref_audio` required) | ✓ (PCM stream) | — | ✓ |
 | OmniVoice | `k2-fsa/OmniVoice` | ✓ | — | — | — |
@@ -202,6 +203,51 @@ python fish_speech/gradio_demo.py --api-base http://localhost:8091  # if server
 
 ---
 
+## Ming-omni-tts
+
+Dense 0.5B two-stage TTS served through `/v1/audio/speech`. Ming uses the standard speech endpoint plus structured controls in `instructions`, `voice`, `language`, `ref_audio`, `ref_text`, and `speaker_embedding`.
+
+### Launch
+```bash
+bash examples/online_serving/text_to_speech/ming_tts/run_server.sh
+```
+Equivalent manual command:
+```bash
+vllm-omni serve inclusionAI/Ming-omni-tts-0.5B \
+    --deploy-config vllm_omni/deploy/ming_tts.yaml \
+    --host 0.0.0.0 --port 8091 \
+    --enforce-eager --omni
+```
+
+### Sending requests
+```bash
+python examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py \
+    --text "你好，这是 Ming 在线语音合成测试。"
+```
+
+Structured dialect control:
+```bash
+python examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py \
+    --text "我觉得社会企业同个人都有责任" \
+    --instruction-json '{"方言":"广粤话"}' \
+    --ref-audio /path/to/yue_prompt.wav
+```
+
+Zero-shot cloning:
+```bash
+python examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py \
+    --text "我们的愿景是构建未来服务业的数字化基础设施，为世界带来更多微小而美好的改变。" \
+    --ref-audio /path/to/10002287-00000094.wav \
+    --ref-text "在此奉劝大家别乱打美白针。"
+```
+
+### Notes
+- `run_curl.sh` keeps a small sanity subset; use the Ming README for the broader request cookbook.
+- Online serving is speech-shaped today; music-only `bgm` and text-to-audio `tta` remain offline examples.
+- Full request details live in [`ming_tts/README.md`](ming_tts/README.md).
+
+---
+
 ## Ming-flash-omni-TTS
 
 Standalone talker-only deployment of Ming-flash-omni-2.0. Voice is controlled through caption text passed via `instructions`.

examples/online_serving/text_to_speech/ming_tts/README.md

@@ -0,0 +1,113 @@
+# Ming-omni-tts Online Serving
+
+Serve the dense `inclusionAI/Ming-omni-tts-0.5B` two-stage TTS model through
+the OpenAI-compatible `/v1/audio/speech` endpoint.
+
+## Start Server
+
+```bash
+vllm-omni serve inclusionAI/Ming-omni-tts-0.5B \
+    --deploy-config vllm_omni/deploy/ming_tts.yaml \
+    --omni \
+    --port 8091 \
+    --enforce-eager
+```
+
+Or:
+
+```bash
+cd examples/online_serving/text_to_speech/ming_tts
+./run_server.sh
+```
+
+The tested ROCm environment is summarized in the
+[Ming recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/inclusionAI/Ming-omni-tts-0.5B.md).
+
+## Send Requests
+
+The Python client targets `http://localhost:8091/v1` with `api_key=EMPTY`; it
+does not call OpenAI's hosted API.
+
+```bash
+python openai_speech_client.py \
+    --text "你好，这是 Ming 在线语音合成测试。" \
+    --max-new-tokens 200
+```
+
+Style or dialect controls can be plain text or Ming JSON. The upstream
+dialect example also uses `yue_prompt.wav` for speaker conditioning:
+
+```bash
+python openai_speech_client.py \
+    --text "我觉得社会企业同个人都有责任" \
+    --instruction-json '{"方言":"广粤话"}' \
+    --ref-audio /path/to/yue_prompt.wav \
+    --max-new-tokens 200
+```
+
+When `--ref-audio` is supplied without `--ref-text`, the server extracts the
+Ming speaker embedding, matching upstream `use_spk_emb=True`, without using the
+audio as a zero-shot prompt.
+
+Reference-audio cloning:
+
+```bash
+python openai_speech_client.py \
+    --text "我们的愿景是构建未来服务业的数字化基础设施，为世界带来更多微小而美好的改变。" \
+    --ref-audio /path/to/10002287-00000094.wav \
+    --ref-text "在此奉劝大家别乱打美白针。" \
+    --max-new-tokens 200
+```
+
+Podcast-style multi-speaker prompt:
+
+```bash
+python openai_speech_client.py \
+    --text " speaker_1:你可以说一下，就大概说一下，可能虽然我也不知道，我看过那部电影没有。
+ speaker_2:就是那个叫什么，变相一节课的嘛。
+ speaker_1:嗯。
+ speaker_2:一部搞笑的电影。
+ speaker_1:一部搞笑的。" \
+    --ref-audio /path/to/CTS-CN-F2F-2019-11-11-423-012-A.wav \
+    --ref-audio /path/to/CTS-CN-F2F-2019-11-11-423-012-B.wav \
+    --ref-text " speaker_1:并且我们还要进行每个月还要考核 笔试的话还要进行笔试，做个，当服务员还要去笔试了
+ speaker_2:对啊，这真的很奇怪，就是 单纯的因，单纯自己工资不高，只是因为可能人家那个店比较出名一点，就对你苛刻要求"
+```
+
+Streaming PCM:
+
+```bash
+python openai_speech_client.py \
+    --text "你好，这是流式输出测试。" \
+    --stream \
+    --output ming_output.pcm
+```
+
+`run_curl.sh` keeps small smoke checks:
+
+```bash
+./run_curl.sh basic
+REF_AUDIO=/path/to/reference.wav REF_TEXT="在此奉劝大家别乱打美白针。" ./run_curl.sh zero_shot
+./run_curl.sh stream
+```
+
+## Request Fields
+
+| Field | Ming meaning |
+|-------|--------------|
+| `input` | target text |
+| `instructions` | plain style text, or JSON object for structured Ming controls |
+| `voice` | Ming IP voice label unless it resolves to an uploaded speaker |
+| `language` | Ming `方言` control |
+| `ref_audio` | speaker reference; with `ref_text`, also supplies the prompt waveform |
+| `ref_text` | transcript enabling zero-shot or podcast prompt-latent conditioning |
+| `speaker_embedding` | 192-d Ming speaker embedding |
+| `max_new_tokens` | Ming `max_decode_steps` |
+
+## Notes
+
+- `ref_audio` accepts local paths through the client, remote URLs, `file://`,
+  or `data:` URLs.
+- Non-streaming responses return WAV bytes; streaming responses return PCM.
+- Music-only `bgm` generation is offline-only until the API exposes Ming
+  prompt-mode selection.

examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py

@@ -0,0 +1,225 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""OpenAI-compatible client for Ming-omni-tts via /v1/audio/speech.
+
+Examples:
+    python openai_speech_client.py --text "你好，世界"
+    python openai_speech_client.py --text "我觉得社会企业同个人都有责任" \
+        --instruction-json '{"方言":"广粤话"}' --ref-audio yue_prompt.wav --max-new-tokens 200
+    python openai_speech_client.py \
+        --text "我们的愿景是构建未来服务业的数字化基础设施，为世界带来更多微小而美好的改变。" \
+        --ref-audio 10002287-00000094.wav --ref-text "在此奉劝大家别乱打美白针。" \
+        --max-new-tokens 200
+    python openai_speech_client.py --text "speaker_1:你好。 speaker_2:你好。" \
+        --ref-audio speaker_1.wav --ref-audio speaker_2.wav --ref-text "speaker_1:你好。 speaker_2:你好。"
+    python openai_speech_client.py --text "你好，这是流式输出测试。" \
+        --stream --output ming_output.pcm
+"""
+
+import argparse
+import base64
+import json
+import os
+
+import httpx
+
+DEFAULT_API_BASE = "http://localhost:8091"
+DEFAULT_API_KEY = "EMPTY"
+DEFAULT_MODEL = "inclusionAI/Ming-omni-tts-0.5B"
+EXPECTED_SPEAKER_EMBEDDING_DIM = 192
+
+
+def encode_audio_to_base64(audio_path: str) -> str:
+    """Encode a local audio file to a base64 data URL."""
+    if not os.path.exists(audio_path):
+        raise FileNotFoundError(f"Audio file not found: {audio_path}")
+    ext = audio_path.lower().rsplit(".", 1)[-1]
+    mime_map = {
+        "wav": "audio/wav",
+        "mp3": "audio/mpeg",
+        "flac": "audio/flac",
+        "ogg": "audio/ogg",
+        "aac": "audio/aac",
+    }
+    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}"
+
+
+def load_speaker_embedding(path: str) -> list[float]:
+    """Load and validate a 192-d Ming speaker embedding JSON file."""
+    with open(path, encoding="utf-8") as f:
+        data = json.load(f)
+
+    if not isinstance(data, list):
+        raise ValueError("speaker_embedding file must contain a JSON list")
+    if len(data) != EXPECTED_SPEAKER_EMBEDDING_DIM:
+        raise ValueError(
+            f"Ming dense speaker_embedding must have {EXPECTED_SPEAKER_EMBEDDING_DIM} values, got {len(data)}"
+        )
+
+    values = []
+    for index, value in enumerate(data):
+        try:
+            values.append(float(value))
+        except (TypeError, ValueError) as exc:
+            raise ValueError(f"speaker_embedding[{index}] must be a number, got {value!r}") from exc
+    return values
+
+
+def build_instruction_payload(args) -> str | None:
+    """Return a string payload for the API `instructions` field."""
+    if args.instructions and args.instruction_json:
+        raise ValueError("Use either --instructions or --instruction-json, not both")
+    if args.instruction_json:
+        parsed = json.loads(args.instruction_json)
+        return json.dumps(parsed, ensure_ascii=False)
+    return args.instructions
+
+
+def validate_args(args) -> None:
+    """Fail fast on invalid combinations before hitting the server."""
+    if args.ref_text and not args.ref_audio:
+        raise ValueError("--ref-audio is required when --ref-text is provided")
+    if args.speaker_embedding and args.ref_audio and len(args.ref_audio) > 1:
+        raise ValueError("--speaker-embedding cannot be combined with multiple --ref-audio values")
+
+
+def run_tts(args) -> None:
+    """Generate speech via the OpenAI-compatible /v1/audio/speech API."""
+    validate_args(args)
+
+    payload = {
+        "model": args.model,
+        "input": args.text,
+        "response_format": args.response_format,
+    }
+
+    if args.voice:
+        payload["voice"] = args.voice
+    if args.task_type:
+        payload["task_type"] = args.task_type
+    if args.dialect:
+        payload["language"] = args.dialect
+
+    instructions = build_instruction_payload(args)
+    if instructions:
+        payload["instructions"] = instructions
+
+    if args.ref_audio:
+        ref_audio = []
+        for audio in args.ref_audio:
+            if audio.startswith(("http://", "https://", "data:")):
+                ref_audio.append(audio)
+            else:
+                ref_audio.append(encode_audio_to_base64(audio))
+        payload["ref_audio"] = ref_audio[0] if len(ref_audio) == 1 else ref_audio
+    if args.ref_text:
+        payload["ref_text"] = args.ref_text
+    if args.speaker_embedding:
+        payload["speaker_embedding"] = load_speaker_embedding(args.speaker_embedding)
+    if args.max_new_tokens:
+        payload["max_new_tokens"] = args.max_new_tokens
+    if args.stream:
+        payload["stream"] = True
+        payload["response_format"] = "pcm"
+
+    api_url = f"{args.api_base}/v1/audio/speech"
+    headers = {
+        "Content-Type": "application/json",
+        "Authorization": f"Bearer {args.api_key}",
+    }
+
+    print(f"Model: {args.model}")
+    print(f"Text: {args.text}")
+    print(f"Payload keys: {sorted(payload)}")
+
+    if args.stream:
+        output_path = args.output or "ming_output.pcm"
+        with httpx.Client(timeout=300.0) as client:
+            with client.stream("POST", api_url, json=payload, headers=headers) as response:
+                if response.status_code != 200:
+                    print(f"Error: {response.status_code}")
+                    print(response.read().decode())
+                    return
+                with open(output_path, "wb") as f:
+                    for chunk in response.iter_bytes():
+                        f.write(chunk)
+        print(f"Streamed PCM audio to: {output_path}")
+        return
+
+    with httpx.Client(timeout=300.0) as client:
+        response = client.post(api_url, json=payload, headers=headers)
+
+    if response.status_code != 200:
+        print(f"Error: {response.status_code}")
+        print(response.text)
+        return
+
+    try:
+        text = response.content.decode("utf-8")
+        if text.startswith('{"error"'):
+            print(f"Error: {text}")
+            return
+    except UnicodeDecodeError:
+        pass
+
+    output_path = args.output or "ming_output.wav"
+    with open(output_path, "wb") as f:
+        f.write(response.content)
+    print(f"Audio saved to: {output_path}")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="OpenAI-compatible client for Ming-omni-tts via /v1/audio/speech")
+    parser.add_argument("--api-base", default=DEFAULT_API_BASE, help="API base URL")
+    parser.add_argument("--api-key", default=DEFAULT_API_KEY, help="API key")
+    parser.add_argument("--model", "-m", default=DEFAULT_MODEL, help="Model name or path")
+    parser.add_argument("--text", required=True, help="Text to synthesize")
+    parser.add_argument(
+        "--task-type",
+        default=None,
+        choices=["CustomVoice", "VoiceDesign", "Base"],
+        help="Optional compatibility task type. Ming accepts the same field but primarily uses prompt metadata.",
+    )
+    parser.add_argument(
+        "--voice",
+        default=None,
+        help="Maps to Ming `IP` when using built-in character voices, or to an uploaded voice sample name",
+    )
+    parser.add_argument("--dialect", default=None, help="Maps to Ming `方言`")
+    parser.add_argument("--instructions", default=None, help="Free-form Ming instruction string")
+    parser.add_argument(
+        "--instruction-json",
+        default=None,
+        help='Structured Ming instruction JSON, for example \'{"情感":"高兴"}\'',
+    )
+    parser.add_argument(
+        "--ref-audio",
+        action="append",
+        default=None,
+        help="Reference audio path, URL, or data URL. Repeat for podcast-style multi-speaker prompts.",
+    )
+    parser.add_argument("--ref-text", default=None, help="Reference transcript for cloning")
+    parser.add_argument(
+        "--speaker-embedding", default=None, help="Path to a JSON file containing a 192-d speaker embedding"
+    )
+    parser.add_argument("--max-new-tokens", type=int, default=None, help="Override ming_max_decode_steps")
+    parser.add_argument("--stream", action="store_true", help="Enable streaming PCM output")
+    parser.add_argument(
+        "--response-format",
+        default="wav",
+        choices=["wav", "mp3", "flac", "pcm", "aac", "opus"],
+        help="Audio format when not streaming",
+    )
+    parser.add_argument("--output", "-o", default=None, help="Output file path")
+    args = parser.parse_args()
+    try:
+        run_tts(args)
+    except (FileNotFoundError, ValueError, json.JSONDecodeError) as exc:
+        raise SystemExit(f"Error: {exc}") from exc
+
+
+if __name__ == "__main__":
+    main()

examples/online_serving/text_to_speech/ming_tts/run_curl.sh

@@ -0,0 +1,76 @@
+#!/bin/bash
+set -euo pipefail
+
+MODE="${1:-basic}"
+HOST="${HOST:-localhost}"
+PORT="${PORT:-8091}"
+MODEL="${MODEL:-inclusionAI/Ming-omni-tts-0.5B}"
+API_URL="http://${HOST}:${PORT}/v1/audio/speech"
+TEXT="${TEXT:-你好，这是 Ming 在线语音合成测试。}"
+OUTPUT="${OUTPUT:-ming_output.wav}"
+STREAM_OUTPUT="${STREAM_OUTPUT:-ming_output.pcm}"
+REF_AUDIO="${REF_AUDIO:-}"
+REF_TEXT="${REF_TEXT:-}"
+
+post_json() {
+    local payload="$1"
+    local output_path="$2"
+    curl -X POST "$API_URL" \
+        -H "Content-Type: application/json" \
+        -H "Authorization: Bearer EMPTY" \
+        -d "$payload" \
+        --output "$output_path"
+}
+
+case "$MODE" in
+    basic)
+        post_json "{
+            \"model\": \"${MODEL}\",
+            \"input\": \"${TEXT}\",
+            \"response_format\": \"wav\"
+        }" "$OUTPUT"
+        ;;
+    zero_shot)
+        if [ -z "$REF_AUDIO" ] || [ -z "$REF_TEXT" ]; then
+            echo "zero_shot requires REF_AUDIO and REF_TEXT" >&2
+            exit 1
+        fi
+        python - <<'PY' > /tmp/ming_zero_shot_payload.json
+import base64
+import json
+import mimetypes
+import os
+from pathlib import Path
+
+path = Path(os.environ["REF_AUDIO"])
+mime_type = mimetypes.guess_type(path.name)[0] or "audio/wav"
+payload = {
+    "model": os.environ["MODEL"],
+    "input": os.environ["TEXT"],
+    "ref_audio": f"data:{mime_type};base64,{base64.b64encode(path.read_bytes()).decode('utf-8')}",
+    "ref_text": os.environ["REF_TEXT"],
+    "response_format": "wav",
+}
+print(json.dumps(payload, ensure_ascii=False))
+PY
+        curl -X POST "$API_URL" \
+            -H "Content-Type: application/json" \
+            -H "Authorization: Bearer EMPTY" \
+            --data-binary @/tmp/ming_zero_shot_payload.json \
+            --output "$OUTPUT"
+        rm -f /tmp/ming_zero_shot_payload.json
+        ;;
+    stream)
+        post_json "{
+            \"model\": \"${MODEL}\",
+            \"input\": \"${TEXT}\",
+            \"stream\": true,
+            \"response_format\": \"pcm\"
+        }" "$STREAM_OUTPUT"
+        ;;
+    *)
+        echo "Unknown mode: $MODE" >&2
+        echo "Supported sanity checks: basic, zero_shot, stream" >&2
+        exit 1
+        ;;
+esac

examples/online_serving/text_to_speech/ming_tts/run_server.sh

@@ -0,0 +1,25 @@
+#!/bin/bash
+# Launch vLLM-Omni server for Ming-omni-tts.
+#
+# Usage:
+#   ./run_server.sh
+#   PORT=8000 ./run_server.sh
+
+set -e
+
+DIR="$(cd "$(dirname "$0")" && pwd)"
+ROOT="$(cd "$DIR/../../../.." && pwd)"
+
+MODEL="${MODEL:-inclusionAI/Ming-omni-tts-0.5B}"
+PORT="${PORT:-8091}"
+DEPLOY_CONFIG="${DEPLOY_CONFIG:-$ROOT/vllm_omni/deploy/ming_tts.yaml}"
+
+echo "Starting Ming-omni-tts server with model: $MODEL"
+echo "Deploy config: $DEPLOY_CONFIG"
+
+vllm-omni serve "$MODEL" \
+    --deploy-config "$DEPLOY_CONFIG" \
+    --host 0.0.0.0 \
+    --port "$PORT" \
+    --enforce-eager \
+    --omni

recipes/README.md

@@ -33,6 +33,7 @@ recipes/
 | [`fishaudio/Fish-Speech-S2-Pro.md`](./fishaudio/Fish-Speech-S2-Pro.md) | Online serving for TTS | 1x A800 80GB |
 | [`Helios/Helios.md`](./Helios/Helios.md) | Text-to-video, image-to-video, and video-to-video generation | 1x NVIDIA H20 |
 | [`inclusionAI/Ming-flash-omni-2.0.md`](./inclusionAI/Ming-flash-omni-2.0.md) | Online serving for multimodal chat + standalone TTS | 4x H100 / 1x H100 80GB |
+| [`inclusionAI/Ming-omni-tts-0.5B.md`](./inclusionAI/Ming-omni-tts-0.5B.md) | Offline + online dense Ming TTS/audio generation | 1x AMD `gfx942` (ROCm 7.2) |
 | [`LTX/LTX-2.md`](./LTX/LTX-2.md) | Text-to-video and image-to-video serving | 1x H200 141GB |
 | [`LTX/LTX-2.3.md`](./LTX/LTX-2.3.md) | Text-to-video with audio generation (22B) | 1x GPU (96GB VRAM) |
 | [`mistralai/Voxtral-TTS.md`](./mistralai/Voxtral-TTS.md) | Online serving for TTS | 1x RTX 4090 24GB |

recipes/inclusionAI/Ming-omni-tts-0.5B.md

@@ -0,0 +1,82 @@
+# Ming-omni-tts 0.5B
+
+> Offline and online dense Ming text-to-speech on ROCm
+
+## Summary
+
+- Vendor: inclusionAI
+- Model: `inclusionAI/Ming-omni-tts-0.5B`
+- Task: Text-to-speech with style, dialect, cloning, and multi-speaker controls
+- Mode: Online serving via the OpenAI-compatible `/v1/audio/speech` API; offline inference
+- Maintainer: Community
+
+## References
+
+- [Model card](https://huggingface.co/inclusionAI/Ming-omni-tts-0.5B)
+- [Upstream repository](https://github.com/inclusionAI/Ming-omni-tts)
+- [Offline example](../../examples/offline_inference/text_to_speech/ming_tts/)
+- [Online example](../../examples/online_serving/text_to_speech/ming_tts/)
+
+## Hardware Support
+
+## ROCm
+
+### 1x AMD `gfx942`
+
+#### Environment
+
+- OS: Ubuntu 22.04.5 LTS, x86_64
+- Python: 3.12.13
+- ROCm / HIP: 7.2.53211
+- vLLM version: 0.22.0
+- vLLM-Omni version or commit: 0.1.dev1873 / `99c35c410`
+- Docker image: `vllm/vllm-omni-rocm:v0.22.0`
+
+#### Command
+
+From the vLLM-Omni repository root:
+
+```bash
+docker run --rm \
+    --group-add=video \
+    --ipc=host \
+    --cap-add=SYS_PTRACE \
+    --security-opt seccomp=unconfined \
+    --device /dev/kfd \
+    --device /dev/dri \
+    -v ~/.cache/huggingface:/root/.cache/huggingface \
+    -v "$PWD":/app/vllm-omni \
+    -w /app/vllm-omni \
+    -e VLLM_ROCM_USE_AITER=0 \
+    -p 8091:8091 \
+    vllm/vllm-omni-rocm:v0.22.0 \
+    --model inclusionAI/Ming-omni-tts-0.5B \
+    --deploy-config vllm_omni/deploy/ming_tts.yaml \
+    --omni \
+    --port 8091 \
+    --enforce-eager
+```
+
+#### Verification
+
+```bash
+python examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py \
+    --text "我觉得社会企业同个人都有责任" \
+    --instruction-json '{"方言":"广粤话"}' \
+    --ref-audio /path/to/yue_prompt.wav \
+    --max-new-tokens 200 \
+    --output dialect.wav
+```
+
+`--ref-audio` matches upstream `use_spk_emb=True`; do not add `--ref-text`
+for the dialect case.
+
+## Notes
+
+- The official ROCm image includes the platform dependencies.
+- See the [ROCm installation guide](../../docs/getting_started/installation/gpu.md)
+  for interactive and source-build workflows.
+- The [offline example](../../examples/offline_inference/text_to_speech/ming_tts/)
+  can be run from an interactive container.
+- The tested environment uses `--enforce-eager`.
+- Output is mono 44.1 kHz audio.

requirements/common.txt

@@ -22,4 +22,3 @@ prettytable>=3.8.0
 aenum==3.1.16
 pyzmq>=25.0.0
 janus>=1.0.0
-pydub

requirements/xpu.txt

@@ -1,2 +1,3 @@
 -r common.txt
 onnxruntime>=1.23.2
+auto-round-lib

tests/core/sched/test_omni_generation_scheduler_update_session.py

@@ -20,6 +20,7 @@ from vllm.sampling_params import SamplingParams
 from vllm.v1.engine import EngineCoreEventType
 from vllm.v1.request import Request, RequestStatus, StreamingUpdate
 from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
+from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
 
 # isort: on
 
@@ -142,3 +143,300 @@ class TestReplaceSessionWithStreamingUpdate:
 
         assert session.events
         assert session.events[-1].type == EngineCoreEventType.QUEUED
+
+
+class _RealignSchedulerStub(OmniSchedulerMixin):
+    """Minimal scheduler surface for ``_realign_request_status_to_queues``.
+
+    Pins ``self.requests`` / ``self.running`` / ``self.waiting`` to the
+    same shapes the upstream ``Scheduler.finish_requests`` reads, so the
+    helper sees realistic state without spinning a real scheduler.
+    """
+
+    def __init__(
+        self,
+        *,
+        requests: dict,
+        running: list,
+        waiting: list,
+    ) -> None:
+        self.requests = requests
+        self.running = running
+        self.waiting = waiting
+
+
+class TestRealignRequestStatusToQueues:
+    """Regression for the residual hang described in
+    https://github.com/vllm-project/vllm-omni/pull/3774 -- chunk-transfer
+    adapter's ``requests_origin_status`` table goes stale on the
+    ``waiting → running`` admit transition, and an abort that lands
+    before the next deque round-trip stomps stale ``WAITING`` onto a
+    request that actually lives in ``self.running``. After
+    ``max_num_seqs`` such aborts every ``input_batch`` slot is leaked and
+    new requests hang at ``chunks=0``.
+    """
+
+    def test_running_with_stale_waiting_status_is_realigned_to_running(
+        self,
+    ) -> None:
+        """The exact race the hang reproduces: a request lives in
+        ``self.running`` but its ``status`` is ``WAITING``. After
+        realign, status must be ``RUNNING`` so upstream
+        ``Scheduler.finish_requests`` removes it from ``self.running``.
+        """
+        req = _make_request(request_id="req-stale")
+        req.status = RequestStatus.WAITING  # stale -- actually in running
+
+        sched = _RealignSchedulerStub(
+            requests={req.request_id: req},
+            running=[req],
+            waiting=[],
+        )
+        sched._realign_request_status_to_queues([req.request_id])
+
+        assert req.status == RequestStatus.RUNNING
+
+    def test_waiting_with_stale_running_status_is_realigned_to_waiting(
+        self,
+    ) -> None:
+        """Symmetric case: request lives in ``self.waiting`` but status
+        is ``RUNNING``. Upstream's else branch should run, so realign to
+        ``WAITING``.
+        """
+        req = _make_request(request_id="req-stale-2")
+        req.status = RequestStatus.RUNNING  # stale -- actually in waiting
+
+        sched = _RealignSchedulerStub(
+            requests={req.request_id: req},
+            running=[],
+            waiting=[req],
+        )
+        sched._realign_request_status_to_queues([req.request_id])
+
+        assert req.status == RequestStatus.WAITING
+
+    def test_already_aligned_status_is_left_unchanged(self) -> None:
+        """Healthy case: status matches actual queue. No status mutation
+        means no spurious side effects.
+        """
+        req_running = _make_request(request_id="req-r")
+        req_running.status = RequestStatus.RUNNING
+        req_waiting = _make_request(request_id="req-w")
+        req_waiting.status = RequestStatus.WAITING
+
+        sched = _RealignSchedulerStub(
+            requests={
+                req_running.request_id: req_running,
+                req_waiting.request_id: req_waiting,
+            },
+            running=[req_running],
+            waiting=[req_waiting],
+        )
+        sched._realign_request_status_to_queues([req_running.request_id, req_waiting.request_id])
+
+        assert req_running.status == RequestStatus.RUNNING
+        assert req_waiting.status == RequestStatus.WAITING
+
+    def test_unknown_request_id_is_skipped_silently(self) -> None:
+        """Aligning ids that aren't in ``self.requests`` (already freed
+        upstream, or never existed) must be a no-op.
+        """
+        sched = _RealignSchedulerStub(requests={}, running=[], waiting=[])
+        sched._realign_request_status_to_queues(["never-existed"])  # no raise
+
+    def test_request_in_neither_queue_is_left_unchanged(self) -> None:
+        """If a tracked request is in neither queue (e.g. parked in a
+        chunk-transfer deque), realign must not invent a status. The
+        adapter / deque purge owns that surface; the realign helper is
+        only for the admit-transition staleness.
+        """
+        req = _make_request(request_id="req-parked")
+        req.status = RequestStatus.WAITING_FOR_CHUNK
+
+        sched = _RealignSchedulerStub(
+            requests={req.request_id: req},
+            running=[],
+            waiting=[],
+        )
+        sched._realign_request_status_to_queues([req.request_id])
+
+        assert req.status == RequestStatus.WAITING_FOR_CHUNK
+
+    def test_request_ids_str_is_treated_as_single_id(self) -> None:
+        """Match upstream ``Scheduler.finish_requests`` resolution: a
+        bare string is treated as one id, not iterated as characters.
+        """
+        req = _make_request(request_id="req-s")
+        req.status = RequestStatus.WAITING
+
+        sched = _RealignSchedulerStub(
+            requests={req.request_id: req},
+            running=[req],
+            waiting=[],
+        )
+        sched._realign_request_status_to_queues(req.request_id)
+
+        assert req.status == RequestStatus.RUNNING
+
+    def test_request_ids_none_aligns_every_known_request(self) -> None:
+        """``request_ids=None`` matches upstream's "all requests" path.
+        Realign must walk every entry in ``self.requests`` and fix any
+        stale status it finds.
+        """
+        stale = _make_request(request_id="req-stale-none")
+        stale.status = RequestStatus.WAITING  # but actually in running
+
+        clean = _make_request(request_id="req-clean-none")
+        clean.status = RequestStatus.RUNNING
+
+        sched = _RealignSchedulerStub(
+            requests={
+                stale.request_id: stale,
+                clean.request_id: clean,
+            },
+            running=[stale, clean],
+            waiting=[],
+        )
+        sched._realign_request_status_to_queues(None)
+
+        assert stale.status == RequestStatus.RUNNING
+        assert clean.status == RequestStatus.RUNNING
+
+    def test_finished_request_is_skipped(self) -> None:
+        """Already-finished requests must not be touched -- they may
+        have legitimate finished statuses (FINISHED_STOPPED etc.) that
+        a status flip would corrupt.
+        """
+        req = _make_request(request_id="req-finished")
+        req.status = RequestStatus.FINISHED_STOPPED
+
+        sched = _RealignSchedulerStub(
+            requests={req.request_id: req},
+            running=[req],  # not realistic, but we want to prove the guard fires
+            waiting=[],
+        )
+        sched._realign_request_status_to_queues([req.request_id])
+
+        assert req.status == RequestStatus.FINISHED_STOPPED
+
+
+class TestPurgeFinishedFromRunning:
+    """Regression for the residual ``self.running`` slot leak surface
+    paired with ``_realign_request_status_to_queues``: even after
+    realign + ``super().finish_requests`` runs, corner cases can leave
+    already-finished or untracked entries in ``self.running`` -- e.g.
+    a connector cleanup that pops from ``self.requests`` without
+    unwinding ``self.running``, or a ``status`` mid-transition when
+    finish ran. The defensive post-finish purge sweeps those residues
+    so the worker's ``input_batch`` slot never pins a freed request.
+
+    See https://github.com/vllm-project/vllm-omni/pull/3774 discussion
+    and the residual-hang reproduction in that PR.
+    """
+
+    def test_finished_request_is_purged_from_running(self) -> None:
+        """``is_finished()`` request lingering in ``self.running`` must
+        be swept so its ``input_batch`` slot is freed."""
+        finished = _make_request(request_id="req-finished")
+        finished.status = RequestStatus.FINISHED_STOPPED  # is_finished() True
+
+        sched = _RealignSchedulerStub(
+            requests={finished.request_id: finished},
+            running=[finished],
+            waiting=[],
+        )
+        sched._purge_finished_from_running()
+
+        assert sched.running == []
+
+    def test_untracked_request_is_purged_from_running(self) -> None:
+        """Request lingering in ``self.running`` but no longer present
+        in ``self.requests`` (already freed by upstream / connector)
+        must be swept."""
+        untracked = _make_request(request_id="req-untracked")
+        untracked.status = RequestStatus.RUNNING
+
+        sched = _RealignSchedulerStub(
+            requests={},  # already deleted from self.requests
+            running=[untracked],
+            waiting=[],
+        )
+        sched._purge_finished_from_running()
+
+        assert sched.running == []
+
+    def test_healthy_running_is_left_unchanged(self) -> None:
+        """Live tracked running requests must not be touched -- the
+        purge is defensive, not aggressive."""
+        alive = _make_request(request_id="req-alive")
+        alive.status = RequestStatus.RUNNING
+
+        sched = _RealignSchedulerStub(
+            requests={alive.request_id: alive},
+            running=[alive],
+            waiting=[],
+        )
+        sched._purge_finished_from_running()
+
+        assert sched.running == [alive]
+
+    def test_empty_running_is_noop(self) -> None:
+        """Empty ``self.running`` must short-circuit cleanly."""
+        sched = _RealignSchedulerStub(requests={}, running=[], waiting=[])
+        sched._purge_finished_from_running()
+
+        assert sched.running == []
+
+    def test_mixed_alive_and_dead_keeps_only_alive_in_order(self) -> None:
+        """Mixed ``self.running`` -- some alive, some finished, some
+        untracked. Sweep keeps only the alive ones and preserves their
+        relative order."""
+        alive_a = _make_request(request_id="req-alive-a")
+        alive_a.status = RequestStatus.RUNNING
+        finished_b = _make_request(request_id="req-finished-b")
+        finished_b.status = RequestStatus.FINISHED_STOPPED
+        untracked_c = _make_request(request_id="req-untracked-c")
+        untracked_c.status = RequestStatus.RUNNING
+        alive_d = _make_request(request_id="req-alive-d")
+        alive_d.status = RequestStatus.RUNNING
+
+        sched = _RealignSchedulerStub(
+            requests={
+                alive_a.request_id: alive_a,
+                finished_b.request_id: finished_b,  # tracked but is_finished
+                # untracked_c is NOT in self.requests
+                alive_d.request_id: alive_d,
+            },
+            running=[alive_a, finished_b, untracked_c, alive_d],
+            waiting=[],
+        )
+        sched._purge_finished_from_running()
+
+        assert sched.running == [alive_a, alive_d]
+
+    def test_in_place_mutation_keeps_helper_local_list_identity(self) -> None:
+        """The slice-assign form (``self.running[:] = ...``) avoids an
+        extra rebind inside this helper. Note that across a full
+        ``finish_requests`` call upstream ``Scheduler.finish_requests``
+        already rebinds ``self.running``, so global identity is not a
+        contract -- this test only pins the local helper behavior so a
+        future refactor away from in-place mutation is intentional.
+        """
+        finished = _make_request(request_id="req-finished")
+        finished.status = RequestStatus.FINISHED_STOPPED
+        alive = _make_request(request_id="req-alive")
+        alive.status = RequestStatus.RUNNING
+
+        sched = _RealignSchedulerStub(
+            requests={
+                finished.request_id: finished,
+                alive.request_id: alive,
+            },
+            running=[finished, alive],
+            waiting=[],
+        )
+        held_ref = sched.running
+        sched._purge_finished_from_running()
+
+        assert held_ref is sched.running
+        assert held_ref == [alive]

tests/core/sched/test_omni_scheduler_finish_requests_purge.py

@@ -0,0 +1,158 @@
+"""Integration tests for the post-finish defensive purge of
+``self.running`` placed in
+``OmniARScheduler.finish_requests`` /
+``OmniGenerationScheduler.finish_requests``.
+
+Pins *placement* — that ``_purge_finished_from_running`` runs AFTER
+``super().finish_requests`` so it catches entries upstream's removal
+missed (e.g. status mid-transition, connector cleanup popping
+``self.requests`` without unwinding ``self.running``). Direct helper
+unit tests live in ``test_omni_scheduler_mixin.py``; these tests
+exercise the full ``finish_requests`` path end-to-end.
+
+Regression for the residual self.running slot leak surface paired
+with ``_realign_request_status_to_queues`` (vllm-omni#3774).
+"""
+
+from __future__ import annotations
+
+import pytest
+from vllm.v1.request import RequestStatus
+
+import vllm_omni.core.sched.omni_ar_scheduler as ar_sched_mod
+import vllm_omni.core.sched.omni_generation_scheduler as gen_sched_mod
+from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
+from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _StubRequest:
+    """Minimal Request stub with the surface the helpers exercise."""
+
+    def __init__(self, request_id: str, status: RequestStatus) -> None:
+        self.request_id = request_id
+        self.status = status
+
+    def is_finished(self) -> bool:
+        return RequestStatus.is_finished(self.status)
+
+
+_SCHEDULER_PARAMS = [
+    pytest.param(OmniARScheduler, ar_sched_mod, id="ar"),
+    pytest.param(OmniGenerationScheduler, gen_sched_mod, id="generation"),
+]
+
+
+def _make_scheduler(scheduler_cls, *, requests, running, waiting):
+    """Bare scheduler with just the surface ``finish_requests`` reads."""
+    scheduler = scheduler_cls.__new__(scheduler_cls)
+    scheduler.chunk_transfer_adapter = None
+    scheduler.input_coordinator = None
+    scheduler.requests = requests
+    scheduler.running = running
+    scheduler.waiting = waiting
+    return scheduler
+
+
+@pytest.mark.parametrize(("scheduler_cls", "scheduler_mod"), _SCHEDULER_PARAMS)
+def test_finish_requests_purges_finished_request_from_running(
+    monkeypatch: pytest.MonkeyPatch,
+    scheduler_cls,
+    scheduler_mod,
+) -> None:
+    """A request whose ``status`` is a finished value but still lives
+    in ``self.running`` after ``super().finish_requests`` returns must
+    be swept by the post-finish purge. Mocks ``super()`` to simulate
+    the corner case where upstream's status-driven removal missed the
+    entry.
+    """
+    finished = _StubRequest("req-finished", RequestStatus.FINISHED_STOPPED)
+    scheduler = _make_scheduler(
+        scheduler_cls,
+        requests={finished.request_id: finished},
+        running=[finished],
+        waiting=[],
+    )
+
+    def fake_finish_requests(self, request_ids, finished_status):
+        # super() returning the abort tuple but leaving self.running untouched
+        return [(finished.request_id, 0)]
+
+    monkeypatch.setattr(scheduler_mod.VLLMScheduler, "finish_requests", fake_finish_requests)
+
+    result = scheduler_cls.finish_requests(
+        scheduler,
+        [finished.request_id],
+        RequestStatus.FINISHED_STOPPED,
+    )
+
+    assert scheduler.running == []
+    assert result == [(finished.request_id, 0)]
+
+
+@pytest.mark.parametrize(("scheduler_cls", "scheduler_mod"), _SCHEDULER_PARAMS)
+def test_finish_requests_purges_untracked_request_from_running(
+    monkeypatch: pytest.MonkeyPatch,
+    scheduler_cls,
+    scheduler_mod,
+) -> None:
+    """A request that has been deleted from ``self.requests`` (e.g.
+    after upstream's ``_free_request`` ran but a status mismatch
+    caused the ``self.running`` removal branch to no-op) must be
+    swept by the post-finish purge.
+    """
+    untracked = _StubRequest("req-untracked", RequestStatus.RUNNING)
+    scheduler = _make_scheduler(
+        scheduler_cls,
+        requests={},  # already deleted from self.requests
+        running=[untracked],
+        waiting=[],
+    )
+
+    def fake_finish_requests(self, request_ids, finished_status):
+        return [(untracked.request_id, 0)]
+
+    monkeypatch.setattr(scheduler_mod.VLLMScheduler, "finish_requests", fake_finish_requests)
+
+    result = scheduler_cls.finish_requests(
+        scheduler,
+        [untracked.request_id],
+        RequestStatus.FINISHED_STOPPED,
+    )
+
+    assert scheduler.running == []
+    assert result == [(untracked.request_id, 0)]
+
+
+@pytest.mark.parametrize(("scheduler_cls", "scheduler_mod"), _SCHEDULER_PARAMS)
+def test_finish_requests_leaves_healthy_running_intact(
+    monkeypatch: pytest.MonkeyPatch,
+    scheduler_cls,
+    scheduler_mod,
+) -> None:
+    """A live ``RUNNING`` request that is not in the finish set must
+    not be swept by the defensive purge -- this pins that the predicate
+    is conservative and won't aggressively drop healthy entries.
+    """
+    alive = _StubRequest("req-alive", RequestStatus.RUNNING)
+    leaving_id = "req-leaving"
+    scheduler = _make_scheduler(
+        scheduler_cls,
+        requests={alive.request_id: alive},
+        running=[alive],  # leaving already removed by super() (typical happy path)
+        waiting=[],
+    )
+
+    def fake_finish_requests(self, request_ids, finished_status):
+        return [(leaving_id, 0)]
+
+    monkeypatch.setattr(scheduler_mod.VLLMScheduler, "finish_requests", fake_finish_requests)
+
+    scheduler_cls.finish_requests(
+        scheduler,
+        [leaving_id],
+        RequestStatus.FINISHED_STOPPED,
+    )
+
+    assert scheduler.running == [alive]

tests/core/sched/test_omni_scheduler_input_coordinator_cleanup.py

@@ -46,6 +46,9 @@ def test_finish_requests_cleans_input_coordinator_for_finished_ids(
     scheduler = scheduler_cls.__new__(scheduler_cls)
     scheduler.chunk_transfer_adapter = None
     scheduler.input_coordinator = coordinator
+    scheduler.requests = {}
+    scheduler.running = []
+    scheduler.waiting = []
 
     def fake_finish_requests(self, request_ids, finished_status):
         assert request_ids == ["req-a", "req-b"]

tests/dfx/perf/tests/test_tts.json

@@ -339,5 +339,66 @@
         }
       }
     ]
+  },
+  {
+    "test_name": "test_voxtral_tts",
+    "server_params": {
+      "model": "mistralai/Voxtral-4B-TTS-2603"
+    },
+    "benchmark_params": [
+      {
+        "task": "default_voice",
+        "eval_phase": "latency",
+        "dataset_name": "seed-tts-text",
+        "backend": "openai-audio-speech",
+        "endpoint": "/v1/audio/speech",
+        "dataset_path": "benchmarks/build_dataset/seed_tts_smoke",
+        "num_prompts": [20],
+        "max_concurrency": [1],
+        "seed_tts_locale": "en",
+        "extra_body": {"voice": "casual_male"},
+        "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
+        "baseline": {
+          "median_audio_ttfp_ms": [150],
+          "median_audio_rtf": [0.15]
+        }
+      },
+      {
+        "task": "default_voice",
+        "eval_phase": "throughput",
+        "dataset_name": "seed-tts-text",
+        "backend": "openai-audio-speech",
+        "endpoint": "/v1/audio/speech",
+        "dataset_path": "benchmarks/build_dataset/seed_tts_smoke",
+        "num_prompts": [80],
+        "max_concurrency": [8],
+        "seed_tts_locale": "en",
+        "extra_body": {"voice": "casual_male"},
+        "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
+        "baseline": {
+          "median_audio_ttfp_ms": [1500],
+          "median_audio_rtf": [0.30],
+          "audio_throughput": [30.0]
+        }
+      },
+      {
+        "task": "default_voice",
+        "eval_phase": "stress",
+        "dataset_name": "seed-tts-text",
+        "backend": "openai-audio-speech",
+        "endpoint": "/v1/audio/speech",
+        "dataset_path": "benchmarks/build_dataset/seed_tts_smoke",
+        "num_prompts": [100],
+        "request_rate": [2.0],
+        "seed_tts_locale": "en",
+        "extra_body": {"voice": "casual_male"},
+        "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
+        "baseline": {
+          "median_audio_ttfp_ms": [3000],
+          "median_audio_rtf": [0.25],
+          "audio_throughput": [4.0]
+        }
+      }
+    ]
   }
 ]

tests/dfx/reliability/test_reliability_qwen3_omni.py

@@ -319,6 +319,7 @@ QWEN_PARAMS = create_reliability_omni_server_params(RELIABILITY_SCENARIOS, DEPLO
 
 
 @pytest.mark.slow
+@pytest.mark.skip(reason="issue#4285")
 @hardware_test(res={"cuda": "H100"}, num_cards=2)
 @pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
 def test_reliability_fault_gpu_oom_error_contract_consistent_chat_speech(
@@ -387,6 +388,7 @@ def test_reliability_fault_gpu_oom_error_contract_consistent_chat_speech(
 
 
 @pytest.mark.slow
+@pytest.mark.skip(reason="issue#4285")
 @hardware_test(res={"cuda": "H100"}, num_cards=2)
 @pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
 def test_reliability_fault_gpu_oom_chat_large_payload_failure(omni_server_function, openai_client_function) -> None:
@@ -429,6 +431,7 @@ def test_reliability_fault_gpu_oom_chat_large_payload_failure(omni_server_functi
 
 
 @pytest.mark.slow
+@pytest.mark.skip(reason="issue#4285")
 @hardware_test(res={"cuda": "H100"}, num_cards=2)
 @pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
 def test_reliability_fault_gpu_oom_concurrent_pressure_failure(omni_server_function, openai_client_function) -> None:
@@ -771,7 +774,7 @@ def test_reliability_fault_process_kill_tree_with_load_fast_fail_and_cleanup(
 
 
 @pytest.mark.slow
-@pytest.mark.skip(reason="OOM injection causes the serving instance to exit")
+@pytest.mark.skip(reason="issue#4285")
 @hardware_test(res={"cuda": "H100"}, num_cards=2)
 @pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
 def test_reliability_fault_gpu_oom_state_converges_after_fault_removed(

tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py

@@ -97,6 +97,260 @@ class TestPipelineIndependence:
         assert modules.resident_names == ["vocoder"]
         assert len(modules.vaes) == 2
 
+    def test_ltx23_pipeline_has_diffusion_pipeline_profiler_mixin(self):
+        """LTX23Pipeline must support lightweight stage timing."""
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
+        from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
+
+        assert issubclass(LTX23Pipeline, DiffusionPipelineProfilerMixin)
+
+
+class TestLTX23VaeDecodeParallel:
+    """Test LTX-2.3 video VAE tiled parallel helpers without loading weights."""
+
+    def test_ltx23_video_vae_is_distributed_tile_only_class(self):
+        from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import (
+            DistributedAutoencoderKLLTX2Video,
+        )
+        from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import DistributedVaeMixin
+
+        assert issubclass(DistributedAutoencoderKLLTX2Video, DistributedVaeMixin)
+        assert not hasattr(DistributedAutoencoderKLLTX2Video, "patch_split")
+
+    def test_ltx23_video_vae_tile_split_uses_native_ltx23_tile_geometry(self):
+        from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import (
+            DistributedAutoencoderKLLTX2Video,
+        )
+
+        vae = SimpleNamespace(
+            spatial_compression_ratio=32,
+            tile_sample_min_height=512,
+            tile_sample_min_width=512,
+            tile_sample_stride_height=448,
+            tile_sample_stride_width=448,
+            temporal_compression_ratio=8,
+            dtype=torch.float32,
+        )
+
+        z = torch.zeros(1, 2, 5, 16, 24)
+        tasks, grid_spec = DistributedAutoencoderKLLTX2Video.tile_split(vae, z)
+
+        assert grid_spec.grid_shape == (2, 2)
+        assert grid_spec.split_dims == (3, 4)
+        assert grid_spec.tile_spec["sample_height"] == 512
+        assert grid_spec.tile_spec["sample_width"] == 768
+        assert grid_spec.tile_spec["blend_height"] == 64
+        assert grid_spec.tile_spec["blend_width"] == 64
+        assert grid_spec.tile_spec["max_tile_output_shape"] == (1, 3, 33, 512, 512)
+        assert grid_spec.tile_spec["tile_output_shapes"] == {
+            0: (1, 3, 33, 512, 512),
+            1: (1, 3, 33, 512, 320),
+            2: (1, 3, 33, 64, 512),
+            3: (1, 3, 33, 64, 320),
+        }
+        assert [task.grid_coord for task in tasks] == [(0, 0), (0, 1), (1, 0), (1, 1)]
+        assert [tuple(task.tensor.shape) for task in tasks] == [
+            (1, 2, 5, 16, 16),
+            (1, 2, 5, 16, 10),
+            (1, 2, 5, 2, 16),
+            (1, 2, 5, 2, 10),
+        ]
+        assert [task.workload for task in tasks] == [5 * 16 * 16, 5 * 16 * 10, 5 * 2 * 16, 5 * 2 * 10]
+
+    def test_ltx23_video_vae_tile_merge_blends_and_crops_like_tiled_decode(self):
+        from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import (
+            DistributedAutoencoderKLLTX2Video,
+        )
+        from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import GridSpec
+
+        class FakeVae:
+            def __init__(self):
+                self.blend_calls = []
+
+            def clear_cache(self):
+                pass
+
+            def blend_v(self, _previous, current, blend_height):
+                self.blend_calls.append(("v", blend_height))
+                return current
+
+            def blend_h(self, _previous, current, blend_width):
+                self.blend_calls.append(("h", blend_width))
+                return current
+
+        fake_vae = FakeVae()
+        grid_spec = GridSpec(
+            split_dims=(3, 4),
+            grid_shape=(2, 2),
+            tile_spec={
+                "sample_height": 10,
+                "sample_width": 10,
+                "blend_height": 1,
+                "blend_width": 2,
+                "tile_sample_stride_height": 5,
+                "tile_sample_stride_width": 5,
+            },
+        )
+        tiles = {
+            (0, 0): torch.full((1, 3, 2, 6, 6), 1.0),
+            (0, 1): torch.full((1, 3, 2, 6, 6), 2.0),
+            (1, 0): torch.full((1, 3, 2, 6, 6), 3.0),
+            (1, 1): torch.full((1, 3, 2, 6, 6), 4.0),
+        }
+
+        merged = DistributedAutoencoderKLLTX2Video.tile_merge(fake_vae, tiles, grid_spec)
+
+        assert merged.shape == (1, 3, 2, 10, 10)
+        assert fake_vae.blend_calls == [("h", 2), ("v", 1), ("v", 1), ("h", 2)]
+        torch.testing.assert_close(merged[:, :, :, :5, :5], torch.ones(1, 3, 2, 5, 5))
+        torch.testing.assert_close(merged[:, :, :, :5, 5:], torch.full((1, 3, 2, 5, 5), 2.0))
+        torch.testing.assert_close(merged[:, :, :, 5:, :5], torch.full((1, 3, 2, 5, 5), 3.0))
+        torch.testing.assert_close(merged[:, :, :, 5:, 5:], torch.full((1, 3, 2, 5, 5), 4.0))
+
+    def test_ltx23_video_vae_tiled_decode_dispatches_to_tile_operator(self):
+        from vllm_omni.diffusion.distributed.autoencoders import autoencoder_kl_ltx2
+
+        z = torch.zeros(1, 2, 1, 16, 24)
+        expected = torch.ones(1, 3, 1, 512, 768)
+        seen = {}
+
+        class FakeExecutor:
+            def execute(self, tensor, operator, broadcast_result=True):
+                seen["tensor"] = tensor
+                seen["operator"] = operator
+                seen["broadcast_result"] = broadcast_result
+                return expected
+
+        vae = SimpleNamespace(distributed_executor=FakeExecutor(), is_distributed_enabled=lambda: True)
+        vae.tile_split = autoencoder_kl_ltx2.DistributedAutoencoderKLLTX2Video.tile_split.__get__(vae)
+        vae.tile_exec = autoencoder_kl_ltx2.DistributedAutoencoderKLLTX2Video.tile_exec.__get__(vae)
+        vae.tile_merge = autoencoder_kl_ltx2.DistributedAutoencoderKLLTX2Video.tile_merge.__get__(vae)
+
+        output = autoencoder_kl_ltx2.DistributedAutoencoderKLLTX2Video.tiled_decode(
+            vae,
+            z,
+            temb=torch.tensor(0.5),
+            return_dict=False,
+        )
+
+        assert len(output) == 1
+        assert output[0] is expected
+        assert seen["tensor"] is z
+        assert seen["broadcast_result"] is False
+        assert seen["operator"].split.__name__ == "tile_split"
+        assert seen["operator"].merge.__name__ == "tile_merge"
+
+    def test_ltx23_vae_executor_gathers_known_tile_shapes_and_returns_empty_on_non_rank0(self):
+        from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import LTX2VaeExecutor
+        from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
+            DistributedOperator,
+            GridSpec,
+            TileTask,
+        )
+
+        z = torch.zeros(1, 1, 1, 1, 1)
+        tile_output_shapes = {
+            0: (1, 1, 1, 2, 2),
+            1: (1, 1, 1, 2, 1),
+            2: (1, 1, 1, 1, 2),
+            3: (1, 1, 1, 1, 1),
+        }
+        tasks = [
+            TileTask(0, (0, 0), z, workload=4),
+            TileTask(1, (0, 1), z, workload=2),
+            TileTask(2, (1, 0), z, workload=2),
+            TileTask(3, (1, 1), z, workload=1),
+        ]
+        grid_spec = GridSpec(
+            split_dims=(3, 4),
+            grid_shape=(2, 2),
+            tile_spec={
+                "max_tile_output_shape": (1, 1, 1, 2, 2),
+                "tile_output_shapes": tile_output_shapes,
+            },
+            output_dtype=torch.float32,
+        )
+        seen = {}
+
+        def exec_tile(task):
+            return torch.full(tile_output_shapes[task.tile_id], float(task.tile_id + 1))
+
+        def merge_tiles(coord_tensor_map, passed_grid_spec):
+            seen["merged_shapes"] = {coord: tuple(tile.shape) for coord, tile in coord_tensor_map.items()}
+            assert passed_grid_spec is grid_spec
+            return torch.stack(
+                [
+                    coord_tensor_map[(0, 0)].flatten()[0],
+                    coord_tensor_map[(0, 1)].flatten()[0],
+                    coord_tensor_map[(1, 0)].flatten()[0],
+                    coord_tensor_map[(1, 1)].flatten()[0],
+                ]
+            )
+
+        operator = DistributedOperator(split=lambda _z: (tasks, grid_spec), exec=exec_tile, merge=merge_tiles)
+
+        rank0_executor = object.__new__(LTX2VaeExecutor)
+        rank0_executor.parallel_size = 2
+        rank0_executor.world_size = 2
+        rank0_executor.rank = 0
+
+        def gather_rank0(local_tile_tensor):
+            assigned = rank0_executor._balance_tasks(tasks, 2)
+            rank1_results = [(task.tile_id, exec_tile(task)) for task in assigned[1]]
+            rank1_tile_tensor = rank0_executor._pack_local_tiles_without_meta(
+                rank1_results,
+                list(local_tile_tensor.shape),
+                z.device,
+                torch.float32,
+            )
+            seen["rank0_gather_shape"] = tuple(local_tile_tensor.shape)
+            return [local_tile_tensor, rank1_tile_tensor]
+
+        def fail_final_sync(*_args, **_kwargs):
+            raise AssertionError("broadcast_result=False should not sync the final result")
+
+        rank0_executor.gather_tensors = gather_rank0
+        rank0_executor._sync_final_result = fail_final_sync
+
+        rank0_result = rank0_executor.execute(z, operator, broadcast_result=False)
+
+        torch.testing.assert_close(rank0_result, torch.tensor([1.0, 2.0, 3.0, 4.0]))
+        assert seen["rank0_gather_shape"] == (2, 1, 1, 1, 2, 2)
+        assert seen["merged_shapes"] == {
+            (0, 0): (1, 1, 1, 2, 2),
+            (0, 1): (1, 1, 1, 2, 1),
+            (1, 0): (1, 1, 1, 1, 2),
+            (1, 1): (1, 1, 1, 1, 1),
+        }
+
+        non_rank0_executor = object.__new__(LTX2VaeExecutor)
+        non_rank0_executor.parallel_size = 2
+        non_rank0_executor.world_size = 2
+        non_rank0_executor.rank = 1
+
+        def gather_rank1(local_tile_tensor):
+            seen["rank1_gather_shape"] = tuple(local_tile_tensor.shape)
+            return None
+
+        def fail_non_rank0_merge(*_args, **_kwargs):
+            raise AssertionError("non-rank0 should not merge gathered tiles")
+
+        non_rank0_executor.gather_tensors = gather_rank1
+        non_rank0_executor._sync_final_result = fail_final_sync
+
+        empty_result = non_rank0_executor.execute(
+            z,
+            DistributedOperator(
+                split=lambda _z: (tasks, grid_spec),
+                exec=exec_tile,
+                merge=fail_non_rank0_merge,
+            ),
+            broadcast_result=False,
+        )
+
+        assert tuple(empty_result.shape) == (0,)
+        assert seen["rank1_gather_shape"] == (2, 1, 1, 1, 2, 2)
+
 
 class TestCFGParallelHelpers:
     """Test LTX-2.3 CFG helper math without loading model weights."""

tests/distributed/omni_connectors/test_chunk_transfer_adapter.py

@@ -387,14 +387,15 @@ def test_process_and_restore_queues(build_adapter):
     running_req = _req("r1", RequestStatus.RUNNING)
     waiting_queue = DummyWaitingQueue([waiting_req])
     running_queue = [running_req]
+    scheduler_requests = {waiting_req.request_id: waiting_req, running_req.request_id: running_req}
 
-    adapter.process_pending_chunks(waiting_queue, running_queue)
+    adapter.process_pending_chunks(waiting_queue, running_queue, scheduler_requests=scheduler_requests)
     assert waiting_req.status == RequestStatus.WAITING_FOR_CHUNK
     assert running_req.status == RequestStatus.WAITING_FOR_CHUNK
     assert waiting_queue == []
     assert running_queue == []
 
-    adapter.restore_queues(waiting_queue, running_queue)
+    adapter.restore_queues(waiting_queue, running_queue, scheduler_requests=scheduler_requests)
     assert waiting_queue == [waiting_req]
     assert running_queue == [running_req]
     assert adapter.waiting_for_chunk_waiting_requests == deque()
@@ -897,6 +898,8 @@ def test_omni_ar_scheduler_finish_requests(mocker: MockerFixture):
     sched = OmniARScheduler.__new__(OmniARScheduler)
     sched.chunk_transfer_adapter = adapter
     sched.requests = {}
+    sched.running = []
+    sched.waiting = []
 
     with patch.object(VLLMScheduler, "finish_requests", _super_finish):
         OmniARScheduler.finish_requests(sched, ["r1"], RequestStatus.FINISHED_ABORTED)
@@ -1170,3 +1173,164 @@ def test_deferred_finish_not_finished_still_emits_output(mocker: MockerFixture):
     assert eco.outputs[0].finish_reason == "stop"
     assert eco.outputs[0].kv_transfer_params is None
     assert scheduler._pending_finish_reqs == []
+
+
+# ---------------------------------------------------------------
+# Zombie purge tests (regression for vllm-project/vllm-omni#3736)
+# ---------------------------------------------------------------
+
+
+def test_process_pending_chunks_purges_zombies_in_running_deque(
+    build_adapter,
+):
+    """A request aborted while parked in ``waiting_for_chunk_running_requests``
+    is no longer tracked by ``scheduler.requests`` once
+    ``Scheduler._free_request`` runs. ``process_pending_chunks`` must drop
+    that zombie *before* ``restore_queues`` would re-inject it onto the
+    running queue.
+
+    Regression for https://github.com/vllm-project/vllm-omni/issues/3736
+    (engine-core ``KeyError`` on aborts under chunk-transfer pressure).
+    """
+    adapter, _ = build_adapter(stage_id=1, max_num_seqs=8)
+
+    live_req = _req("live-1", RequestStatus.WAITING_FOR_CHUNK)
+    zombie_req = _req("zombie-1", RequestStatus.WAITING_FOR_CHUNK)
+    adapter.waiting_for_chunk_running_requests.append(live_req)
+    adapter.waiting_for_chunk_running_requests.append(zombie_req)
+    # Mirror state the adapter would carry for an in-flight request so we can
+    # later assert ``cleanup_receiver`` actually fired against the zombie.
+    adapter.requests_origin_status[zombie_req.request_id] = RequestStatus.RUNNING
+    adapter.requests_origin_status[live_req.request_id] = RequestStatus.RUNNING
+
+    waiting_queue = DummyWaitingQueue()
+    running_queue: list = []
+    # Simulate the post-abort scheduler: only ``live-1`` is still tracked.
+    scheduler_requests = {live_req.request_id: live_req}
+
+    adapter.process_pending_chunks(waiting_queue, running_queue, scheduler_requests=scheduler_requests)
+
+    # 1. Zombie was popped from the deque.
+    assert zombie_req not in adapter.waiting_for_chunk_running_requests
+    # 2. Live request is still in the deque.
+    assert live_req in adapter.waiting_for_chunk_running_requests
+    # 3. ``cleanup_receiver`` ran for the zombie (drops origin-status mapping
+    #    and registers the id as cancelled so a late load/poll is dropped too).
+    assert zombie_req.request_id not in adapter.requests_origin_status
+    assert zombie_req.request_id in adapter._cancelled_load_reqs
+    # 4. Live request's bookkeeping is untouched.
+    assert adapter.requests_origin_status[live_req.request_id] == RequestStatus.RUNNING
+    assert live_req.request_id not in adapter._cancelled_load_reqs
+
+    # 5. ``restore_queues`` (which the scheduler runs in its ``finally``
+    #    clause) now only re-injects the live request -- the zombie is
+    #    gone, so the worker's ``_update_states`` cannot crash on it.
+    adapter.restore_queues(waiting_queue, running_queue, scheduler_requests=scheduler_requests)
+    assert running_queue == [live_req]
+
+
+def test_process_pending_chunks_purges_zombies_in_waiting_deque(build_adapter):
+    """Zombie purge applies symmetrically to the waiting-side deque.
+
+    Regression for https://github.com/vllm-project/vllm-omni/issues/3736 --
+    aborted requests that landed in ``waiting_for_chunk_waiting_requests``
+    must also be removed before ``restore_queues`` re-injects them, since
+    ``restore_queues`` uses ``waiting_queue.add_request`` rather than
+    ``running_queue.extend`` for that path.
+    """
+    adapter, _ = build_adapter(stage_id=1, max_num_seqs=8)
+
+    live_req = _req("live-w", RequestStatus.WAITING_FOR_CHUNK)
+    zombie_req = _req("zombie-w", RequestStatus.WAITING_FOR_CHUNK)
+    adapter.waiting_for_chunk_waiting_requests.append(live_req)
+    adapter.waiting_for_chunk_waiting_requests.append(zombie_req)
+
+    waiting_queue = DummyWaitingQueue()
+    running_queue: list = []
+    scheduler_requests = {live_req.request_id: live_req}
+
+    adapter.process_pending_chunks(waiting_queue, running_queue, scheduler_requests=scheduler_requests)
+
+    assert live_req in adapter.waiting_for_chunk_waiting_requests
+    assert zombie_req not in adapter.waiting_for_chunk_waiting_requests
+    assert zombie_req.request_id in adapter._cancelled_load_reqs
+
+    adapter.restore_queues(waiting_queue, running_queue, scheduler_requests=scheduler_requests)
+    assert waiting_queue == [live_req]
+
+
+def test_purge_preserves_live_order_with_interleaved_zombies(build_adapter):
+    """Interleaved live/zombie ordering -- pin the ``popleft`` + append-live
+    in-place filter at chunk_transfer_adapter._purge_untracked_chunk_requests:
+    survivor order must match insertion order, and ``cleanup_receiver`` must
+    run exactly once per zombie.
+    """
+    adapter, _ = build_adapter(stage_id=1, max_num_seqs=8)
+
+    live1 = _req("live-1", RequestStatus.WAITING_FOR_CHUNK)
+    zombie1 = _req("zombie-1", RequestStatus.WAITING_FOR_CHUNK)
+    live2 = _req("live-2", RequestStatus.WAITING_FOR_CHUNK)
+    zombie2 = _req("zombie-2", RequestStatus.WAITING_FOR_CHUNK)
+    for req in (live1, zombie1, live2, zombie2):
+        adapter.waiting_for_chunk_running_requests.append(req)
+
+    waiting_queue = DummyWaitingQueue()
+    running_queue: list = []
+    scheduler_requests = {live1.request_id: live1, live2.request_id: live2}
+
+    adapter.process_pending_chunks(waiting_queue, running_queue, scheduler_requests=scheduler_requests)
+
+    assert list(adapter.waiting_for_chunk_running_requests) == [live1, live2]
+    assert zombie1.request_id in adapter._cancelled_load_reqs
+    assert zombie2.request_id in adapter._cancelled_load_reqs
+
+
+def test_restore_queues_purges_late_aborts_after_process_pending_chunks(
+    build_adapter,
+):
+    """Race window guard: an abort can fire between ``process_pending_chunks``
+    and the scheduler's ``finally``-clause ``restore_queues`` call. The
+    second purge inside ``restore_queues`` must drop the now-untracked
+    request so it does not get re-injected onto ``running_queue``.
+    """
+    adapter, _ = build_adapter(stage_id=1, max_num_seqs=8)
+
+    req = _req("late-abort", RequestStatus.WAITING_FOR_CHUNK)
+    adapter.waiting_for_chunk_running_requests.append(req)
+
+    waiting_queue = DummyWaitingQueue()
+    running_queue: list = []
+
+    # Tick N: process_pending_chunks runs while req is still tracked.
+    scheduler_requests: dict = {req.request_id: req}
+    adapter.process_pending_chunks(waiting_queue, running_queue, scheduler_requests=scheduler_requests)
+    assert req in adapter.waiting_for_chunk_running_requests
+
+    # Mid-tick: abort fires and the scheduler's free path deletes the entry.
+    del scheduler_requests[req.request_id]
+
+    # finally: restore_queues sees the now-untracked req and must drop it
+    # instead of blindly extending it onto running_queue.
+    adapter.restore_queues(waiting_queue, running_queue, scheduler_requests=scheduler_requests)
+    assert running_queue == []
+    assert req.request_id in adapter._cancelled_load_reqs
+
+
+def test_purge_is_noop_on_empty_deques(build_adapter):
+    """Empty deques short-circuit -- guards against any accidental
+    ``IndexError`` from ``popleft`` on an empty deque if a future caller
+    shadows the empty check.
+    """
+    adapter, _ = build_adapter(stage_id=1, max_num_seqs=8)
+    assert len(adapter.waiting_for_chunk_waiting_requests) == 0
+    assert len(adapter.waiting_for_chunk_running_requests) == 0
+
+    waiting_queue = DummyWaitingQueue()
+    running_queue: list = []
+
+    adapter.process_pending_chunks(waiting_queue, running_queue, scheduler_requests={})
+    assert len(adapter.waiting_for_chunk_waiting_requests) == 0
+    assert len(adapter.waiting_for_chunk_running_requests) == 0
+    adapter.restore_queues(waiting_queue, running_queue, scheduler_requests={})
+    assert running_queue == []
+    assert waiting_queue == []

tests/e2e/accuracy/test_hunyuan_image3.py

@@ -29,7 +29,46 @@ from tests.e2e.accuracy.helpers import (
 )
 from tests.helpers.mark import hardware_test
 from tests.helpers.runtime import OmniRunner, OmniServer
-from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import build_prompt_tokens, resolve_stop_token_ids
+from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import (
+    build_prompt_tokens,
+    resolve_stop_token_ids,
+)
+
+# CoT structural markers that MUST appear in AR output text.
+# The AR trajectory under think_recaption mode is:
+#   <think> ... </think> <recaption> ... <relation_1> ... </relation_1>
+#   <relation_2> ... </relation_2> </recaption>
+# Test deploy config sets include_stop_str_in_output=True and
+# skip_special_tokens=False, so ALL opening and closing markers
+# should be present in the detokenized text.
+_THINK_CLOSE = "</" + "think>"  # avoid Python slash-t-as-TAB pitfall
+_COT_REQUIRED_MARKERS = (
+    "<recaption>",
+    "</recaption>",
+    "<relation_1>",
+    "</relation_1>",
+    "<relation_2>",
+    "</relation_2>",
+    _THINK_CLOSE,
+)
+
+
+def _assert_cot_structural_markers(cot_text: str, label: str) -> None:
+    """Assert that the AR CoT output contains all required structural tags.
+
+    Guards against regressions where detokenizer stop-token stripping
+    or truncation logic loses CoT content.  Under the test deploy config
+    (include_stop_str_in_output=True, skip_special_tokens=False),
+    ALL markers should be present in the detokenized text.
+    """
+    missing = [m for m in _COT_REQUIRED_MARKERS if m not in cot_text]
+    assert not missing, (
+        f"[{label}] CoT text missing required structural markers: {missing}. "
+        f"This typically indicates that stop-token stripping or "
+        f"_truncate_at_cot_end truncation removed CoT content. "
+        f"CoT text length={len(cot_text)}, first 100 chars: {repr(cot_text[:100])}"
+    )
+
 
 os.environ["DIFFUSION_ATTENTION_BACKEND"] = "TORCH_SDPA"
 
@@ -129,6 +168,7 @@ _DEPLOY_CONFIG = {
                 "stop_token_ids": [128025],
                 "detokenize": True,
                 "skip_special_tokens": False,
+                "include_stop_str_in_output": True,
             },
         },
         {
@@ -400,6 +440,8 @@ def test_image_to_image_alignment_online(accuracy_artifact_root: Path, accuracy_
     ]
     print("[ONLINE] " + tabulate(table, headers=["Metric", "Value", "L20x Reference"], tablefmt="grid"))
 
+    _assert_cot_structural_markers(online_cot, "ONLINE")
+
     assert cot_results["cot_semantic_sim"] >= THRESHOLDS["cot_semantic_sim"], (
         f"[ONLINE] COT semantic similarity {cot_results['cot_semantic_sim']:.4f} below threshold {THRESHOLDS['cot_semantic_sim']}"
     )
@@ -508,6 +550,8 @@ def test_image_to_image_alignment(accuracy_artifact_root: Path, accuracy_assets_
 
     print(tabulate(table, headers=["Metric", "Value", "L20x Reference"], tablefmt="grid"))
 
+    _assert_cot_structural_markers(omni_cot, "OFFLINE")
+
     assert cot_results["cot_semantic_sim"] >= THRESHOLDS["cot_semantic_sim"], (
         f"COT semantic similarity {cot_results['cot_semantic_sim']:.4f} is below threshold {THRESHOLDS['cot_semantic_sim']}"
     )

tests/e2e/offline_inference/test_ming_flash_omni_expansion.py

@@ -51,7 +51,7 @@ test_params = [(m, c) for m in models for c in [get_eager_config()]]
 
 @pytest.mark.core_model
 @pytest.mark.omni
-@hardware_test(res={"cuda": "H100"}, num_cards=4)
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4)
 @pytest.mark.parametrize("omni_runner", test_params_thinker, indirect=True)
 def test_text_to_text(omni_runner, omni_runner_handler) -> None:
     """
@@ -67,7 +67,7 @@ def test_text_to_text(omni_runner, omni_runner_handler) -> None:
 
 @pytest.mark.core_model
 @pytest.mark.omni
-@hardware_test(res={"cuda": "H100"}, num_cards=4)
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4)
 @pytest.mark.parametrize("omni_runner", test_params_thinker, indirect=True)
 def test_image_to_text(omni_runner, omni_runner_handler) -> None:
     """
@@ -84,7 +84,7 @@ def test_image_to_text(omni_runner, omni_runner_handler) -> None:
 
 @pytest.mark.core_model
 @pytest.mark.omni
-@hardware_test(res={"cuda": "H100"}, num_cards=4)
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4)
 @pytest.mark.parametrize("omni_runner", test_params_thinker, indirect=True)
 def test_audio_to_text(omni_runner, omni_runner_handler) -> None:
     """
@@ -103,7 +103,7 @@ def test_audio_to_text(omni_runner, omni_runner_handler) -> None:
 
 @pytest.mark.core_model
 @pytest.mark.omni
-@hardware_test(res={"cuda": "H100"}, num_cards=4)
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4)
 @pytest.mark.parametrize("omni_runner", test_params_thinker, indirect=True)
 def test_video_to_text(omni_runner, omni_runner_handler) -> None:
     """
@@ -120,7 +120,7 @@ def test_video_to_text(omni_runner, omni_runner_handler) -> None:
 
 @pytest.mark.core_model
 @pytest.mark.omni
-@hardware_test(res={"cuda": "H100"}, num_cards=4)
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4)
 @pytest.mark.parametrize("omni_runner", test_params_thinker, indirect=True)
 def test_mixed_to_text(omni_runner, omni_runner_handler) -> None:
     """
@@ -140,7 +140,7 @@ def test_mixed_to_text(omni_runner, omni_runner_handler) -> None:
 
 @pytest.mark.advanced_model
 @pytest.mark.omni
-@hardware_test(res={"cuda": "H100"}, num_cards=4)
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4)
 @pytest.mark.parametrize("omni_runner", test_params, indirect=True)
 def test_text_to_audio(omni_runner, omni_runner_handler) -> None:
     """
@@ -156,7 +156,7 @@ def test_text_to_audio(omni_runner, omni_runner_handler) -> None:
 
 @pytest.mark.advanced_model
 @pytest.mark.omni
-@hardware_test(res={"cuda": "H100"}, num_cards=4)
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4)
 @pytest.mark.parametrize("omni_runner", test_params, indirect=True)
 def test_image_to_audio(omni_runner, omni_runner_handler) -> None:
     """

tests/e2e/offline_inference/test_ming_tts.py

@@ -0,0 +1,295 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""End-to-end offline inference tests for Ming-omni-tts."""
+
+import asyncio
+import uuid
+
+import numpy as np
+import pytest
+import torch
+from transformers import AutoTokenizer
+from vllm import SamplingParams
+from vllm.inputs import tokens_input
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
+from tests.helpers.stage_config import get_deploy_config_path
+from vllm_omni import AsyncOmni
+from vllm_omni.model_executor.models.ming_tts.config_ming_tts import (
+    KEY_MAX_DECODE_STEPS,
+    SAMPLE_RATE,
+    TEXT_EOS_TOKEN_ID,
+)
+from vllm_omni.model_executor.models.ming_tts.prompt_assembly import DEFAULT_PROMPT, build_ming_dense_prompt
+
+MODEL = "inclusionAI/Ming-omni-tts-0.5B"
+DEPLOY_CONFIG = get_deploy_config_path("ming_tts.yaml")
+TEST_TEXT = "我会一直在这里陪着你，直到你慢慢地沉入那个最温柔的梦里。"
+TEST_INSTRUCTION = "轻柔的ASMR耳语，慢速，贴近麦克风"
+MIN_AUDIO_SAMPLES = 1000
+
+
+@pytest.fixture(scope="module")
+def ming_tokenizer():
+    return AutoTokenizer.from_pretrained(MODEL, trust_remote_code=False)
+
+
+@pytest.fixture
+def ming_engine():
+    with OmniRunner(
+        MODEL,
+        deploy_config=DEPLOY_CONFIG,
+        stage_init_timeout=300,
+        enforce_eager=True,
+    ) as runner:
+        yield runner.omni
+
+
+@pytest.fixture
+def async_omni_engine():
+    engine = AsyncOmni(
+        model=MODEL,
+        deploy_config=DEPLOY_CONFIG,
+        stage_init_timeout=300,
+        enforce_eager=True,
+    )
+    yield engine
+    engine.shutdown()
+
+
+def _build_prompt(
+    tokenizer,
+    *,
+    text: str = TEST_TEXT,
+    instruction=TEST_INSTRUCTION,
+    use_zero_spk_emb: bool = True,
+) -> dict:
+    prompt_dict = build_ming_dense_prompt(
+        tokenizer,
+        prompt=DEFAULT_PROMPT,
+        text=text,
+        instruction=instruction,
+        runtime_controls={KEY_MAX_DECODE_STEPS: 200},
+        use_zero_spk_emb=use_zero_spk_emb,
+    )
+    prompt = tokens_input(prompt_token_ids=prompt_dict["prompt_token_ids"])
+    prompt["prompt"] = prompt_dict["prompt"]
+    prompt["text"] = prompt_dict["text"]
+    prompt["additional_information"] = prompt_dict["additional_information"]
+    return prompt
+
+
+def _sampling_params_list() -> list[SamplingParams]:
+    return [
+        SamplingParams(
+            temperature=0.0,
+            max_tokens=201,
+            stop_token_ids=[int(TEXT_EOS_TOKEN_ID)],
+        ),
+        SamplingParams(temperature=0.0, max_tokens=1),
+    ]
+
+
+def _flatten_audio(audio) -> torch.Tensor:
+    if isinstance(audio, list):
+        parts = [torch.as_tensor(item, dtype=torch.float32).reshape(-1).cpu() for item in audio]
+        parts = [item for item in parts if item.numel() > 0]
+        if not parts:
+            return torch.zeros((0,), dtype=torch.float32)
+        return torch.cat(parts, dim=0)
+    return torch.as_tensor(audio, dtype=torch.float32).reshape(-1).cpu()
+
+
+def _extract_audio(multimodal_output: dict) -> torch.Tensor:
+    assert isinstance(multimodal_output, dict), f"Expected dict, got {type(multimodal_output)}"
+    audio = multimodal_output.get("audio", multimodal_output.get("model_outputs"))
+    assert audio is not None, f"No audio output found, keys={list(multimodal_output.keys())}"
+    waveform = _flatten_audio(audio)
+    if waveform.numel() == 0:
+        raise RuntimeError("Generated audio waveform is empty")
+    return waveform
+
+
+def _extract_sample_rate(multimodal_output: dict) -> int:
+    sample_rate = multimodal_output.get("sr")
+    if sample_rate is None:
+        raise RuntimeError("Expected multimodal_output['sr']")
+    if isinstance(sample_rate, list):
+        sample_rate = sample_rate[-1]
+    if hasattr(sample_rate, "item"):
+        sample_rate = sample_rate.item()
+    return int(sample_rate)
+
+
+def _extract_final_audio_outputs(outputs):
+    final_outputs = []
+    for item in outputs:
+        if getattr(item, "final_output_type", None) == "audio":
+            final_outputs.append(item)
+            continue
+        request_output = getattr(item, "request_output", None)
+        if request_output is None:
+            continue
+        multimodal_output = getattr(request_output, "multimodal_output", None)
+        if isinstance(multimodal_output, dict):
+            final_outputs.append(item)
+            continue
+        completions = getattr(request_output, "outputs", None) or []
+        if any(isinstance(getattr(completion, "multimodal_output", None), dict) for completion in completions):
+            final_outputs.append(item)
+    return final_outputs
+
+
+def _extract_multimodal_output(output) -> dict:
+    multimodal_output = getattr(output, "multimodal_output", None)
+    if isinstance(multimodal_output, dict):
+        return multimodal_output
+
+    request_output = getattr(output, "request_output", None)
+    if request_output is not None:
+        multimodal_output = getattr(request_output, "multimodal_output", None)
+        if isinstance(multimodal_output, dict):
+            return multimodal_output
+        completions = getattr(request_output, "outputs", None) or []
+        for completion in completions:
+            multimodal_output = getattr(completion, "multimodal_output", None)
+            if isinstance(multimodal_output, dict):
+                return multimodal_output
+
+    raise AssertionError("No multimodal audio output found in Ming generate results")
+
+
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_ming_tts_offline_basic(ming_engine, ming_tokenizer) -> None:
+    """Test blocking Ming generation through Omni."""
+    outputs = ming_engine.generate(
+        prompts=[_build_prompt(ming_tokenizer)],
+        sampling_params_list=_sampling_params_list(),
+        py_generator=False,
+    )
+    final_outputs = _extract_final_audio_outputs(outputs)
+    assert len(final_outputs) == 1, f"Expected one final audio output, got {len(final_outputs)}"
+    multimodal_output = _extract_multimodal_output(final_outputs[0])
+    waveform = _extract_audio(multimodal_output)
+    sample_rate = _extract_sample_rate(multimodal_output)
+    assert waveform.ndim == 1
+    assert waveform.shape[0] == waveform.numel()
+    assert waveform.numel() > MIN_AUDIO_SAMPLES
+    assert np.max(np.abs(waveform.numpy())) > 0.01, "Audio appears silent"
+    assert sample_rate == SAMPLE_RATE, f"Expected Ming output sample rate {SAMPLE_RATE}, got {sample_rate}"
+
+
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_ming_tts_speaker_conditioning_differs(ming_engine, ming_tokenizer) -> None:
+    """Test that different Ming speaker controls produce different waveform outputs."""
+    style_outputs = ming_engine.generate(
+        prompts=[_build_prompt(ming_tokenizer)],
+        sampling_params_list=_sampling_params_list(),
+        py_generator=False,
+    )
+    ip_outputs = ming_engine.generate(
+        prompts=[_build_prompt(ming_tokenizer, text=TEST_TEXT, instruction={"IP": "灵小甄"}, use_zero_spk_emb=True)],
+        sampling_params_list=_sampling_params_list(),
+        py_generator=False,
+    )
+
+    style_final_outputs = _extract_final_audio_outputs(style_outputs)
+    ip_final_outputs = _extract_final_audio_outputs(ip_outputs)
+    assert len(style_final_outputs) == 1, "No style audio output produced"
+    assert len(ip_final_outputs) == 1, "No IP audio output produced"
+
+    style_waveform = _extract_audio(_extract_multimodal_output(style_final_outputs[0]))
+    ip_waveform = _extract_audio(_extract_multimodal_output(ip_final_outputs[0]))
+    assert style_waveform.numel() > MIN_AUDIO_SAMPLES
+    assert ip_waveform.numel() > MIN_AUDIO_SAMPLES
+    assert np.max(np.abs(style_waveform.numpy())) > 0.01, "Style audio appears silent"
+    assert np.max(np.abs(ip_waveform.numpy())) > 0.01, "IP audio appears silent"
+
+    overlap = min(int(style_waveform.numel()), int(ip_waveform.numel()))
+    mean_abs_diff = torch.mean(torch.abs(style_waveform[:overlap] - ip_waveform[:overlap])).item()
+    assert style_waveform.shape != ip_waveform.shape or mean_abs_diff > 1e-4, (
+        "Speaker-conditioned outputs should differ, but style and IP waveforms were effectively identical"
+    )
+
+
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_ming_tts_multiple_prompts_queued(ming_engine, ming_tokenizer) -> None:
+    """Regression: supported max_num_seqs=1 config must still drain queued prompts."""
+    prompts = [
+        _build_prompt(
+            ming_tokenizer,
+            text="第一条语音用于验证 Ming dense 队列中的长请求可以完成。",
+            instruction="平静自然的旁白，语速中等",
+        ),
+        _build_prompt(
+            ming_tokenizer,
+            text="第二条短语音也必须生成完成。",
+            instruction="清晰明亮的女声",
+        ),
+    ]
+    outputs = ming_engine.generate(
+        prompts=prompts,
+        sampling_params_list=_sampling_params_list(),
+        py_generator=False,
+    )
+    final_outputs = _extract_final_audio_outputs(outputs)
+    assert len(final_outputs) == len(prompts), f"Expected {len(prompts)} audio outputs, got {len(final_outputs)}"
+    for i, output in enumerate(final_outputs):
+        waveform = _extract_audio(_extract_multimodal_output(output))
+        duration_s = waveform.numel() / SAMPLE_RATE
+        assert 0.1 < duration_s < 30.0, f"Request {i} audio duration out of range: {duration_s:.2f}s"
+        assert np.max(np.abs(waveform.numpy())) > 0.01, f"Request {i} audio appears silent"
+
+
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_ming_tts_offline_streaming(async_omni_engine, ming_tokenizer) -> None:
+    """Test async_chunk streaming Ming generation through AsyncOmni."""
+
+    async def _run() -> None:
+        all_audio_chunks = []
+        accumulated_samples = 0
+        chunk_idx = 0
+        sample_rate = None
+        async for stage_output in async_omni_engine.generate(
+            prompt=_build_prompt(ming_tokenizer),
+            request_id=str(uuid.uuid4()),
+            sampling_params_list=_sampling_params_list(),
+        ):
+            multimodal_output = stage_output.multimodal_output or {}
+            audio = multimodal_output.get("audio")
+            if "sr" in multimodal_output:
+                sample_rate = _extract_sample_rate(multimodal_output)
+            if audio is None:
+                continue
+            finished = stage_output.finished
+            if isinstance(audio, torch.Tensor):
+                if finished:
+                    audio_chunk = audio[accumulated_samples:].float().detach().cpu()
+                else:
+                    audio_chunk = audio.float().detach().cpu()
+            elif isinstance(audio, list):
+                audio_chunk = _flatten_audio(audio)
+            else:
+                audio_chunk = torch.as_tensor(audio, dtype=torch.float32).reshape(-1).cpu()
+            accumulated_samples += int(audio_chunk.numel())
+            chunk_idx += 1
+            if audio_chunk.numel() > 0:
+                all_audio_chunks.append(audio_chunk)
+        assert all_audio_chunks, "No streaming audio chunks received"
+        waveform = torch.cat(all_audio_chunks, dim=0)
+        assert waveform.numel() > MIN_AUDIO_SAMPLES
+        assert np.max(np.abs(waveform.numpy())) > 0.01, "Audio appears silent"
+        assert sample_rate is not None, "Streaming path did not return a sample rate"
+        assert sample_rate == SAMPLE_RATE, f"Expected Ming output sample rate {SAMPLE_RATE}, got {sample_rate}"
+
+    asyncio.run(_run())

tests/e2e/offline_inference/test_sensenova_u1_expansion.py

@@ -0,0 +1,160 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+Expanded end-to-end test for SenseNova-U1 in offline mode.
+
+This test file primarily covers end-to-end tests for LoRA support.
+
+Validates that LoRA adapters are correctly loaded, applied with controllable
+scale, and cleanly deactivated. Uses a synthetic rank-4 adapter targeting the
+first decoder layer's gen-path (``_mot_gen``) output projection, so the
+perturbation is exercised during the denoising loop.
+
+Assertions:
+  (a) LoRA at scale=1.0 visibly changes the output  (diff > 0.5)
+  (b) scale=2.0 produces a different delta than scale=1.0  (scale sensitivity)
+  (c) The delta is bounded  (not corrupted)
+  (d) Deactivating LoRA exactly restores the baseline  (diff == 0)
+
+Observed on H100 (seed=42, 8 steps, 512x512): diff_1x~62.2, diff_2x~76.9,
+diff_restored=0.0. Thresholds match test_bagel_expansion.py and hold with margin.
+"""
+
+import json
+import os
+from pathlib import Path
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+import numpy as np
+import pytest
+import torch
+from PIL import Image
+from safetensors.torch import save_file
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
+from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.lora.request import LoRARequest
+from vllm_omni.lora.utils import stable_lora_int_id
+
+MODEL = "SenseNova/SenseNova-U1-8B-MoT"
+DEFAULT_PROMPT = "a photo of a cat sitting on a laptop keyboard"
+SIZE = (512, 512)  # (width, height); small for fast E2E
+
+# SenseNova-U1-8B-MoT hidden_size (Qwen3-based). o_proj is square (hidden x hidden)
+# regardless of GQA, so this is the robust LoRA target. Verify against config.json.
+_LORA_DIM = 4096
+_LORA_MODULE = "language_model.model.layers.0.self_attn.o_proj_mot_gen"
+_LORA_RANK = 4
+
+
+# ---------------------------------------------------------------------------
+# Helpers (reused from test_sensenova_u1_text2img.py patterns)
+# ---------------------------------------------------------------------------
+
+
+def _build_sampling_params(
+    lora_request: LoRARequest | None = None,
+    lora_scale: float = 1.0,
+) -> OmniDiffusionSamplingParams:
+    return OmniDiffusionSamplingParams(
+        height=SIZE[1],
+        width=SIZE[0],
+        seed=42,
+        num_inference_steps=8,
+        extra_args={
+            "cfg_scale": 4.0,
+            "cfg_norm": "none",
+            "timestep_shift": 3.0,
+            "cfg_interval": (0.0, 1.0),
+            "batch_size": 1,
+            "think": False,
+            "t_eps": 0.02,
+        },
+        lora_request=lora_request,
+        lora_scale=lora_scale,
+    )
+
+
+def _extract_generated_image(omni_outputs: list) -> Image.Image | None:
+    for req_output in omni_outputs:
+        if images := getattr(req_output, "images", None):
+            return images[0]
+    return None
+
+
+def _generate(omni: Omni, lora_request: LoRARequest | None = None, lora_scale: float = 1.0) -> Image.Image:
+    outputs = list(
+        omni.generate(
+            prompts={"prompt": DEFAULT_PROMPT, "modalities": ["image"]},
+            sampling_params_list=_build_sampling_params(lora_request, lora_scale),
+        )
+    )
+    img = _extract_generated_image(outputs)
+    assert img is not None, "No image generated"
+    return img
+
+
+def _make_file_lora_request(adapter_dir: Path) -> LoRARequest:
+    """Write synthetic adapter to disk and return a file-backed LoRARequest."""
+    adapter_dir.mkdir(parents=True, exist_ok=True)
+    gen = torch.Generator().manual_seed(42)
+    lora_a = torch.randn((_LORA_RANK, _LORA_DIM), dtype=torch.float32, generator=gen) * 0.1
+    lora_b = torch.randn((_LORA_DIM, _LORA_RANK), dtype=torch.float32, generator=gen) * 0.5
+    save_file(
+        {
+            f"base_model.model.{_LORA_MODULE}.lora_A.weight": lora_a,
+            f"base_model.model.{_LORA_MODULE}.lora_B.weight": lora_b,
+        },
+        str(adapter_dir / "adapter_model.safetensors"),
+    )
+    (adapter_dir / "adapter_config.json").write_text(
+        json.dumps({"r": _LORA_RANK, "lora_alpha": _LORA_RANK, "target_modules": [_LORA_MODULE]}),
+        encoding="utf-8",
+    )
+    lora_dir = str(adapter_dir)
+    return LoRARequest(lora_name="test_file", lora_int_id=stable_lora_int_id(lora_dir), lora_path=lora_dir)
+
+
+# ---------------------------------------------------------------------------
+# Test
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.full_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "H100"})
+def test_sensenova_u1_lora_scale_and_deactivation(tmp_path) -> None:
+    """Validate LoRA effect, scale sensitivity, bounded perturbation, and clean deactivation."""
+    lora_request = _make_file_lora_request(tmp_path / "sensenova_lora")
+
+    with OmniRunner(MODEL, stage_configs_path=None) as runner:
+        omni = runner.omni
+        baseline = _generate(omni)
+        img_1x = _generate(omni, lora_request, lora_scale=1.0)
+        img_2x = _generate(omni, lora_request, lora_scale=2.0)
+        restored = _generate(omni)
+
+    baseline_arr = np.array(baseline, dtype=np.int16)
+    diff_1x = np.abs(baseline_arr - np.array(img_1x, dtype=np.int16)).mean()
+    diff_2x = np.abs(baseline_arr - np.array(img_2x, dtype=np.int16)).mean()
+    diff_restored = np.abs(baseline_arr - np.array(restored, dtype=np.int16)).mean()
+
+    # (a) Adapter has visible effect at both scales
+    assert diff_1x > 0.5, f"LoRA scale=1.0 had no visible effect: diff={diff_1x}"
+    assert diff_2x > 0.5, f"LoRA scale=2.0 had no visible effect: diff={diff_2x}"
+
+    # (b) Different scales produce different outputs
+    assert not np.isclose(diff_1x, diff_2x, atol=1.0), (
+        f"LoRA scale has no effect: diff_1x={diff_1x:.2f}, diff_2x={diff_2x:.2f}"
+    )
+
+    # (c) Output is not corrupted
+    assert diff_1x < 80, f"LoRA output looks corrupted: diff_1x={diff_1x}"
+    assert diff_2x < 120, f"LoRA output looks corrupted: diff_2x={diff_2x}"
+
+    # (d) Deactivation fully restores base model
+    assert diff_restored == 0.0, f"Base model not restored after LoRA deactivation: diff={diff_restored}"

tests/e2e/online_serving/test_ming_tts.py

@@ -0,0 +1,92 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+E2E online serving tests for Ming-omni-tts (dense) model.
+Tests text-to-audio via /v1/audio/speech endpoint.
+"""
+
+import os
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+import pytest
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniServerParams
+from tests.helpers.stage_config import get_deploy_config_path
+
+pytestmark = [pytest.mark.advanced_model, pytest.mark.omni]
+
+MODEL = "inclusionAI/Ming-omni-tts-0.5B"
+DEPLOY_CONFIG = get_deploy_config_path("ming_tts.yaml")
+
+no_async_chunk_params = [
+    pytest.param(
+        OmniServerParams(
+            model=MODEL,
+            stage_config_path=DEPLOY_CONFIG,
+            server_args=["--enforce-eager", "--no-async-chunk"],
+        ),
+        id="no_async_chunk",
+    )
+]
+
+async_chunk_params = [
+    pytest.param(
+        OmniServerParams(
+            model=MODEL,
+            stage_config_path=DEPLOY_CONFIG,
+            server_args=["--enforce-eager"],
+        ),
+        id="async_chunk",
+    )
+]
+
+
+def get_prompt(prompt_type="zh"):
+    prompts = {
+        "zh": "今天天气真不错，适合出去散散步。",
+        "zh_short": "这款产品的名字，叫变态坑爹牛肉丸。",
+    }
+    return prompts.get(prompt_type, prompts["zh"])
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", no_async_chunk_params, indirect=True)
+def test_text_to_audio_non_streaming_001(omni_server, openai_client) -> None:
+    """
+    Deploy Setting: ming_tts.yaml with --no-async-chunk
+    Input Modal: text
+    Output Modal: audio
+    Input Setting: stream=False
+    Datasets: two concurrent requests
+    """
+    request_config = {
+        "model": omni_server.model,
+        "input": get_prompt("zh"),
+        "stream": False,
+        "response_format": "wav",
+        "timeout": 300.0,
+    }
+    openai_client.send_audio_speech_request(request_config, request_num=2)
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", async_chunk_params, indirect=True)
+def test_text_to_audio_streaming_001(omni_server, openai_client) -> None:
+    """
+    Deploy Setting: ming_tts.yaml (async_chunk=true)
+    Input Modal: text + voice
+    Output Modal: audio (streamed)
+    Input Setting: stream=True
+    Datasets: single request
+    """
+    request_config = {
+        "model": omni_server.model,
+        "input": get_prompt("zh_short"),
+        "voice": "灵小甄",
+        "stream": True,
+        "response_format": "wav",
+        "timeout": 300.0,
+    }
+    openai_client.send_audio_speech_request(request_config)

tests/e2e/online_serving/test_voxtral_tts.py

@@ -91,3 +91,90 @@ class TestVoxtralTTSFixedVoice:
                     "timeout": 120.0,
                 }
             )
+
+    @pytest.mark.advanced_model
+    @pytest.mark.tts
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_speech_speed(self, omni_server, openai_client) -> None:
+        """Request with speed parameters"""
+        speeds = [0.5, 1, 1.5, 2, 2.5]
+        for speed in speeds:
+            openai_client.send_audio_speech_request(
+                {
+                    "model": omni_server.model,
+                    "input": "The boy was there when the sun rose.",
+                    "voice": "casual_female",
+                    "language": "English",
+                    "response_format": "wav",
+                    "timeout": 120.0,
+                    "speed": speed,
+                }
+            )
+
+    @pytest.mark.advanced_model
+    @pytest.mark.tts
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_speech_instructions(self, omni_server, openai_client) -> None:
+        """Request with instructions parameters"""
+        instructions = [
+            "Speak formally",
+            "Speak angrily",
+            "Deliver with a sad voice",
+            "Speak with a chirpy happy voice",
+        ]
+        for instruction in instructions:
+            openai_client.send_audio_speech_request(
+                {
+                    "model": omni_server.model,
+                    "input": "The boy was there when the sun rose.",
+                    "voice": "casual_female",
+                    "language": "English",
+                    "response_format": "wav",
+                    "timeout": 120.0,
+                    "instructions": instruction,
+                }
+            )
+
+    @pytest.mark.advanced_model
+    @pytest.mark.tts
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_speech_response_formats(self, omni_server, openai_client) -> None:
+        """Test TTS with different response formats"""
+        response_formats = ["wav", "mp3"]
+        for response_format in response_formats:
+            openai_client.send_audio_speech_request(
+                {
+                    "model": omni_server.model,
+                    "input": "Testing various response formats.",
+                    "voice": "casual_male",
+                    "language": "English",
+                    "response_format": response_format,
+                    "timeout": 120.0,
+                }
+            )
+
+    @pytest.mark.advanced_model
+    @pytest.mark.tts
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_speech_batches(self, omni_server, openai_client) -> None:
+        """Test TTS batches"""
+        items = [
+            {"input": "The birch canoe slid on the smooth planks."},
+            {"input": "Glue the sheet to the dark blue background."},
+            {"input": "It's easy to tell the depth of a well."},
+            {"input": "These days a chicken leg is a rare dish."},
+            {"input": "Rice is often served in round bowls."},
+        ]
+
+        openai_client.send_audio_speech_batch_http_request(
+            {
+                "json": {
+                    "model": omni_server.model,
+                    "items": items,
+                    "voice": "casual_male",
+                    "language": "English",
+                    "response_format": "wav",
+                },
+                "timeout": 120.0,
+            }
+        )

tests/helpers/media.py

@@ -588,15 +588,32 @@ def cosine_similarity_text(text1, text2, n: int = 3):
     return cosine * length_harmony
 
 
-def _merge_base64_audio_to_segment(base64_list: list[str]):
-    from pydub import AudioSegment
+class _AudioBuffer:
+    """Minimal replacement for pydub.AudioSegment used by test helpers."""
 
-    merged = None
+    def __init__(self, data: np.ndarray, sample_rate: int):
+        self.data = data
+        self.sample_rate = sample_rate
+
+    def export(self, buf: io.BytesIO, format: str = "wav"):
+        sf.write(buf, self.data, self.sample_rate, format=format.upper())
+        buf.seek(0)
+
+
+def _merge_base64_audio_to_segment(base64_list: list[str]) -> _AudioBuffer:
+    from vllm.multimodal.media.audio import AudioMediaIO
+
+    io_ = AudioMediaIO()
+    chunks: list[np.ndarray] = []
+    sample_rate: int | None = None
     for b64 in base64_list:
         raw = base64.b64decode(b64.split(",", 1)[-1])
-        seg = AudioSegment.from_file(io.BytesIO(raw))
-        merged = seg if merged is None else merged + seg
-    return merged
+        waveform, sr = io_.load_bytes(raw)
+        if sample_rate is None:
+            sample_rate = int(sr)
+        chunks.append(waveform)
+    merged = np.concatenate(chunks) if chunks else np.array([], dtype=np.float32)
+    return _AudioBuffer(merged, sample_rate or 16000)
 
 
 @contextmanager

tests/helpers/runtime.py

@@ -1767,7 +1767,16 @@ class OpenAIClientHandler:
         # Qwen3-TTS custom fields, forwarded via extra_body.
         extra_body: dict[str, Any] = {}
         # Keep this list aligned with vllm_omni.entrypoints.openai.protocol.audio params.
-        for key in ("task_type", "ref_text", "ref_audio", "language", "max_new_tokens", "seed"):
+        for key in (
+            "task_type",
+            "ref_text",
+            "ref_audio",
+            "language",
+            "max_new_tokens",
+            "seed",
+            "instructions",
+            "speed",
+        ):
             if key in request_config:
                 extra_body[key] = request_config[key]
 

tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py

@@ -84,6 +84,11 @@ def _make_code2wav_model(
         token_mel_ratio=2 if with_stride_cfg else 0,
     )
     model.code2wav = _DummyCode2Wav(vocab_size=4, num_samples=num_samples, outputs=outputs)
+    # Short-circuit the lazy TensorRT estimator swap: these tests exercise the
+    # forward audio logic, not the TRT path. On a GPU CI runner the swap would
+    # otherwise run and dereference ``self.model_dir`` (only set in __init__,
+    # which this fixture bypasses via object.__new__).
+    model._code2wav_trt_done = True
     model.source_cache_len = 4
     model.speech_window = torch.hamming_window(8, periodic=False)
     model._stream_audio_cache_by_req = {}

tests/model_executor/models/ming_tts/test_ming_shared_modules.py

@@ -0,0 +1,185 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech
+from vllm_omni.model_executor.models.common.ming.audio_vae import AudioVAEConfig
+from vllm_omni.model_executor.models.common.ming.fm import Solver
+from vllm_omni.model_executor.models.ming_tts.constants import (
+    AGGREGATOR_HIDDEN_SIZE,
+    HISTORY_PATCH_SIZE,
+    LATENT_DIM,
+    LLM_HIDDEN_SIZE,
+    LLM_VOCAB_SIZE,
+    PATCH_SIZE,
+    SAMPLE_RATE,
+    VAE_PATCH_SIZE,
+)
+from vllm_omni.model_executor.models.ming_tts.flowloss_head import FlowLoss
+from vllm_omni.model_executor.models.ming_tts.ming_tts_llm import MingLLMModel
+from vllm_omni.model_executor.models.ming_tts.validation import validate_ming_tts_config
+from vllm_omni.model_executor.models.output_templates import OmniOutput
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.tts]
+
+
+def test_ming_tts_audio_vae_uses_common_config():
+    """AudioVAEConfig is shared by Ming dense and Ming flash modules."""
+    cfg = AudioVAEConfig(sample_rate=16000, patch_size=-1)
+
+    assert cfg.sample_rate == 16000
+    assert cfg.patch_size == -1
+
+
+def test_ming_tts_cfm_solver_uses_common_implementation():
+    """Ming dense imports the shared solver implementation directly."""
+    assert Solver.__module__ == "vllm_omni.model_executor.models.common.ming.fm"
+
+
+def test_ming_tts_flowloss_preserves_checkpoint_prefix():
+    flowloss = FlowLoss(z_channels=4, llm_cond_dim=8, hidden_size=16, depth=1, num_heads=2)
+
+    assert any(name.startswith("cfm.model.") for name in flowloss.state_dict())
+
+
+def test_ming_dense_validation_rejects_semantic_audio_vae_config():
+    """Dense 0.5B validation rejects semantic AudioVAE configs."""
+    cfg = SimpleNamespace(
+        audio_dummy_token_id=151705,
+        audio_eos_token_id=151704,
+        text_eos_token_id=151669,
+        audio_tokenizer_config=AudioVAEConfig(
+            sample_rate=SAMPLE_RATE,
+            patch_size=VAE_PATCH_SIZE,
+            semantic_module_kwargs={"whisper_encoder": {}},
+            enc_kwargs={"latent_dim": LATENT_DIM, "input_dim": 882, "hop_size": 882},
+            dec_kwargs={"latent_dim": LATENT_DIM, "output_dim": 882},
+        ),
+        latent_dim=LATENT_DIM,
+        patch_size=PATCH_SIZE,
+        history_patch_size=HISTORY_PATCH_SIZE,
+        llm_hidden_size=LLM_HIDDEN_SIZE,
+        llm_vocab_size=LLM_VOCAB_SIZE,
+        sample_rate=SAMPLE_RATE,
+        vae_patch_size=VAE_PATCH_SIZE,
+        llm_config={"hidden_size": LLM_HIDDEN_SIZE},
+        aggregator_config={"hidden_size": AGGREGATOR_HIDDEN_SIZE},
+        ditar_config={"hidden_size": AGGREGATOR_HIDDEN_SIZE},
+        latent_chunk_size=1,
+        latent_left_context=0,
+        max_decode_steps=1,
+        stop_head_threshold=0.5,
+        stop_head_min_steps=0,
+    )
+
+    with pytest.raises(ValueError, match="semantic_module_kwargs"):
+        validate_ming_tts_config(cfg)
+
+
+def test_ming_instruction_parser_preserves_dense_and_flash_defaults():
+    """Ming dense and Ming flash keep distinct instruction defaults."""
+    serving = object.__new__(OmniOpenAIServingSpeech)
+    serving.uploaded_speakers = {"uploaded": {}}
+
+    dense_plain = serving._parse_ming_instruction(SimpleNamespace(instructions="calm", language=None, voice=None))
+    assert dense_plain == "calm"
+
+    dense_with_fields = serving._parse_ming_instruction(
+        SimpleNamespace(instructions="calm", language="Auto", voice="灵小甄")
+    )
+    assert dense_with_fields == {"IP": "灵小甄", "风格": "calm"}
+
+    flash_fields = serving._parse_ming_instruction_fields(
+        SimpleNamespace(instructions="calm", language="粤语", voice="灵小甄")
+    )
+    assert flash_fields == {"风格": "calm"}
+
+
+def _make_ming_logits_model(vocab_size=8):
+    model = object.__new__(MingLLMModel)
+    model.ming_config = SimpleNamespace(
+        llm_vocab_size=vocab_size,
+        max_decode_steps=1,
+        stop_head_min_steps=0,
+        text_eos_token_id=7,
+    )
+    model._last_text_mode = False
+    model._last_ming_next_token_ids = None
+    return model
+
+
+def test_ming_compute_logits_uses_cached_forced_next_token_ids():
+    model = _make_ming_logits_model()
+    model._last_ming_next_token_ids = [2, 5]
+
+    logits = MingLLMModel.compute_logits(model, torch.zeros((2, 4)), SimpleNamespace())
+
+    assert logits.shape == (2, 8)
+    assert logits[0, 2].item() == 0.0
+    assert logits[1, 5].item() == 0.0
+    assert torch.isneginf(logits[0, [0, 1, 3, 4, 5, 6, 7]]).all()
+    assert torch.isneginf(logits[1, [0, 1, 2, 3, 4, 6, 7]]).all()
+    assert model._last_ming_next_token_ids is None
+
+
+def test_ming_compute_logits_falls_back_to_dummy_token_id():
+    model = _make_ming_logits_model()
+
+    logits = MingLLMModel.compute_logits(model, torch.zeros((1, 4)), SimpleNamespace())
+    assert logits.shape == (1, 8)
+    assert logits[0, 7].item() == 0.0
+    assert torch.isneginf(logits[0, [0, 1, 2, 3, 4, 5, 6]]).all()
+
+
+def test_ming_forward_non_decode_return_clears_cached_forced_next_token_ids():
+    class FakeBackbone:
+        def __call__(self, **kwargs):
+            return kwargs["inputs_embeds"]
+
+    model = _make_ming_logits_model()
+    model.model = FakeBackbone()
+    model._last_ming_next_token_ids = [2]
+
+    output = MingLLMModel.forward(
+        model,
+        input_ids=torch.tensor([1]),
+        positions=torch.tensor([0]),
+        inputs_embeds=torch.zeros((1, 4)),
+        model_intermediate_buffer=[],
+    )
+
+    assert isinstance(output, OmniOutput)
+    assert output.multimodal_outputs is None
+    assert model._last_ming_next_token_ids is None
+
+
+def test_ming_compute_logits_rejects_forced_token_batch_mismatch():
+    model = _make_ming_logits_model()
+    model._last_ming_next_token_ids = [2]
+
+    with pytest.raises(RuntimeError, match="batch mismatch"):
+        MingLLMModel.compute_logits(model, torch.zeros((2, 4)), SimpleNamespace())
+
+
+def test_ming_compute_logits_text_mode_delegates_to_backbone():
+    class FakeBackbone:
+        def __init__(self):
+            self.hidden_states = None
+
+        def compute_logits(self, hidden_states):
+            self.hidden_states = hidden_states
+            return torch.ones((hidden_states.shape[0], 3))
+
+    model = _make_ming_logits_model(vocab_size=3)
+    model._last_text_mode = True
+    model.model = FakeBackbone()
+    hidden_states = torch.zeros((2, 4))
+
+    logits = MingLLMModel.compute_logits(model, hidden_states, SimpleNamespace())
+
+    assert torch.equal(model.model.hidden_states, hidden_states)
+    assert torch.equal(logits, torch.ones((2, 3)))

tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py

@@ -921,6 +921,7 @@ def test_cosyvoice3_full_payload_replace_keys_present() -> None:
 def test_ming_flash_omni_thinker2talker_token_only_smoke() -> None:
     """Smoke: ming_flash_omni token-only carries voice metadata."""
     from vllm_omni.model_executor.stage_input_processors.ming_flash_omni import (
+        thinker2talker,
         thinker2talker_token_only,
     )
 
@@ -938,14 +939,15 @@ def test_ming_flash_omni_thinker2talker_token_only_smoke() -> None:
 
     src = [_Wrap("hello world")]
     prompt = _Prompt({"voice_name": "ZH_FEMALE", "prompt_text": "ref text"})
-    out = thinker2talker_token_only(src, prompt=prompt)
-    assert len(out) == 1
-    assert out[0]["prompt_token_ids"] == [0]  # talker self-tokenizes; dummy id
-    info = out[0]["additional_information"]
-    assert info["text"] == "hello world"
-    assert info["voice_name"] == "ZH_FEMALE"
-    assert info["prompt_text"] == "ref text"
-    assert info["ming_task"] == "omni"
+    for func in (thinker2talker, thinker2talker_token_only):
+        out = func(src, prompt=prompt)
+        assert len(out) == 1
+        assert out[0]["prompt_token_ids"] == [0]
+        info = out[0]["additional_information"]
+        assert info["text"] == "hello world"
+        assert info["voice_name"] == "ZH_FEMALE"
+        assert info["prompt_text"] == "ref text"
+        assert info["ming_task"] == "omni"
 
 
 def test_qwen2_5_omni_thinker2talker_token_only_smoke() -> None:

tests/test_pydub_removal.py

@@ -0,0 +1,182 @@
+"""Functional tests for pydub removal.
+
+Verifies that the new av/numpy/torchaudio replacements produce
+identical results to the old pydub-based code for all three paths:
+  1. WAV PCM → duration & frame count (pure arithmetic)
+  2. Non-WAV (mp3/ogg) → duration & frame count (av.open)
+  3. TTS WER PCM export: resample to 24000Hz mono s16
+"""
+
+import io
+import wave
+
+import av
+import numpy as np
+import soundfile as sf
+
+
+def _make_wav_bytes(sample_rate=16000, channels=1, sample_width=2, duration_sec=1.0):
+    """Generate a WAV file with a sine wave, return (wav_bytes, pcm_bytes, params)."""
+    n_samples = int(sample_rate * duration_sec)
+    t = np.linspace(0, duration_sec, n_samples, endpoint=False)
+    samples = (np.sin(2 * np.pi * 440 * t) * 32767 * 0.5).astype(np.int16)
+    if channels == 2:
+        samples = np.column_stack([samples, samples]).flatten()
+
+    buf = io.BytesIO()
+    with wave.open(buf, "wb") as wf:
+        wf.setnchannels(channels)
+        wf.setsampwidth(sample_width)
+        wf.setframerate(sample_rate)
+        wf.writeframes(samples.tobytes())
+    wav_bytes = buf.getvalue()
+
+    # Extract raw PCM (skip WAV header)
+    with wave.open(io.BytesIO(wav_bytes), "rb") as wr:
+        pcm_bytes = wr.readframes(wr.getnframes())
+        params = (wr.getnchannels(), wr.getsampwidth(), wr.getframerate())
+
+    return wav_bytes, pcm_bytes, params
+
+
+def _make_ogg_bytes(sample_rate=16000, channels=1, duration_sec=1.0):
+    """Generate an OGG Vorbis file with a sine wave, return bytes."""
+    n_samples = int(sample_rate * duration_sec)
+    t = np.linspace(0, duration_sec, n_samples, endpoint=False)
+    samples = (np.sin(2 * np.pi * 440 * t) * 0.5).astype(np.float32)
+    buf = io.BytesIO()
+    sf.write(buf, samples, sample_rate, format="OGG", subtype="VORBIS")
+    return buf.getvalue()
+
+
+def test_wav_path_duration_and_frames():
+    """Path 1: WAV PCM buffer → duration & frames via arithmetic."""
+    for sr in [16000, 24000, 44100]:
+        for ch in [1, 2]:
+            for dur in [0.5, 1.0, 2.5]:
+                _, pcm_bytes, (channels, sample_width, frame_rate) = _make_wav_bytes(
+                    sample_rate=sr, channels=ch, duration_sec=dur
+                )
+
+                # New code logic (from patch.py)
+                frame_width = sample_width * channels
+                audio_frames = len(pcm_bytes) // frame_width
+                audio_duration_sec = audio_frames / frame_rate
+
+                assert abs(audio_duration_sec - dur) < 0.001, f"sr={sr} ch={ch} dur={dur}: got {audio_duration_sec}"
+                expected_frames = int(sr * dur)
+                assert audio_frames == expected_frames, (
+                    f"sr={sr} ch={ch} dur={dur}: frames {audio_frames} != {expected_frames}"
+                )
+
+    print("[PASS] test_wav_path_duration_and_frames")
+
+
+def test_non_wav_path_duration_av():
+    """Path 2: non-WAV audio → duration & frames via av."""
+    for dur in [0.5, 1.0, 2.0]:
+        ogg_bytes = _make_ogg_bytes(sample_rate=16000, duration_sec=dur)
+
+        container = av.open(io.BytesIO(ogg_bytes))
+        stream = container.streams.audio[0]
+        audio_frames = sum(f.samples for f in container.decode(stream))
+        audio_duration_sec = audio_frames / stream.rate if stream.rate else 0.0
+        container.close()
+
+        # OGG encoding may add padding, allow small tolerance
+        assert abs(audio_duration_sec - dur) < 0.05, f"dur={dur}: got {audio_duration_sec}"
+
+    print("[PASS] test_non_wav_path_duration_av")
+
+
+def test_tts_pcm_export_wav_path():
+    """Path 3a: WAV PCM → resample to 24000Hz mono s16."""
+    import torch
+    import torchaudio
+
+    _, pcm_bytes, (channels, sample_width, frame_rate) = _make_wav_bytes(
+        sample_rate=16000, channels=2, sample_width=2, duration_sec=1.0
+    )
+
+    # New code logic
+    pcm_array = np.frombuffer(pcm_bytes, dtype=np.int16)
+    if channels > 1:
+        pcm_array = pcm_array.reshape(-1, channels).mean(axis=1).astype(np.int16)
+    if frame_rate != 24000:
+        t = torch.from_numpy(pcm_array.astype(np.float32)).unsqueeze(0)
+        t = torchaudio.functional.resample(t, frame_rate, 24000)
+        pcm_array = t.squeeze(0).to(torch.int16).numpy()
+
+    result_bytes = pcm_array.tobytes()
+    # Should be mono, 24000Hz, 2 bytes per sample
+    expected_samples = int(24000 * 1.0)
+    actual_samples = len(result_bytes) // 2
+    assert abs(actual_samples - expected_samples) < 10, f"expected ~{expected_samples} samples, got {actual_samples}"
+    assert len(pcm_array.shape) == 1, "should be mono (1D array)"
+
+    print("[PASS] test_tts_pcm_export_wav_path")
+
+
+def test_tts_pcm_export_non_wav_path():
+    """Path 3b: non-WAV → resample to 24000Hz mono s16 via av.AudioResampler."""
+    ogg_bytes = _make_ogg_bytes(sample_rate=16000, channels=1, duration_sec=1.0)
+
+    container = av.open(io.BytesIO(ogg_bytes))
+    resampler = av.AudioResampler(format="s16", layout="mono", rate=24000)
+    frames = []
+    for frame in container.decode(audio=0):
+        for rf in resampler.resample(frame):
+            frames.append(rf.to_ndarray().flatten())
+    container.close()
+
+    result = np.concatenate(frames)
+    result_bytes = result.tobytes()
+    expected_samples = int(24000 * 1.0)
+    actual_samples = len(result_bytes) // 2
+    assert abs(actual_samples - expected_samples) < 500, f"expected ~{expected_samples} samples, got {actual_samples}"
+
+    print("[PASS] test_tts_pcm_export_non_wav_path")
+
+
+def test_media_merge_and_export():
+    """media.py: _merge_base64_audio_to_segment decode + concat + export WAV."""
+    import base64 as b64mod
+    import sys
+
+    sys.path.insert(0, ".")
+    from tests.helpers.media import _merge_base64_audio_to_segment
+
+    # Create two WAV chunks as base64
+    chunks_b64 = []
+    for dur in [0.5, 0.3]:
+        wav_bytes, _, _ = _make_wav_bytes(sample_rate=16000, channels=1, sample_width=2, duration_sec=dur)
+        chunks_b64.append(b64mod.b64encode(wav_bytes).decode())
+
+    merged = _merge_base64_audio_to_segment(chunks_b64)
+
+    # Test export to WAV
+    buf = io.BytesIO()
+    merged.export(buf, format="wav")
+    wav_out = buf.getvalue()
+    assert len(wav_out) > 44, f"WAV output too small: {len(wav_out)} bytes"
+
+    # Verify merged duration (~0.8s)
+    out_data, out_sr = sf.read(io.BytesIO(wav_out))
+    actual_dur = len(out_data) / out_sr
+    assert abs(actual_dur - 0.8) < 0.05, f"expected ~0.8s merged duration, got {actual_dur}"
+
+    print(f"[PASS] test_media_merge_and_export (duration={actual_dur:.3f}s, wav_size={len(wav_out)} bytes)")
+
+
+if __name__ == "__main__":
+    test_wav_path_duration_and_frames()
+    test_non_wav_path_duration_av()
+    import importlib.util
+
+    if importlib.util.find_spec("torch") and importlib.util.find_spec("torchaudio"):
+        test_tts_pcm_export_wav_path()
+    else:
+        print("[SKIP] test_tts_pcm_export_wav_path (torch/torchaudio not available)")
+    test_tts_pcm_export_non_wav_path()
+    test_media_merge_and_export()
+    print("\nAll tests passed.")

tests/worker/test_omni_gpu_model_runner.py

@@ -22,6 +22,7 @@ class DummyInputBatch:
 
     def __init__(self, req_ids):
         self.req_ids = req_ids
+        self.req_id_to_index = {r: i for i, r in enumerate(req_ids)}
 
 
 class DummyReqState:

vllm_omni/benchmarks/patch/patch.py

@@ -16,8 +16,8 @@ from datetime import datetime
 from typing import Any, Literal
 
 import aiohttp
+import numpy as np
 import pybase64 as base64
-from pydub import AudioSegment
 from tqdm.asyncio import tqdm
 from vllm.benchmarks import datasets
 from vllm.benchmarks.datasets import SampleRequest
@@ -653,9 +653,8 @@ async def async_request_openai_chat_omni_completions(
         # Reset per-attempt state so that retries do not mix partial
         # outputs or metrics from previous attempts.
         generated_text = ""
-        generated_audio = None
         # For wav responses, accumulate decoded PCM bytes per chunk
-        # to avoid repeated AudioSegment decode/concat.
+        # to avoid repeated decode/concat.
         wav_pcm_buffer = bytearray()
         wav_audio_params: tuple[int, int, int] | None = None
         wav_inconsistent_chunk_count = 0
@@ -795,30 +794,33 @@ async def async_request_openai_chat_omni_completions(
 
                     output.latency = timestamp - st
                     output.generated_text = generated_text
+                    audio_duration_sec = 0.0
+                    audio_frames = 0
                     if response_format == "wav" and wav_pcm_buffer and wav_audio_params is not None:
                         channels, sample_width, frame_rate = wav_audio_params
-                        generated_audio = AudioSegment(
-                            data=bytes(wav_pcm_buffer),
-                            sample_width=sample_width,
-                            frame_rate=frame_rate,
-                            channels=channels,
-                        )
+                        frame_width = sample_width * channels
+                        if frame_width > 0:
+                            audio_frames = len(wav_pcm_buffer) // frame_width
+                            audio_duration_sec = audio_frames / frame_rate
+                        else:
+                            logger.warning("Audio frame width is zero")
                     elif audio_bytes_buffer:
                         try:
-                            generated_audio = AudioSegment.from_file(
+                            from vllm.multimodal.audio import get_audio_duration
+                            from vllm.multimodal.media.audio import load_audio
+
+                            waveform, sr = load_audio(
                                 io.BytesIO(bytes(audio_bytes_buffer)),
-                                format=response_format,
+                                sr=None,
+                                mono=False,
                             )
+                            audio_duration_sec = get_audio_duration(y=waveform, sr=sr)
+                            audio_frames = int(audio_duration_sec * sr)
                         except Exception as ex:
                             logger.warning("Failed to decode accumulated audio bytes: %s", ex)
-                    if generated_audio is not None:
-                        output.audio_duration = len(generated_audio) / 1000.0
-                        frame_width = generated_audio.frame_width
-                        if frame_width > 0:
-                            output.audio_frames = len(generated_audio.raw_data) // frame_width
-                        else:
-                            output.audio_frames = 0
-                            logger.warning("Audio frame width is zero")
+                    if audio_duration_sec > 0 or audio_frames > 0:
+                        output.audio_duration = audio_duration_sec
+                        output.audio_frames = audio_frames
                         audio_duration = output.audio_duration
                         if audio_duration > 0:
                             output.audio_rtf = audio_generate_time / output.audio_duration
@@ -827,8 +829,29 @@ async def async_request_openai_chat_omni_completions(
                             logger.warning("Audio duration is zero")
                         if _seed_tts_capture_pcm_for_wer() and getattr(request_func_input, "seed_tts_row", False):
                             try:
-                                seg = generated_audio.set_frame_rate(24000).set_channels(1).set_sample_width(2)
-                                output.tts_output_pcm_bytes = bytes(seg.raw_data)
+                                if response_format == "wav" and wav_pcm_buffer and wav_audio_params is not None:
+                                    from vllm.multimodal.audio import AudioResampler
+
+                                    pcm_channels, pcm_sw, pcm_rate = wav_audio_params
+                                    pcm = np.frombuffer(
+                                        bytes(wav_pcm_buffer), dtype=np.int16 if pcm_sw == 2 else np.float32
+                                    )
+                                    if pcm_channels > 1:
+                                        pcm = pcm.reshape(-1, pcm_channels).mean(axis=1).astype(pcm.dtype)
+                                    pcm_f32 = pcm.astype(np.float32) / 32767.0 if pcm.dtype == np.int16 else pcm
+                                    if pcm_rate != 24000:
+                                        resampler = AudioResampler(target_sr=24000)
+                                        pcm_f32 = resampler.resample(pcm_f32, orig_sr=pcm_rate)
+                                    output.tts_output_pcm_bytes = (pcm_f32 * 32767).astype(np.int16).tobytes()
+                                elif audio_bytes_buffer:
+                                    from vllm.multimodal.media.audio import load_audio
+
+                                    waveform, _ = load_audio(
+                                        io.BytesIO(bytes(audio_bytes_buffer)),
+                                        sr=24000,
+                                        mono=True,
+                                    )
+                                    output.tts_output_pcm_bytes = (waveform * 32767).astype(np.int16).tobytes()
                             except Exception as ex:
                                 logger.warning("seed_tts WER PCM export failed: %s", ex)
                     output.success = True

vllm_omni/config/pipeline_registry.py

@@ -101,6 +101,10 @@ _OMNI_PIPELINES: dict[str, tuple[str, str]] = {
         "vllm_omni.model_executor.models.mimo_audio.pipeline",
         "MIMO_AUDIO_PIPELINE",
     ),
+    "ming_tts": (
+        "vllm_omni.model_executor.models.ming_tts.pipeline",
+        "MING_TTS_PIPELINE",
+    ),
     "voxtral_tts": (
         "vllm_omni.model_executor.models.voxtral_tts.pipeline",
         "VOXTRAL_TTS_PIPELINE",

vllm_omni/config/stage_config.py

@@ -1128,6 +1128,19 @@ class StageConfigFactory:
                             continue
                     return cls._create_from_registry(registered.model_type, cli_overrides, deploy_config_path)
 
+        # An explicit deploy config can select a pipeline when the model's HF
+        # metadata is too generic for registry auto-detection.
+        if deploy_config_path is not None:
+            deploy_path = Path(deploy_config_path)
+            if deploy_path.exists():
+                deploy_config = load_deploy_config(deploy_path)
+                if deploy_config.pipeline in _PIPELINE_REGISTRY:
+                    return cls._create_from_registry(
+                        deploy_config.pipeline,
+                        cli_overrides,
+                        deploy_config_path,
+                    )
+
         # Not in the pipeline registry — let the caller fall back to the
         # legacy ``stage_configs/*.yaml`` path (resolve_model_config_path).
         return None

vllm_omni/core/sched/omni_ar_scheduler.py

@@ -218,7 +218,9 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         self._process_pending_input_timeouts()
 
         if self.chunk_transfer_adapter:
-            self.chunk_transfer_adapter.process_pending_chunks(self.waiting, self.running)
+            self.chunk_transfer_adapter.process_pending_chunks(
+                self.waiting, self.running, scheduler_requests=self.requests
+            )
 
         try:
             scheduler_output = super().schedule()
@@ -626,7 +628,34 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         if self.chunk_transfer_adapter:
             self.chunk_transfer_adapter.finish_requests(request_ids, finished_status, self.requests)
 
+        # Realign stale ``request.status`` (chunk-transfer-adapter's
+        # ``requests_origin_status`` table doesn't follow the
+        # ``waiting → running`` admit transition; without this, an abort
+        # arriving between admit and the next deque round-trip leaves
+        # the request in ``self.running`` with ``status=WAITING`` and
+        # upstream ``Scheduler.finish_requests`` silently fails to
+        # release the worker's ``input_batch`` slot -- after
+        # ``max_num_seqs`` such aborts new requests hang at
+        # ``chunks=0``). Only the ``async_chunk`` path triggers the
+        # staleness; with ``async_chunk`` disabled this is a cheap O(n)
+        # no-op over an already-aligned set, kept unconditional so the
+        # abort path stays uniform across configurations. See
+        # ``OmniSchedulerMixin._realign_request_status_to_queues`` and
+        # #3774 discussion.
+        self._realign_request_status_to_queues(request_ids)
+
         finished = super().finish_requests(request_ids, finished_status)
+
+        # Defensive post-finish purge: belt-and-suspenders to the
+        # realignment above. Even after realign + ``super()``, corner
+        # cases (mid-transition status, connector cleanups that pop
+        # from ``self.requests`` without unwinding ``self.running``)
+        # can leave already-finished or untracked entries in
+        # ``self.running``. Sweep them now so the worker's
+        # ``input_batch`` slot never pins a freed request and starves
+        # new admissions. See ``OmniSchedulerMixin._purge_finished_from_running``.
+        self._purge_finished_from_running()
+
         input_coordinator = getattr(self, "input_coordinator", None)
         if input_coordinator is not None:
             for request_id, _ in finished:

vllm_omni/core/sched/omni_generation_scheduler.py

@@ -87,7 +87,9 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
         self._consume_pending_connector_output(model_mode="generation")
         self._process_pending_input_timeouts()
         if self.chunk_transfer_adapter:
-            self.chunk_transfer_adapter.process_pending_chunks(self.waiting, self.running)
+            self.chunk_transfer_adapter.process_pending_chunks(
+                self.waiting, self.running, scheduler_requests=self.requests
+            )
 
         # OMNI: Track requests that are already finished (e.g., marked by connector)
         # These should be removed from running and not scheduled
@@ -364,7 +366,23 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
         if self.chunk_transfer_adapter:
             self.chunk_transfer_adapter.finish_requests(request_ids, finished_status, self.requests)
 
+        # See ``OmniSchedulerMixin._realign_request_status_to_queues`` --
+        # closes the residual hang on the
+        # ``waiting → running → abort-before-next-deque-round-trip`` race
+        # surfaced by #3774's reproduction matrix. Only the
+        # ``async_chunk`` path actually triggers the staleness; with
+        # ``async_chunk`` disabled this is a cheap O(n) no-op, kept
+        # unconditional so the abort path stays uniform.
+        self._realign_request_status_to_queues(request_ids)
+
         finished = super().finish_requests(request_ids, finished_status)
+
+        # See ``OmniSchedulerMixin._purge_finished_from_running`` --
+        # defensive belt-and-suspenders sweep paired with the realign
+        # above. Closes residual ``self.running`` slot leaks if any
+        # corner case slips past upstream's status-driven removal.
+        self._purge_finished_from_running()
+
         if self.input_coordinator is not None:
             for request_id, _ in finished:
                 self._free_input_coordinator_request(request_id)

vllm_omni/core/sched/omni_scheduler_mixin.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import os
 import time
+from collections.abc import Iterable
 from typing import Any
 
 from vllm.logger import init_logger
@@ -156,3 +157,122 @@ class OmniSchedulerMixin:
             return None
         self._last_stats_time = now
         return super().make_stats(*args, **kwargs)
+
+    def _realign_request_status_to_queues(
+        self,
+        request_ids: str | Iterable[str] | None,
+    ) -> None:
+        """Realign ``request.status`` to actual queue membership.
+
+        ``OmniChunkTransferAdapter._process_chunk_queue`` stamps
+        ``requests_origin_status[req.id] = WAITING`` (or ``RUNNING``) when
+        first parking a request in a chunk-transfer deque. On the next
+        tick, when the chunk arrives, ``_process_chunk_queue`` sets
+        ``request.status = target_status`` and continues, but
+        ``requests_origin_status`` is left at its first-park value -- no
+        hook updates it on the ``waiting → running`` admit transition
+        that ``super().schedule()`` later performs. The table stays
+        stale until the request makes another deque round-trip.
+
+        If an abort lands in the gap between admit and the next deque
+        round-trip, ``chunk_transfer_adapter.finish_requests`` reads the
+        stale ``WAITING`` from ``requests_origin_status``, stomps it
+        onto ``request.status``, and the upstream
+        ``Scheduler.finish_requests`` else branch silently fails to
+        remove from ``self.running`` -- the request stays alive in
+        ``self.running`` and the worker's ``input_batch`` slot leaks.
+        After ``max_num_seqs`` such aborts every new request hangs at
+        ``chunks=0`` until the client times out.
+
+        Realign here: if a request lives in ``self.running`` but its
+        status is not ``RUNNING``, set it to ``RUNNING``; symmetrically
+        flip ``RUNNING → WAITING`` when the request is actually in
+        ``self.waiting``. This is a localized safety net for
+        ``requests_origin_status`` staleness on the admit transition;
+        it does not touch the adapter's invariants and is complementary
+        to the chunk-transfer-adapter deque purge that already runs
+        inside ``process_pending_chunks`` / ``restore_queues``.
+
+        Note on scope: only the ``async_chunk`` path actually triggers
+        the ``requests_origin_status`` staleness this helper repairs.
+        When ``async_chunk`` is disabled, no chunk-transfer round-trip
+        occurs between admit and finish, so the realignment walk is a
+        cheap O(n) no-op over an already-aligned set. The call is kept
+        unconditional in ``finish_requests`` to (a) keep the abort path
+        uniform and (b) defend any future configuration that re-enables
+        chunk transfer from rediscovering the same regression.
+
+        See https://github.com/vllm-project/vllm-omni/pull/3774 and the
+        residual-hang reproduction discussed in that PR.
+        """
+        # Mirror the upstream Scheduler.finish_requests resolution of
+        # ``request_ids`` so realignment touches exactly the set that
+        # ``super().finish_requests`` will then walk.
+        if isinstance(request_ids, str):
+            ids_to_align: Iterable[str] = (request_ids,)
+        elif request_ids is None:
+            ids_to_align = list(self.requests.keys())
+        else:
+            ids_to_align = list(request_ids)
+
+        if not ids_to_align:
+            return
+
+        running_ids = {r.request_id for r in self.running}
+        waiting_ids = {r.request_id for r in self.waiting}
+
+        for rid in ids_to_align:
+            req = self.requests.get(rid)
+            if req is None or req.is_finished():
+                continue
+            if rid in running_ids and req.status != RequestStatus.RUNNING:
+                req.status = RequestStatus.RUNNING
+            elif rid in waiting_ids and req.status == RequestStatus.RUNNING:
+                req.status = RequestStatus.WAITING
+
+    def _purge_finished_from_running(self) -> None:
+        """Defensive post-finish sweep of ``self.running``.
+
+        Belt-and-suspenders to ``_realign_request_status_to_queues``:
+        even after status realignment lets upstream
+        ``Scheduler.finish_requests`` pick the right removal branch,
+        a future regression or an unexpected ``status`` mid-transition
+        could still leave already-finished entries in ``self.running``.
+        Sweeping here guarantees the worker's ``input_batch`` slot is
+        not pinned by a freed request.
+
+        Complementary to ``_realign_request_status_to_queues``: realign
+        is preventive (fix ``status`` before ``super().finish_requests``
+        so the right branch fires); this purge is defensive (sweep the
+        residue after ``super().finish_requests`` so any stale entries
+        are reclaimed).
+
+        Scope of the predicate. ``is_finished()`` covers entries the
+        upstream ``finish_requests`` already drained from ``self.requests``
+        but failed to remove from ``self.running``; the
+        ``request_id not in self.requests`` arm catches the same surface
+        from a different angle and is the post-cleanup mirror of the
+        deque purge ``_purge_untracked_chunk_requests`` already runs at
+        the chunk-transfer-adapter layer. It does **not** by itself make
+        arbitrary direct deletions of ``self.requests`` safe -- callers
+        that pop ``self.requests`` outside the standard finish path
+        still have to go through ``_free_request`` (or equivalent) for
+        block / connector / coordinator cleanup. This sweep only
+        reclaims the ``self.running`` slot reference.
+
+        In-place via ``self.running[:] = ...`` for minor consistency
+        with idiomatic vLLM scheduler mutation; upstream
+        ``Scheduler.finish_requests`` itself rebinds ``self.running``,
+        so list identity across the whole call is not preserved -- the
+        slice form is just to avoid an extra rebind inside this helper.
+
+        Assumes the upstream V1 invariant that scheduler ticks are
+        serialized on a single thread; in-place mutation here is no more
+        racy than the rest of the scheduler under that assumption.
+
+        See https://github.com/vllm-project/vllm-omni/pull/3774
+        discussion.
+        """
+        if not self.running:
+            return
+        self.running[:] = [req for req in self.running if not req.is_finished() and req.request_id in self.requests]

vllm_omni/data_entry_keys.py

@@ -128,6 +128,7 @@ class EmbeddingsStruct(_StructBase):
     voice: torch.Tensor | None = None
     speech_feat: torch.Tensor | None = None
     speech_token: torch.Tensor | None = None
+    speech_token_len: torch.Tensor | None = None
     embedding: torch.Tensor | None = None
     thinker_reply: torch.Tensor | None = None
 

vllm_omni/deploy/cosyvoice3.yaml

@@ -7,9 +7,22 @@
 # references stay in the yaml unconditionally; the runtime only activates
 # them when ``async_chunk: true``.
 #
-# enforce_eager=true everywhere:
-#   * Stage 0 (talker) — cudagraph not verified on this checkpoint.
-#   * Stage 1 (code2wav) — CUDA graphs don't work with dynamic conv shapes.
+# enforce_eager:
+#   * Stage 0 (talker) — enforce_eager: false (CUDA graphs enabled).
+#   * Stage 1 (code2wav) — enforce_eager: true (CUDA graphs don't work with
+#     dynamic conv shapes).
+#
+# Streaming continuity is GPU-class sensitive — TUNE PER GPU.
+#   The defaults below (max_num_seqs: 8, codec_chunk_frames: 15) are tuned for
+#   H100, where streaming continuity at concurrency 4 stays ~100%. They push
+#   more audio per scheduler step through the SharedMemoryConnector than a
+#   slower GPU can keep up with: on H20-3e the same config regressed continuity
+#   to ~78% at c=4 (vs ~86% with smaller batch/chunk), because code2wav can't
+#   produce mel chunks fast enough to feed the streaming output → buffer
+#   underruns → audible gaps.
+#   On a slower GPU class, lower max_num_seqs (e.g. 4) and/or codec_chunk_frames
+#   first. Raising connector_get_sleep_s also relaxes the consumer's polling
+#   cadence, but trades against TTFP — keep it as a last resort.
 async_chunk: true
 dtype: float32
 
@@ -21,16 +34,16 @@ connectors:
       connector_get_sleep_s: 0.01
       connector_get_max_wait_first_chunk: 3000
       connector_get_max_wait: 300
-      codec_chunk_frames: 25
+      codec_chunk_frames: 15
       codec_left_context_frames: 25
       codec_vocab_size: 6561
 
 stages:
   - stage_id: 0
     max_num_batched_tokens: 32768
-    max_num_seqs: 1
+    max_num_seqs: 8
     gpu_memory_utilization: 0.4
-    enforce_eager: true
+    enforce_eager: false
     trust_remote_code: true
     enable_prefix_caching: false
     devices: "0"
@@ -49,7 +62,7 @@ stages:
 
   - stage_id: 1
     max_num_batched_tokens: 32768
-    max_num_seqs: 1
+    max_num_seqs: 8
     gpu_memory_utilization: 0.2
     enforce_eager: true
     trust_remote_code: true

vllm_omni/deploy/hunyuan_image_3_moe.yaml

@@ -62,6 +62,7 @@ stages:
       max_tokens: 8192
       detokenize: true
       skip_special_tokens: false
+      include_stop_str_in_output: true,
 
   - stage_id: 1
     max_num_seqs: 1

vllm_omni/deploy/ming_tts.yaml

@@ -0,0 +1,55 @@
+# Ming-omni-tts deploy: stage-0 LLM+flow -> stage-1 audio VAE.
+# Verified legacy settings migrated from ming_tts.yaml and ming_tts_async_chunk.yaml.
+# Default mode is async-chunk streaming on a single GPU.
+# pipeline: required — HF config reports model_type="dense" with no architectures,
+# so auto-detection falls through. This key routes the deploy-config probe in
+# StageConfigFactory.create_from_model to the correct registry entry.
+pipeline: ming_tts
+async_chunk: true
+trust_remote_code: false
+dtype: bfloat16
+enable_prefix_caching: false
+
+connectors:
+  connector_of_shared_memory:
+    name: SharedMemoryConnector
+    extra:
+      latent_chunk_size: 25
+      latent_left_context: 0
+
+stages:
+  - stage_id: 0
+    max_num_seqs: 1
+    gpu_memory_utilization: 0.45
+    enforce_eager: true
+    enable_prefix_caching: false
+    async_scheduling: false
+    max_num_batched_tokens: 8192
+    max_model_len: 8192
+    devices: "0"
+    output_connectors:
+      to_stage_1: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 512
+      detokenize: true
+
+  - stage_id: 1
+    max_num_seqs: 1
+    gpu_memory_utilization: 0.25
+    enforce_eager: true
+    enable_prefix_caching: false
+    async_scheduling: false
+    max_num_batched_tokens: 8192
+    max_model_len: 8192
+    devices: "0"
+    input_connectors:
+      from_stage_0: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 1
+      detokenize: false

vllm_omni/deploy/qwen3_omni_moe_mori_intranode.yaml

@@ -0,0 +1,107 @@
+# Qwen3-Omni-MoE intra-node deploy using Mori XGMI transfer.
+#
+# Targets a single AMD Instinct MI300X OAM node (8x GPU, 192GB HBM3,
+# gfx942).  Each stage occupies one dedicated GPU (3 GPUs total) and
+# stage-to-stage streaming chunks are shipped via Mori's RDMA IOEngine
+# (with ``backend_type: xgmi`` selecting the AMD Infinity Fabric
+# GPU-to-GPU direct path for intranode transfers).
+#
+# Pipeline wiring (stage_id, execution_type, model_stage, input_sources,
+# custom processors, async-chunk builders) lives in
+# ``vllm_omni/model_executor/models/qwen3_omni/pipeline.py`` and is
+# merged into the stages below by the deploy loader.
+#
+# Why this yaml exists:
+#   * On CUDA/H100 Qwen3-Omni-MoE ships via ``SharedMemoryConnector``
+#     (see ``qwen3_omni_moe.yaml``).  On MI300X the SHM path still works
+#     but XGMI RDMA can substantially reduce stage-transition latency
+#     for the thinker->talker hidden-state stream and the talker->code2wav
+#     codec-frame stream, both of which run in async-chunk mode.
+#   * This config swaps the connector definition to
+#     ``MoriTransferEngineConnector`` while keeping the pipeline topology,
+#     sampling params, and async_chunk=true gating identical to the CUDA
+#     deploy, so the chunk_transfer_adapter path (scheduler-owned
+#     background put/get threads) is fully exercised.
+#
+# Usage:
+#   vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --log-stats \
+#     --deploy-config vllm_omni/deploy/qwen3_omni_moe_mori_intranode.yaml
+#
+# Dependencies: mori (``pip install mori``), torch-rocm>=2.9, vllm-omni
+# with #2383 deploy schema, ibverbs/RDMA NICs exposed via /dev/infiniband.
+#
+# Limitation: TP=1-only.  Connector raises NotImplementedError on
+# tp_world_size > 1 (chunk-path endpoint derivation isn't rank-aware yet).
+async_chunk: true
+
+connectors:
+  mori_connector:
+    name: MoriTransferEngineConnector
+    extra:
+      host: "auto"
+      zmq_port: 50051
+      backend_type: "xgmi"            # AMD Infinity Fabric GPU-to-GPU direct
+      device_name: ""                 # leave empty to honour $MORI_RDMA_DEVICES
+      memory_pool_size: 536870912     # 512 MB, matches Qwen2.5-Omni Mori deploy
+      memory_pool_device: "cuda"      # GPU memory required for XGMI RDMA transfers
+      xgmi_num_streams: 64
+      xgmi_num_events: 64
+      qp_per_transfer: 1
+      num_worker_threads: 1
+      post_batch_size: -1
+      # Chunk-transfer knobs consumed by talker2code2wav_async_chunk:
+      # 25 decoded codec frames per connector.put (~40 ms at 625 Hz
+      # codec rate) with a 25-frame left context to satisfy the
+      # code2wav receptive field.  Kept identical to the SHM deploy so
+      # accuracy stays comparable when swapping backends.
+      codec_chunk_frames: 25
+      codec_left_context_frames: 25
+
+stages:
+  - stage_id: 0
+    gpu_memory_utilization: 0.9
+    enforce_eager: true               # ROCm override: avoid flashinfer autotune path
+    devices: "0"
+    output_connectors:
+      to_stage_1: mori_connector
+    default_sampling_params:
+      temperature: 0.4
+      top_p: 0.9
+      top_k: 1
+      max_tokens: 2048
+      seed: 42
+      repetition_penalty: 1.05
+
+  - stage_id: 1
+    gpu_memory_utilization: 0.9
+    enforce_eager: true
+    devices: "1"
+    input_connectors:
+      from_stage_0: mori_connector
+    output_connectors:
+      to_stage_2: mori_connector
+    default_sampling_params:
+      temperature: 0.9
+      top_k: 50
+      max_tokens: 4096
+      seed: 42
+      repetition_penalty: 1.05
+
+  - stage_id: 2
+    gpu_memory_utilization: 0.3
+    max_num_seqs: 1
+    enforce_eager: true
+    async_scheduling: false
+    # Codec prefill length (Q * num_frames) exceeds 32k default;
+    # matches qwen3_omni_moe.yaml code2wav sizing.
+    max_num_batched_tokens: 51200
+    devices: "2"
+    input_connectors:
+      from_stage_1: mori_connector
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 65536
+      seed: 42
+      repetition_penalty: 1.1

vllm_omni/deploy/voxcpm2.yaml

@@ -1,5 +1,6 @@
 # VoxCPM2 deploy: single-stage AR pipeline with per-request state batching.
-# Verified on 1x H20 141GB.
+# Runs on large (H20 141GB) and small (24GB L4) cards alike: the KV cache is
+# pinned to a right-sized budget instead of ballooning to fill the card.
 #
 # VoxCPM2-local runtime knobs are passed through hf_overrides.
 async_chunk: false
@@ -9,6 +10,16 @@ stages:
   - stage_id: 0
     max_num_seqs: 32
     gpu_memory_utilization: 0.9
+    # Right-size the KV cache instead of letting it fill gpu_memory_utilization.
+    # max_num_seqs(32) * max_model_len(4096) ~= 131k tokens; 6 GiB holds ~175k
+    # (42x concurrency, above the 32-seq scheduler cap, so no throughput loss).
+    # The remaining VRAM stays free for the VoxCPM2 diffusion side-path
+    # (LocDiT / AudioVAE / code2wav + batch CUDA-graph pools), which allocates
+    # outside vLLM's KV accounting during inference. Keeps peak ~13 GiB on any
+    # card and avoids per-rebase OOM tuning on small CI GPUs.
+    # Assumes prefix caching off and max_num_seqs*max_model_len <= ~175k tokens;
+    # raise this cap if either grows or enable_prefix_caching is turned on.
+    kv_cache_memory_bytes: 6442450944  # 6 GiB
     enforce_eager: true
     trust_remote_code: true
     enable_prefix_caching: false

vllm_omni/diffusion/attention/backends/sage_attn.py

@@ -1,6 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import os
+
 import torch
 from vllm.logger import init_logger
 
@@ -9,17 +11,43 @@ from vllm_omni.diffusion.attention.backends.abstract import (
     AttentionImpl,
     AttentionMetadata,
 )
+from vllm_omni.platforms import current_omni_platform
 
 logger = init_logger(__name__)
 
-try:
-    from sageattention import sageattn
-except ImportError:
-    logger.warning(
-        "SageAttentionBackend is not available. You may install sage-attention"
-        " by pip install git+https://github.com/thu-ml/SageAttention.git"
-    )
-    raise ImportError
+_SAGE_ATTN_TENSOR_LAYOUT = os.environ.get("SAGE_ATTN_TENSOR_LAYOUT", "NHD").upper()
+assert _SAGE_ATTN_TENSOR_LAYOUT in ("NHD", "HND"), (
+    f"SAGE_ATTN_TENSOR_LAYOUT must be 'NHD' or 'HND', got '{_SAGE_ATTN_TENSOR_LAYOUT}'"
+)
+
+if current_omni_platform.is_xpu():
+    try:
+        import inspect
+
+        from auto_round_kernel import ARK
+
+        _ark = ARK()
+        xpu_sageattn = _ark.sagev1
+        _sagev1_params = inspect.signature(xpu_sageattn).parameters
+        _sagev1_has_tensor_layout = "tensor_layout" in _sagev1_params
+        _sagev1_scale_param = "sm_scale" if "sm_scale" in _sagev1_params else "scale"
+    except ImportError:
+        logger.warning(
+            "XPU SageAttention (auto_round_kernel.ARK.sagev1) is not available. "
+            "Install auto-round-lib for XPU sage attention support."
+        )
+        xpu_sageattn = None
+        _sagev1_has_tensor_layout = False
+        _sagev1_scale_param = "scale"
+else:
+    try:
+        from sageattention import sageattn
+    except ImportError:
+        logger.warning(
+            "SageAttentionBackend is not available. You may install sage-attention"
+            " by pip install git+https://github.com/thu-ml/SageAttention.git"
+        )
+        raise ImportError
 
 # TODO add sage3 attention backend
 
@@ -73,3 +101,50 @@ class SageAttentionImpl(AttentionImpl):
             sm_scale=self.softmax_scale,
         )
         return output
+
+    def forward_xpu(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata = None,
+    ) -> torch.Tensor:
+        if xpu_sageattn is None:
+            raise ImportError("XPU SageAttention requires auto-round-lib. Install with: pip install auto-round-lib")
+        orig_dtype = query.dtype
+        q = query.to(torch.float16) if orig_dtype != torch.float16 else query
+        k = key.to(torch.float16) if orig_dtype != torch.float16 else key
+        v = value.to(torch.float16) if orig_dtype != torch.float16 else value
+
+        if _sagev1_has_tensor_layout:
+            if _SAGE_ATTN_TENSOR_LAYOUT == "HND":
+                q = q.transpose(1, 2).contiguous()
+                k = k.transpose(1, 2).contiguous()
+                v = v.transpose(1, 2).contiguous()
+            output = xpu_sageattn(
+                q,
+                k,
+                v,
+                tensor_layout=_SAGE_ATTN_TENSOR_LAYOUT,
+                is_causal=self.causal,
+                **{_sagev1_scale_param: self.softmax_scale},
+            )
+            if _SAGE_ATTN_TENSOR_LAYOUT == "HND":
+                output = output.transpose(1, 2).contiguous()
+        else:
+            # No tensor_layout support: kernel expects HND [B, H, S, D]
+            q = q.transpose(1, 2).contiguous()
+            k = k.transpose(1, 2).contiguous()
+            v = v.transpose(1, 2).contiguous()
+            output = xpu_sageattn(
+                q,
+                k,
+                v,
+                is_causal=self.causal,
+                **{_sagev1_scale_param: self.softmax_scale},
+            )
+            output = output.transpose(1, 2).contiguous()
+
+        if orig_dtype != torch.float16:
+            output = output.to(orig_dtype)
+        return output

vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_ltx2.py

@@ -0,0 +1,210 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from typing import Any
+
+import torch
+from diffusers import AutoencoderKLLTX2Video
+from diffusers.models.autoencoders.vae import DecoderOutput
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
+    DistributedOperator,
+    DistributedVaeExecutor,
+    DistributedVaeMixin,
+    GridSpec,
+    TileTask,
+)
+
+logger = init_logger(__name__)
+
+
+class LTX2VaeExecutor(DistributedVaeExecutor):
+    def _pack_local_tiles_without_meta(self, local_results, global_padding_shape, device, dtype):
+        tile_tensor = torch.zeros(global_padding_shape, device=device, dtype=dtype)
+        for idx, (_, t_tensor) in enumerate(local_results):
+            slices = tuple(slice(0, s) for s in t_tensor.shape)
+            tile_tensor[idx][slices] = t_tensor
+        return tile_tensor
+
+    def _unpack_tiles_with_known_metadata(self, tile_gather, assigned, tile_output_shapes):
+        coord_tensor_map = {}
+        for r, rank_tasks in enumerate(assigned):
+            tiles_src = tile_gather[r]
+            for idx, task in enumerate(rank_tasks):
+                output_shape = tile_output_shapes[task.tile_id]
+                slices = tuple(slice(0, dim) for dim in output_shape)
+                coord_tensor_map[task.grid_coord] = tiles_src[idx][slices]
+
+        return coord_tensor_map
+
+    def execute(self, z: torch.Tensor, operator: DistributedOperator, broadcast_result: bool = True):
+        pp_size = min(self.parallel_size, self.world_size)
+
+        tiletask_list, grid_spec = operator.split(z)
+        max_tile_output_shape = grid_spec.tile_spec.get("max_tile_output_shape")
+        tile_output_shapes = grid_spec.tile_spec.get("tile_output_shapes")
+        if max_tile_output_shape is None or tile_output_shapes is None:
+            return super().execute(z, operator, broadcast_result=broadcast_result)
+
+        assigned = self._balance_tasks(tiletask_list, pp_size)
+        local_tasks = assigned[self.rank] if self.rank < pp_size else []
+        local_results = [(t.tile_id, operator.exec(t)) for t in local_tasks]
+
+        output_dtype = grid_spec.output_dtype if grid_spec.output_dtype is not None else z.dtype
+        max_local_tile_count = max(len(tasks) for tasks in assigned)
+        global_padding_shape = [max_local_tile_count, *max_tile_output_shape]
+        local_tile_tensor = self._pack_local_tiles_without_meta(
+            local_results,
+            global_padding_shape,
+            z.device,
+            output_dtype,
+        )
+        tile_gather = self.gather_tensors(local_tile_tensor)
+
+        if self.rank != 0:
+            result = torch.empty(0, device=z.device)
+        else:
+            coord_tensor_map = self._unpack_tiles_with_known_metadata(tile_gather, assigned, tile_output_shapes)
+            result = operator.merge(coord_tensor_map, grid_spec)
+
+        if broadcast_result:
+            result = self._sync_final_result(result, z.ndim, z.device, output_dtype)
+        return result
+
+
+class DistributedAutoencoderKLLTX2Video(AutoencoderKLLTX2Video, DistributedVaeMixin):
+    @classmethod
+    def from_pretrained(cls, *args: Any, **kwargs: Any):
+        model = super().from_pretrained(*args, **kwargs)
+        model.init_distributed()
+        return model
+
+    def init_distributed(self):
+        self.distributed_executor = LTX2VaeExecutor()
+
+    def tile_split(self, z: torch.Tensor) -> tuple[list[TileTask], GridSpec]:
+        _, _, num_frames, height, width = z.shape
+        sample_height = height * self.spatial_compression_ratio
+        sample_width = width * self.spatial_compression_ratio
+
+        tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
+        tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
+        tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio
+        tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio
+        sample_num_frames = (num_frames - 1) * self.temporal_compression_ratio + 1
+        output_channels = int(getattr(getattr(self, "config", None), "out_channels", 3))
+
+        tiletask_list = []
+        tile_output_shapes = {}
+        for i in range(0, height, tile_latent_stride_height):
+            for j in range(0, width, tile_latent_stride_width):
+                tile = z[:, :, :num_frames, i : i + tile_latent_min_height, j : j + tile_latent_min_width]
+                tile_id = len(tiletask_list)
+                tile_output_shapes[tile_id] = (
+                    z.shape[0],
+                    output_channels,
+                    sample_num_frames,
+                    tile.shape[3] * self.spatial_compression_ratio,
+                    tile.shape[4] * self.spatial_compression_ratio,
+                )
+                tiletask_list.append(
+                    TileTask(
+                        tile_id,
+                        (i // tile_latent_stride_height, j // tile_latent_stride_width),
+                        tile,
+                        workload=tile.shape[2] * tile.shape[3] * tile.shape[4],
+                    )
+                )
+
+        tile_spec = {
+            "sample_height": sample_height,
+            "sample_width": sample_width,
+            "blend_height": self.tile_sample_min_height - self.tile_sample_stride_height,
+            "blend_width": self.tile_sample_min_width - self.tile_sample_stride_width,
+            "tile_sample_stride_height": self.tile_sample_stride_height,
+            "tile_sample_stride_width": self.tile_sample_stride_width,
+            "max_tile_output_shape": (
+                z.shape[0],
+                output_channels,
+                sample_num_frames,
+                self.tile_sample_min_height,
+                self.tile_sample_min_width,
+            ),
+            "tile_output_shapes": tile_output_shapes,
+        }
+        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=tile_spec,
+            output_dtype=self.dtype,
+        )
+        return tiletask_list, grid_spec
+
+    def tile_exec(
+        self,
+        task: TileTask,
+        temb: torch.Tensor | None = None,
+        causal: bool | None = None,
+    ) -> torch.Tensor:
+        if hasattr(self, "clear_cache"):
+            self.clear_cache()
+        return self.decoder(task.tensor, temb, causal=causal)
+
+    def tile_merge(self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid_spec: GridSpec) -> torch.Tensor:
+        grid_h, grid_w = grid_spec.grid_shape
+        result_rows = []
+
+        if hasattr(self, "clear_cache"):
+            self.clear_cache()
+
+        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_sample_stride_height"],
+                        : grid_spec.tile_spec["tile_sample_stride_width"],
+                    ]
+                )
+            result_rows.append(torch.cat(result_row, dim=-1))
+
+        dec = torch.cat(result_rows, dim=3)[
+            :, :, :, : grid_spec.tile_spec["sample_height"], : grid_spec.tile_spec["sample_width"]
+        ]
+        return dec
+
+    def tiled_decode(
+        self,
+        z: torch.Tensor,
+        temb: torch.Tensor | None = None,
+        causal: bool | None = None,
+        return_dict: bool = True,
+    ):
+        if not self.is_distributed_enabled():
+            if causal is None:
+                return super().tiled_decode(z, temb, return_dict=return_dict)
+            return super().tiled_decode(z, temb, causal=causal, return_dict=return_dict)
+
+        logger.debug("LTX2 video VAE decode running with distributed tiled executor")
+        result = self.distributed_executor.execute(
+            z,
+            DistributedOperator(
+                split=self.tile_split,
+                exec=lambda task: self.tile_exec(task, temb=temb, causal=causal),
+                merge=self.tile_merge,
+            ),
+            broadcast_result=False,
+        )
+        if not return_dict:
+            return (result,)
+
+        return DecoderOutput(sample=result)

vllm_omni/diffusion/lora/manager.py

@@ -366,7 +366,11 @@ class DiffusionLoRAManager:
             fully_sharded_loras=False,
         )
 
-        for component_name in ("transformer", "transformer_2", "dit", "bagel"):
+        # Default denoising components plus any a pipeline opts into via
+        # ``_lora_components`` (opt-in; other pipelines unchanged).
+        default_components = ("transformer", "transformer_2", "dit", "bagel")
+        extra_components = tuple(getattr(self.pipeline, "_lora_components", ()) or ())
+        for component_name in (*default_components, *extra_components):
             if not hasattr(self.pipeline, component_name):
                 continue
             component = getattr(self.pipeline, component_name)

vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py

@@ -23,7 +23,7 @@ from typing import Any, ClassVar
 
 import numpy as np
 import torch
-from diffusers import AutoencoderKLLTX2Audio, AutoencoderKLLTX2Video, FlowMatchEulerDiscreteScheduler
+from diffusers import AutoencoderKLLTX2Audio, FlowMatchEulerDiscreteScheduler
 from diffusers.pipelines.ltx2 import LTX2TextConnectors
 from diffusers.pipelines.ltx2.vocoder import LTX2Vocoder
 from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import retrieve_timesteps
@@ -36,6 +36,7 @@ from vllm.logger import init_logger
 from vllm.model_executor.models.utils import AutoWeightsLoader
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import DistributedAutoencoderKLLTX2Video
 from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.parallel_state import (
     get_cfg_group,
@@ -48,6 +49,7 @@ from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_p
 from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
 from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
 from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
+from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 
 from .pipeline_ltx2 import (
@@ -113,7 +115,13 @@ def get_ltx2_post_process_func(od_config: OmniDiffusionConfig):
     return post_process_func
 
 
-class LTX23Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, SupportsComponentDiscovery):
+class LTX23Pipeline(
+    nn.Module,
+    CFGParallelMixin,
+    ProgressBarMixin,
+    SupportsComponentDiscovery,
+    DiffusionPipelineProfilerMixin,
+):
     """Fully independent LTX-2.3 pipeline.
 
     Key differences from LTX2Pipeline:
@@ -193,7 +201,7 @@ class LTX23Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, SupportsCompo
 
         # --- VAE, Audio VAE ---
         self.vae = from_pretrained_with_prefetch(
-            AutoencoderKLLTX2Video.from_pretrained,
+            DistributedAutoencoderKLLTX2Video.from_pretrained,
             model,
             subfolder="vae",
             prefetch_list=ltx2_subfolders,
@@ -266,6 +274,10 @@ class LTX23Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, SupportsCompo
         self._num_timesteps = None
         self._current_timestep = None
 
+        self.setup_diffusion_pipeline_profiler(
+            enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler
+        )
+
     def _place_aux_components(self) -> None:
         parallel_config = getattr(self.od_config, "parallel_config", None)
         use_managed_placement = bool(
@@ -1253,14 +1265,18 @@ class LTX23Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, SupportsCompo
 
             latents = latents.to(self.vae.dtype)
             video = self.vae.decode(latents, timestep_decode, return_dict=False)[0]
-            video = self.video_processor.postprocess_video(video, output_type=output_type)
+            if video.numel() > 0:
+                video = self.video_processor.postprocess_video(video, output_type=output_type)
 
             audio_latents = audio_latents.to(self.audio_vae.dtype)
             generated_mel_spectrograms = self.audio_vae.decode(audio_latents, return_dict=False)[0]
 
             audio = self.vocoder(generated_mel_spectrograms)
 
-        return DiffusionOutput(output=(video, audio))
+        return DiffusionOutput(
+            output=(video, audio),
+            stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
+        )
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         loader = AutoWeightsLoader(self)

vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py

@@ -536,6 +536,10 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli
     _vae_modules: ClassVar[list[str]] = []
     _resident_modules: ClassVar[list[str]] = ["fm_modules"]
 
+    # Top-level module(s) the diffusion LoRA manager scans (both MoT branches).
+    # TODO: promote to a shared LoRA contract/mixin instead of per-pipeline opt-in.
+    _lora_components: ClassVar[list[str]] = ["language_model"]
+
     def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
         super().__init__()
         self.od_config = od_config
@@ -1454,6 +1458,7 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli
             (".gate_up_proj", ".gate_proj", 0),
             (".gate_up_proj", ".up_proj", 1),
         ]
+        self.stacked_params_mapping = stacked_params_mapping
 
         params_dict = dict(self.named_parameters())
         loaded_params: set[str] = set()

vllm_omni/distributed/__init__.py

@@ -6,6 +6,7 @@ from .omni_connectors import (
     MooncakeConnector,
     MooncakeStoreConnector,
     MooncakeTransferEngineConnector,
+    MoriTransferEngineConnector,
     OmniConnectorBase,
     OmniConnectorFactory,
     OmniTransferConfig,
@@ -25,6 +26,7 @@ __all__ = [
     "MooncakeConnector",  # compat alias
     "MooncakeStoreConnector",
     "MooncakeTransferEngineConnector",
+    "MoriTransferEngineConnector",
     "SharedMemoryConnector",
     "YuanrongConnector",
     "YuanrongTransferEngineConnector",

vllm_omni/distributed/omni_connectors/__init__.py

@@ -17,6 +17,11 @@ try:
     from .connectors.mooncake_transfer_engine_connector import MooncakeTransferEngineConnector
 except ImportError:
     MooncakeTransferEngineConnector = None  # RDMA deps (msgspec/zmq/mooncake) not installed
+
+try:
+    from .connectors.mori_transfer_engine_connector import MoriTransferEngineConnector
+except ImportError:
+    MoriTransferEngineConnector = None  # RDMA deps (msgspec/zmq/mori) not installed
 from .factory import OmniConnectorFactory
 from .utils.config import ConnectorSpec, OmniTransferConfig
 from .utils.initialization import (
@@ -44,6 +49,7 @@ __all__ = [
     "MooncakeConnector",  # compat alias → MooncakeStoreConnector
     "MooncakeStoreConnector",
     "MooncakeTransferEngineConnector",
+    "MoriTransferEngineConnector",
     "SharedMemoryConnector",
     "YuanrongConnector",
     "YuanrongTransferEngineConnector",

vllm_omni/distributed/omni_connectors/connectors/mori_transfer_engine_connector.py

@@ -0,0 +1,962 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""OmniConnector backed by the Mori RDMA transfer engine.
+
+Implements the ``OmniConnector`` interface using Mori's ``IOEngine`` /
+``MemoryDesc`` API to perform zero-copy RDMA transfers between
+disaggregated prefill and decode workers.
+
+Notable design points:
+  * Remote peers must be registered via ``register_remote_engine()``
+    with the peer's ``EngineDesc`` before any data transfer can proceed.
+  * Memory regions are described by ``MemoryDesc`` objects (serialisable
+    via pack/unpack) rather than raw virtual addresses.
+  * Transfers are dispatched through ``IOEngine.batch_write()`` and
+    tracked asynchronously via ``TransferStatus`` objects.
+
+ZMQ handshake protocol:
+  * **Pull request** – receiver sends ``MoriPullRequest`` (msgspec-encoded)
+    containing its ``EngineDesc`` and pool ``MemoryDesc`` so the sender can
+    register the remote engine and RDMA-write directly into the receiver's
+    pool at the specified offset.
+  * **Query request** – receiver sends ``QUERY_INFO`` prefix + ``QueryRequest``
+    to check whether data is ready on the sender side (metadata-less get path).
+"""
+
+import os
+import queue
+import socket
+import threading
+import time as _time_mod
+import uuid
+from concurrent.futures import ThreadPoolExecutor
+from typing import Any
+
+import msgspec
+import torch
+import zmq
+
+from ..utils.logging import get_connector_logger
+from ..utils.memory_pool import BufferAllocator, ManagedBuffer
+from ..utils.serialization import OmniSerializer
+from .base import OmniConnectorBase
+
+logger = get_connector_logger(__name__)
+
+try:
+    from mori.cpp import TransferStatus  # noqa: F401
+    from mori.io import (
+        BackendType,
+        EngineDesc,
+        IOEngine,
+        IOEngineConfig,
+        MemoryDesc,
+        PollCqMode,
+        RdmaBackendConfig,
+        XgmiBackendConfig,
+    )
+except ImportError:
+    IOEngine = None
+
+# Supported backend types for Mori. Kept as string constants so configuration
+# (YAML / CLI / dict) stays transport-agnostic. ``rdma`` uses NIC-based RDMA
+# (RoCE / IB, GDR-capable); ``xgmi`` uses AMD Infinity Fabric GPU-to-GPU
+# direct links and therefore requires a CUDA pool.
+_SUPPORTED_BACKENDS = ("rdma", "xgmi")
+
+_BUFFER_TTL_SECONDS = 300
+
+TRANS_DONE = b"trans_done"
+TRANS_ERROR = b"trans_error"
+QUERY_INFO = b"query_info"
+INFO_NOT_FOUND = b"info_not_found"
+
+
+# ---------------------------------------------------------------------------
+# ZMQ message types
+# ---------------------------------------------------------------------------
+
+
+class MoriPullRequest(msgspec.Struct):
+    """Receiver → sender: request an RDMA write into the receiver's pool."""
+
+    request_id: str
+    engine_desc_packed: bytes
+    mem_desc_packed: bytes
+    dst_offset: int
+    length: int
+
+
+class QueryRequest(msgspec.Struct):
+    """Receiver → sender: query whether data for *request_id* is ready."""
+
+    request_id: str
+
+
+class QueryResponse(msgspec.Struct):
+    """Sender → receiver: metadata about a ready buffer."""
+
+    request_id: str
+    data_size: int
+    is_fast_path: bool
+
+
+# ---------------------------------------------------------------------------
+# Connector
+# ---------------------------------------------------------------------------
+
+
+class MoriTransferEngineConnector(OmniConnectorBase):
+    """OmniConnector backed by the Mori RDMA IOEngine.
+
+    Topology: 1 sender ↔ 1 receiver per key (same as
+    ``MooncakeTransferEngineConnector``).
+    """
+
+    supports_raw_data: bool = True
+
+    # ------------------------------------------------------------------ init
+    def __init__(self, config: dict[str, Any]):
+        if IOEngine is None:
+            raise ImportError("Mori is not available. Install via: pip install mori")
+
+        self._closed = False
+        self._bind_error: Exception | None = None
+
+        # Teardown-safe defaults
+        self._stop_event = threading.Event()
+        self._sender_executor: ThreadPoolExecutor | None = None
+        self._listener_thread: threading.Thread | None = None
+        self._listener_ready = threading.Event()
+        self._local_buffers: dict[str, Any] = {}
+        self._local_buffers_lock = threading.Lock()
+        self._req_local = threading.local()
+        self._worker_local = threading.local()
+        self._last_ttl_check: float = _time_mod.monotonic()
+
+        # Track remote engines already registered with IOEngine.
+        # ``mori.io.IOEngine.deregister_remote_engine`` requires the original
+        # ``EngineDesc`` object (its pybind signature is typed
+        # ``deregister_remote_engine(self, engine_desc: EngineDesc)``), not
+        # just the key string -- so we keep the desc alongside the key so
+        # ``close()`` can pass a valid object back at teardown.
+        self._registered_engines: dict[str, EngineDesc] = {}
+        self._registered_engines_lock = threading.Lock()
+
+        self._metrics = {
+            "puts": 0,
+            "gets": 0,
+            "bytes_transferred": 0,
+            "errors": 0,
+            "timeouts": 0,
+        }
+
+        self.config = config
+
+        # Stage id is read by OmniChunkTransferAdapter (process_pending_chunks /
+        # _send_single_request / _poll_single_request / ...) to decide whether
+        # this connector sits at the sender or receiver end of a stage pair.
+        # Kept parallel to ``SharedMemoryConnector`` and other OmniConnector
+        # implementations. Populated by ``build_stage_connectors`` via the
+        # ``stage_id`` key that ``get_connectors_config_for_stage`` injects.
+        self.stage_id = config.get("stage_id", -1)
+
+        # ---- Host / ZMQ ----
+        host_cfg = config.get("host", "127.0.0.1")
+        if host_cfg.lower() == "auto":
+            self.host = self._get_local_ip()
+            logger.info(f"Auto-detected local IP: {self.host}")
+        else:
+            self.host = host_cfg
+        self.zmq_port = config.get("zmq_port", 50051)
+
+        # ---- Backend selection ----
+        # Defaults to "rdma" so existing deployments (pre-XGMI support) keep
+        # working unchanged. ``xgmi`` selects AMD Infinity Fabric GPU-to-GPU
+        # direct links; see ``mori.io.XgmiBackendConfig``.
+        backend_type_cfg = str(config.get("backend_type", "rdma")).lower()
+        if backend_type_cfg not in _SUPPORTED_BACKENDS:
+            raise ValueError(f"Invalid backend_type={backend_type_cfg!r}. Supported: {list(_SUPPORTED_BACKENDS)}.")
+        self.backend_type = backend_type_cfg
+
+        # ---- RDMA device (RDMA backend only) ----
+        self.device_name = ""
+        if self.backend_type == "rdma":
+            self.device_name = config.get("device_name", "")
+            if not self.device_name:
+                env_dev = os.environ.get("MORI_RDMA_DEVICES", "")
+                if env_dev:
+                    self.device_name = env_dev
+                    logger.info(f"Using MORI_RDMA_DEVICES from env: {self.device_name}")
+        elif config.get("device_name"):
+            logger.warning(
+                "device_name=%r is ignored for backend_type='xgmi' (XGMI does not use an RDMA NIC).",
+                config.get("device_name"),
+            )
+
+        # ---- Pool config ----
+        self.pool_size = config.get("memory_pool_size", 1024**3)
+        self.pool_device = config.get("memory_pool_device", "cpu")
+        if self.backend_type == "xgmi" and self.pool_device == "cpu":
+            raise ValueError(
+                "backend_type='xgmi' requires memory_pool_device='cuda': "
+                "XGMI is a GPU-to-GPU fabric and cannot address CPU memory."
+            )
+
+        # ---- Sender info (receiver uses this when metadata=None) ----
+        self.sender_host = config.get("sender_host", None)
+        self.sender_zmq_port = config.get("sender_zmq_port", None)
+
+        # ---- Role ----
+        # ``sender``   : binds the ZMQ listener, accepts ``put()``.
+        # ``receiver`` : no listener, only ``get()`` by querying an
+        #                upstream sender whose endpoint lives in
+        #                ``sender_host`` / ``sender_zmq_port``.
+        role = str(config.get("role", "sender")).lower()
+        if role not in {"sender", "receiver"}:
+            raise ValueError(f"Invalid role={role!r} for MoriTransferEngineConnector. Expected 'sender' or 'receiver'.")
+        self.role = role
+        self.can_put = role == "sender"
+
+        # ---- Mori IOEngine ----
+        if self.device_name:
+            os.environ["MORI_RDMA_DEVICES"] = self.device_name
+
+        engine_config = IOEngineConfig(host=self.host, port=0)
+        self.engine_key = f"omni-{role}-{uuid.uuid4().hex[:8]}-pid{os.getpid()}-{self.host}"
+        self.engine = IOEngine(self.engine_key, engine_config)
+
+        # ---- Backend creation (per backend_type) ----
+        # ``IOEngine.batch_write`` is backend-agnostic, so only construction
+        # differs; the data-plane / ZMQ handshake paths below are identical.
+        if self.backend_type == "xgmi":
+            xgmi_cfg = XgmiBackendConfig()
+            xgmi_cfg.num_streams = config.get("xgmi_num_streams", 64)
+            xgmi_cfg.num_events = config.get("xgmi_num_events", 64)
+            self.engine.create_backend(BackendType.XGMI, xgmi_cfg)
+            logger.info(f"Mori backend: XGMI (num_streams={xgmi_cfg.num_streams}, num_events={xgmi_cfg.num_events})")
+        else:  # rdma
+            qp_per_transfer = config.get("qp_per_transfer", 1)
+            post_batch_size = config.get("post_batch_size", -1)
+            num_workers = config.get("num_worker_threads", 1)
+            rdma_cfg = RdmaBackendConfig(
+                qp_per_transfer,
+                post_batch_size,
+                num_workers,
+                PollCqMode.POLLING,
+                False,
+            )
+            self.engine.create_backend(BackendType.RDMA, rdma_cfg)
+            logger.info(
+                f"Mori backend: RDMA (qp_per_transfer={qp_per_transfer}, "
+                f"num_workers={num_workers}, device={self.device_name or 'auto'})"
+            )
+
+        self.engine_desc: EngineDesc = self.engine.get_engine_desc()
+        self.engine_desc_packed: bytes = self.engine_desc.pack()
+        logger.info(f"Mori IOEngine ready: key={self.engine_key} at {self.engine_desc.host}:{self.engine_desc.port}")
+
+        # ---- Pool allocation & Mori memory registration ----
+        logger.info(
+            f"Allocating {self.backend_type.upper()} pool: {self.pool_size / 1024**2:.2f} MB on {self.pool_device}"
+        )
+        try:
+            if self.pool_device == "cpu":
+                self.pool = torch.empty(self.pool_size, dtype=torch.uint8).pin_memory()
+            else:
+                self.pool = torch.empty(
+                    self.pool_size,
+                    dtype=torch.uint8,
+                    device=self.pool_device,
+                )
+            self.pool_mem_desc: MemoryDesc = self.engine.register_torch_tensor(self.pool)
+            self.pool_mem_desc_packed: bytes = self.pool_mem_desc.pack()
+            self.base_ptr = self.pool.data_ptr()
+        except Exception as e:
+            logger.error(f"Failed to allocate/register pool: {e}")
+            raise
+
+        self.allocator = BufferAllocator(self.pool_size, alignment=4096)
+
+        # ---- ZMQ / threading ----
+        self.zmq_ctx = zmq.Context()
+        self._sender_executor = ThreadPoolExecutor(max_workers=4, thread_name_prefix="mori-sender")
+
+        logger.info(
+            f"MoriTransferEngineConnector config:\n"
+            f"  Local: host={self.host}, zmq_port={self.zmq_port}\n"
+            f"  Remote: sender_host={self.sender_host}, "
+            f"sender_zmq_port={self.sender_zmq_port}\n"
+            f"  Role: can_put={self.can_put}, "
+            f"configured_role={config.get('role', 'sender')}"
+        )
+
+        if self.can_put:
+            self._last_ttl_check = _time_mod.monotonic()
+            self._listener_thread = threading.Thread(target=self._zmq_listener_loop, daemon=True)
+            self._listener_thread.start()
+            self._listener_ready.wait(timeout=1.0)
+            if self._bind_error is not None:
+                raise RuntimeError(
+                    f"MoriTransferEngineConnector failed to bind ZMQ on {self.host}:{self.zmq_port}: {self._bind_error}"
+                ) from self._bind_error
+            logger.info(f"MoriTransferEngineConnector SENDER ready (ZMQ on {self.host}:{self.zmq_port})")
+        else:
+            if not self.sender_host or self.sender_host.lower() == "auto":
+                logger.info("MoriTransferEngineConnector RECEIVER: awaiting sender info via update_sender_info().")
+            else:
+                logger.info(
+                    f"MoriTransferEngineConnector RECEIVER ready (sender at {self.sender_host}:{self.sender_zmq_port})"
+                )
+
+    # -------------------------------------------------------- public helpers
+    def get_connection_info(self) -> dict[str, Any]:
+        return {
+            "host": self.host,
+            "zmq_port": self.zmq_port,
+            "engine_key": self.engine_key,
+            "can_put": self.can_put,
+        }
+
+    def update_sender_info(self, sender_host: str, sender_zmq_port: int) -> None:
+        """Inject the sender's ZMQ endpoint into the receiver connector."""
+        self.sender_host = sender_host
+        self.sender_zmq_port = sender_zmq_port
+        logger.info(f"Sender info updated: host={sender_host!r}, zmq_port={sender_zmq_port}")
+
+    # -------------------------------------------------- internal helpers
+    @staticmethod
+    def _get_local_ip() -> str:
+        try:
+            with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s:
+                s.connect(("8.8.8.8", 80))
+                return s.getsockname()[0]
+        except Exception:
+            try:
+                return socket.gethostbyname(socket.gethostname())
+            except Exception:
+                return "127.0.0.1"
+
+    def _get_req_socket(self, zmq_addr: str, timeout_ms: int) -> zmq.Socket:
+        cache: dict[str, zmq.Socket] | None = getattr(self._req_local, "cache", None)
+        if cache is None:
+            cache = {}
+            self._req_local.cache = cache
+
+        sock = cache.get(zmq_addr)
+        if sock is None:
+            sock = self.zmq_ctx.socket(zmq.REQ)
+            sock.connect(zmq_addr)
+            cache[zmq_addr] = sock
+        sock.setsockopt(zmq.SNDTIMEO, timeout_ms)
+        sock.setsockopt(zmq.RCVTIMEO, timeout_ms)
+        return sock
+
+    def _invalidate_req_socket(self, zmq_addr: str) -> None:
+        cache: dict[str, zmq.Socket] | None = getattr(self._req_local, "cache", None)
+        if cache is None:
+            return
+        sock = cache.pop(zmq_addr, None)
+        if sock is not None:
+            try:
+                sock.close(linger=0)
+            except Exception:
+                pass
+
+    def _ensure_remote_registered(self, engine_desc_packed: bytes) -> str:
+        """Register a remote engine with the local IOEngine (idempotent)."""
+        desc = EngineDesc.unpack(engine_desc_packed)
+        key = desc.key
+        with self._registered_engines_lock:
+            if key not in self._registered_engines:
+                self.engine.register_remote_engine(desc)
+                self._registered_engines[key] = desc
+                logger.debug(f"Registered remote Mori engine: {key}")
+        return key
+
+    # ----------------------------------------------------------------- put()
+    def put(
+        self,
+        from_stage: str,
+        to_stage: str,
+        put_key: str,
+        data: Any,
+    ) -> tuple[bool, int, dict[str, Any] | None]:
+        if self._closed:
+            raise RuntimeError("Cannot put: MoriTransferEngineConnector is closed")
+        if not self.can_put:
+            logger.warning(f"Rejecting put for {put_key}: connector is receiver-only")
+            return False, 0, None
+
+        put_key = self._make_key(put_key, from_stage, to_stage)
+        try:
+            src_offset = 0
+            size = 0
+            holder = None
+            is_fast_path = True
+            should_release = False
+
+            # Reject empty payloads early
+            if isinstance(data, bytes) and len(data) == 0:
+                return False, 0, None
+            if isinstance(data, torch.Tensor) and data.nbytes == 0:
+                return False, 0, None
+
+            # Serialize non-raw types
+            if not isinstance(data, (ManagedBuffer, torch.Tensor, bytes)):
+                data = OmniSerializer.serialize(data)
+                is_fast_path = False
+
+            if isinstance(data, ManagedBuffer):
+                if data.pool_tensor.data_ptr() != self.pool.data_ptr():
+                    data = data.tensor.contiguous()
+                else:
+                    src_offset = data.offset
+                    size = data.size
+                    holder = data
+                    should_release = False
+
+            if isinstance(data, (torch.Tensor, bytes)):
+                if isinstance(data, torch.Tensor):
+                    size = data.nbytes
+                    tensor_data = data
+                else:
+                    size = len(data)
+                    tensor_data = torch.frombuffer(data, dtype=torch.uint8)
+
+                try:
+                    offset = self.allocator.alloc(size)
+                    holder = ManagedBuffer(self.allocator, offset, size, self.pool)
+                    should_release = True
+                except MemoryError:
+                    logger.error(f"Pool exhausted for {size} bytes")
+                    return False, 0, None
+
+                try:
+                    dst_t = holder.tensor
+                    if isinstance(data, torch.Tensor):
+                        if not data.is_contiguous():
+                            data = data.contiguous()
+                        src_view = data.view(torch.uint8).flatten()
+                        if src_view.device != dst_t.device:
+                            dst_t.copy_(src_view, non_blocking=True)
+                            if src_view.is_cuda:
+                                with torch.cuda.device(src_view.device):
+                                    torch.cuda.current_stream().synchronize()
+                            elif dst_t.is_cuda:
+                                with torch.cuda.device(dst_t.device):
+                                    torch.cuda.current_stream().synchronize()
+                        else:
+                            dst_t.copy_(src_view)
+                            if dst_t.is_cuda:
+                                with torch.cuda.device(dst_t.device):
+                                    torch.cuda.current_stream().synchronize()
+                    else:
+                        dst_t.copy_(tensor_data)
+                        if dst_t.is_cuda:
+                            with torch.cuda.device(dst_t.device):
+                                torch.cuda.current_stream().synchronize()
+                except Exception as e:
+                    holder.release()
+                    logger.error(f"Failed to copy data to pool: {e}")
+                    return False, 0, None
+
+                src_offset = offset
+
+            if size <= 0:
+                if should_release and isinstance(holder, ManagedBuffer):
+                    holder.release()
+                return False, 0, None
+
+            with self._local_buffers_lock:
+                old = self._local_buffers.pop(put_key, None)
+                if old:
+                    _, _, oh, osr, _, _ = old
+                    if osr and isinstance(oh, ManagedBuffer):
+                        oh.release()
+                        logger.warning(f"Released stale buffer for duplicate key: {put_key}")
+                # (offset, size, holder, should_release, is_fast_path, ts)
+                self._local_buffers[put_key] = (
+                    src_offset,
+                    size,
+                    holder,
+                    should_release,
+                    is_fast_path,
+                    _time_mod.monotonic(),
+                )
+
+            metadata = {
+                "source_host": self.host,
+                "source_port": self.zmq_port,
+                "data_size": size,
+                "is_fast_path": is_fast_path,
+            }
+            self._metrics["puts"] += 1
+            self._metrics["bytes_transferred"] += size
+            return True, size, metadata
+
+        except Exception as e:
+            self._metrics["errors"] += 1
+            logger.error(f"Put failed for {put_key}: {e}", exc_info=True)
+            return False, 0, None
+
+    # ------------------------------------------------ query (no-metadata get)
+    def _query_metadata_from_sender(self, get_key: str) -> dict[str, Any] | None:
+        zmq_addr = f"tcp://{self.sender_host}:{self.sender_zmq_port}"
+        req_socket = self._get_req_socket(zmq_addr, timeout_ms=5000)
+        try:
+            q = QueryRequest(request_id=get_key)
+            req_socket.send(QUERY_INFO + msgspec.msgpack.encode(q))
+            resp = req_socket.recv()
+            if resp == INFO_NOT_FOUND:
+                return None
+            qr = msgspec.msgpack.decode(resp, type=QueryResponse)
+            return {
+                "source_host": self.sender_host,
+                "source_port": self.sender_zmq_port,
+                "data_size": qr.data_size,
+                "is_fast_path": qr.is_fast_path,
+            }
+        except Exception as e:
+            self._invalidate_req_socket(zmq_addr)
+            logger.debug(f"Metadata query failed for {get_key}: {e}")
+            return None
+
+    # ----------------------------------------------------------------- get()
+    def get(
+        self,
+        from_stage: str,
+        to_stage: str,
+        get_key: str,
+        metadata: dict[str, Any] | None = None,
+    ) -> tuple[Any, int] | None:
+        if self._closed:
+            raise RuntimeError("Cannot get: MoriTransferEngineConnector is closed")
+
+        get_key = self._make_key(get_key, from_stage, to_stage)
+        _t0 = _time_mod.perf_counter()
+
+        # Resolve metadata
+        if not metadata:
+            if not self.sender_host or not self.sender_zmq_port or str(self.sender_host).lower() == "auto":
+                raise RuntimeError("get(metadata=None) requires sender info. Call update_sender_info() first.")
+            metadata = self._query_metadata_from_sender(get_key)
+            if not metadata:
+                return None
+
+        _t1 = _time_mod.perf_counter()
+        _query_ms = (_t1 - _t0) * 1000
+
+        src_host = metadata.get("source_host")
+        src_port = metadata.get("source_port")
+        data_size = metadata.get("data_size", 0)
+        is_fast_path = metadata.get("is_fast_path", False)
+
+        if not src_host or not src_port or str(src_host).lower() == "auto":
+            logger.error(f"Invalid metadata for {get_key}")
+            return None
+        if data_size == 0:
+            logger.warning(f"Skipping get for {get_key}: data_size is 0")
+            return None
+
+        # Allocate destination buffer
+        try:
+            offset = self.allocator.alloc(data_size)
+            recv_buf = ManagedBuffer(self.allocator, offset, data_size, self.pool)
+        except MemoryError:
+            logger.error(f"Failed to allocate {data_size} bytes for get")
+            return None
+
+        _t2 = _time_mod.perf_counter()
+        _alloc_ms = (_t2 - _t1) * 1000
+
+        # Build pull request with Mori-specific EngineDesc + MemoryDesc
+        pull_req = MoriPullRequest(
+            request_id=get_key,
+            engine_desc_packed=self.engine_desc_packed,
+            mem_desc_packed=self.pool_mem_desc_packed,
+            dst_offset=offset,
+            length=data_size,
+        )
+
+        _base_timeout_ms = 30000
+        _size_timeout_ms = max(0, (data_size // (100 * 1024 * 1024))) * 5000
+        _total_timeout_ms = _base_timeout_ms + _size_timeout_ms
+        zmq_addr = f"tcp://{src_host}:{src_port}"
+        req_socket = self._get_req_socket(zmq_addr, timeout_ms=_total_timeout_ms)
+
+        try:
+            req_socket.send(msgspec.msgpack.encode(pull_req))
+            resp = req_socket.recv()
+
+            _t3 = _time_mod.perf_counter()
+            _rdma_ms = (_t3 - _t2) * 1000
+
+            if resp == TRANS_DONE:
+                if self.pool.is_cuda:
+                    with torch.cuda.device(self.pool.device):
+                        torch.cuda.current_stream().synchronize()
+
+                _t4 = _time_mod.perf_counter()
+                _sync_ms = (_t4 - _t3) * 1000
+
+                if is_fast_path:
+                    _total_ms = (_time_mod.perf_counter() - _t0) * 1000
+                    _mbps = (data_size / 1024 / 1024) / (_total_ms / 1000) if _total_ms > 0 else 0
+                    logger.info(
+                        f"[MORI GET] {get_key}: query={_query_ms:.1f}ms, "
+                        f"alloc={_alloc_ms:.1f}ms, rdma={_rdma_ms:.1f}ms, "
+                        f"sync={_sync_ms:.1f}ms, total={_total_ms:.1f}ms, "
+                        f"{_mbps:.1f} MB/s (fast_path)"
+                    )
+                    self._metrics["gets"] += 1
+                    self._metrics["bytes_transferred"] += data_size
+                    return recv_buf, data_size
+                else:
+                    try:
+                        raw = recv_buf.to_bytes()
+                        val = OmniSerializer.deserialize(raw)
+                        _total_ms = (_time_mod.perf_counter() - _t0) * 1000
+                        _mbps = (data_size / 1024 / 1024) / (_total_ms / 1000) if _total_ms > 0 else 0
+                        logger.info(f"[MORI GET] {get_key}: total={_total_ms:.1f}ms, {_mbps:.1f} MB/s (deserialized)")
+                        self._metrics["gets"] += 1
+                        self._metrics["bytes_transferred"] += data_size
+                        return val, data_size
+                    finally:
+                        recv_buf.release()
+            else:
+                self._metrics["errors"] += 1
+                logger.error(f"MORI get failed: received {resp} instead of TRANS_DONE")
+                recv_buf.release()
+                return None
+        except Exception as e:
+            self._invalidate_req_socket(zmq_addr)
+            self._metrics["timeouts"] += 1
+            logger.error(f"MORI get error: {e}", exc_info=True)
+            recv_buf.release()
+            return None
+
+    # ------------------------------------------------------------- cleanup()
+    def cleanup(
+        self,
+        request_id: str,
+        from_stage: str | None = None,
+        to_stage: str | None = None,
+    ) -> None:
+        if (from_stage is None) != (to_stage is None):
+            raise ValueError("cleanup() requires both from_stage and to_stage, or neither.")
+        if from_stage is not None and to_stage is not None:
+            request_id = self._make_key(request_id, from_stage, to_stage)
+        with self._local_buffers_lock:
+            item = self._local_buffers.pop(request_id, None)
+            if item:
+                _, _, holder, sr, _, _ = item
+                if sr and isinstance(holder, ManagedBuffer):
+                    holder.release()
+
+    # -------------------------------------------------------------- health()
+    def health(self) -> dict[str, Any]:
+        if self._closed:
+            return {"status": "unhealthy", "error": "Connector is closed"}
+        return {
+            "status": "healthy",
+            "host": self.host,
+            "pool_device": self.pool_device,
+            "pool_size": self.pool_size,
+            "engine_key": self.engine_key,
+            **self._metrics,
+        }
+
+    # --------------------------------------------------------------- close()
+    def close(self) -> None:
+        if getattr(self, "_closed", True):
+            return
+        self._closed = True
+        logger.info("Closing MoriTransferEngineConnector...")
+
+        self._stop_event.set()
+
+        if self._listener_thread is not None and self._listener_thread.is_alive():
+            self._listener_thread.join(timeout=2.0)
+            if self._listener_thread.is_alive():
+                logger.warning("Listener thread did not stop gracefully")
+
+        if self._sender_executor is not None:
+            self._sender_executor.shutdown(wait=True, cancel_futures=False)
+
+        with self._local_buffers_lock:
+            for _k, item in list(self._local_buffers.items()):
+                _, _, holder, sr, _, _ = item
+                if sr and isinstance(holder, ManagedBuffer):
+                    holder.release()
+            self._local_buffers.clear()
+
+        cache: dict[str, zmq.Socket] | None = getattr(self._req_local, "cache", None)
+        if cache:
+            for _addr, sock in cache.items():
+                try:
+                    sock.close(linger=0)
+                except Exception:
+                    pass
+            cache.clear()
+
+        try:
+            if hasattr(self, "zmq_ctx"):
+                self.zmq_ctx.term()
+        except Exception as e:
+            logger.warning(f"Failed to terminate ZMQ context: {e}")
+
+        # ---- Release Mori IOEngine resources ----
+        # ``mori.io.IOEngine`` does not expose an explicit ``shutdown()`` /
+        # ``close()`` entry-point, so we best-effort deregister every resource
+        # we registered with it and then drop the only reference we hold so
+        # Python GC triggers the C++ destructor.
+        engine = getattr(self, "engine", None)
+        if engine is not None:
+            pool_desc = getattr(self, "pool_mem_desc", None)
+            if pool_desc is not None:
+                try:
+                    engine.deregister_memory(pool_desc)
+                except Exception as e:
+                    logger.warning(f"Mori deregister_memory failed: {e}")
+                self.pool_mem_desc = None  # type: ignore[assignment]
+
+            with self._registered_engines_lock:
+                remote_descs = list(self._registered_engines.values())
+                self._registered_engines.clear()
+            for desc in remote_descs:
+                try:
+                    engine.deregister_remote_engine(desc)
+                except Exception as e:
+                    logger.debug(f"Mori deregister_remote_engine({desc.key!r}) failed: {e}")
+
+            # Drop the only reference to IOEngine; the C++ destructor runs
+            # under Python GC and unwinds backend / RDMA fabric state.
+            self.engine = None  # type: ignore[assignment]
+
+        self.pool = None  # type: ignore[assignment]
+        logger.info("MoriTransferEngineConnector closed.")
+
+    # ============================================================== LISTENER
+    def _cleanup_stale_buffers(self) -> None:
+        """Reclaim buffers older than ``_BUFFER_TTL_SECONDS``.
+
+        Prevents permanent memory leaks when a receiver crashes or times out
+        without ever pulling the data.
+
+        TODO(zejwang): In extreme rare case, long transfer time, there might
+        exist TTL cleanup vs in-flight RDMA transfer conflict, which will be
+        handled in a follow-up PR. Same race is acknowledged in
+        ``MooncakeTransferEngineConnector._cleanup_stale_buffers``.
+        """
+        now = _time_mod.monotonic()
+        with self._local_buffers_lock:
+            stale = [k for k, v in self._local_buffers.items() if now - v[5] > _BUFFER_TTL_SECONDS]
+            for k in stale:
+                item = self._local_buffers.pop(k)
+                _, _, holder, sr, _, _ = item
+                if sr and isinstance(holder, ManagedBuffer):
+                    holder.release()
+                logger.warning(f"TTL expired ({_BUFFER_TTL_SECONDS}s): reclaimed buffer for {k}")
+
+    def _zmq_listener_loop(self) -> None:
+        sock = self.zmq_ctx.socket(zmq.ROUTER)
+
+        try:
+            sock.bind(f"tcp://{self.host}:{self.zmq_port}")
+        except zmq.ZMQError as exc:
+            # Any bind failure (EADDRINUSE, EADDRNOTAVAIL, EACCES, etc.) is
+            # fatal for a sender — fail fast so __init__ propagates the error.
+            # There is no silent receiver fallback; roles are explicitly
+            # assigned (matches MooncakeTransferEngineConnector).
+            logger.error(f"ZMQ bind failed on {self.host}:{self.zmq_port}: {exc} (errno={exc.errno})")
+            self.can_put = False
+            self._bind_error = exc
+            self._listener_ready.set()
+            return
+
+        self._listener_ready.set()
+
+        notify_addr = f"inproc://mori-notify-{id(self)}"
+        notify_recv = self.zmq_ctx.socket(zmq.PULL)
+        notify_recv.bind(notify_addr)
+
+        poller = zmq.Poller()
+        poller.register(sock, zmq.POLLIN)
+        poller.register(notify_recv, zmq.POLLIN)
+
+        response_queue: queue.Queue = queue.Queue()
+
+        try:
+            while not self._stop_event.is_set():
+                try:
+                    events = dict(poller.poll(1000))
+
+                    if notify_recv in events:
+                        while True:
+                            try:
+                                notify_recv.recv(zmq.NOBLOCK)
+                            except zmq.Again:
+                                break
+
+                    while True:
+                        try:
+                            identity, response = response_queue.get_nowait()
+                            sock.send_multipart([identity, b"", response])
+                        except queue.Empty:
+                            break
+
+                    now = _time_mod.monotonic()
+                    if now - self._last_ttl_check >= 10.0:
+                        self._last_ttl_check = now
+                        self._cleanup_stale_buffers()
+
+                    if sock in events:
+                        frames = sock.recv_multipart()
+                        if len(frames) >= 2:
+                            self._sender_executor.submit(
+                                self._handle_pull_request,
+                                response_queue,
+                                notify_addr,
+                                frames[0],
+                                frames[-1],
+                            )
+                except zmq.ContextTerminated:
+                    break
+                except Exception:
+                    logger.debug("Listener loop error", exc_info=True)
+        finally:
+            try:
+                notify_recv.close(linger=0)
+                sock.close(linger=0)
+            except Exception:
+                pass
+
+    def _handle_pull_request(
+        self,
+        response_queue: queue.Queue,
+        notify_addr: str,
+        identity: bytes,
+        payload: bytes,
+    ) -> None:
+        try:
+            if payload.startswith(QUERY_INFO):
+                self._handle_query_request(
+                    response_queue,
+                    notify_addr,
+                    identity,
+                    payload[len(QUERY_INFO) :],
+                )
+                return
+
+            pull = msgspec.msgpack.decode(payload, type=MoriPullRequest)
+
+            with self._local_buffers_lock:
+                item = self._local_buffers.get(pull.request_id)
+
+            if not item:
+                response_queue.put((identity, TRANS_ERROR))
+                self._notify_listener(notify_addr)
+                return
+
+            src_offset, src_size, _, _, _, _ = item
+
+            # Validate against sender's own src_size before RDMA write:
+            # length < src_size silently truncates the payload; length >
+            # src_size reads past this allocation into adjacent in-flight
+            # buffers because pool_mem_desc covers the whole pool, not just
+            # this slot, so Mori does not bound-check. Refuse on mismatch
+            # rather than corrupt or leak data silently.
+            if pull.length != src_size:
+                logger.error(
+                    "Length mismatch for %s: sender src_size=%d != receiver length=%d; refusing transfer",
+                    pull.request_id,
+                    src_size,
+                    pull.length,
+                )
+                response_queue.put((identity, TRANS_ERROR))
+                self._notify_listener(notify_addr)
+                return
+
+            # Register the receiver's IOEngine (idempotent)
+            self._ensure_remote_registered(pull.engine_desc_packed)
+
+            # Reconstruct the receiver's pool MemoryDesc
+            remote_mem = MemoryDesc.unpack(pull.mem_desc_packed)
+
+            # RDMA write from local pool → remote pool. Use src_size (sender's
+            # own record); it has just been validated to equal pull.length.
+            transfer_uid = self.engine.allocate_transfer_uid()
+            statuses = self.engine.batch_write(
+                [self.pool_mem_desc],
+                [[src_offset]],
+                [remote_mem],
+                [[pull.dst_offset]],
+                [[src_size]],
+                [transfer_uid],
+            )
+
+            success = True
+            for st in statuses:
+                st.Wait()
+                if st.Failed():
+                    logger.error(f"RDMA write failed for {pull.request_id}: {st.Message()}")
+                    success = False
+
+            if success:
+                self.cleanup(pull.request_id)
+                response_queue.put((identity, TRANS_DONE))
+            else:
+                logger.warning(f"RDMA write failed for {pull.request_id}. Buffer retained for retry.")
+                response_queue.put((identity, TRANS_ERROR))
+
+        except Exception as e:
+            logger.error(f"Pull request handler error: {e}", exc_info=True)
+            response_queue.put((identity, TRANS_ERROR))
+
+        self._notify_listener(notify_addr)
+
+    def _handle_query_request(
+        self,
+        response_queue: queue.Queue,
+        notify_addr: str,
+        identity: bytes,
+        payload: bytes,
+    ) -> None:
+        try:
+            q = msgspec.msgpack.decode(payload, type=QueryRequest)
+            with self._local_buffers_lock:
+                item = self._local_buffers.get(q.request_id)
+            if not item:
+                response_queue.put((identity, INFO_NOT_FOUND))
+            else:
+                _, data_size, _, _, is_fast, _ = item
+                resp = QueryResponse(
+                    request_id=q.request_id,
+                    data_size=data_size,
+                    is_fast_path=is_fast,
+                )
+                response_queue.put((identity, msgspec.msgpack.encode(resp)))
+        except Exception as e:
+            logger.error(f"Query handler error: {e}")
+            response_queue.put((identity, INFO_NOT_FOUND))
+
+        self._notify_listener(notify_addr)
+
+    def _notify_listener(self, notify_addr: str) -> None:
+        try:
+            local = self._worker_local
+            sock = getattr(local, "notify_socket", None)
+            cached_addr = getattr(local, "notify_addr", None)
+            if sock is None or cached_addr != notify_addr:
+                if sock is not None:
+                    sock.close(linger=0)
+                sock = self.zmq_ctx.socket(zmq.PUSH)
+                sock.connect(notify_addr)
+                local.notify_socket = sock
+                local.notify_addr = notify_addr
+            sock.send(b"", zmq.NOBLOCK)
+        except Exception:
+            local.notify_socket = None
+            local.notify_addr = None

vllm_omni/distributed/omni_connectors/factory.py

@@ -114,11 +114,23 @@ def _create_mooncake_transfer_engine_connector(config: dict[str, Any]) -> OmniCo
     return MooncakeTransferEngineConnector(config)
 
 
+def _create_mori_transfer_engine_connector(config: dict[str, Any]) -> OmniConnectorBase:
+    try:
+        from .connectors.mori_transfer_engine_connector import MoriTransferEngineConnector
+    except ImportError:
+        import sys
+
+        sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
+        from omni_connectors.connectors.mori_transfer_engine_connector import MoriTransferEngineConnector
+    return MoriTransferEngineConnector(config)
+
+
 # Register connectors
 OmniConnectorFactory.register_connector("MooncakeStoreConnector", _create_mooncake_store_connector)
 OmniConnectorFactory.register_connector("MooncakeTransferEngineConnector", _create_mooncake_transfer_engine_connector)
 OmniConnectorFactory.register_connector("SharedMemoryConnector", _create_shm_connector)
 OmniConnectorFactory.register_connector("YuanrongConnector", _create_yuanrong_connector)
 OmniConnectorFactory.register_connector("YuanrongTransferEngineConnector", _create_yuanrong_transfer_engine_connector)
+OmniConnectorFactory.register_connector("MoriTransferEngineConnector", _create_mori_transfer_engine_connector)
 # Backward-compatible aliases – will be removed in the future
 OmniConnectorFactory.register_connector("MooncakeConnector", _create_mooncake_store_connector)

vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py

@@ -444,12 +444,37 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
         self,
         waiting_queue: Any,
         running_queue: list[Request],
+        *,
+        scheduler_requests: dict[str, Request] | None = None,
     ) -> None:
         """
         Process pending chunks for waiting and running queues.
+
+        When ``scheduler_requests`` is provided, purges any
+        ``waiting_for_chunk_*_requests`` deque entries whose
+        ``request_id`` is no longer tracked by it (e.g. after a
+        mid-flight abort that ran ``Scheduler._free_request``) before
+        processing chunks. Without this purge, ``restore_queues`` would
+        later re-inject the freed ``Request`` onto ``running_queue`` and
+        the worker's ``_update_states`` would crash with ``KeyError``
+        reading ``self.requests[req_id]``. See vllm-project/vllm-omni#3736.
+
+        ``scheduler_requests`` is keyword-only and optional; production
+        schedulers always pass their live request map, while legacy
+        callers that don't track aborts may omit it to keep the prior
+        (unguarded) behaviour.
         """
         if self.connector.stage_id == 0:
             return
+
+        # Purge deque entries whose request was freed mid-flight (abort →
+        # Scheduler._free_request) before any chunk processing, so neither
+        # the legacy nor the active-stream path can re-inject a zombie
+        # Request onto the queues. See vllm-project/vllm-omni#3736.
+        if scheduler_requests is not None:
+            self._purge_untracked_chunk_requests(self.waiting_for_chunk_waiting_requests, scheduler_requests)
+            self._purge_untracked_chunk_requests(self.waiting_for_chunk_running_requests, scheduler_requests)
+
         if self._active_window <= 0:
             self._process_chunk_queue_legacy(
                 waiting_queue, self.waiting_for_chunk_waiting_requests, RequestStatus.WAITING, self._finished_load_reqs
@@ -558,6 +583,29 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
             self.requests_origin_status[request.request_id] = target_status
             waiting_for_chunk_list.append(request)
 
+    def _purge_untracked_chunk_requests(
+        self,
+        deque_list: deque[Any],
+        scheduler_requests: dict[str, Request],
+    ) -> None:
+        """Drop deque entries whose ``request_id`` is not in
+        ``scheduler_requests`` and reclaim their receiver-side state.
+
+        Handles requests that were aborted mid-flight while parked in a
+        chunk-transfer deque: ``Scheduler._free_request`` deleted the
+        entry from ``scheduler.requests`` but the deque still holds a
+        reference to the now-freed ``Request``. Order of survivors is
+        preserved.
+        """
+        if not deque_list:
+            return
+        for _ in range(len(deque_list)):
+            request = deque_list.popleft()
+            if request.request_id in scheduler_requests:
+                deque_list.append(request)
+            else:
+                self.cleanup_receiver(request.request_id)
+
     def restore_queues(
         self,
         waiting_queue: Any,
@@ -566,7 +614,23 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
     ) -> None:
         """
         Restore requests waiting for chunk to the waiting and running queues.
+
+        Re-runs the zombie purge first to close the race window where an
+        abort fires *between* ``process_pending_chunks`` and the
+        ``finally``-clause ``restore_queues`` call. Without the second
+        purge, ``running_queue.extend(...)`` would still re-inject a
+        freed ``Request`` and crash the worker on the next tick.
+
+        ``scheduler_requests`` is optional for back-compat with legacy
+        callers (older tests pass only the two queue arguments). When
+        provided, it gates both the deque purge and the per-request
+        admit checks below; when ``None``, the purge is skipped and
+        every parked request is restored unconditionally (the
+        pre-purge behavior).
         """
+        if scheduler_requests is not None:
+            self._purge_untracked_chunk_requests(self.waiting_for_chunk_waiting_requests, scheduler_requests)
+            self._purge_untracked_chunk_requests(self.waiting_for_chunk_running_requests, scheduler_requests)
         # Add request waiting for chunk to the waiting and running queue
         for request in self.waiting_for_chunk_waiting_requests:
             if scheduler_requests is None or request.request_id in scheduler_requests:

vllm_omni/distributed/omni_connectors/utils/config.py

@@ -11,6 +11,7 @@ logger = get_connector_logger(__name__)
 TRANSFER_ENGINE_CONNECTOR_NAMES = frozenset(
     {
         "MooncakeTransferEngineConnector",
+        "MoriTransferEngineConnector",
         "YuanrongTransferEngineConnector",
     }
 )

vllm_omni/engine/arg_utils.py

@@ -35,6 +35,7 @@ def _register_omni_hf_configs() -> None:
     try:
         from transformers import AutoConfig
 
+        from vllm_omni.model_executor.models.ming_tts.config_ming_tts import MingDenseConfig
         from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import (
             Qwen3TTSConfig,
         )
@@ -55,6 +56,7 @@ def _register_omni_hf_configs() -> None:
         _CONFIG_REGISTRY = None
 
     for model_type, config_cls in [
+        ("dense", MingDenseConfig),
         ("qwen3_tts", Qwen3TTSConfig),
         ("cosyvoice3", CosyVoice3Config),
         ("glm_tts", GLMTTSConfig),

vllm_omni/engine/async_omni_engine.py

@@ -1363,6 +1363,7 @@ class AsyncOmniEngine:
 
         async def _run_orchestrator() -> None:
             self._initialize_stages(stage_init_timeout)
+
             pd_config = self._detect_pd_config()
             coordinator_pub_address: str | None = None
             load_balancer_factory: Callable[[], LoadBalancer] | None = None

vllm_omni/engine/orchestrator.py

@@ -689,7 +689,7 @@ class Orchestrator:
                             await self._handle_kv_ready_raw_outputs(stage_id, raw_outputs)
                             for eco in raw_outputs.outputs:
                                 req_state = self.request_states.get(getattr(eco, "request_id", None))
-                                if req_state is None:
+                                if req_state is None or not req_state.streaming.enabled:
                                     continue
                                 req_state.streaming.segment_finished = bool(getattr(eco, "is_segment_finished", False))
                                 req_state.streaming.new_prompt_len_snapshot = getattr(

vllm_omni/entrypoints/openai/protocol/audio.py

@@ -7,6 +7,43 @@ from pydantic import AliasChoices, BaseModel, Field, field_validator, model_vali
 _MAX_EMBEDDING_DIM = 8192
 
 
+def _normalize_ref_audio_value(value):
+    if value is None:
+        return None
+    if isinstance(value, str):
+        return value
+    if isinstance(value, (list, tuple)):
+        items = []
+        for item in value:
+            if not isinstance(item, str):
+                raise TypeError("'ref_audio' list entries must be strings")
+            items.append(item)
+        if not items:
+            raise ValueError("'ref_audio' list cannot be empty")
+        return items
+    raise TypeError("'ref_audio' must be a string or list of strings")
+
+
+def _normalize_speaker_embedding_value(value):
+    if value is None:
+        return None
+    if not isinstance(value, (list, tuple)):
+        raise TypeError("'speaker_embedding' must be a list of numbers or list of embedding vectors")
+    if not value:
+        return []
+
+    first = value[0]
+    if isinstance(first, (list, tuple)):
+        embeddings = []
+        for item in value:
+            if not isinstance(item, (list, tuple)):
+                raise TypeError("'speaker_embedding' must not mix flat and nested values")
+            embeddings.append([float(x) for x in item])
+        return embeddings
+
+    return [float(x) for x in value]
+
+
 class OpenAICreateSpeechRequest(BaseModel):
     input: str
     model: str | None = None
@@ -46,7 +83,7 @@ class OpenAICreateSpeechRequest(BaseModel):
         default=None,
         description="Language code (e.g., 'Chinese', 'English', 'Auto')",
     )
-    ref_audio: str | None = Field(
+    ref_audio: str | list[str] | None = Field(
         default=None,
         description="Reference audio for voice cloning (Base task). URL, base64, or file URI.",
     )
@@ -73,7 +110,7 @@ class OpenAICreateSpeechRequest(BaseModel):
         default=None,
         description="Use speaker embedding only without in-context learning (Base task)",
     )
-    speaker_embedding: list[float] | None = Field(
+    speaker_embedding: list[float] | list[list[float]] | None = Field(
         default=None,
         max_length=_MAX_EMBEDDING_DIM,
         description="Pre-computed speaker embedding vector (1024-dim for 0.6B, "
@@ -108,10 +145,29 @@ class OpenAICreateSpeechRequest(BaseModel):
             raise ValueError("'sse' is not a supported stream_format yet. Please use 'audio'.")
         return v
 
+    @field_validator("ref_audio", mode="before")
+    @classmethod
+    def normalize_ref_audio(cls, v):
+        return _normalize_ref_audio_value(v)
+
     @field_validator("speaker_embedding")
     @classmethod
-    def validate_speaker_embedding(cls, v: list[float] | None) -> list[float] | None:
-        if v is not None and not all(math.isfinite(x) for x in v):
+    def validate_speaker_embedding(
+        cls, v: list[float] | list[list[float]] | None
+    ) -> list[float] | list[list[float]] | None:
+        v = _normalize_speaker_embedding_value(v)
+        if v is None:
+            return None
+        if not v:
+            return []
+        if isinstance(v[0], list):
+            for item in v:
+                if not item:
+                    raise ValueError("'speaker_embedding' nested vectors must be non-empty")
+                if not all(math.isfinite(x) for x in item):
+                    raise ValueError("'speaker_embedding' values must be finite (no NaN or Inf)")
+            return v
+        if not all(math.isfinite(x) for x in v):
             raise ValueError("'speaker_embedding' values must be finite (no NaN or Inf)")
         return v
 
@@ -210,7 +266,7 @@ class OpenAICreateSpeechRequest(BaseModel):
     @model_validator(mode="after")
     def validate_embedding_constraints(self) -> "OpenAICreateSpeechRequest":
         if self.speaker_embedding is not None:
-            if self.ref_audio is not None:
+            if self.ref_audio is not None and not isinstance(self.ref_audio, list):
                 raise ValueError("'speaker_embedding' and 'ref_audio' are mutually exclusive")
         return self
 

vllm_omni/entrypoints/openai/serving_speech.py

@@ -66,6 +66,7 @@ from vllm_omni.utils.speaker_cache import (
 logger = init_logger(__name__)
 
 # TTS Configuration
+_MING_TTS_MODEL_ARCHS = {"MingTTSForConditionalGeneration"}
 _VOXTRAL_TTS_MODEL_STAGES = {"audio_generation"}
 _QWEN3_TTS_MODEL_STAGES = {"qwen3_tts"}
 _FISH_TTS_MODEL_STAGES = {"fish_speech_slow_ar"}
@@ -126,6 +127,7 @@ _QWEN3_TTS_REF_AUDIO_CACHE_KEY = "_qwen3_tts_ref_audio_cache_key"
 _TTS_MAX_INSTRUCTIONS_LENGTH = 500
 _TTS_MAX_NEW_TOKENS_MIN = 1
 _TTS_MAX_NEW_TOKENS_MAX = 4096
+_MING_DEFAULT_PROMPT = MING_DEFAULT_PROMPT
 
 
 def _create_wav_header(sample_rate: int, num_channels: int = 1, bits_per_sample: int = 16) -> bytes:
@@ -410,6 +412,11 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         )
         self._fish_speech_tokenizer = None
         self._covo_audio_tokenizer = None
+        # Cached per process: the CosyVoice3 Qwen tokenizer + resolved model
+        # path used for dynamic-token sizing. Without this, every request
+        # re-ran snapshot_download + reloaded the tokenizer (~100 ms on the
+        # TTFP critical path) in _apply_cosyvoice3_dynamic_tokens.
+        self._cosyvoice3_tokenizer = None
 
         self._is_cosyvoice3 = (
             self._tts_stage is not None
@@ -523,6 +530,32 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
 
     def _load_codec_frame_rate(self) -> float | None:
         """Load codec frame rate from speech tokenizer config for prompt length estimation."""
+        if self._tts_model_type == "ming_tts":
+            try:
+                from vllm_omni.model_executor.models.ming_tts.config_ming_tts import MingTTSConfig
+
+                hf_config = self.engine_client.model_config.hf_config
+                ming_cfg = MingTTSConfig.from_hf_config(hf_config)
+                patch_size = int(ming_cfg.patch_size)
+                audio_frame_hop = int(ming_cfg.audio_frame_hop)
+                sample_rate = int(ming_cfg.sample_rate)
+                if patch_size <= 0 or audio_frame_hop <= 0 or sample_rate <= 0:
+                    raise ValueError(
+                        "Ming config has invalid tokenizer timing values: "
+                        f"patch_size={patch_size}, audio_frame_hop={audio_frame_hop}, sample_rate={sample_rate}"
+                    )
+                rate = float(sample_rate) / float(audio_frame_hop * patch_size)
+                logger.info(
+                    "Derived Ming codec frame rate: %.1f Hz (sample_rate=%s, audio_frame_hop=%s, patch_size=%s)",
+                    rate,
+                    sample_rate,
+                    audio_frame_hop,
+                    patch_size,
+                )
+                return rate
+            except Exception as e:
+                logger.warning(f"Failed to derive Ming codec frame rate from hf_config: {e}")
+
         try:
             model_path = self.engine_client.model_config.model
             st_config_path = os.path.join(model_path, "speech_tokenizer", "config.json")
@@ -567,7 +600,15 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
     def _find_tts_stage(self):
         """Find and return the TTS stage config, or None if not found."""
         for stage in self.engine_client.stage_configs:
-            if stage.engine_args.model_stage in _TTS_MODEL_STAGES:
+            engine_args = stage.engine_args
+            model_stage = engine_args.model_stage
+            model_arch = getattr(engine_args, "model_arch", None)
+            worker_type = getattr(engine_args, "worker_type", None)
+            if model_stage in _TTS_MODEL_STAGES:
+                return stage
+            # Ming dense identifies its AR entry stage by architecture because
+            # it does not use a dedicated TTS model_stage value.
+            if model_arch in _MING_TTS_MODEL_ARCHS and worker_type == "ar":
                 return stage
         return None
 
@@ -597,6 +638,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             return "voxcpm2"
         if model_stage in _MING_TTS_MODEL_STAGES:
             return "ming_flash_omni_tts"
+        if model_arch in _MING_TTS_MODEL_ARCHS:
+            return "ming_tts"
         if model_stage in _MOSS_TTS_MODEL_STAGES:
             return "moss_tts_nano"
         if model_stage in _MOSS_TTS_FULL_MODEL_STAGES:
@@ -677,6 +720,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         try:
             if self._tts_model_type == "glm_tts":
                 return set()
+            if self._tts_model_type == "ming_tts":
+                return set()
             if self._tts_model_type == "voxcpm2":
                 return {"default"}
             if self._tts_model_type == "voxtral_tts":
@@ -1209,9 +1254,13 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             raise ValueError("'speaker_embedding' values must be finite (no NaN or Inf)")
 
         emb_dim = len(embedding)
-        dim_err = self._validate_qwen_tts_speaker_embedding_dim(emb_dim)
-        if dim_err is not None:
-            raise ValueError(dim_err)
+        if self._tts_model_type == "ming_tts":
+            if emb_dim != 192:
+                logger.warning("speaker_embedding has %d dimensions; Ming dense expects 192", emb_dim)
+        else:
+            dim_err = self._validate_qwen_tts_speaker_embedding_dim(emb_dim)
+            if dim_err is not None:
+                raise ValueError(dim_err)
 
         async with self._upload_lock:
             voice_name_lower = name.lower()
@@ -1298,7 +1347,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
 
     def _is_tts_model(self) -> bool:
         """Check if the current model is a supported TTS model."""
-        return any(stage.engine_args.model_stage in _TTS_MODEL_STAGES for stage in self.engine_client.stage_configs)
+        return self._find_tts_stage() is not None
 
     def _validate_tts_request(self, request: OpenAICreateSpeechRequest) -> str | None:
         """Validate TTS request parameters. Returns error message or None."""
@@ -1311,6 +1360,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         if self._tts_model_type == "voxcpm2":
             return self._validate_voxcpm2_request(request)
         if self._tts_model_type == "ming_flash_omni_tts":
+            return self._validate_ming_flash_omni_tts_request(request)
+        if self._tts_model_type == "ming_tts":
             return self._validate_ming_tts_request(request)
         if self._tts_model_type in ("moss_tts_nano", "moss_tts"):
             return self._validate_moss_tts_request(request)
@@ -1357,7 +1408,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         ids = self._voxcpm2_tokenizer.encode(text, add_special_tokens=True)
         return split_multichar_chinese(ids, self._voxcpm2_split_map)
 
-    def _validate_ming_tts_request(self, request: OpenAICreateSpeechRequest) -> str | None:
+    def _validate_ming_flash_omni_tts_request(self, request: OpenAICreateSpeechRequest) -> str | None:
         """Validate Ming-flash-omni standalone-talker request parameters."""
         if not request.input or not request.input.strip():
             return "Input text cannot be empty"
@@ -1394,6 +1445,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
 
     def _validate_ref_audio_format(self, ref_audio: str) -> str | None:
         """Validate ref_audio is a supported URI format. Returns error or None."""
+        if not isinstance(ref_audio, str):
+            return "ref_audio must be a URL (http/https), base64 data URL (data:...), or file URI (file://...)"
         if not (
             ref_audio.startswith(("http://", "https://"))
             or ref_audio.startswith("data:")
@@ -2061,6 +2114,99 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
 
         return None
 
+    def _validate_ming_tts_request(self, request: OpenAICreateSpeechRequest) -> str | None:
+        """Validate Ming TTS request parameters. Returns error message or None."""
+        if not request.input or not request.input.strip():
+            return "Input text cannot be empty"
+
+        if isinstance(request.ref_audio, list):
+            return self._validate_ming_tts_podcast_request(request)
+        return self._validate_ming_tts_single_speaker_request(request)
+
+    def _validate_ming_tts_single_speaker_request(self, request: OpenAICreateSpeechRequest) -> str | None:
+        if request.ref_audio is not None:
+            fmt_err = self._validate_ref_audio_format(request.ref_audio)
+            if fmt_err:
+                return fmt_err
+
+        if request.speaker_embedding is not None:
+            if not request.speaker_embedding:
+                return "'speaker_embedding' must be a non-empty list of floats"
+            emb_len = len(request.speaker_embedding)
+            if emb_len != 192:
+                logger.warning(
+                    "speaker_embedding has %d dimensions; Ming dense expects 192. "
+                    "Wrong dimensions will likely fail or degrade output.",
+                    emb_len,
+                )
+
+        voice_lower = request.voice.lower() if isinstance(request.voice, str) else None
+        uploaded_voice = bool(voice_lower and voice_lower in self.uploaded_speakers)
+        clone_source_present = request.ref_audio is not None or request.speaker_embedding is not None or uploaded_voice
+
+        if request.task_type == "Base" and not clone_source_present:
+            return "Base task requires 'ref_audio', 'speaker_embedding', or an uploaded voice sample"
+
+        if request.ref_audio is not None and request.ref_text is not None and not request.ref_text.strip():
+            return "'ref_text' must be non-empty when provided with 'ref_audio'"
+
+        # Ming offline ref-audio cases use prompt_waveform without prompt_text;
+        # keep the transcript requirement for other TTS models.
+        if request.ref_audio is not None and request.speaker_embedding is None and not self._is_ming_tts_model():
+            uploaded_ref_text = self.uploaded_speakers[voice_lower].get("ref_text") if uploaded_voice else None
+            if not (request.ref_text and request.ref_text.strip()) and not uploaded_ref_text:
+                return "Reference-audio cloning requires non-empty 'ref_text'"
+
+        if request.ref_text is not None and request.ref_audio is None and not uploaded_voice:
+            return "'ref_text' requires 'ref_audio' or an uploaded voice sample"
+
+        if request.instructions and len(request.instructions) > self._max_instructions_length:
+            return f"Instructions too long (max {self._max_instructions_length} characters)"
+
+        if request.max_new_tokens is not None:
+            if request.max_new_tokens < _TTS_MAX_NEW_TOKENS_MIN:
+                return f"max_new_tokens must be at least {_TTS_MAX_NEW_TOKENS_MIN}"
+            if request.max_new_tokens > _TTS_MAX_NEW_TOKENS_MAX:
+                return f"max_new_tokens cannot exceed {_TTS_MAX_NEW_TOKENS_MAX}"
+
+        return None
+
+    def _validate_ming_tts_podcast_request(self, request: OpenAICreateSpeechRequest) -> str | None:
+        if len(request.ref_audio) < 2:
+            return "Podcast-style Ming requests require at least two 'ref_audio' clips"
+
+        for ref_audio in request.ref_audio:
+            fmt_err = self._validate_ref_audio_format(ref_audio)
+            if fmt_err:
+                return fmt_err
+
+        if not request.ref_text or not request.ref_text.strip():
+            return "Podcast-style Ming requests require non-empty 'ref_text'"
+
+        if request.speaker_embedding is not None:
+            embeddings = request.speaker_embedding
+            embedding_count = len(embeddings) if embeddings and isinstance(embeddings[0], list) else 1
+            if embedding_count != len(request.ref_audio):
+                return (
+                    "Podcast-style Ming requests require one speaker embedding per ref_audio clip; "
+                    f"got {embedding_count} embeddings for {len(request.ref_audio)} clips"
+                )
+            if embeddings and isinstance(embeddings[0], list):
+                for item in embeddings:
+                    if len(item) != 192:
+                        return "Podcast-style Ming speaker embeddings must each have 192 dimensions"
+
+        if request.instructions and len(request.instructions) > self._max_instructions_length:
+            return f"Instructions too long (max {self._max_instructions_length} characters)"
+
+        if request.max_new_tokens is not None:
+            if request.max_new_tokens < _TTS_MAX_NEW_TOKENS_MIN:
+                return f"max_new_tokens must be at least {_TTS_MAX_NEW_TOKENS_MIN}"
+            if request.max_new_tokens > _TTS_MAX_NEW_TOKENS_MAX:
+                return f"max_new_tokens cannot exceed {_TTS_MAX_NEW_TOKENS_MAX}"
+
+        return None
+
     async def _resolve_ref_audio(self, ref_audio_str: str) -> tuple[list[float], int]:
         """Resolve ref_audio to (wav_samples, sample_rate).
 
@@ -2187,6 +2333,142 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
     def _discard_ref_audio_artifact_warmup(self, request_id: str) -> None:
         self._request_ref_audio_artifact_keys.pop(request_id, None)
 
+    async def _resolve_ref_audio_many(self, ref_audio_list: list[str]) -> list[tuple[list[float], int]]:
+        resolved = []
+        for ref_audio in ref_audio_list:
+            resolved.append(await self._resolve_ref_audio(ref_audio))
+        return resolved
+
+    # ---- Ming TTS helpers ----
+
+    def _is_ming_tts_model(self) -> bool:
+        return self._tts_model_type == "ming_tts"
+
+    def _coerce_ming_prompt_waveform(self, wav_samples, sample_rate):
+        from torchaudio.functional import resample as resample_audio
+
+        from vllm_omni.model_executor.models.ming_tts.config_ming_tts import SAMPLE_RATE
+
+        waveform = torch.as_tensor(wav_samples, dtype=torch.float32).reshape(1, -1)
+        if int(sample_rate) != SAMPLE_RATE:
+            waveform = resample_audio(waveform, int(sample_rate), SAMPLE_RATE)
+        return waveform
+
+    def _build_ming_prompt_waveform(
+        self,
+        ref_audio_data: tuple[list[float], int] | list[tuple[list[float], int]] | None,
+    ):
+        if isinstance(ref_audio_data, list):
+            return torch.cat(
+                [self._coerce_ming_prompt_waveform(item[0], item[1]) for item in ref_audio_data],
+                dim=-1,
+            )
+        if ref_audio_data is not None:
+            return self._coerce_ming_prompt_waveform(ref_audio_data[0], ref_audio_data[1])
+        return None
+
+    def _extract_ming_speaker_embeddings_from_ref_audio(
+        self,
+        ref_audio_data_list: list[tuple[list[float], int]],
+    ) -> list[list[float]]:
+        from vllm_omni.model_executor.models.ming_tts.speaker_extractor import MingSpeakerEmbeddingExtractor
+
+        extractor = MingSpeakerEmbeddingExtractor(self.engine_client.model_config.model, target_sr=16000)
+        embeddings = []
+        for wav_samples, sr in ref_audio_data_list:
+            waveform = torch.as_tensor(wav_samples, dtype=torch.float32).reshape(1, -1)
+            embedding = extractor.extract_from_waveform(waveform, int(sr))
+            flat = embedding.detach().reshape(-1).to(torch.float32).cpu()
+            if int(flat.numel()) != 192:
+                raise ValueError(f"Ming speaker extractor returned {int(flat.numel())} dims; expected 192")
+            embeddings.append(flat.tolist())
+        return embeddings
+
+    def _parse_ming_instruction_fields(
+        self,
+        request,
+        *,
+        include_language=False,
+        include_voice=False,
+        plain_text_passthrough=False,
+    ):
+        instruction_text = request.instructions.strip() if isinstance(request.instructions, str) else None
+        instruction_dict: dict[str, Any] = {}
+
+        if include_language and request.language not in (None, "", "Auto"):
+            instruction_dict["方言"] = request.language
+
+        voice_lower = request.voice.lower() if isinstance(request.voice, str) else None
+        if include_voice and request.voice and not (voice_lower and voice_lower in self.uploaded_speakers):
+            instruction_dict["IP"] = request.voice
+
+        if instruction_text:
+            try:
+                parsed = json.loads(instruction_text)
+            except json.JSONDecodeError:
+                parsed = None
+            if isinstance(parsed, dict):
+                instruction_dict.update(parsed)
+            elif instruction_dict or not plain_text_passthrough:
+                instruction_dict["风格"] = instruction_text
+            else:
+                return instruction_text
+
+        return instruction_dict or None
+
+    def _parse_ming_instruction(self, request: OpenAICreateSpeechRequest) -> Any:
+        """Build a Ming instruction payload from OpenAI speech fields."""
+        return self._parse_ming_instruction_fields(
+            request,
+            include_language=True,
+            include_voice=True,
+            plain_text_passthrough=True,
+        )
+
+    def _build_ming_dense_prompt(
+        self,
+        request: OpenAICreateSpeechRequest,
+        *,
+        ref_audio_data: tuple[list[float], int] | list[tuple[list[float], int]] | None = None,
+    ) -> dict[str, Any]:
+        """Build a Ming dense prompt directly from the OpenAI speech request."""
+        from transformers import AutoTokenizer
+
+        from vllm_omni.model_executor.models.ming_tts.config_ming_tts import KEY_MAX_DECODE_STEPS
+        from vllm_omni.model_executor.models.ming_tts.prompt_assembly import build_ming_dense_prompt
+
+        if self._tts_tokenizer is None:
+            model_name = self.engine_client.model_config.model
+            self._tts_tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=False)
+
+        ref_text = request.ref_text
+        prompt_waveform = self._build_ming_prompt_waveform(ref_audio_data) if ref_text is not None else None
+        speaker_embedding = request.speaker_embedding
+        use_zero_spk_emb = prompt_waveform is None and speaker_embedding is None
+
+        runtime_controls = {}
+        if request.max_new_tokens is not None:
+            runtime_controls[KEY_MAX_DECODE_STEPS] = request.max_new_tokens
+
+        prompt_dict = build_ming_dense_prompt(
+            self._tts_tokenizer,
+            # bgm / music-prompt mode not supported online;
+            # requires prompt_mode API extension (deferred).
+            prompt=_MING_DEFAULT_PROMPT,
+            text=request.input,
+            runtime_controls=runtime_controls or None,
+            instruction=self._parse_ming_instruction(request),
+            prompt_text=ref_text,
+            prompt_waveform=prompt_waveform,
+            speaker_embedding=speaker_embedding,
+            use_zero_spk_emb=use_zero_spk_emb,
+        )
+        prompt = tokens_input(prompt_token_ids=prompt_dict["prompt_token_ids"])
+        prompt["prompt"] = prompt_dict["prompt"]
+        prompt["text"] = prompt_dict["text"]
+        prompt["additional_information"] = prompt_dict["additional_information"]
+        return prompt
+
     async def _generate_audio_chunks(
         self,
         generator,
@@ -2658,16 +2940,21 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
 
         sampling_params_list = copy.deepcopy(sampling_params_list)
         hf_cfg = self.model_config.hf_config
-        model_path = self.engine_client.model_config.model
-        if not os.path.isdir(model_path):
-            from huggingface_hub import snapshot_download
+        # Build the Qwen tokenizer once per process (resolving the model dir via
+        # snapshot_download at most once) and reuse it across requests.
+        tokenizer = self._cosyvoice3_tokenizer
+        if tokenizer is None:
+            model_path = self.engine_client.model_config.model
+            if not os.path.isdir(model_path):
+                from huggingface_hub import snapshot_download
 
-            model_path = snapshot_download(model_path)
-        tokenizer = get_qwen_tokenizer(
-            token_path=os.path.join(model_path, hf_cfg.qwen_pretrain_path),
-            skip_special_tokens=hf_cfg.skip_special_tokens,
-            version=hf_cfg.version,
-        )
+                model_path = snapshot_download(model_path)
+            tokenizer = get_qwen_tokenizer(
+                token_path=os.path.join(model_path, hf_cfg.qwen_pretrain_path),
+                skip_special_tokens=hf_cfg.skip_special_tokens,
+                version=hf_cfg.version,
+            )
+            self._cosyvoice3_tokenizer = tokenizer
         _, text_token_len = extract_text_token(
             request.input,
             tokenizer,
@@ -2786,25 +3073,12 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
 
     # ---- Ming-flash-omni standalone-talker (TTS) helpers ----
 
-    def _build_ming_prompt(self, request: OpenAICreateSpeechRequest) -> dict[str, Any]:
+    def _build_ming_flash_omni_prompt(self, request: OpenAICreateSpeechRequest) -> dict[str, Any]:
         # request.instructions accepts two forms:
         # 1. Plain text: mapped to the caption's 风格 (style) field
         # 2. JSON object: parsed and splatted into the caption. Unlocks
         #       Unknown keys are dropped by `ming_create_instruction`.
-        caption_fields: dict[str, Any] = {}
-        if request.instructions:
-            stripped = request.instructions.strip()
-            if stripped.startswith("{"):
-                try:
-                    parsed = json.loads(stripped)
-                except json.JSONDecodeError:
-                    parsed = None
-                if isinstance(parsed, dict):
-                    caption_fields.update(parsed)
-                else:
-                    caption_fields["风格"] = request.instructions
-            else:
-                caption_fields["风格"] = request.instructions
+        caption_fields = self._parse_ming_instruction_fields(request) or {}
 
         has_spk_emb = request.speaker_embedding is not None
 
@@ -2869,6 +3143,11 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                 request, ref_audio_data=ref_audio_data, has_inline_ref_audio=has_inline_ref_audio
             )
             tts_params = {}
+            # Fish Speech resolves conditioning entirely from raw request fields
+            # (text / ref_audio / ref_text rebuilt into inputs_embeds); unlike
+            # uploaded-voice models it has no resolved tts_params to fold in, so
+            # the salt is derived from the request alone.
+            prompt["cache_salt"] = _conditioning_cache_salt(request, tts_params)
         elif self._tts_model_type == "omnivoice":
             if not request.input or not request.input.strip():
                 raise ValueError("Input text cannot be empty")
@@ -2967,8 +3246,29 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                 prompt = await self._build_glm_tts_prompt(request, has_inline_ref_audio=has_inline_ref_audio)
                 tts_params = {}
             elif self._tts_model_type == "ming_flash_omni_tts":
-                prompt = self._build_ming_prompt(request)
+                prompt = self._build_ming_flash_omni_prompt(request)
                 tts_params = {}
+            elif self._tts_model_type == "ming_tts":
+                ref_audio_source = request.ref_audio
+                voice_lower = request.voice.lower() if isinstance(request.voice, str) else None
+                if ref_audio_source is None and voice_lower in self.uploaded_speakers:
+                    ref_audio_source = self._get_uploaded_audio_data(request.voice)
+                    if request.ref_text is None:
+                        request.ref_text = self.uploaded_speakers[voice_lower].get("ref_text")
+                ref_audio_data = None
+                if isinstance(ref_audio_source, list):
+                    ref_audio_data = await self._resolve_ref_audio_many(ref_audio_source)
+                    if request.speaker_embedding is None:
+                        request.speaker_embedding = self._extract_ming_speaker_embeddings_from_ref_audio(ref_audio_data)
+                elif ref_audio_source is not None and isinstance(ref_audio_source, str):
+                    wav_list, sr = await self._resolve_ref_audio(ref_audio_source)
+                    ref_audio_data = (wav_list, sr)
+                    if request.speaker_embedding is None:
+                        request.speaker_embedding = self._extract_ming_speaker_embeddings_from_ref_audio(
+                            [ref_audio_data]
+                        )[0]
+                prompt = self._build_ming_dense_prompt(request, ref_audio_data=ref_audio_data)
+                tts_params = prompt.get("additional_information", {})
             elif self._tts_model_type in ("moss_tts_nano", "moss_tts"):
                 tts_params = await self._build_moss_tts_params(request)
                 if request.voice:
@@ -3054,6 +3354,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             model_type = "voxcpm2"
         elif self._tts_model_type == "ming_flash_omni_tts":
             model_type = "ming_flash_omni_tts"
+        elif self._tts_model_type == "ming_tts":
+            model_type = "ming_tts"
         elif self._tts_model_type == "moss_tts_nano":
             model_type = "moss_tts_nano"
         elif self._tts_model_type == "moss_tts":
@@ -3162,6 +3464,17 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                     getattr(sampling_params_list[0], "min_tokens", 0),
                     request.max_new_tokens,
                 )
+        elif self._tts_model_type == "ming_tts" and sampling_params_list:
+            import copy
+
+            from vllm_omni.model_executor.models.ming_tts.config_ming_tts import TEXT_EOS_TOKEN_ID
+
+            sampling_params_list = copy.deepcopy(sampling_params_list)
+            sampling_params_list[0].stop_token_ids = [int(TEXT_EOS_TOKEN_ID)]
+            if request.max_new_tokens is not None:
+                # Ming emits TEXT_EOS after the latent decode budget is exhausted, so
+                # Stage-0 needs one extra token beyond ming_max_decode_steps.
+                sampling_params_list[0].max_tokens = int(request.max_new_tokens) + 1
 
         if request.seed is not None and sampling_params_list:
             if sampling_params_list is self.engine_client.default_sampling_params_list:

vllm_omni/model_executor/models/common/ming/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm_omni/model_executor/models/common/ming/audio_dsp.py

@@ -0,0 +1,209 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Adopted from https://github.com/inclusionAI/Ming-omni-tts/blob/main/audio_tokenizer/istft.py
+
+import torch
+import torch.nn as nn
+
+
+class ISTFT(nn.Module):
+    """
+    Custom implementation of ISTFT since torch.istft doesn't allow custom padding (other than `center=True`) with
+    windowing. This is because the NOLA (Nonzero Overlap Add) check fails at the edges.
+    See issue: https://github.com/pytorch/pytorch/issues/62323
+    Specifically, in the context of neural vocoding we are interested in "same" padding analogous to CNNs.
+    The NOLA constraint is met as we trim padded samples anyway.
+
+    Args:
+        n_fft (int): Size of Fourier transform.
+        hop_length (int): The distance between neighboring sliding window frames.
+        win_length (int): The size of window frame and STFT filter.
+        padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same".
+    """
+
+    def __init__(self, n_fft: int, hop_length: int, win_length: int, padding: str = "same"):
+        super().__init__()
+        if padding not in ["center", "same"]:
+            raise ValueError("Padding must be 'center' or 'same'.")
+        self.padding = padding
+        self.n_fft = n_fft
+        self.hop_length = hop_length
+        self.win_length = win_length
+        window = torch.hann_window(win_length)
+        self.register_buffer("window", window)
+
+        self.buffer_len = self.win_length - self.hop_length
+
+    def _buffer_process(
+        self,
+        x: torch.Tensor,
+        buffer: torch.Tensor | None,
+        pad: int,
+        last_chunk: bool = False,
+        streaming: bool = False,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        if streaming:
+            if buffer is None:
+                # first chunk
+                x = x[:, pad:]
+            if buffer is not None:
+                # next chunk
+                x[:, : self.buffer_len] += buffer
+            buffer = x[:, -self.buffer_len :]
+            if not last_chunk:
+                x = x[:, : -self.buffer_len]
+            else:
+                x = x[:, :-pad]
+        else:
+            x = x[:, pad:-pad]
+
+        return x, buffer
+
+    def forward(
+        self,
+        spec: torch.Tensor,
+        audio_buffer: torch.Tensor | None = None,
+        window_buffer: torch.Tensor | None = None,
+        streaming: bool = False,
+        last_chunk: bool = False,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]:
+        """
+        Compute the Inverse Short Time Fourier Transform (ISTFT) of a complex spectrogram.
+
+        Args:
+            spec (Tensor): Input complex spectrogram of shape (B, N, T), where B is the batch size,
+                            N is the number of frequency bins, and T is the number of time frames.
+            audio_buffer (Tensor): [Streaming Input/State] The audio overlap buffer from the previous chunk.
+                            Shape: (B, win_length - hop_length)
+            window_buffer (Tensor): [Streaming Input/State] The window overlap buffer from the previous chunk.
+            streaming: If `True`, the function operates in streaming mode, processing `spec` as a single chunk.
+            last_chunk: When `streaming=True` and `last_chunk=True`, the function can perform final "flush" operations
+
+        Returns:
+            Reconstructed signal, plus streaming buffers when `padding="same"`.
+        """
+        if self.padding == "center":
+            # Fallback to pytorch native implementation
+            return torch.istft(spec, self.n_fft, self.hop_length, self.win_length, self.window, center=True)
+        elif self.padding == "same":
+            pad = (self.win_length - self.hop_length) // 2
+        else:
+            raise ValueError("Padding must be 'center' or 'same'.")
+
+        if spec.dim() != 3:
+            raise ValueError(f"Expected spec rank-3 [Batch, Freq, Time], got {tuple(spec.shape)}")
+        _, _, T = spec.shape
+
+        # Inverse FFT
+        ifft = torch.fft.irfft(spec, self.n_fft, dim=1, norm="backward")
+        ifft = ifft * self.window[None, :, None]
+
+        # Overlap and Add
+        output_size = (T - 1) * self.hop_length + self.win_length
+        y = torch.nn.functional.fold(
+            ifft,
+            output_size=(1, output_size),
+            kernel_size=(1, self.win_length),
+            stride=(1, self.hop_length),
+        )[:, 0, 0, :]
+
+        y, audio_buffer = self._buffer_process(y, audio_buffer, pad, last_chunk=last_chunk, streaming=streaming)
+
+        # Window envelope
+        window_sq = self.window.square().expand(1, T, -1).transpose(1, 2)
+        window_envelope = (
+            torch.nn.functional.fold(
+                window_sq,
+                output_size=(1, output_size),
+                kernel_size=(1, self.win_length),
+                stride=(1, self.hop_length),
+            )
+            .squeeze(0)
+            .squeeze(0)
+        )
+
+        window_envelope, window_buffer = self._buffer_process(
+            window_envelope, window_buffer, pad, last_chunk=last_chunk, streaming=streaming
+        )
+        window_envelope = window_envelope.squeeze()
+
+        # Normalize
+        if not (window_envelope > 1e-11).all():
+            raise RuntimeError("ISTFT window envelope underflowed; invalid overlap-add state.")
+        y = y / window_envelope
+
+        return y, audio_buffer, window_buffer
+
+
+class FourierHead(nn.Module):
+    """Base class for inverse fourier modules."""
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x (Tensor): Input tensor of shape (B, L, H), where B is the batch size,
+                        L is the sequence length, and H denotes the model dimension.
+
+        Returns:
+            Tensor: Reconstructed time-domain audio signal of shape (B, T), where T is the length of the output signal.
+        """
+        raise NotImplementedError("Subclasses must implement the forward method.")
+
+
+class ISTFTHead(FourierHead):
+    """
+    ISTFT Head module for predicting STFT complex coefficients.
+
+    Args:
+        dim (int): Hidden dimension of the model.
+        n_fft (int): Size of Fourier transform.
+        hop_length (int): The distance between neighboring sliding window frames, which should align with
+                          the resolution of the input features.
+        padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same".
+    """
+
+    def __init__(self, dim: int, n_fft: int, hop_length: int, padding: str = "same"):
+        super().__init__()
+        out_dim = n_fft + 2
+        self.out = torch.nn.Linear(dim, out_dim)
+        self.istft = ISTFT(n_fft=n_fft, hop_length=hop_length, win_length=n_fft, padding=padding)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        audio_buffer: torch.Tensor | None = None,
+        window_buffer: torch.Tensor | None = None,
+        streaming: bool = False,
+        last_chunk: bool = False,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None]:
+        """
+        Forward pass of the ISTFTHead module.
+
+        Args:
+            x (Tensor): Input tensor of shape (B, L, H), where B is the batch size,
+                        L is the sequence length, and H denotes the model dimension.
+
+        Returns:
+            Audio, predicted spectrogram coefficients, and streaming buffers.
+        """
+        x_pred = self.out(x)
+        x_pred = x_pred.transpose(1, 2)
+        mag, p = x_pred.chunk(2, dim=1)
+        mag = torch.exp(mag)
+        mag = torch.clip(mag, max=1e2)  # safeguard to prevent excessively large magnitudes
+        # wrapping happens here. These two lines produce real and imaginary value
+        x = torch.cos(p)
+        y = torch.sin(p)
+        # recalculating phase here does not produce anything new
+        # only costs time
+        # phase = torch.atan2(y, x)
+        # S = mag * torch.exp(phase * 1j)
+        # better directly produce the complex value
+        S = mag * (x + 1j * y)
+        audio, audio_buffer, window_buffer = self.istft(
+            S, audio_buffer=audio_buffer, window_buffer=window_buffer, streaming=streaming, last_chunk=last_chunk
+        )
+        return audio.unsqueeze(1), x_pred, audio_buffer, window_buffer
+
+
+__all__ = ["FourierHead", "ISTFT", "ISTFTHead"]

vllm_omni/model_executor/models/common/ming/audio_vae.py

@@ -1,10 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# Copyright 2025 The vLLM-Omni team.
-# Copyright (c) Ant Group. All rights reserved.
-# Adapted from:
-# https://github.com/inclusionAI/Ming/tree/e58533db227031990c5a6864dcf5f08fb53ed0d2/AudioVAE
-
-from __future__ import annotations
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import torch
 import torch.nn as nn
@@ -13,6 +8,8 @@ from transformers import PretrainedConfig, PreTrainedModel, Qwen2Config, Qwen2Mo
 from transformers.utils import is_flash_attn_2_available
 from vllm.logger import init_logger
 
+from .audio_dsp import ISTFTHead
+
 logger = init_logger(__name__)
 
 
@@ -21,114 +18,50 @@ class AudioVAEConfig(PretrainedConfig):
 
     def __init__(
         self,
-        sample_rate: int = 44100,
-        enc_kwargs: dict | None = None,
-        dec_kwargs: dict | None = None,
-        init_method: str = "kaiming",
-        patch_size: int = 4,
+        sample_rate=44100,
+        enc_kwargs=None,
+        semantic_module_kwargs=None,
+        dec_kwargs=None,
+        hifi_gan_disc_kwargs=None,
+        spec_disc_kwargs=None,
+        lambda_disc=1.0,
+        lambda_mel_loss=15,
+        lambda_adv=1.0,
+        lambda_feat_match_loss=1.0,
+        lambda_semantic=5.0,
+        init_method="kaiming",
+        patch_size=4,
         **kwargs,
     ):
         self.sample_rate = sample_rate
         self.enc_kwargs = enc_kwargs or {}
+        self.semantic_module_kwargs = semantic_module_kwargs
         self.dec_kwargs = dec_kwargs or {}
+        self.hifi_gan_disc_kwargs = hifi_gan_disc_kwargs
+        self.spec_disc_kwargs = spec_disc_kwargs
+        self.lambda_disc = lambda_disc
+        self.lambda_mel_loss = lambda_mel_loss
+        self.lambda_adv = lambda_adv
+        self.lambda_feat_match_loss = lambda_feat_match_loss
+        self.lambda_semantic = lambda_semantic
         self.init_method = init_method
         self.patch_size = patch_size
         super().__init__(**kwargs)
 
 
-class ISTFT(nn.Module):
-    def __init__(self, n_fft: int, hop_length: int, win_length: int, padding: str = "same"):
-        super().__init__()
-        if padding not in ["center", "same"]:
-            raise ValueError("Padding must be 'center' or 'same'.")
-        self.padding = padding
-        self.n_fft = n_fft
-        self.hop_length = hop_length
-        self.win_length = win_length
-        window = torch.hann_window(win_length)
-        self.register_buffer("window", window)
-        self.buffer_len = self.win_length - self.hop_length
-
-    def _buffer_process(self, x, buffer, pad, last_chunk=False, streaming=False):
-        if streaming:
-            if buffer is None:
-                x = x[:, pad:]
-            if buffer is not None:
-                x[:, : self.buffer_len] += buffer
-            buffer = x[:, -self.buffer_len :]
-            if not last_chunk:
-                x = x[:, : -self.buffer_len]
-            else:
-                x = x[:, :-pad]
-        else:
-            x = x[:, pad:-pad]
-        return x, buffer
-
-    def forward(self, spec, audio_buffer=None, window_buffer=None, streaming=False, last_chunk=False):
-        if self.padding == "center":
-            return torch.istft(spec, self.n_fft, self.hop_length, self.win_length, self.window, center=True)
-        elif self.padding == "same":
-            pad = (self.win_length - self.hop_length) // 2
-        else:
-            raise ValueError("Padding must be 'center' or 'same'.")
-
-        B, N, T = spec.shape
-        ifft = torch.fft.irfft(spec, self.n_fft, dim=1, norm="backward")
-        ifft = ifft * self.window[None, :, None]
-
-        output_size = (T - 1) * self.hop_length + self.win_length
-        y = torch.nn.functional.fold(
-            ifft,
-            output_size=(1, output_size),
-            kernel_size=(1, self.win_length),
-            stride=(1, self.hop_length),
-        )[:, 0, 0, :]
-
-        y, audio_buffer = self._buffer_process(y, audio_buffer, pad, last_chunk=last_chunk, streaming=streaming)
-
-        window_sq = self.window.square().expand(1, T, -1).transpose(1, 2)
-        window_envelope = (
-            torch.nn.functional.fold(
-                window_sq,
-                output_size=(1, output_size),
-                kernel_size=(1, self.win_length),
-                stride=(1, self.hop_length),
-            )
-            .squeeze(0)
-            .squeeze(0)
-        )
-
-        window_envelope, window_buffer = self._buffer_process(
-            window_envelope, window_buffer, pad, last_chunk=last_chunk, streaming=streaming
-        )
-        window_envelope = window_envelope.squeeze()
-
-        assert (window_envelope > 1e-11).all()
-        y = y / window_envelope
-
-        return y, audio_buffer, window_buffer
-
-
-class ISTFTHead(nn.Module):
-    def __init__(self, dim: int, n_fft: int, hop_length: int, padding: str = "same"):
-        super().__init__()
-        out_dim = n_fft + 2
-        self.out = nn.Linear(dim, out_dim)
-        self.istft = ISTFT(n_fft=n_fft, hop_length=hop_length, win_length=n_fft, padding=padding)
-
-    def forward(self, x, audio_buffer=None, window_buffer=None, streaming=False, last_chunk=False):
-        x_pred = self.out(x)
-        x_pred = x_pred.transpose(1, 2)
-        mag, p = x_pred.chunk(2, dim=1)
-        mag = torch.exp(mag)
-        mag = torch.clip(mag, max=1e2)
-        x = torch.cos(p)
-        y = torch.sin(p)
-        S = mag * (x + 1j * y)
-        audio, audio_buffer, window_buffer = self.istft(
-            S, audio_buffer=audio_buffer, window_buffer=window_buffer, streaming=streaming, last_chunk=last_chunk
-        )
-        return audio.unsqueeze(1), x_pred, audio_buffer, window_buffer
+def _qwen2_config(backbone):
+    config = Qwen2Config.from_dict(config_dict=backbone)
+    if (getattr(config, "_attn_implementation", None) or getattr(config, "attn_implementation", None)) not in (
+        None,
+        "flash_attention_2",
+    ):
+        return config
+    if is_flash_attn_2_available():
+        config._attn_implementation_autoset = True
+        config._attn_implementation = "flash_attention_2"
+    else:
+        config._attn_implementation = "sdpa"
+    return config
 
 
 class StreamingLinearUpsample(nn.Module):
@@ -138,6 +71,8 @@ class StreamingLinearUpsample(nn.Module):
         self.upsampler = nn.Upsample(scale_factor=scale_factor, mode="linear", align_corners=False)
 
     def forward(self, x, state=None, is_last=False):
+        if x is None and is_last and (state is None or state.get("prev_chunk") is None):
+            raise ValueError("Received end-of-stream without any latent chunk to upsample.")
         if state is None:
             state = {"prev_chunk": None, "history_last": None, "is_first": True}
 
@@ -188,16 +123,72 @@ class StreamingLinearUpsample(nn.Module):
         return final_out, state
 
 
+class Encoder(nn.Module):
+    def __init__(self, encoder_args, input_dim=320, hop_size=320, latent_dim=64, patch_size=-1):
+        super().__init__()
+        config = _qwen2_config(encoder_args)
+        logger.info("AudioVAE Encoder: using attn_implementation=%r", config._attn_implementation)
+        self.encoder = Qwen2Model(config)
+        self.input_dim = input_dim
+        self.hop_size = hop_size
+        self.latent_dim = latent_dim
+        self.fc1 = nn.Linear(input_dim, config.hidden_size, bias=False)
+        self.fc2 = nn.Linear(config.hidden_size, config.hidden_size)
+        self.fc3 = nn.Linear(config.hidden_size, latent_dim * 2)
+        self.norm = nn.LayerNorm(config.hidden_size)
+        self.patch_size = patch_size
+        if patch_size != -1:
+            aggregator_config = _qwen2_config({**encoder_args, "num_hidden_layers": 4})
+            self.aggregator = Qwen2Model(aggregator_config)
+            self.cls_embed = nn.Parameter(torch.rand(1, 1, config.hidden_size))
+            self.cls_embed.data.normal_(0, 0.02)
+
+    def get_frames(self, x):
+        num_frames_total = (x.size(-1) + self.hop_size - 1) // self.hop_size
+        expected_len = (num_frames_total - 1) * self.hop_size + self.input_dim
+        padding_needed = expected_len - x.size(-1)
+        waveform = F.pad(x, (0, padding_needed), value=0.0)
+        frames = waveform.unfold(dimension=-1, size=self.input_dim, step=self.hop_size)
+        return frames
+
+    def pad_patch_insert_cls(self, x):
+        bsz, _, dim = x.size()
+        num_frame = x.size(1)
+        r = num_frame % self.patch_size
+        pad_num = self.patch_size - r if r else 0
+        x = F.pad(x, (0, 0, 0, pad_num), value=0.0)
+        # [Batch, Time, Dimension] -> [Batch*PatchGroups, Patch, Dimension].
+        x = x.reshape(-1, self.patch_size, dim)
+        x = torch.cat((x, self.cls_embed.expand(x.size(0), -1, -1)), dim=1)
+        # [Batch*PatchGroups, Patch+1, Dimension] -> [Batch, Time, Dimension].
+        x = x.reshape(bsz, -1, dim)
+        return x
+
+    def forward(self, waveform):
+        x = self.get_frames(waveform)
+        x = self.fc1(x)
+        x = self.fc2(x)
+        x = self.encoder(inputs_embeds=x)
+        x = x.last_hidden_state
+
+        if self.patch_size != -1:
+            x = self.pad_patch_insert_cls(x)
+            x = self.aggregator(inputs_embeds=x)
+            x = x.last_hidden_state
+            bsz, _, dim = x.size()
+            # [Batch, Time, Dimension] -> [Batch*PatchGroups, Patch+1, Dimension].
+            x = x.reshape(-1, self.patch_size + 1, dim)
+            # [Batch*PatchGroups, 1, Dimension] -> [Batch, PatchGroups, Dimension].
+            x = x[:, -1:, :].reshape(bsz, -1, dim)
+
+        x = self.fc3(x)
+        return x, waveform.unsqueeze(1)
+
+
 class Decoder(nn.Module):
     def __init__(self, decoder_args, output_dim=320, latent_dim=64, patch_size=-1):
         super().__init__()
-        config = Qwen2Config.from_dict(config_dict=decoder_args)
-        if is_flash_attn_2_available():
-            config._attn_implementation_autoset = True
-            config._attn_implementation = "flash_attention_2"
-        else:
-            config._attn_implementation = "sdpa"
-
+        config = _qwen2_config(decoder_args)
         logger.info("AudioVAE Decoder: using attn_implementation=%r", config._attn_implementation)
         self.decoder = Qwen2Model(config)
         self.output_dim = output_dim
@@ -211,7 +202,21 @@ class Decoder(nn.Module):
         if self.patch_size != -1:
             self.upsampling = StreamingLinearUpsample(scale_factor=patch_size)
 
+    def forward(self, x):
+        x = self.fc1(x)
+
+        if self.patch_size != -1:
+            # [Batch, Time, Dimension] -> [Batch, Dimension, Time] -> [Batch, Time, Dimension].
+            x = self.upsampling.upsampler(x.transpose(1, 2)).transpose(1, 2)
+
+        x = self.decoder(inputs_embeds=x)
+        x = x.last_hidden_state
+        x, _ = self.head(x)
+        return x, None
+
     def low_level_reconstruct(self, x, past_key_values=None, use_cache=False, stream_state=None, last_chunk=False):
+        if stream_state is None:
+            stream_state = (None, None, None)
         upsample_state, audio_buffer, window_buffer = stream_state
         bsz, device, dtype = x.size(0), x.device, x.dtype
         x = self.fc1(x)
@@ -222,6 +227,7 @@ class Decoder(nn.Module):
                     stream_state = (upsample_state, audio_buffer, window_buffer)
                     return torch.empty(bsz, 1, 0, device=device, dtype=dtype), stream_state, past_key_values
             else:
+                # [Batch, Time, Dimension] -> [Batch, Dimension, Time] -> [Batch, Time, Dimension].
                 x = self.upsampling.upsampler(x.transpose(1, 2)).transpose(1, 2)
 
         hidden_states_list = []
@@ -239,7 +245,6 @@ class Decoder(nn.Module):
                 past_len = 0
 
             curr_len = x.shape[1]
-
             if past_len < target_len and (past_len + curr_len) >= sw_size:
                 fill_len = target_len - past_len
                 x_fill = x[:, :fill_len, :]
@@ -269,86 +274,39 @@ class Decoder(nn.Module):
         return x_out, stream_state, past_key_values
 
 
-class Encoder(nn.Module):
-    def __init__(self, encoder_args, input_dim=320, hop_size=320, latent_dim=64, patch_size=-1):
-        super().__init__()
-        config = Qwen2Config.from_dict(config_dict=encoder_args)
-        if is_flash_attn_2_available():
-            config._attn_implementation_autoset = True
-            config._attn_implementation = "flash_attention_2"
-        else:
-            config._attn_implementation = "sdpa"
-
-        logger.info("AudioVAE Encoder: using attn_implementation=%r", config._attn_implementation)
-        self.encoder = Qwen2Model(config)
-        self.input_dim = input_dim
-        self.hop_size = hop_size
-        self.latent_dim = latent_dim
-        self.fc1 = nn.Linear(input_dim, config.hidden_size, bias=False)
-        self.fc2 = nn.Linear(config.hidden_size, config.hidden_size)
-        self.fc3 = nn.Linear(config.hidden_size, latent_dim * 2)
-        self.norm = nn.LayerNorm(config.hidden_size)
-        self.patch_size = patch_size
-        if patch_size != -1:
-            config.num_hidden_layers = 4
-            self.aggregator = Qwen2Model(config)
-            self.cls_embed = nn.Parameter(torch.rand(1, 1, config.hidden_size))
-            self.cls_embed.data.normal_(0, 0.02)
-
-    def get_frames(self, x):
-        num_frames_total = (x.size(-1) + self.hop_size - 1) // self.hop_size
-        expected_len = (num_frames_total - 1) * self.hop_size + self.input_dim
-        padding_needed = expected_len - x.size(-1)
-        waveform = F.pad(x, (0, padding_needed), value=0.0)
-        frames = waveform.unfold(dimension=-1, size=self.input_dim, step=self.hop_size)
-        return frames
-
-    def pad_patch_insert_cls(self, x):
-        bsz, _, dim = x.size()
-        num_frame = x.size(1)
-        r = num_frame % self.patch_size
-        pad_num = self.patch_size - r if r else 0
-        x = F.pad(x, (0, 0, 0, pad_num), value=0.0)
-        x = x.reshape(-1, self.patch_size, dim)
-        x = torch.cat((x, self.cls_embed.expand(x.size(0), -1, -1)), dim=1)
-        x = x.reshape(bsz, -1, dim)
-        return x
-
-    def forward(self, waveform):
-        x = self.get_frames(waveform)
-        x = self.fc1(x)
-        x = self.fc2(x)
-        x = self.encoder(inputs_embeds=x)
-        x = x.last_hidden_state
-
-        if self.patch_size != -1:
-            x = self.pad_patch_insert_cls(x)
-            x = self.aggregator(inputs_embeds=x)
-            x = x.last_hidden_state
-            bsz, _, dim = x.size()
-            x = x.reshape(-1, self.patch_size + 1, dim)
-            x = x[:, -1:, :].reshape(bsz, -1, dim)
-
-        x = self.fc3(x)
-        return x, waveform.unsqueeze(1)
-
-
 class AudioVAE(PreTrainedModel):
     config_class = AudioVAEConfig
 
-    def __init__(self, config: AudioVAEConfig):
+    def __init__(self, config):
         super().__init__(config)
+        self.config = config
+        enc_kwargs = config.enc_kwargs
+        dec_kwargs = config.dec_kwargs
+        for key in ("backbone", "input_dim", "latent_dim"):
+            if key not in enc_kwargs:
+                raise ValueError(f"AudioVAE.enc_kwargs missing required key: {key}")
+        for key in ("backbone", "output_dim", "latent_dim"):
+            if key not in dec_kwargs:
+                raise ValueError(f"AudioVAE.dec_kwargs missing required key: {key}")
+
+        hop_size = enc_kwargs.get("hop_size", enc_kwargs["input_dim"])
+        if enc_kwargs["input_dim"] != hop_size:
+            raise ValueError(f"AudioVAE encoder input_dim ({enc_kwargs['input_dim']}) != hop_size ({hop_size}).")
+        if hop_size != dec_kwargs["output_dim"]:
+            raise ValueError(
+                f"AudioVAE encoder hop_size ({hop_size}) != decoder output_dim ({dec_kwargs['output_dim']})."
+            )
         self.encoder = Encoder(
-            encoder_args=config.enc_kwargs["backbone"],
-            input_dim=config.enc_kwargs["input_dim"],
-            hop_size=config.enc_kwargs.get("hop_size", 320),
-            latent_dim=config.enc_kwargs["latent_dim"],
+            encoder_args=enc_kwargs["backbone"],
+            input_dim=enc_kwargs["input_dim"],
+            hop_size=hop_size,
+            latent_dim=enc_kwargs["latent_dim"],
             patch_size=config.patch_size,
         )
         self.decoder = Decoder(
-            decoder_args=config.dec_kwargs["backbone"],
-            output_dim=config.dec_kwargs["output_dim"],
-            latent_dim=config.dec_kwargs["latent_dim"],
+            decoder_args=dec_kwargs["backbone"],
+            output_dim=dec_kwargs["output_dim"],
+            latent_dim=dec_kwargs["latent_dim"],
             patch_size=config.patch_size,
         )
         self.post_init()
@@ -367,21 +325,50 @@ class AudioVAE(PreTrainedModel):
             if module.padding_idx is not None:
                 module.weight.data[module.padding_idx].zero_()
 
+    @torch.inference_mode()
     def encode_latent(self, waveform, waveform_length):
-        from diffusers.models.autoencoders.autoencoder_oobleck import OobleckDiagonalGaussianDistribution
+        if waveform.ndim != 2:
+            raise ValueError(f"Expected waveform rank-2 [Batch, Time], got {tuple(waveform.shape)}")
+        if waveform_length.ndim != 1:
+            raise ValueError(f"Expected waveform_length rank-1 [Batch], got {tuple(waveform_length.shape)}")
+        if waveform.shape[0] != waveform_length.shape[0]:
+            raise ValueError(
+                "Batch mismatch: "
+                f"waveform batch={waveform.shape[0]} vs "
+                f"waveform_length batch={waveform_length.shape[0]}"
+            )
+        if torch.any(waveform_length <= 0):
+            raise ValueError("waveform_length must be strictly positive.")
 
         frame_num = torch.ceil(waveform_length / self.config.enc_kwargs["input_dim"]).to(torch.int32)
         if self.config.patch_size != -1:
             frame_num = torch.ceil(frame_num / self.config.patch_size)
-        h, y = self.encoder(waveform)
+        h, _ = self.encoder(waveform)
         h = h.transpose(1, 2)
 
-        posterior = OobleckDiagonalGaussianDistribution(h)
-        latent = posterior.sample()
+        mean, logvar = torch.chunk(h, 2, dim=1)
+        logvar = torch.clamp(logvar, -30.0, 20.0)
+        std = torch.exp(0.5 * logvar)
+        latent = mean + std * torch.randn_like(mean)
         latent = latent.transpose(1, 2)
         return latent, frame_num
 
+    @torch.inference_mode()
     def decode(self, latent, past_key_values=None, use_cache=False, stream_state=(None, None, None), last_chunk=False):
+        if latent.dim() != 3:
+            raise ValueError(f"Expected latent rank-3 [B,T,D], got shape={tuple(latent.shape)}")
+        if latent.shape[0] <= 0:
+            raise ValueError("latent batch size must be positive.")
+
+        target_dtype = next(self.decoder.parameters()).dtype
+        target_device = next(self.decoder.parameters()).device
+        if latent.dtype != target_dtype or latent.device != target_device:
+            latent = latent.to(device=target_device, dtype=target_dtype)
+
+        expected_latent_dim = self.config.dec_kwargs["latent_dim"]
+        if latent.shape[-1] != expected_latent_dim:
+            raise ValueError(f"Latent dim mismatch in decode(): got {latent.shape[-1]}, expected {expected_latent_dim}")
+
         waveform, stream_state, past_key_values = self.decoder.low_level_reconstruct(
             latent,
             past_key_values=past_key_values,

vllm_omni/model_executor/models/common/ming/dit.py

@@ -0,0 +1,193 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import torch
+import torch.nn.functional as F
+from torch import nn
+from x_transformers.x_transformers import apply_rotary_pos_emb
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim, eps=1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        if self.weight.dtype in [torch.float16, torch.bfloat16]:
+            x = x.to(self.weight.dtype)
+        x = F.rms_norm(x, normalized_shape=(x.shape[-1],), weight=self.weight, eps=self.eps)
+        return x
+
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, dim_out=None, mult=4, dropout=0.0, approximate="none"):
+        super().__init__()
+        inner_dim = int(dim * mult)
+        dim_out = dim_out if dim_out is not None else dim
+
+        activation = nn.GELU(approximate=approximate)
+        project_in = nn.Sequential(nn.Linear(dim, inner_dim), activation)
+        self.ff = nn.Sequential(project_in, nn.Dropout(dropout), nn.Linear(inner_dim, dim_out))
+
+    def forward(self, x):
+        return self.ff(x)
+
+
+class Attention(nn.Module):
+    def __init__(
+        self,
+        dim,
+        heads=8,
+        dim_head=64,
+        dropout=0.0,
+        qk_norm=None,
+        pe_attn_head=None,
+        attn_mask_enabled=True,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.heads = heads
+        self.inner_dim = dim_head * heads
+        self.dropout = dropout
+        self.to_q = nn.Linear(dim, self.inner_dim)
+        self.to_k = nn.Linear(dim, self.inner_dim)
+        self.to_v = nn.Linear(dim, self.inner_dim)
+        if qk_norm is None:
+            self.q_norm = None
+            self.k_norm = None
+        elif qk_norm == "rms_norm":
+            self.q_norm = RMSNorm(dim_head)
+            self.k_norm = RMSNorm(dim_head)
+        else:
+            raise ValueError(f"Unimplemented qk_norm: {qk_norm}")
+
+        self.to_out = nn.ModuleList([])
+        self.to_out.append(nn.Linear(self.inner_dim, dim))
+        self.to_out.append(nn.Dropout(dropout))
+        self.pe_attn_head = pe_attn_head
+        self.attn_mask_enabled = attn_mask_enabled
+
+    def forward(self, x, mask=None, rope=None):
+        batch_size = x.shape[0]
+        query = self.to_q(x)
+        key = self.to_k(x)
+        value = self.to_v(x)
+
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // self.heads
+        query = query.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
+        key = key.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
+
+        if self.q_norm is not None:
+            query = self.q_norm(query)
+        if self.k_norm is not None:
+            key = self.k_norm(key)
+
+        if rope is not None:
+            freqs, xpos_scale = rope
+            q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0)
+            if self.pe_attn_head is not None:
+                on = self.pe_attn_head
+                query[:, :on, :, :] = apply_rotary_pos_emb(query[:, :on, :, :], freqs, q_xpos_scale)
+                key[:, :on, :, :] = apply_rotary_pos_emb(key[:, :on, :, :], freqs, k_xpos_scale)
+            else:
+                query = apply_rotary_pos_emb(query, freqs, q_xpos_scale)
+                key = apply_rotary_pos_emb(key, freqs, k_xpos_scale)
+
+        if self.attn_mask_enabled and mask is not None:
+            valid_sample_indices = mask.any(dim=1)
+            final_output = torch.zeros_like(query).to(query.device)
+            attn_mask = mask[valid_sample_indices]
+            query = query[valid_sample_indices]
+            key = key[valid_sample_indices]
+            value = value[valid_sample_indices]
+            attn_mask = attn_mask.unsqueeze(1).unsqueeze(1)
+            attn_mask = attn_mask.expand(valid_sample_indices.sum().item(), self.heads, query.shape[-2], key.shape[-2])
+        else:
+            attn_mask = None
+
+        x = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=False)
+        if self.attn_mask_enabled and mask is not None:
+            final_output[valid_sample_indices] = x
+            x = final_output
+
+        x = x.transpose(1, 2).reshape(batch_size, -1, self.heads * head_dim)
+        x = x.to(query.dtype)
+        x = self.to_out[0](x)
+        x = self.to_out[1](x)
+
+        if mask is not None:
+            mask = mask.unsqueeze(-1)
+            x = x.masked_fill(~mask, 0.0)
+
+        return x
+
+
+class DiTBlock(nn.Module):
+    def __init__(
+        self,
+        hidden_size,
+        num_heads,
+        mlp_ratio=4.0,
+        dropout=0.1,
+        qk_norm=None,
+        pe_attn_head=None,
+        attn_mask_enabled=True,
+        **kwargs,
+    ):
+        super().__init__()
+        self.norm1 = RMSNorm(hidden_size)
+        self.attn = Attention(
+            dim=hidden_size,
+            heads=num_heads,
+            dim_head=hidden_size // num_heads,
+            dropout=dropout,
+            qk_norm=qk_norm,
+            pe_attn_head=pe_attn_head,
+            attn_mask_enabled=attn_mask_enabled,
+        )
+        self.norm2 = RMSNorm(hidden_size)
+        self.mlp = FeedForward(dim=hidden_size, mult=mlp_ratio, dropout=dropout, approximate="tanh")
+
+    def forward(self, x, mask, rope):
+        x = x + self.attn(self.norm1(x), mask=mask, rope=rope)
+        x = x + self.mlp(self.norm2(x))
+        return x
+
+
+class FinalLayer(nn.Module):
+    def __init__(self, hidden_size, out_channels):
+        super().__init__()
+        self.norm_final = RMSNorm(hidden_size)
+        self.linear = nn.Linear(hidden_size, out_channels, bias=True)
+
+    def forward(self, x):
+        x = self.norm_final(x)
+        x = self.linear(x)
+        return x
+
+
+class CondEmbedder(nn.Module):
+    def __init__(self, input_feature_size, hidden_size):
+        super().__init__()
+        self.cond_embedder = nn.Linear(input_feature_size, hidden_size)
+
+    def forward(self, llm_cond):
+        return self.cond_embedder(llm_cond)
+
+
+def get_epss_timesteps(n, device, dtype):
+    dt = 1 / 32
+    predefined_timesteps = {
+        5: [0, 2, 4, 8, 16, 32],
+        6: [0, 2, 4, 6, 8, 16, 32],
+        7: [0, 2, 4, 6, 8, 16, 24, 32],
+        10: [0, 2, 4, 6, 8, 12, 16, 20, 24, 28, 32],
+        12: [0, 2, 4, 6, 8, 10, 12, 14, 16, 20, 24, 28, 32],
+        16: [0, 1, 2, 3, 4, 5, 6, 7, 8, 10, 12, 14, 16, 20, 24, 28, 32],
+    }
+    t = predefined_timesteps.get(n, [])
+    if not t:
+        return torch.linspace(0, 1, n + 1, device=device, dtype=dtype)
+    return dt * torch.tensor(t, device=device, dtype=dtype)

vllm_omni/model_executor/models/common/ming/fm.py

@@ -0,0 +1,68 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+
+from .dit import get_epss_timesteps
+
+
+def apply_sway_sampling(t, sway_sampling_coef):
+    if sway_sampling_coef is None:
+        return t
+    return t + sway_sampling_coef * (torch.cos(torch.pi / 2 * t) - 1 + t)
+
+
+def build_timesteps(steps, device, dtype, use_epss=True, sway_sampling_coef=-1.0):
+    if steps <= 0:
+        raise ValueError(f"steps must be positive, got {steps}")
+    if use_epss:
+        t = get_epss_timesteps(steps, device=device, dtype=dtype)
+    else:
+        t = torch.linspace(0, 1, steps + 1, device=device, dtype=dtype)
+    return apply_sway_sampling(t, sway_sampling_coef)
+
+
+class Solver:
+    def __init__(self, func, y0, sigma=0.25, temperature=1.5) -> None:
+        self.func = func
+        self.y0 = y0
+        self.sigma = sigma
+        self.temperature = temperature
+
+    def integrate(self, t):
+        solution = torch.empty(len(t), *self.y0.shape, dtype=self.y0.dtype, device=self.y0.device)
+        solution[0] = self.y0
+
+        j = 1
+        y0 = self.y0
+        for t0, t1 in zip(t[:-1], t[1:]):
+            dt = t1 - t0
+            f0 = self.func(t0, y0)
+            dy = dt * f0
+            y1 = y0 + dy
+
+            while j < len(t) and t1 >= t[j]:
+                solution[j] = self._linear_interp(t0, t1, y0, y1, t[j])
+                j += 1
+
+            noise = torch.randn_like(y0)
+            shift = self.sigma * (self.temperature**0.5) * (abs(dt) ** 0.5) * noise
+            y0 = y1 + shift
+
+        return solution
+
+    def _linear_interp(self, t0, t1, y0, y1, t):
+        if t == t0:
+            return y0
+        if t == t1:
+            return y1
+        slope = (t - t0) / (t1 - t0)
+        return y0 + slope * (y1 - y0)
+
+
+def integrate_cfm_steps(fn, y0, t, sde_args, sde_rnd, steps):
+    for step in range(steps):
+        dt = t[step + 1] - t[step]
+        y0 = y0 + fn(t[step], y0) * dt
+        y0 = y0 + sde_args[1] * (sde_args[2] ** 0.5) * (dt.abs() ** 0.5) * sde_rnd[step]
+    return y0

vllm_omni/model_executor/models/cosyvoice3/code2wav_core/cfm.py

@@ -141,24 +141,37 @@ class ConditionalCFM(BASECFM):
         if isinstance(self.estimator, torch.nn.Module):
             return self.estimator(x, mask, mu, t, spks, cond)
         else:
+            # TensorRT estimator: bind raw device pointers. The flow runs in
+            # fp32 but the engine may have fp16 I/O (strongly-typed fp16 engine),
+            # so cast inputs/output to the engine's dtype at the boundary. Keep
+            # references to the cast buffers alive until execute completes (a bare
+            # ``.contiguous().data_ptr()`` could free the temp -> dangling ptr).
+            io_dtype = getattr(self.estimator, "io_dtype", x.dtype)
             [estimator, stream], trt_engine = self.estimator.acquire_estimator()
             # NOTE need to synchronize when switching stream
             torch.cuda.current_stream().synchronize()
             with stream:
-                estimator.set_input_shape("x", (2, 80, x.size(2)))
-                estimator.set_input_shape("mask", (2, 1, x.size(2)))
-                estimator.set_input_shape("mu", (2, 80, x.size(2)))
+                x_e = x.to(io_dtype).contiguous()
+                mask_e = mask.to(io_dtype).contiguous()
+                mu_e = mu.to(io_dtype).contiguous()
+                t_e = t.to(io_dtype).contiguous()
+                spks_e = spks.to(io_dtype).contiguous()
+                cond_e = cond.to(io_dtype).contiguous()
+                out_e = torch.empty_like(x_e)
+                estimator.set_input_shape("x", (2, 80, x_e.size(2)))
+                estimator.set_input_shape("mask", (2, 1, x_e.size(2)))
+                estimator.set_input_shape("mu", (2, 80, x_e.size(2)))
                 estimator.set_input_shape("t", (2,))
                 estimator.set_input_shape("spks", (2, 80))
-                estimator.set_input_shape("cond", (2, 80, x.size(2)))
+                estimator.set_input_shape("cond", (2, 80, x_e.size(2)))
                 data_ptrs = [
-                    x.contiguous().data_ptr(),
-                    mask.contiguous().data_ptr(),
-                    mu.contiguous().data_ptr(),
-                    t.contiguous().data_ptr(),
-                    spks.contiguous().data_ptr(),
-                    cond.contiguous().data_ptr(),
-                    x.data_ptr(),
+                    x_e.data_ptr(),
+                    mask_e.data_ptr(),
+                    mu_e.data_ptr(),
+                    t_e.data_ptr(),
+                    spks_e.data_ptr(),
+                    cond_e.data_ptr(),
+                    out_e.data_ptr(),
                 ]
                 for i, j in enumerate(data_ptrs):
                     estimator.set_tensor_address(trt_engine.get_tensor_name(i), j)
@@ -166,7 +179,7 @@ class ConditionalCFM(BASECFM):
                 assert estimator.execute_async_v3(torch.cuda.current_stream().cuda_stream) is True
                 torch.cuda.current_stream().synchronize()
             self.estimator.release_estimator(estimator, stream)
-            return x
+            return out_e.to(x.dtype)
 
 
 class CausalConditionalCFM(ConditionalCFM):

vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py

@@ -3,10 +3,16 @@
 import os
 from collections.abc import Iterable, Mapping, Sequence
 from functools import partial
+from math import gcd
 from threading import Lock
 
+import numpy as np
+import onnxruntime
 import torch
 import torch.nn as nn
+from huggingface_hub import snapshot_download
+from scipy.signal import resample_poly
+from transformers import Qwen2Config
 from transformers.feature_extraction_utils import BatchFeature
 from vllm.config import VllmConfig
 from vllm.config.multimodal import BaseDummyOptions
@@ -32,12 +38,15 @@ from vllm.v1.sample.metadata import SamplingMetadata
 from vllm.v1.sample.sampler import Sampler
 
 from vllm_omni.data_entry_keys import EmbeddingsStruct, OmniPayloadStruct, to_dict, to_struct
+from vllm_omni.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokenizer
 from vllm_omni.model_executor.models.cosyvoice3.utils import (
     concat_text_with_prompt_ids,
     extract_speech_feat,
-    extract_speech_token,
     extract_spk_embedding,
+    extract_spk_embedding_trt,
     extract_text_token,
+    mel_spectrogram,
+    unpad_prompt_conditioning,
 )
 from vllm_omni.model_executor.models.output_templates import OmniOutput
 from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config
@@ -45,6 +54,23 @@ from vllm_omni.utils.speaker_cache import get_speaker_cache
 
 logger = init_logger(__name__)
 
+# Process-wide cache of per-model mm-processor runtime components (tokenizer,
+# feat_extractor, campplus session/engine). The mm processor is re-created per
+# request (mm_processor_cache_gb: 0), so this avoids rebuilding them every time.
+_RUNTIME_COMPONENTS_CACHE: dict[str, dict] = {}
+
+
+def _cosyvoice3_trt_enabled() -> bool:
+    """COSYVOICE3_TRT env toggle (default on) for the optional TensorRT paths.
+
+    Gates both the talker speaker-embedding (campplus) engine and the code2wav
+    flow-decoder estimator engine. Env-var toggle (cf. mimo_audio's
+    ``MIMO_AUDIO_TOKENIZER_CUDA_GRAPH``) because these run outside the stage
+    worker, so deploy-yaml ``hf_overrides`` / per-stage ``env`` do not reach
+    them — export ``COSYVOICE3_TRT=0`` in the launching shell to disable.
+    """
+    return os.environ.get("COSYVOICE3_TRT", "1") not in ("0", "false", "False", "")
+
 
 class CosyVoice3MultiModalProcessingInfo(BaseProcessingInfo):
     def get_hf_config(self):
@@ -73,39 +99,136 @@ class CosyVoice3MultiModalProcessor(BaseMultiModalProcessor[CosyVoice3MultiModal
         if cached_model_dir == model_dir:
             return
 
-        # If model_dir is an HF repo ID (not a local path), resolve to cache
-        if not os.path.isdir(model_dir):
-            from huggingface_hub import snapshot_download
+        # The mm processor is re-created per request (deploy config sets
+        # ``mm_processor_cache_gb: 0``), so the instance-level guard above never
+        # hits across requests. Without a process-wide cache, every request
+        # would re-run ``snapshot_download``, rebuild the Qwen tokenizer and
+        # create a fresh ONNX campplus session — ~hundreds of ms of pure
+        # overhead on the TTFP critical path. Build the heavy components once
+        # per process, keyed by model_dir, and reuse them.
+        comps = _RUNTIME_COMPONENTS_CACHE.get(model_dir)
+        if comps is None:
+            comps = self._build_runtime_components(model_dir, config)
+            _RUNTIME_COMPONENTS_CACHE[model_dir] = comps
 
+        self.tokenizer = comps["tokenizer"]
+        self.feat_extractor = comps["feat_extractor"]
+        self.campplus_session = comps["campplus_session"]
+        self.campplus_trt = comps["campplus_trt"]
+        self._cached_model_dir = model_dir
+        self._speaker_cache = get_speaker_cache()
+
+    def _build_runtime_components(self, model_dir: str, config: CosyVoice3Config) -> dict:
+        """Build the per-model runtime components once (cached process-wide)."""
+        # If model_dir is an HF repo ID (not a local path), resolve to cache.
+        if not os.path.isdir(model_dir):
             model_dir = snapshot_download(model_dir)
 
-        import onnxruntime
-
-        from vllm_omni.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokenizer
-        from vllm_omni.model_executor.models.cosyvoice3.utils import mel_spectrogram
-
-        option = onnxruntime.SessionOptions()
-        option.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
-        option.intra_op_num_threads = 1
-
-        self.tokenizer = get_qwen_tokenizer(
+        tokenizer = get_qwen_tokenizer(
             token_path=os.path.join(model_dir, config.qwen_pretrain_path),
             skip_special_tokens=config.skip_special_tokens,
             version=config.version,
         )
-        self.speech_tokenizer = onnxruntime.InferenceSession(
-            os.path.join(model_dir, config.speech_tokenizer_path),
-            sess_options=option,
-            providers=["CUDAExecutionProvider" if torch.cuda.is_available() else "CPUExecutionProvider"],
-        )
-        self.feat_extractor = partial(mel_spectrogram, **getattr(config, "feat_extractor", {}))
-        self.campplus_session = onnxruntime.InferenceSession(
-            os.path.join(model_dir, config.campplus_onxx_path),
-            sess_options=option,
-            providers=["CPUExecutionProvider"],
-        )
-        self._cached_model_dir = model_dir
-        self._speaker_cache = get_speaker_cache()
+        feat_extractor = partial(mel_spectrogram, **getattr(config, "feat_extractor", {}))
+
+        campplus_onnx_path = os.path.join(model_dir, config.campplus_onxx_path)
+        # TensorRT speaker-embedding path (default on); a prebuilt TRT engine
+        # runs campplus on GPU. The engine itself is cached process-wide by
+        # ``get_campplus_trt``.
+        campplus_trt = None
+        if self._speaker_embedding_trt_enabled() and torch.cuda.is_available():
+            try:
+                from vllm_omni.model_executor.models.cosyvoice3.speaker_embedding_trt import (
+                    get_campplus_trt,
+                )
+
+                campplus_trt = get_campplus_trt(campplus_onnx_path, device="cuda")
+                logger.info("CosyVoice3: using TensorRT campplus speaker embedding")
+            except Exception as exc:  # pragma: no cover - defensive fallback
+                logger.warning(
+                    "CosyVoice3: TensorRT campplus build failed (%s); falling back to ONNX-Runtime speaker embedding",
+                    exc,
+                )
+                campplus_trt = None
+
+        # Only build the CPU ONNX campplus session when TRT is unavailable —
+        # otherwise it is never used and creating it (ORT_ENABLE_ALL graph
+        # optimization) costs ~hundreds of ms.
+        campplus_session = None
+        if campplus_trt is None:
+            option = onnxruntime.SessionOptions()
+            option.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
+            option.intra_op_num_threads = 1
+            campplus_session = onnxruntime.InferenceSession(
+                campplus_onnx_path,
+                sess_options=option,
+                providers=["CPUExecutionProvider"],
+            )
+
+        return {
+            "tokenizer": tokenizer,
+            "feat_extractor": feat_extractor,
+            "campplus_session": campplus_session,
+            "campplus_trt": campplus_trt,
+        }
+
+    @staticmethod
+    def _speaker_embedding_trt_enabled() -> bool:
+        """Whether to use the TensorRT campplus speaker-embedding path (default on)."""
+        return _cosyvoice3_trt_enabled()
+
+    # Class-level cached s3tokenizer model — loaded once per process on first
+    # call to ``_extract_speech_token_via_s3`` and shared across all
+    # processor instances.
+    _s3_model = None
+
+    @classmethod
+    def _ensure_s3_model(cls):
+        if cls._s3_model is not None:
+            return cls._s3_model
+
+        # s3tokenizer is imported lazily (kept off the module top level) so
+        # callers that don't use the CosyVoice3 talker need not install it.
+        try:
+            import s3tokenizer as _s3
+        except ImportError as e:
+            raise ImportError(
+                "CosyVoice3 speech-token extraction requires the 's3tokenizer' "
+                "package; install it with `pip install s3tokenizer`."
+            ) from e
+
+        model = _s3.load_model("speech_tokenizer_v3_25hz")
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        model = model.to(device).eval()
+        cls._s3_model = (model, _s3, device)
+        return cls._s3_model
+
+    def _extract_speech_token_via_s3(self, audio, return_device):
+        """Drop-in replacement for ``extract_speech_token`` that uses the
+        S3Tokenizer PyTorch model on GPU. Returns the same
+        ``(speech_token[1, T], speech_token_len[1])`` int32 tensors as the
+        ONNX path so the rest of ``_call_hf_processor`` is unchanged.
+        """
+        model, _s3, dev = self._ensure_s3_model()
+
+        # audio is a (waveform_ndarray, sr) tuple; resample to 16 kHz mono float32.
+        wav, sr = audio
+        wav = np.asarray(wav, dtype=np.float32)
+        if wav.ndim == 2:
+            wav = wav.mean(axis=1)
+        if int(sr) != 16000:
+            g = gcd(int(sr), 16000)
+            wav = resample_poly(wav, 16000 // g, int(sr) // g).astype(np.float32)
+        audio_t = torch.from_numpy(wav)
+
+        mel = _s3.log_mel_spectrogram(audio_t)
+        mels_p, mels_lens = _s3.padding([mel])
+        with torch.inference_mode():
+            codes, codes_lens = model.quantize(mels_p.to(dev), mels_lens.to(dev))
+        n = int(codes_lens[0].item())
+        speech_token = codes[:1, :n].to(dtype=torch.int32, device=return_device)
+        speech_token_len = torch.tensor([n], dtype=torch.int32, device=return_device)
+        return speech_token, speech_token_len
 
     def _call_hf_processor(
         self,
@@ -185,7 +308,11 @@ class CosyVoice3MultiModalProcessor(BaseMultiModalProcessor[CosyVoice3MultiModal
                 )
                 return ft
 
-        speech_token, speech_token_len = extract_speech_token(audio, self.speech_tokenizer, device)
+        # Speech-token extraction via the S3Tokenizer PyTorch model on GPU
+        # (~30x faster than the bundled ``speech_tokenizer_v3.onnx`` CPU ONNX
+        # path in this venv).
+        speech_token, speech_token_len = self._extract_speech_token_via_s3(audio, device)
+
         speech_feat, speech_feat_len = extract_speech_feat(audio, self.feat_extractor, device)
 
         if config.sample_rate == 24000:
@@ -193,7 +320,10 @@ class CosyVoice3MultiModalProcessor(BaseMultiModalProcessor[CosyVoice3MultiModal
             speech_feat, speech_feat_len[:] = speech_feat[:, : 2 * token_len], 2 * token_len
             speech_token, speech_token_len[:] = speech_token[:, :token_len], token_len
 
-        embedding = extract_spk_embedding(audio, self.campplus_session, device)
+        if self.campplus_trt is not None:
+            embedding = extract_spk_embedding_trt(audio, self.campplus_trt, device)
+        else:
+            embedding = extract_spk_embedding(audio, self.campplus_session, device)
 
         # Cache the extracted artifacts for named speakers
         if cache_key is not None:
@@ -318,8 +448,6 @@ class CosyVoice3Model(
         self.model_stage = vllm_config.model_config.model_stage
         model_dir = vllm_config.model_config.model
         if not os.path.isdir(model_dir):
-            from huggingface_hub import snapshot_download
-
             model_dir = snapshot_download(model_dir)
         self.model_dir = model_dir
         self.model = None
@@ -377,8 +505,6 @@ class CosyVoice3Model(
         the pretrained model directory. The cache config is inherited from the parent
         to enable PagedAttention with the same memory configuration.
         """
-        from transformers import Qwen2Config
-
         qwen_config_path = os.path.join(self.model_dir, self.config.llm["llm"]["pretrain_path"])
         qwen_hf_config = Qwen2Config.from_pretrained(qwen_config_path)
 
@@ -469,26 +595,32 @@ class CosyVoice3Model(
         top_k: int,
         generator: torch.Generator | None,
     ) -> int:
+        """Vectorized nucleus + top-k sampling.
+
+        Distribution-equivalent to the reference iterative implementation: the
+        keep-set is identical (token i is kept iff
+        ``cumsum(sorted_probs)[i] - sorted_probs[i] < top_p`` AND ``i < top_k``)
+        and the renormalized sampling distribution matches, but the exact token
+        drawn for a given seed is NOT guaranteed to match. The reference draws
+        via ``multinomial`` over the stacked kept subset while this draws over
+        the full sorted vector (zeroed outside the keep-set), so the generator
+        advances over different-sized inputs and may yield a different sample.
+        The win: no per-token ``.item()`` D2H syncs from the Python loop —
+        those dominated the sampler CPU time in profiling.
+        """
         probs = weighted_scores.softmax(dim=0)
         sorted_prob, sorted_idx = probs.sort(descending=True, stable=True)
-        kept_probs: list[torch.Tensor] = []
-        kept_indices: list[torch.Tensor] = []
-        cum_prob = 0.0
-        max_keep = len(sorted_idx) if top_k <= 0 else min(int(top_k), len(sorted_idx))
-        for i in range(len(sorted_idx)):
-            if cum_prob < top_p and len(kept_probs) < max_keep:
-                cum_prob += float(sorted_prob[i].item())
-                kept_probs.append(sorted_prob[i])
-                kept_indices.append(sorted_idx[i])
-            else:
-                break
-
-        if not kept_probs:
-            return int(sorted_idx[0].item())
-
-        sample_probs = torch.stack(kept_probs)
-        sample_idx = cls._multinomial_sample(sample_probs, generator=generator)
-        return int(torch.stack(kept_indices)[int(sample_idx.item())].item())
+        cum_before = sorted_prob.cumsum(dim=0) - sorted_prob
+        mask = cum_before < top_p
+        if top_k > 0:
+            n = sorted_prob.shape[0]
+            mask = mask & (torch.arange(n, device=mask.device) < min(int(top_k), n))
+        weights = sorted_prob * mask.to(sorted_prob.dtype)
+        # First token always passes (cum_before[0] = 0 < top_p for any top_p > 0),
+        # so ``weights`` is guaranteed to have at least one nonzero entry. The
+        # final ``.item()`` is the ONLY D2H sync per call.
+        sample_idx = torch.multinomial(weights, 1, replacement=True, generator=generator)
+        return int(sorted_idx[sample_idx].item())
 
     @classmethod
     def _ras_sample_one(
@@ -627,11 +759,21 @@ class CosyVoice3Model(
         else:
             raise RuntimeError(f"compute_logits is only valid for {self.model_stage}.")
 
-    def embed_multimodal(self, **kwargs: object) -> torch.Tensor:
+    def embed_multimodal(self, **kwargs: object):
         if self.model_stage == "cosyvoice3_talker":
             speech_token = kwargs["speech_token"]
-            speech_token_emb = self.model.speech_embedding(speech_token)
-            return speech_token_emb
+            # vLLM's _execute_mm_encoder batches ALL multimodal items scheduled
+            # in one engine step into a single call, expecting one embedding
+            # tensor per item back. When >=2 requests prefill in the same step
+            # (likely once mm preprocessing is fast, e.g. TensorRT speaker
+            # embedding), speech_token arrives as a list of per-item tensors
+            # (variable length), so embed it per item and return a list of
+            # [T_i, emb] tensors. The single-item path keeps returning a
+            # [1, T, emb] tensor (the caller's extend() iterates dim 0).
+            if isinstance(speech_token, (list, tuple)):
+                emb_dim = self.model.speech_embedding.weight.shape[1]
+                return [self.model.speech_embedding(t).reshape(-1, emb_dim) for t in speech_token]
+            return self.model.speech_embedding(speech_token)
         else:
             raise RuntimeError(f"embed_multimodal is only valid for {self.model_stage}.")
 
@@ -720,6 +862,121 @@ class CosyVoice3Model(
         else:
             raise RuntimeError(f"embed_input_ids is not valid for {self.model_stage}.")
 
+    @staticmethod
+    def _split_prompt_conditioning(speech_token, speech_feat, embedding, speech_token_len):
+        """Split collated prompt conditioning into per-request lists.
+
+        Returns ``(speech_token_list, speech_feat_list, embedding_list,
+        speech_token_len_list)``, one rank-correct tensor per request:
+        ``speech_token`` ``[1, T]`` (possibly right-padded), ``speech_feat``
+        ``[1, 2T, F]``, ``embedding`` ``[1, D]``, ``speech_token_len`` ``[1, 1]``.
+
+        This runs inside the talker ``forward`` which may be CUDA-graph
+        captured, so it must NOT trigger any host<->device sync (no ``.item()``
+        / ``.tolist()`` / Python-int slicing by tensor value). Splitting is pure
+        on-device indexing; right-padding is dropped later in the (eager)
+        code2wav stage using the per-request ``speech_token_len``. Robust to
+        inputs arriving as padded ``[B, ...]`` batch tensors or per-request
+        lists.
+        """
+
+        def _rows(x, want_dim):
+            if isinstance(x, (list, tuple)):
+                items = list(x)
+            elif isinstance(x, torch.Tensor):
+                # ``x.shape[0]`` is a static Python int — no device sync.
+                items = [x[i] for i in range(x.shape[0])]
+            else:
+                return None
+            out = []
+            for t in items:
+                if isinstance(t, torch.Tensor):
+                    while t.dim() < want_dim:
+                        t = t.unsqueeze(0)
+                    if t.shape[0] != 1:
+                        t = t[:1]
+                    t = t.contiguous()
+                out.append(t)
+            return out
+
+        st_out = _rows(speech_token, 2)
+        sf_out = _rows(speech_feat, 3)
+        emb_out = _rows(embedding, 2)
+        stl_out = _rows(speech_token_len, 2)
+        return st_out, sf_out, emb_out, stl_out
+
+    def _resolve_flow_estimator_onnx(self) -> str | None:
+        """Locate the flow-decoder estimator ONNX for the TensorRT engine.
+
+        Prefers the fp16 (strongly-typed) ONNX → fp16 engine. Order: env
+        ``COSYVOICE3_ESTIMATOR_ONNX`` → ``<model_dir>/<fp16 name>`` → fetched
+        from ``flow_estimator_onnx_repo`` → bundled fp32 ONNX (fp32+TF32). Each
+        is checked for existence; returns ``None`` if nothing is found.
+        """
+        env_path = os.environ.get("COSYVOICE3_ESTIMATOR_ONNX")
+        if env_path and os.path.exists(env_path):
+            return env_path
+
+        fp16_name = getattr(self.config, "flow_estimator_onnx_path", "flow.decoder.estimator.autocast_fp16.onnx")
+        local_fp16 = os.path.join(self.model_dir, fp16_name)
+        if os.path.exists(local_fp16):
+            return local_fp16
+
+        repo = getattr(self.config, "flow_estimator_onnx_repo", None)
+        if repo:
+            try:
+                fetched_dir = snapshot_download(repo, allow_patterns=[fp16_name])
+                fetched = os.path.join(fetched_dir, fp16_name)
+                if os.path.exists(fetched):
+                    return fetched
+            except Exception as exc:  # pragma: no cover - network/repo issues
+                logger.warning("CosyVoice3 code2wav: could not fetch fp16 estimator ONNX from %s (%s)", repo, exc)
+
+        fp32_name = getattr(self.config, "flow_estimator_onnx_path_fp32", "flow.decoder.estimator.fp32.onnx")
+        local_fp32 = os.path.join(self.model_dir, fp32_name)
+        if os.path.exists(local_fp32):
+            logger.info("CosyVoice3 code2wav: fp16 estimator ONNX unavailable, falling back to fp32 (%s)", local_fp32)
+            return local_fp32
+        return None
+
+    def _maybe_enable_code2wav_trt(self) -> None:
+        """Swap the flow-decoder estimator to a TensorRT engine once (lazy).
+
+        Runs on the first code2wav step — after weights are loaded — so the
+        torch estimator is fully built first and then dropped. No-op unless
+        ``COSYVOICE3_TRT`` is on (default), CUDA is available, and the estimator
+        ONNX ships with the model. Falls back to the torch estimator on any
+        failure. The upstream ``CausalConditionalCFM.forward_estimator`` detects
+        the non-``nn.Module`` estimator and drives the TRT engine.
+        """
+        if getattr(self, "_code2wav_trt_done", False):
+            return
+        self._code2wav_trt_done = True
+        if not (_cosyvoice3_trt_enabled() and torch.cuda.is_available()):
+            return
+        onnx_path = self._resolve_flow_estimator_onnx()
+        if onnx_path is None:
+            logger.warning("CosyVoice3 code2wav: no flow-estimator ONNX available; keeping torch estimator")
+            return
+        try:
+            from vllm_omni.model_executor.models.cosyvoice3.flow_estimator_trt import (
+                build_flow_estimator_trt,
+            )
+
+            wrapper = build_flow_estimator_trt(onnx_path, device="cuda")
+            # ``estimator`` is a registered nn.Module submodule; delete it first
+            # (frees the torch estimator weights) so the TRT wrapper can be set
+            # as a plain attribute — nn.Module.__setattr__ rejects non-Modules.
+            decoder = self.code2wav.flow_model.decoder
+            del decoder.estimator
+            decoder.estimator = wrapper
+            logger.info("CosyVoice3: using TensorRT flow-decoder estimator (code2wav)")
+        except Exception as exc:  # pragma: no cover - defensive fallback
+            logger.warning(
+                "CosyVoice3 code2wav: TensorRT estimator build failed (%s); keeping torch estimator",
+                exc,
+            )
+
     def forward(
         self,
         input_ids: torch.Tensor,
@@ -741,18 +998,43 @@ class CosyVoice3Model(
             if "speech_token" in kwargs:
                 # Prompt conditioning tensors for code2wav: live under
                 # ``embed.*`` per OmniPayloadStruct schema.
+                #
+                # vLLM hands these mm fields to forward() as collated, padded
+                # batch tensors (speech_token [B, maxT], speech_feat
+                # [B, 2*maxT, F], embedding [B, D]). Emitting them raw makes the
+                # downstream per-request payload split fragile at batch>1: it
+                # intermittently de-batches speech_token to 1-D and leaks the
+                # whole [B, D] embedding to every request, which corrupts voice
+                # conditioning and crashes code2wav (`prompt_token.shape[1]`).
+                # Instead, split into an explicit per-request list of unpadded,
+                # correctly-ranked tensors so ``to_payload_element`` splits them
+                # deterministically by request index (list[idx]).
+                speech_token_list, speech_feat_list, embedding_list, speech_token_len_list = (
+                    self._split_prompt_conditioning(
+                        kwargs.get("speech_token"),
+                        kwargs.get("speech_feat"),
+                        kwargs.get("embedding"),
+                        kwargs.get("speech_token_len"),
+                    )
+                )
                 multimodal_outputs = to_dict(
                     OmniPayloadStruct(
                         embed=EmbeddingsStruct(
-                            speech_token=kwargs.get("speech_token"),
-                            speech_feat=kwargs.get("speech_feat"),
-                            embedding=kwargs.get("embedding"),
+                            speech_token=speech_token_list,
+                            speech_feat=speech_feat_list,
+                            speech_token_len=speech_token_len_list,
+                            embedding=embedding_list,
                         ),
                     )
                 )
 
             return OmniOutput(text_hidden_states=hidden_states, multimodal_outputs=multimodal_outputs)
         elif self.model_stage == "cosyvoice3_code2wav":
+            # Lazily swap the flow-decoder estimator to a TensorRT engine on the
+            # first code2wav step (after weights are loaded), gated by the same
+            # COSYVOICE3_TRT env toggle as the talker speaker embedding.
+            self._maybe_enable_code2wav_trt()
+
             runtime_info = kwargs.get("model_intermediate_buffer")
             if runtime_info is None:
                 runtime_info = kwargs.get("runtime_additional_information", [])
@@ -784,6 +1066,11 @@ class CosyVoice3Model(
                 speech_token = embed.speech_token if embed else None
                 speech_feat = embed.speech_feat if embed else None
                 embedding = embed.embedding if embed else None
+                # Drop any right-padding carried from batched talker emission.
+                if speech_token is not None and speech_feat is not None:
+                    speech_token, speech_feat = unpad_prompt_conditioning(
+                        speech_token, speech_feat, embed.speech_token_len if embed else None
+                    )
                 if speech_token is None or speech_feat is None or embedding is None:
                     if stream_finished and req_id is not None and hasattr(self, "_stream_vocoder_cache_by_req"):
                         with self._stream_audio_cache_lock:

vllm_omni/model_executor/models/cosyvoice3/flow_estimator_trt.py

@@ -0,0 +1,154 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""TensorRT engine for the CosyVoice3 flow-decoder (CFM) DiT estimator.
+
+The estimator is the per-step network the conditional flow-matching ODE solver
+calls during code2wav (token -> mel). It dominates code2wav latency; the
+upstream ``CausalConditionalCFM.forward_estimator`` already supports running it
+through a TensorRT engine (it switches on ``self.estimator`` not being an
+``nn.Module`` and drives it via ``acquire_estimator`` / ``execute_async_v3``).
+This module builds that engine from the bundled ``flow.decoder.estimator*.onnx``
+and wraps it so it can be dropped in for the torch estimator.
+
+The estimator engine has 6 inputs ``x, mask, mu, t, spks, cond`` and one output.
+``x/mask/mu/cond`` carry a dynamic time dim; ``t``/``spks`` are fixed, so only
+the former need an optimization profile (TRT infers the latter from the ONNX).
+Shapes mirror the CosyVoice runtime.
+
+Precision: TensorRT >= 11 dropped the weakly-typed FP16/INT8 builder flags, so
+fp16 only comes from a STRONGLY_TYPED network built from an fp16 ONNX
+(``*autocast_fp16*``, fp16 I/O). An fp32 ONNX is built fp32 + the TF32 matmul
+flag. ``EXPLICIT_BATCH`` is implicit (no flag) and ``ITensor.dtype`` is
+read-only, so neither is set here.
+"""
+
+from __future__ import annotations
+
+import os
+import queue
+
+import torch
+from vllm.logger import init_logger
+
+from vllm_omni.model_executor.models.cosyvoice3.speaker_embedding_trt import (
+    _resolve_plan_path,
+    _trt_logger,
+)
+
+logger = init_logger(__name__)
+
+# Optimization-profile shapes for the dynamic-length inputs (CFG batch dim = 2,
+# 80 mel channels, time dim min/opt/max). Matches CosyVoice's token2wav.
+_DYNAMIC_INPUTS = ("x", "mask", "mu", "cond")
+_MIN_SHAPES = ((2, 80, 4), (2, 1, 4), (2, 80, 4), (2, 80, 4))
+_OPT_SHAPES = ((2, 80, 500), (2, 1, 500), (2, 80, 500), (2, 80, 500))
+_MAX_SHAPES = ((2, 80, 3000), (2, 1, 3000), (2, 80, 3000), (2, 80, 3000))
+
+
+def _is_fp16_onnx(onnx_path: str) -> bool:
+    """Heuristic: the project's fp16 estimator ONNX is exported strongly-typed
+    and named ``*autocast_fp16*`` / ``*fp16*`` (vs ``*fp32*``)."""
+    name = os.path.basename(onnx_path).lower()
+    return "fp16" in name or "autocast" in name
+
+
+def _convert_onnx_to_trt(onnx_path: str, plan_path: str, strongly_typed: bool) -> None:
+    import tensorrt as trt
+
+    logger.info(
+        "Building flow-estimator TensorRT engine from %s (%s) ...",
+        onnx_path,
+        "strongly-typed/fp16" if strongly_typed else "fp32+TF32",
+    )
+    trt_logger = _trt_logger()
+    builder = trt.Builder(trt_logger)
+    # STRONGLY_TYPED takes precision from the ONNX graph (fp16 engine from an
+    # fp16 ONNX) — this is the only way to get fp16 in TRT>=11, which dropped
+    # the weakly-typed FP16 BuilderFlag. Otherwise EXPLICIT_BATCH is implicit,
+    # so create the network with no flags.
+    if strongly_typed:
+        network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED))
+    else:
+        network = builder.create_network(0)
+    parser = trt.OnnxParser(network, trt_logger)
+    with open(onnx_path, "rb") as f:
+        if not parser.parse(f.read()):
+            errs = "; ".join(str(parser.get_error(i)) for i in range(parser.num_errors))
+            raise ValueError(f"Failed to parse {onnx_path}: {errs}")
+
+    config = builder.create_builder_config()
+    config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 32)  # 4 GiB
+    if not strongly_typed:
+        # fp32 ONNX: enable the best available reduced-precision matmul flag
+        # (FP16 on older TRT, else TF32 on Ampere+/Hopper). For a strongly-typed
+        # network the precision is fixed by the graph, so no flag is set.
+        for _flag_name in ("FP16", "TF32"):
+            _flag = getattr(trt.BuilderFlag, _flag_name, None)
+            if _flag is not None:
+                config.set_flag(_flag)
+                break
+
+    profile = builder.create_optimization_profile()
+    for name, mn, op, mx in zip(_DYNAMIC_INPUTS, _MIN_SHAPES, _OPT_SHAPES, _MAX_SHAPES):
+        profile.set_shape(name, mn, op, mx)
+    config.add_optimization_profile(profile)
+
+    engine_bytes = builder.build_serialized_network(network, config)
+    if engine_bytes is None:
+        raise RuntimeError(f"TensorRT failed to build flow-estimator engine from {onnx_path}")
+    tmp = plan_path + ".tmp"
+    with open(tmp, "wb") as f:
+        f.write(engine_bytes)
+    os.replace(tmp, plan_path)
+    logger.info("Wrote flow-estimator TensorRT engine to %s", plan_path)
+
+
+class TrtContextWrapper:
+    """Pool of TensorRT execution contexts for the flow estimator.
+
+    Exposes the ``acquire_estimator`` / ``release_estimator`` contract that
+    ``CausalConditionalCFM.forward_estimator`` expects.
+    """
+
+    def __init__(
+        self, engine, device: str | torch.device, io_dtype: torch.dtype = torch.float32, trt_concurrent: int = 1
+    ):
+        self.trt_engine = engine
+        # Engine I/O dtype (fp16 for a strongly-typed fp16 engine). The flow runs
+        # in fp32, so forward_estimator casts to/from this at the boundary.
+        self.io_dtype = io_dtype
+        self._pool: queue.Queue = queue.Queue(maxsize=trt_concurrent)
+        for _ in range(trt_concurrent):
+            ctx = engine.create_execution_context()
+            assert ctx is not None, "failed to create TRT execution context (out of memory?)"
+            stream = torch.cuda.stream(torch.cuda.Stream(torch.device(device)))
+            self._pool.put([ctx, stream])
+
+    def acquire_estimator(self):
+        return self._pool.get(), self.trt_engine
+
+    def release_estimator(self, context, stream):
+        self._pool.put([context, stream])
+
+
+def build_flow_estimator_trt(onnx_path: str, device: str | torch.device) -> TrtContextWrapper:
+    """Build/load the flow-estimator TRT engine and return a context-pool wrapper.
+
+    An fp16 ONNX (``*autocast_fp16*``) is built as a strongly-typed network (the
+    only way to get fp16 in TRT>=11); an fp32 ONNX is built fp32 + TF32.
+    """
+    import tensorrt as trt
+
+    strongly_typed = _is_fp16_onnx(onnx_path)
+    plan_path = _resolve_plan_path(onnx_path, prefix="flow_estimator")
+    if not os.path.exists(plan_path) or os.path.getsize(plan_path) == 0:
+        _convert_onnx_to_trt(onnx_path, plan_path, strongly_typed=strongly_typed)
+
+    runtime = trt.Runtime(_trt_logger())
+    with open(plan_path, "rb") as f:
+        engine = runtime.deserialize_cuda_engine(f.read())
+    if engine is None:
+        raise RuntimeError(f"Failed to deserialize flow-estimator TensorRT engine {plan_path}")
+    logger.info("Loaded flow-estimator TensorRT engine (%s)", plan_path)
+    io_dtype = torch.float16 if strongly_typed else torch.float32
+    return TrtContextWrapper(engine, device=device, io_dtype=io_dtype)

vllm_omni/model_executor/models/cosyvoice3/speaker_embedding_trt.py

@@ -0,0 +1,157 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""TensorRT speaker-embedding (campplus) extractor for CosyVoice3.
+
+Drop-in accelerator for the ONNX-Runtime campplus session used in the
+multimodal processor. ``campplus.onnx`` maps an fbank feature
+``[B, T, 80]`` to a 192-d speaker embedding ``[B, 192]``; on this stack the
+ONNX path only has a CPU execution provider, so it dominates per-request
+prompt-processing latency. Building a TensorRT engine once and running it on
+GPU removes that bottleneck.
+
+The engine is built lazily from ``campplus.onnx`` on first use and cached as
+a ``.plan`` file (engines are device/TRT-version specific, so the cache key
+includes both). Set ``COSYVOICE3_TRT_CACHE`` to override the cache dir.
+
+Written against TensorRT >= 10/11: ``EXPLICIT_BATCH`` is the implicit default
+(create the network with no flags) and ``ITensor.dtype`` is read-only (I/O
+dtypes come from the ONNX graph), so neither is set here.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import os
+import threading
+
+import torch
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+# campplus fbank sequence length bounds for the optimization profile.
+_MIN_T = 4
+_OPT_T = 500
+_MAX_T = 3000
+_EMBED_DIM = 192
+
+
+def _trt_logger():
+    import tensorrt as trt
+
+    return trt.Logger(trt.Logger.WARNING)
+
+
+def _convert_onnx_to_trt(onnx_path: str, plan_path: str) -> None:
+    import tensorrt as trt
+
+    logger.info("Building campplus TensorRT engine from %s ...", onnx_path)
+    trt_logger = _trt_logger()
+    builder = trt.Builder(trt_logger)
+    # EXPLICIT_BATCH is implicit in TRT>=10; create the network with no flags.
+    network = builder.create_network(0)
+    parser = trt.OnnxParser(network, trt_logger)
+    with open(onnx_path, "rb") as f:
+        if not parser.parse(f.read()):
+            errs = "; ".join(str(parser.get_error(i)) for i in range(parser.num_errors))
+            raise ValueError(f"Failed to parse {onnx_path}: {errs}")
+
+    config = builder.create_builder_config()
+    config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 31)  # 2 GiB
+
+    input_name = network.get_input(0).name
+    profile = builder.create_optimization_profile()
+    profile.set_shape(
+        input_name,
+        (1, _MIN_T, 80),
+        (1, _OPT_T, 80),
+        (1, _MAX_T, 80),
+    )
+    config.add_optimization_profile(profile)
+
+    engine_bytes = builder.build_serialized_network(network, config)
+    if engine_bytes is None:
+        raise RuntimeError(f"TensorRT failed to build campplus engine from {onnx_path}")
+    tmp = plan_path + ".tmp"
+    with open(tmp, "wb") as f:
+        f.write(engine_bytes)
+    os.replace(tmp, plan_path)
+    logger.info("Wrote campplus TensorRT engine to %s", plan_path)
+
+
+def _resolve_plan_path(onnx_path: str, prefix: str = "campplus") -> str:
+    cache_dir = os.environ.get("COSYVOICE3_TRT_CACHE") or os.path.join(
+        os.path.expanduser("~"), ".cache", "vllm_omni", "cosyvoice3_trt"
+    )
+    os.makedirs(cache_dir, exist_ok=True)
+    try:
+        # No-arg form targets the current device (avoids the banned
+        # torch.cuda.current_device per ruff TID251).
+        dev_name = torch.cuda.get_device_name()
+    except Exception:
+        dev_name = "unknown"
+    import tensorrt as trt
+
+    st = os.stat(onnx_path)
+    key = f"{os.path.abspath(onnx_path)}|{st.st_size}|{int(st.st_mtime)}|{dev_name}|trt{trt.__version__}"
+    digest = hashlib.sha1(key.encode()).hexdigest()[:16]
+    return os.path.join(cache_dir, f"{prefix}_{digest}.plan")
+
+
+class CampplusTRT:
+    """Lazily-built TensorRT campplus speaker-embedding extractor."""
+
+    def __init__(self, onnx_path: str, device: str | torch.device):
+        import tensorrt as trt
+
+        self.device = torch.device(device)
+        plan_path = _resolve_plan_path(onnx_path)
+        if not os.path.exists(plan_path) or os.path.getsize(plan_path) == 0:
+            _convert_onnx_to_trt(onnx_path, plan_path)
+
+        runtime = trt.Runtime(_trt_logger())
+        with open(plan_path, "rb") as f:
+            self.engine = runtime.deserialize_cuda_engine(f.read())
+        if self.engine is None:
+            raise RuntimeError(f"Failed to deserialize campplus TensorRT engine {plan_path}")
+        self.context = self.engine.create_execution_context()
+        self.input_name = self.engine.get_tensor_name(0)
+        self.output_name = self.engine.get_tensor_name(1)
+        # TRT execution contexts are not safe for concurrent enqueue; serialize.
+        self._lock = threading.Lock()
+        logger.info("Loaded campplus TensorRT engine (%s)", plan_path)
+
+    @torch.inference_mode()
+    def __call__(self, feat: torch.Tensor) -> torch.Tensor:
+        """feat: ``[T, 80]`` fbank features. Returns ``[1, 192]`` embedding."""
+        x = feat.unsqueeze(0).to(self.device, dtype=torch.float32).contiguous()
+        out = torch.empty((x.shape[0], _EMBED_DIM), device=self.device, dtype=torch.float32)
+        stream = torch.cuda.current_stream(self.device)
+        with self._lock:
+            self.context.set_input_shape(self.input_name, tuple(x.shape))
+            self.context.set_tensor_address(self.input_name, x.data_ptr())
+            self.context.set_tensor_address(self.output_name, out.data_ptr())
+            ok = self.context.execute_async_v3(stream.cuda_stream)
+            if not ok:
+                raise RuntimeError("campplus TensorRT execute_async_v3 failed")
+            stream.synchronize()
+        return out
+
+
+# Process-wide cache: the mm processor is re-created per request (the deploy
+# config disables the mm-processor object cache), so building a fresh
+# CampplusTRT each time would re-deserialize the ~33 MB engine and allocate a
+# new execution context on every request — inflating TTFP. Cache by
+# (engine path, device) and reuse; the context serializes concurrent enqueues
+# via its own lock, so a single shared instance is safe across requests.
+_CAMPPLUS_CACHE: dict[tuple[str, str], CampplusTRT] = {}
+
+
+def get_campplus_trt(onnx_path: str, device: str | torch.device) -> CampplusTRT:
+    """Return a process-wide cached :class:`CampplusTRT` (engine + context)."""
+    key = (os.path.abspath(onnx_path), str(device))
+    inst = _CAMPPLUS_CACHE.get(key)
+    if inst is None:
+        inst = CampplusTRT(onnx_path, device)
+        _CAMPPLUS_CACHE[key] = inst
+    return inst

vllm_omni/model_executor/models/cosyvoice3/utils.py

@@ -215,6 +215,44 @@ def extract_spk_embedding(prompt_wav, campplus_session, device):
     return embedding
 
 
+def unpad_prompt_conditioning(speech_token, speech_feat, speech_token_len):
+    """Drop right-padding from per-request prompt conditioning.
+
+    The talker emits ``speech_token`` ``[1, T]`` possibly right-padded to the
+    batch max, plus the true ``speech_token_len``; ``speech_feat`` is at the 2:1
+    mel:token ratio (``[1, 2T, F]``). This trims both to the real length. Called
+    in the (eager) code2wav stage where a host read of the length is allowed.
+    Returns ``(speech_token, speech_feat)`` unchanged when length is unavailable.
+    """
+    if not isinstance(speech_token, torch.Tensor) or speech_token_len is None:
+        return speech_token, speech_feat
+    try:
+        token_len = int(torch.as_tensor(speech_token_len).reshape(-1)[0].item())
+    except (RuntimeError, ValueError, IndexError):
+        return speech_token, speech_feat
+    if token_len <= 0 or speech_token.dim() < 2:
+        return speech_token, speech_feat
+    token_len = min(token_len, int(speech_token.shape[1]))
+    speech_token = speech_token[:, :token_len]
+    if isinstance(speech_feat, torch.Tensor) and speech_feat.dim() >= 2:
+        speech_feat = speech_feat[:, : 2 * token_len]
+    return speech_token, speech_feat
+
+
+def extract_spk_embedding_trt(prompt_wav, campplus_trt, device):
+    """TensorRT counterpart of ``extract_spk_embedding``.
+
+    Identical fbank front-end; the campplus forward runs on a prebuilt
+    TensorRT engine (GPU) instead of the CPU ONNX-Runtime session. Returns the
+    same ``[1, 192]`` embedding tensor on ``device``.
+    """
+    speech = load_wav(prompt_wav, 16000)
+    feat = kaldi.fbank(speech, num_mel_bins=80, dither=0, sample_frequency=16000)
+    feat = feat - feat.mean(dim=0, keepdim=True)
+    embedding = campplus_trt(feat)
+    return embedding.to(device)
+
+
 def extract_text_token(text, tokenizer, allowed_special):
     text_token = tokenizer.encode(text, allowed_special=allowed_special)
     text_token = torch.tensor([text_token], dtype=torch.int32)

vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py

@@ -26,10 +26,10 @@ from vllm.sequence import IntermediateTensors
 
 from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin
 from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
+from vllm_omni.model_executor.models.common.ming.audio_vae import AudioVAE, AudioVAEConfig
 from vllm_omni.model_executor.models.output_templates import OmniOutput
 from vllm_omni.transformers_utils.configs.ming_flash_omni import MingFlashOmniTalkerConfig
 
-from .audio_vae import AudioVAE, AudioVAEConfig
 from .prompt_utils import DEFAULT_PROMPT as MING_DEFAULT_PROMPT
 from .talker_module import CFM, Aggregator, DiT, MingAudioGenerator, build_tts_input
 from .text_processing import segment_and_normalize

vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_thinker.py

@@ -697,6 +697,9 @@ class MingFlashOmniThinkerForConditionalGeneration(
             f"image_gen_scales={list(self.query_tokens_dict.keys())}"
         )
 
+    def get_language_model(self):
+        return self.language_model
+
     def _load_image_gen_query_tokens(self, model_path: str) -> None:
         """Load learnable query tokens used for image generation.
 

vllm_omni/model_executor/models/ming_flash_omni/prompt_utils.py

@@ -110,12 +110,18 @@ BASE_CAPTION_TEMPLATE: dict[str, Any] = {
 }
 
 
-def create_instruction(user_input: dict[str, Any]) -> str:
+def create_instruction(user_input: Any) -> str | None:
     """Return a JSON caption string for ``audio_sequence[0]``.
 
     Only keys already present on the base template are merged in; unknown
     keys are silently ignored to keep the output schema stable.
     """
+    if user_input is None:
+        return None
+    if isinstance(user_input, str):
+        return user_input
+    if not isinstance(user_input, dict):
+        raise ValueError(f"Ming instruction must be str, dict, or None; got {type(user_input).__name__}")
     caption = copy.deepcopy(BASE_CAPTION_TEMPLATE)
     item = caption["audio_sequence"][0]
     for key, value in user_input.items():

vllm_omni/model_executor/models/ming_flash_omni/talker_module.py

@@ -27,223 +27,19 @@ from threading import Lock
 
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
 from transformers import PreTrainedTokenizerBase, Qwen2Config, Qwen2Model, StaticCache
 from vllm.logger import init_logger
 from vllm.platforms import current_platform
-from x_transformers.x_transformers import RotaryEmbedding, apply_rotary_pos_emb
+from x_transformers.x_transformers import RotaryEmbedding
 
 from vllm_omni.model_executor.layers.timestep_embedding import DiTTimestepEmbedding
-
-from .audio_vae import AudioVAE
+from vllm_omni.model_executor.models.common.ming.audio_vae import AudioVAE
+from vllm_omni.model_executor.models.common.ming.dit import CondEmbedder, DiTBlock, FinalLayer, get_epss_timesteps
+from vllm_omni.model_executor.models.common.ming.fm import apply_sway_sampling, integrate_cfm_steps
 
 logger = init_logger(__name__)
 
 
-########################################################################
-# DiT Modules
-# Ported from:
-# https://github.com/inclusionAI/Ming/blob/e58533db227031990c5a6864dcf5f08fb53ed0d2/talker_module/modules.py
-# Ported from:
-# https://github.com/inclusionAI/Ming/blob/e58533db227031990c5a6864dcf5f08fb53ed0d2/talker_module/dit.py
-########################################################################
-
-
-class RMSNorm(nn.Module):
-    def __init__(self, dim: int, eps: float = 1e-6):
-        super().__init__()
-        self.eps = eps
-        self.weight = nn.Parameter(torch.ones(dim))
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        if self.weight.dtype in [torch.float16, torch.bfloat16]:
-            x = x.to(self.weight.dtype)
-        x = F.rms_norm(x, normalized_shape=(x.shape[-1],), weight=self.weight, eps=self.eps)
-        return x
-
-
-class FeedForward(nn.Module):
-    def __init__(
-        self, dim: int, dim_out: int | None = None, mult: float = 4, dropout: float = 0.0, approximate: str = "none"
-    ):
-        super().__init__()
-        inner_dim = int(dim * mult)
-        dim_out = dim_out if dim_out is not None else dim
-
-        activation = nn.GELU(approximate=approximate)
-        project_in = nn.Sequential(nn.Linear(dim, inner_dim), activation)
-        self.ff = nn.Sequential(project_in, nn.Dropout(dropout), nn.Linear(inner_dim, dim_out))
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return self.ff(x)
-
-
-class Attention(nn.Module):
-    def __init__(
-        self,
-        dim: int,
-        heads: int = 8,
-        dim_head: int = 64,
-        dropout: float = 0.0,
-        qk_norm: str | None = None,
-        pe_attn_head: int | None = None,
-        attn_mask_enabled: bool = True,
-    ):
-        super().__init__()
-        self.dim = dim
-        self.heads = heads
-        self.inner_dim = dim_head * heads
-        self.dropout = dropout
-
-        self.to_q = nn.Linear(dim, self.inner_dim)
-        self.to_k = nn.Linear(dim, self.inner_dim)
-        self.to_v = nn.Linear(dim, self.inner_dim)
-        if qk_norm is None:
-            self.q_norm = None
-            self.k_norm = None
-        elif qk_norm == "rms_norm":
-            self.q_norm = RMSNorm(dim_head)
-            self.k_norm = RMSNorm(dim_head)
-        else:
-            raise ValueError(f"Unimplemented qk_norm: {qk_norm}")
-
-        self.to_out = nn.ModuleList([])
-        self.to_out.append(nn.Linear(self.inner_dim, dim))
-        self.to_out.append(nn.Dropout(dropout))
-
-        self.pe_attn_head = pe_attn_head
-        self.attn_mask_enabled = attn_mask_enabled
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        mask: torch.Tensor | None = None,
-        rope: tuple[torch.Tensor, torch.Tensor | None] | None = None,
-    ) -> torch.Tensor:
-        batch_size = x.shape[0]
-
-        query = self.to_q(x)
-        key = self.to_k(x)
-        value = self.to_v(x)
-
-        inner_dim = key.shape[-1]
-        head_dim = inner_dim // self.heads
-        query = query.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
-        key = key.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
-        value = value.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
-
-        if self.q_norm is not None:
-            query = self.q_norm(query)
-        if self.k_norm is not None:
-            key = self.k_norm(key)
-
-        if rope is not None:
-            freqs, xpos_scale = rope
-            q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0)
-
-            if self.pe_attn_head is not None:
-                on = self.pe_attn_head
-                query[:, :on, :, :] = apply_rotary_pos_emb(query[:, :on, :, :], freqs, q_xpos_scale)
-                key[:, :on, :, :] = apply_rotary_pos_emb(key[:, :on, :, :], freqs, k_xpos_scale)
-            else:
-                query = apply_rotary_pos_emb(query, freqs, q_xpos_scale)
-                key = apply_rotary_pos_emb(key, freqs, k_xpos_scale)
-
-        if self.attn_mask_enabled and mask is not None:
-            valid_sample_indices = mask.any(dim=1)
-            final_output = torch.zeros_like(query).to(query.device)
-
-            attn_mask = mask[valid_sample_indices]
-            query = query[valid_sample_indices]
-            key = key[valid_sample_indices]
-            value = value[valid_sample_indices]
-            attn_mask = attn_mask.unsqueeze(1).unsqueeze(1)
-            attn_mask = attn_mask.expand(valid_sample_indices.sum().item(), self.heads, query.shape[-2], key.shape[-2])
-        else:
-            attn_mask = None
-
-        x = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=False)
-        if self.attn_mask_enabled and mask is not None:
-            final_output[valid_sample_indices] = x
-            x = final_output
-
-        x = x.transpose(1, 2).reshape(batch_size, -1, self.heads * head_dim)
-        x = x.to(query.dtype)
-
-        x = self.to_out[0](x)
-        x = self.to_out[1](x)
-
-        if mask is not None:
-            mask = mask.unsqueeze(-1)
-            x = x.masked_fill(~mask, 0.0)
-
-        return x
-
-
-class DiTBlock(nn.Module):
-    """A DiT block with pre-norm and residual connections."""
-
-    def __init__(
-        self,
-        hidden_size: int,
-        num_heads: int,
-        mlp_ratio: float = 4.0,
-        dropout: float = 0.1,
-        qk_norm: str | None = None,
-        pe_attn_head: int | None = None,
-        attn_mask_enabled: bool = True,
-        **kwargs,
-    ):
-        super().__init__()
-        self.norm1 = RMSNorm(hidden_size)
-        self.attn = Attention(
-            dim=hidden_size,
-            heads=num_heads,
-            dim_head=hidden_size // num_heads,
-            dropout=dropout,
-            qk_norm=qk_norm,
-            pe_attn_head=pe_attn_head,
-            attn_mask_enabled=attn_mask_enabled,
-        )
-        self.norm2 = RMSNorm(hidden_size)
-        self.mlp = FeedForward(dim=hidden_size, mult=mlp_ratio, dropout=dropout, approximate="tanh")
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        mask: torch.Tensor | None,
-        rope: tuple[torch.Tensor, torch.Tensor | None] | None,
-    ) -> torch.Tensor:
-        x = x + self.attn(self.norm1(x), mask=mask, rope=rope)
-        x = x + self.mlp(self.norm2(x))
-        return x
-
-
-class FinalLayer(nn.Module):
-    """The final layer of DiT."""
-
-    def __init__(self, hidden_size: int, out_channels: int):
-        super().__init__()
-        self.norm_final = RMSNorm(hidden_size)
-        self.linear = nn.Linear(hidden_size, out_channels, bias=True)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x = self.norm_final(x)
-        x = self.linear(x)
-        return x
-
-
-class CondEmbedder(nn.Module):
-    """Embeds LLM hidden states with optional CFG dropout."""
-
-    def __init__(self, input_feature_size: int, hidden_size: int):
-        super().__init__()
-        self.cond_embedder = nn.Linear(input_feature_size, hidden_size)
-
-    def forward(self, llm_cond: torch.Tensor) -> torch.Tensor:
-        return self.cond_embedder(llm_cond)
-
-
 class DiT(nn.Module):
     """Diffusion model with a Transformer backbone for audio latent generation."""
 
@@ -324,29 +120,6 @@ class DiT(nn.Module):
         return model_out[:, -x.shape[1] :, :]
 
 
-#########################################################################################
-# CFM
-# Ported from:
-# https://github.com/inclusionAI/Ming/blob/e58533db227031990c5a6864dcf5f08fb53ed0d2/talker_module/cfm.py
-#########################################################################################
-
-
-def get_epss_timesteps(n, device, dtype):
-    dt = 1 / 32
-    predefined_timesteps = {
-        5: [0, 2, 4, 8, 16, 32],
-        6: [0, 2, 4, 6, 8, 16, 32],
-        7: [0, 2, 4, 6, 8, 16, 24, 32],
-        10: [0, 2, 4, 6, 8, 12, 16, 20, 24, 28, 32],
-        12: [0, 2, 4, 6, 8, 10, 12, 14, 16, 20, 24, 28, 32],
-        16: [0, 1, 2, 3, 4, 5, 6, 7, 8, 10, 12, 14, 16, 20, 24, 28, 32],
-    }
-    t = predefined_timesteps.get(n, [])
-    if not t:
-        return torch.linspace(0, 1, n + 1, device=device, dtype=dtype)
-    return dt * torch.tensor(t, device=device, dtype=dtype)
-
-
 class CFM(nn.Module):
     """Conditional Flow Matching module for audio latent generation."""
 
@@ -391,15 +164,8 @@ class CFM(nn.Module):
             pred, null_pred = torch.chunk(pred_cfg, 2, dim=0)
             return pred + (pred - null_pred) * sde_args[0]
 
-        if self.sway_sampling_coef is not None:
-            t = t + self.sway_sampling_coef * (torch.cos(torch.pi / 2 * t) - 1 + t)
-
-        for step in range(self.steps):
-            dt = t[step + 1] - t[step]
-            y0 = y0 + fn(t[step], y0) * dt
-            y0 = y0 + sde_args[1] * (sde_args[2] ** 0.5) * (dt.abs() ** 0.5) * sde_rnd[step]
-
-        return y0
+        t = apply_sway_sampling(t, self.sway_sampling_coef)
+        return integrate_cfm_steps(fn, y0, t, sde_args, sde_rnd, self.steps)
 
 
 class CFMGraphExecutor:

vllm_omni/model_executor/models/ming_flash_omni/voice_presets.py

@@ -22,7 +22,8 @@ from .spk_embedding import SpkembExtractor
 from .talker_module import resample
 
 if TYPE_CHECKING:
-    from .audio_vae import AudioVAE
+    from vllm_omni.model_executor.models.common.ming.audio_vae import AudioVAE
+
     from .talker_module import Aggregator
 
 logger = init_logger(__name__)

vllm_omni/model_executor/models/ming_tts/__init__.py

@@ -0,0 +1,13 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from .config_ming_tts import MingDenseConfig
+from .ming_tts import MingTTSForConditionalGeneration
+from .ming_tts_audio_vae import MingAudioVAEModel
+from .ming_tts_llm import MingLLMModel
+
+__all__ = [
+    "MingDenseConfig",
+    "MingTTSForConditionalGeneration",
+    "MingLLMModel",
+    "MingAudioVAEModel",
+]

vllm_omni/model_executor/models/ming_tts/aggregator.py

@@ -0,0 +1,67 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Adopted from https://github.com/inclusionAI/Ming-omni-tts/blob/main/fm/dit.py
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+from x_transformers.x_transformers import RotaryEmbedding
+
+from vllm_omni.model_executor.models.common.ming.dit import DiTBlock, FinalLayer
+
+
+class Aggregator(nn.Module):
+    def __init__(
+        self,
+        in_channels=4,
+        hidden_size=1152,
+        depth=28,
+        num_heads=16,
+        mlp_ratio=4.0,
+        llm_input_dim=896,
+        **kwargs,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = in_channels
+        self.num_heads = num_heads
+
+        self.word_embedder = nn.Embedding(1, hidden_size)
+        self.x_embedder = nn.Linear(in_channels, hidden_size)
+        self.hidden_size = hidden_size
+        self.rotary_embed = RotaryEmbedding(hidden_size // num_heads)
+        self.blocks = nn.ModuleList(
+            [DiTBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio, **kwargs) for _ in range(depth)]
+        )
+        self.final_layer = FinalLayer(hidden_size, llm_input_dim)
+
+    def forward(self, x, mask=None):
+        if x.ndim != 3:
+            raise ValueError(f"Expected x rank-3 [Batch, Time, Dimension], got {tuple(x.shape)}")
+        if x.shape[-1] != self.in_channels:
+            raise ValueError(f"x feature dim mismatch: got {x.shape[-1]}, expected {self.in_channels}")
+
+        # [Batch, Time, Dimension] -> [Batch, Time, Hidden].
+        x = self.x_embedder(x)
+        cls_embed = self.word_embedder(torch.zeros((x.shape[0], 1), dtype=torch.long, device=x.device))
+        # Prepend a learned CLS token: [Batch, Time, Hidden] -> [Batch, Time + 1, Hidden].
+        x = torch.cat([cls_embed, x], dim=1)
+
+        rope = self.rotary_embed.forward_from_seq_len(x.shape[1])
+        if mask is not None:
+            if mask.ndim != 2:
+                raise ValueError(f"Expected mask rank-2 [Batch, Time], got {tuple(mask.shape)}")
+            if mask.shape[0] != x.shape[0] or mask.shape[1] != x.shape[1] - 1:
+                raise ValueError(
+                    f"Mask shape mismatch: got {tuple(mask.shape)}, expected {(x.shape[0], x.shape[1] - 1)}"
+                )
+            mask_pad = mask.clone().detach()[:, :1]
+            mask = torch.cat([mask_pad, mask], dim=-1)
+        for block in self.blocks:
+            x = block(x, mask, rope)
+        x = self.final_layer(x)
+        # Keep the CLS projection only: [Batch, Time + 1, Hidden] -> [Batch, 1, Hidden].
+        return x[:, :1, :]
+
+
+__all__ = ["Aggregator"]