docs: address recipe review feedback

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

.buildkite/pipeline.yml

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

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

@@ -19,7 +19,7 @@ has_new_python=$($PYTHON -c "print(1 if __import__('sys').version_info >= (3,12)
 if [[ "$has_new_python" -eq 0 ]]; then
     # use new python from docker
     docker pull python:3-slim
-    PYTHON="docker run --rm -v $(pwd):/app -w /app python:3-slim python3"
+    PYTHON="docker run --rm --user $(id -u):$(id -g) -v $(pwd):/app -w /app python:3-slim python3"
 fi
 
 echo "Using python interpreter: $PYTHON"
@@ -36,7 +36,7 @@ mkdir -p "$INDICES_OUTPUT_DIR"
 
 # HACK: we do not need regex module here, but it is required by pre-commit hook
 # To avoid any external dependency, we simply replace it back to the stdlib re module
-sed -i 's/import regex as re/import re/g' .buildkite/scripts/generate-nightly-index.py
+sed -i.bak 's/import regex as re/import re/g' .buildkite/scripts/generate-nightly-index.py && rm -f .buildkite/scripts/generate-nightly-index.py.bak
 
 # Generate indices -- the version is just the commit hash (not omni/{commit})
 # because relative paths are computed between the index and wheel directories,
@@ -73,15 +73,16 @@ echo "Pure version (without variant): $pure_version"
 
 # re-generate and copy to /omni/{version}/ only if it does not have "dev" in the version
 if [[ "$version" != *"dev"* ]]; then
-    echo "Re-generating indices for /omni/$pure_version/"
+    s3_version="v$pure_version"
+    echo "Re-generating indices for /omni/$s3_version/"
     rm -rf "${INDICES_OUTPUT_DIR:?}"
     mkdir -p "$INDICES_OUTPUT_DIR"
     # wheel-dir is overridden to be the commit directory, so that the indices point to the correct wheel path
     $PYTHON .buildkite/scripts/generate-nightly-index.py \
-        --version "$pure_version" \
+        --version "$s3_version" \
         --wheel-dir "$BUILDKITE_COMMIT" \
         --current-objects "$obj_json" \
         --output-dir "$INDICES_OUTPUT_DIR" \
         --comment "version $pure_version"
-    aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "s3://$BUCKET/omni/$pure_version/"
+    aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "s3://$BUCKET/omni/$s3_version/"
 fi

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

@@ -4,6 +4,7 @@
 
 import argparse
 import json
+import re
 import sys
 from dataclasses import asdict, dataclass
 from datetime import datetime
@@ -11,8 +12,6 @@ from pathlib import Path
 from typing import Any
 from urllib.parse import quote
 
-import regex as re
-
 
 def normalize_package_name(name: str) -> str:
     """Normalize package name per PEP 503."""

.buildkite/scripts/upload_pipeline_with_skip_ci.sh

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

@@ -390,16 +390,6 @@ steps:
           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"
         '
-
-  - label: "CosyVoice3-TTS E2E Test"
-    timeout_in_minutes: 20
-    depends_on: upload-merge-pipeline
-    commands:
-      - |
-        timeout 20m bash -c '
-          export VLLM_WORKER_MULTIPROC_METHOD=spawn
-          pytest -s -v tests/e2e/online_serving/test_cosyvoice3_tts.py -m "advanced_model" --run-level "advanced_model"
-        '
     agents:
       queue: "mithril-h100-pool"
     plugins:

.buildkite/test-nightly-diffusion.yml

@@ -0,0 +1,367 @@
+# Nightly diffusion GPU tests — appended to the main nightly build via
+#   buildkite-agent pipeline upload .buildkite/test-nightly-diffusion.yml
+# from test-nightly.yml (step key: nightly-diffusion-model-test). Top-level groups are
+# foldable in the Buildkite UI (Other / Wan / Qwen-Image).
+steps:
+  - group: ":card_index_dividers: Other Model Test"
+    key: nightly-other-model-test-group
+    steps:
+      - label: ":full_moon: Diffusion · Other · Function Test with H100"
+        timeout_in_minutes: 120
+        if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
+        commands:
+          - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          - pytest -s -v tests/e2e/online_serving/test_*_expansion.py -k "not test_wan22_expansion and not test_wan_2_1_vace_expansion and not test_qwen_image" -m "advanced_model and diffusion and H100" --run-level "advanced_model"
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 2
+                    volumeMounts:
+                      - name: devshm
+                        mountPath: /dev/shm
+                      - name: hf-cache
+                        mountPath: /root/.cache/huggingface
+                    env:
+                      - name: HF_HOME
+                        value: /root/.cache/huggingface
+                      - name: HF_TOKEN
+                        valueFrom:
+                          secretKeyRef:
+                            name: hf-token-secret
+                            key: token
+                nodeSelector:
+                  node.kubernetes.io/instance-type: gpu-h100-sxm
+                volumes:
+                  - name: devshm
+                    emptyDir:
+                      medium: Memory
+                  - name: hf-cache
+                    hostPath:
+                      path: /mnt/hf-cache
+                      type: DirectoryOrCreate
+
+      - label: ":full_moon: Diffusion · Other · Function Test with L4"
+        timeout_in_minutes: 60
+        if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
+        commands:
+          - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          - pytest -s -v tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and diffusion and L4" --run-level "advanced_model"
+        agents:
+          queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
+        plugins:
+          - docker#v5.2.0:
+              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+              always-pull: true
+              propagate-environment: true
+              shm-size: "8gb"
+              environment:
+                - "HF_HOME=/fsx/hf_cache"
+                - "HF_TOKEN"
+              volumes:
+                - "/fsx/hf_cache:/fsx/hf_cache"
+
+      - label: ":full_moon: Diffusion · Other · Doc Test"
+        timeout_in_minutes: 60
+        if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
+        commands:
+          - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          - export VLLM_TEST_CLEAN_GPU_MEMORY="1"
+          - pytest -s -v tests/examples/online_serving/test_text_to_image.py tests/examples/offline_inference/test_text_to_image.py -m "advanced_model and example and H100" --run-level "advanced_model"
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 2
+                    volumeMounts:
+                      - name: devshm
+                        mountPath: /dev/shm
+                      - name: hf-cache
+                        mountPath: /root/.cache/huggingface
+                    env:
+                      - name: HF_HOME
+                        value: /root/.cache/huggingface
+                      - name: HF_TOKEN
+                        valueFrom:
+                          secretKeyRef:
+                            name: hf-token-secret
+                            key: token
+                nodeSelector:
+                  node.kubernetes.io/instance-type: gpu-h100-sxm
+                volumes:
+                  - name: devshm
+                    emptyDir:
+                      medium: Memory
+                  - name: hf-cache
+                    hostPath:
+                      path: /mnt/hf-cache
+                      type: DirectoryOrCreate
+
+  - group: ":card_index_dividers: Wan Series Model Test"
+    key: nightly-wan-model-test-group
+    steps:
+      - label: ":full_moon: Diffusion · Wan · Function Test"
+        timeout_in_minutes: 90
+        if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
+        commands:
+          - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py tests/e2e/online_serving/test_wan_2_1_vace_expansion.py -m "advanced_model" --run-level "advanced_model"
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 2
+                    volumeMounts:
+                      - name: devshm
+                        mountPath: /dev/shm
+                      - name: hf-cache
+                        mountPath: /root/.cache/huggingface
+                    env:
+                      - name: HF_HOME
+                        value: /root/.cache/huggingface
+                      - name: HF_TOKEN
+                        valueFrom:
+                          secretKeyRef:
+                            name: hf-token-secret
+                            key: token
+                nodeSelector:
+                  node.kubernetes.io/instance-type: gpu-h100-sxm
+                volumes:
+                  - name: devshm
+                    emptyDir:
+                      medium: Memory
+                  - name: hf-cache
+                    hostPath:
+                      path: /mnt/hf-cache
+                      type: DirectoryOrCreate
+
+      - label: ":full_moon: Diffusion · Wan · Accuracy Test"
+        key: nightly-wan22-i2v-accuracy
+        timeout_in_minutes: 180
+        if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
+        commands:
+          - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          - pytest -s -v tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py --run-level advanced_model
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 2
+                    volumeMounts:
+                      - name: devshm
+                        mountPath: /dev/shm
+                      - name: hf-cache
+                        mountPath: /root/.cache/huggingface
+                    env:
+                      - name: HF_HOME
+                        value: /root/.cache/huggingface
+                      - name: HF_TOKEN
+                        valueFrom:
+                          secretKeyRef:
+                            name: hf-token-secret
+                            key: token
+                nodeSelector:
+                  node.kubernetes.io/instance-type: gpu-h100-sxm
+                volumes:
+                  - name: devshm
+                    emptyDir:
+                      medium: Memory
+                  - name: hf-cache
+                    hostPath:
+                      path: /mnt/hf-cache
+                      type: DirectoryOrCreate
+
+  - group: ":card_index_dividers: Qwen-Image Series Model Test"
+    key: nightly-qwen-image-edit-group
+    steps:
+      - label: ":full_moon: Diffusion · Qwen-Image · Function Test with H100"
+        timeout_in_minutes: 120
+        if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
+        commands:
+          - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          - pytest -s -v tests/e2e/online_serving/test_qwen_image*_expansion.py -m "advanced_model and diffusion and H100" --run-level "advanced_model"
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 2
+                    volumeMounts:
+                      - name: devshm
+                        mountPath: /dev/shm
+                      - name: hf-cache
+                        mountPath: /root/.cache/huggingface
+                    env:
+                      - name: HF_HOME
+                        value: /root/.cache/huggingface
+                      - name: HF_TOKEN
+                        valueFrom:
+                          secretKeyRef:
+                            name: hf-token-secret
+                            key: token
+                nodeSelector:
+                  node.kubernetes.io/instance-type: gpu-h100-sxm
+                volumes:
+                  - name: devshm
+                    emptyDir:
+                      medium: Memory
+                  - name: hf-cache
+                    hostPath:
+                      path: /mnt/hf-cache
+                      type: DirectoryOrCreate
+
+      - label: ":full_moon: Diffusion · Qwen-Image · GEBench Accuracy Test"
+        key: nightly-gebench-accuracy
+        timeout_in_minutes: 60
+        if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
+        commands:
+          - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          - pytest -s -v tests/e2e/accuracy/test_gebench_h100_smoke.py --run-level advanced_model --gebench-model Qwen/Qwen-Image-2512 --accuracy-judge-model QuantTrio/Qwen3-VL-30B-A3B-Instruct-AWQ --accuracy-gpu 0 --gebench-port 8093 --accuracy-workers 1
+          - buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gebench_qwen-image-2512/summary*.json"
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 1
+                    volumeMounts:
+                      - name: devshm
+                        mountPath: /dev/shm
+                      - name: hf-cache
+                        mountPath: /root/.cache/huggingface
+                    env:
+                      - name: HF_HOME
+                        value: /root/.cache/huggingface
+                      - name: HF_TOKEN
+                        valueFrom:
+                          secretKeyRef:
+                            name: hf-token-secret
+                            key: token
+                nodeSelector:
+                  node.kubernetes.io/instance-type: gpu-h100-sxm
+                volumes:
+                  - name: devshm
+                    emptyDir:
+                      medium: Memory
+                  - name: hf-cache
+                    hostPath:
+                      path: /mnt/hf-cache
+                      type: DirectoryOrCreate
+
+      - label: ":full_moon: Diffusion · Qwen-Image · GEdit-Bench Accuracy Test"
+        key: nightly-gedit-bench-accuracy
+        timeout_in_minutes: 60
+        if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
+        commands:
+          - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          - pytest -s -v tests/e2e/accuracy/test_gedit_bench_h100_smoke.py --run-level advanced_model --gedit-model Qwen/Qwen-Image-Edit --accuracy-judge-model QuantTrio/Qwen3-VL-30B-A3B-Instruct-AWQ --accuracy-gpu 0 --gedit-port 8093 --gedit-samples-per-group 20 --accuracy-workers 1
+          - buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gedit_scores_qwen-image-edit/qwen-image-edit_all_all_vie_score_*.csv"
+          - buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gedit_scores_qwen-image-edit/qwen-image-edit_all_all_summary_*.json"
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 1
+                    volumeMounts:
+                      - name: devshm
+                        mountPath: /dev/shm
+                      - name: hf-cache
+                        mountPath: /root/.cache/huggingface
+                    env:
+                      - name: HF_HOME
+                        value: /root/.cache/huggingface
+                      - name: VLLM_HTTP_TIMEOUT_KEEP_ALIVE
+                        value: "120"
+                      - name: HF_TOKEN
+                        valueFrom:
+                          secretKeyRef:
+                            name: hf-token-secret
+                            key: token
+                nodeSelector:
+                  node.kubernetes.io/instance-type: gpu-h100-sxm
+                volumes:
+                  - name: devshm
+                    emptyDir:
+                      medium: Memory
+                  - name: hf-cache
+                    hostPath:
+                      path: /mnt/hf-cache
+                      type: DirectoryOrCreate
+
+      - label: ":full_moon: Diffusion · Qwen-Image · Perf Test"
+        key: nightly-qwen-image-performance
+        timeout_in_minutes: 180
+        if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
+        commands:
+          - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
+          - export CACHE_DIT_VERSION=1.3.0
+          - pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json
+          - buildkite-agent artifact upload "tests/dfx/perf/results/benchmark_results_*.json"
+          - buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 4
+                    volumeMounts:
+                      - name: devshm
+                        mountPath: /dev/shm
+                      - name: hf-cache
+                        mountPath: /root/.cache/huggingface
+                    env:
+                      - name: HF_HOME
+                        value: /root/.cache/huggingface
+                      - name: HF_TOKEN
+                        valueFrom:
+                          secretKeyRef:
+                            name: hf-token-secret
+                            key: token
+                nodeSelector:
+                  node.kubernetes.io/instance-type: gpu-h100-sxm
+                volumes:
+                  - name: devshm
+                    emptyDir:
+                      medium: Memory
+                  - name: hf-cache
+                    hostPath:
+                      path: /mnt/hf-cache
+                      type: DirectoryOrCreate

.buildkite/test-nightly.yml

@@ -1,228 +1,199 @@
 steps:
-  - label: ":full_moon: Omni Model Test with H100"
-    timeout_in_minutes: 90
-    depends_on: upload-nightly-pipeline
-    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
-    commands:
-      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - |
-        pytest -s -v \
-          tests/examples/ \
-          tests/e2e/online_serving/test_*_expansion.py \
-          -m "advanced_model and H100 and omni" --run-level "advanced_model"
-    agents:
-      queue: "mithril-h100-pool"
-    plugins:
-      - kubernetes:
-          podSpec:
-            containers:
-              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-                resources:
-                  limits:
-                    nvidia.com/gpu: 2
-                volumeMounts:
+  # Group: collapses under one heading in the Buildkite UI; child steps still run in parallel.
+  - group: ":card_index_dividers: Omni Model Test"
+    key: nightly-omni-test-group
+    steps:
+      - label: ":full_moon: Omni · Function Test with H100"
+        timeout_in_minutes: 90
+        depends_on: upload-nightly-pipeline
+        if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
+        commands:
+          - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          - pytest -s -v tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and H100 and omni" --run-level "advanced_model"
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 2
+                    volumeMounts:
+                      - name: devshm
+                        mountPath: /dev/shm
+                      - name: hf-cache
+                        mountPath: /root/.cache/huggingface
+                    env:
+                      - name: HF_HOME
+                        value: /root/.cache/huggingface
+                      - name: HF_TOKEN
+                        valueFrom:
+                          secretKeyRef:
+                            name: hf-token-secret
+                            key: token
+                nodeSelector:
+                  node.kubernetes.io/instance-type: gpu-h100-sxm
+                volumes:
                   - name: devshm
-                    mountPath: /dev/shm
+                    emptyDir:
+                      medium: Memory
                   - name: hf-cache
-                    mountPath: /root/.cache/huggingface
-                env:
-                  - name: HF_HOME
-                    value: /root/.cache/huggingface
-                  - name: HF_TOKEN
-                    valueFrom:
-                      secretKeyRef:
-                        name: hf-token-secret
-                        key: token
-            nodeSelector:
-              node.kubernetes.io/instance-type: gpu-h100-sxm
-            volumes:
-              - name: devshm
-                emptyDir:
-                  medium: Memory
-              - name: hf-cache
-                hostPath:
-                  path: /mnt/hf-cache
-                  type: DirectoryOrCreate
+                    hostPath:
+                      path: /mnt/hf-cache
+                      type: DirectoryOrCreate
 
-  - label: ":full_moon: Omni Model Test with L4"
-    timeout_in_minutes: 90
+      - label: ":full_moon: Omni · Function Test with L4"
+        timeout_in_minutes: 90
+        depends_on: upload-nightly-pipeline
+        if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
+        commands:
+          - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
+          - pytest -s -v tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and L4 and omni" --run-level "advanced_model"
+        agents:
+          queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
+        plugins:
+          - docker#v5.2.0:
+              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+              always-pull: true
+              propagate-environment: true
+              shm-size: "8gb"
+              environment:
+                - "HF_HOME=/fsx/hf_cache"
+                - "HF_TOKEN"
+              volumes:
+                - "/fsx/hf_cache:/fsx/hf_cache"
+
+      - label: ":full_moon: Omni · Doc Test with L4"
+        timeout_in_minutes: 90
+        depends_on: upload-nightly-pipeline
+        if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
+        commands:
+          - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
+          - pytest -s -v tests/examples/ -m "advanced_model and omni and L4" --run-level "advanced_model"
+        agents:
+          queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
+        plugins:
+          - docker#v5.2.0:
+              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+              always-pull: true
+              propagate-environment: true
+              shm-size: "8gb"
+              environment:
+                - "HF_HOME=/fsx/hf_cache"
+                - "HF_TOKEN"
+              volumes:
+                - "/fsx/hf_cache:/fsx/hf_cache"
+
+      - label: ":full_moon: Omni · Doc Test with H100"
+        timeout_in_minutes: 90
+        depends_on: upload-nightly-pipeline
+        if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
+        commands:
+          - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          - pytest -s -v tests/examples/ -m "advanced_model and omni and H100" --run-level "advanced_model"
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 2
+                    volumeMounts:
+                      - name: devshm
+                        mountPath: /dev/shm
+                      - name: hf-cache
+                        mountPath: /root/.cache/huggingface
+                    env:
+                      - name: HF_HOME
+                        value: /root/.cache/huggingface
+                      - name: HF_TOKEN
+                        valueFrom:
+                          secretKeyRef:
+                            name: hf-token-secret
+                            key: token
+                nodeSelector:
+                  node.kubernetes.io/instance-type: gpu-h100-sxm
+                volumes:
+                  - name: devshm
+                    emptyDir:
+                      medium: Memory
+                  - name: hf-cache
+                    hostPath:
+                      path: /mnt/hf-cache
+                      type: DirectoryOrCreate
+
+      - label: ":full_moon: Omni · Perf Test"
+        key: nightly-omni-performance
+        timeout_in_minutes: 180
+        depends_on: upload-nightly-pipeline
+        if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
+        commands:
+          - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          - export BENCHMARK_DIR=tests/dfx/perf/results
+          - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
+          - pytest -s -v tests/dfx/perf/scripts/run_benchmark.py
+          - buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
+          - buildkite-agent artifact upload "tests/dfx/perf/results/*.html"
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 2
+                    volumeMounts:
+                      - name: devshm
+                        mountPath: /dev/shm
+                      - name: hf-cache
+                        mountPath: /root/.cache/huggingface
+                    env:
+                      - name: HF_HOME
+                        value: /root/.cache/huggingface
+                      - name: HF_TOKEN
+                        valueFrom:
+                          secretKeyRef:
+                            name: hf-token-secret
+                            key: token
+                nodeSelector:
+                  node.kubernetes.io/instance-type: gpu-h100-sxm
+                volumes:
+                  - name: devshm
+                    emptyDir:
+                      medium: Memory
+                  - name: hf-cache
+                    hostPath:
+                      path: /mnt/hf-cache
+                      type: DirectoryOrCreate
+
+  # Dynamically appends steps from test-nightly-diffusion.yml into this build (same mechanism as
+  # pipeline.yml → test-ready.yml / test-merge.yml / test-nightly.yml). Foldable groups stay in the
+  # uploaded YAML (Other / Wan / Qwen-Image).
+  - label: ":card_index_dividers: Diffusion Model Test"
+    key: nightly-diffusion-model-test
     depends_on: upload-nightly-pipeline
     if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
     commands:
-      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
-      - pytest -s -v tests/examples/ tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and L4 and omni" --run-level "advanced_model"
+      - buildkite-agent pipeline upload .buildkite/test-nightly-diffusion.yml
     agents:
-      queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
-    plugins:
-      - docker#v5.2.0:
-          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-          always-pull: true
-          propagate-environment: true
-          shm-size: "8gb"
-          environment:
-            - "HF_HOME=/fsx/hf_cache"
-            - "HF_TOKEN"
-          volumes:
-            - "/fsx/hf_cache:/fsx/hf_cache"
+      queue: "cpu_queue_premerge"
 
-  - label: ":full_moon: Diffusion Model Test with H100"
+  - label: ":bar_chart: Testcase Statistics"
+    key: nightly-testcase-statistics
     timeout_in_minutes: 120
     depends_on: upload-nightly-pipeline
     if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
     commands:
       - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - pytest -s -v tests/e2e/online_serving/test_*_expansion.py -k "not test_wan22_expansion" -m "advanced_model and diffusion and H100" --run-level "advanced_model"
-    agents:
-      queue: "mithril-h100-pool"
-    plugins:
-      - kubernetes:
-          podSpec:
-            containers:
-              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-                resources:
-                  limits:
-                    nvidia.com/gpu: 2
-                volumeMounts:
-                  - name: devshm
-                    mountPath: /dev/shm
-                  - name: hf-cache
-                    mountPath: /root/.cache/huggingface
-                env:
-                  - name: HF_HOME
-                    value: /root/.cache/huggingface
-                  - name: HF_TOKEN
-                    valueFrom:
-                      secretKeyRef:
-                        name: hf-token-secret
-                        key: token
-            nodeSelector:
-              node.kubernetes.io/instance-type: gpu-h100-sxm
-            volumes:
-              - name: devshm
-                emptyDir:
-                  medium: Memory
-              - name: hf-cache
-                hostPath:
-                  path: /mnt/hf-cache
-                  type: DirectoryOrCreate
-
-  - label: ":full_moon: Diffusion Model (Wan2.2) Test with H100"
-    timeout_in_minutes: 90
-    depends_on: upload-nightly-pipeline
-    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
-    commands:
-      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -m "advanced_model" --run-level "advanced_model"
-    agents:
-      queue: "mithril-h100-pool"
-    plugins:
-      - kubernetes:
-          podSpec:
-            containers:
-              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-                resources:
-                  limits:
-                    nvidia.com/gpu: 2
-                volumeMounts:
-                  - name: devshm
-                    mountPath: /dev/shm
-                  - name: hf-cache
-                    mountPath: /root/.cache/huggingface
-                env:
-                  - name: HF_HOME
-                    value: /root/.cache/huggingface
-                  - name: HF_TOKEN
-                    valueFrom:
-                      secretKeyRef:
-                        name: hf-token-secret
-                        key: token
-            nodeSelector:
-              node.kubernetes.io/instance-type: gpu-h100-sxm
-            volumes:
-              - name: devshm
-                emptyDir:
-                  medium: Memory
-              - name: hf-cache
-                hostPath:
-                  path: /mnt/hf-cache
-                  type: DirectoryOrCreate
-
-  - label: ":full_moon: Diffusion Model Test"
-    timeout_in_minutes: 60
-    depends_on: upload-nightly-pipeline
-    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
-    commands:
-      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - pytest -s -v tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and diffusion and L4" --run-level "advanced_model"
-    agents:
-      queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
-    plugins:
-      - docker#v5.2.0:
-          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-          always-pull: true
-          propagate-environment: true
-          shm-size: "8gb"
-          environment:
-            - "HF_HOME=/fsx/hf_cache"
-            - "HF_TOKEN"
-          volumes:
-            - "/fsx/hf_cache:/fsx/hf_cache"
-
-
-  - label: ":full_moon: Doc Example Code Test with H100"
-    timeout_in_minutes: 60
-    depends_on: upload-nightly-pipeline
-    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
-    commands:
-      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - export VLLM_TEST_CLEAN_GPU_MEMORY="1"
-      - pytest -s -v tests/examples/online_serving/test_text_to_image.py tests/examples/offline_inference/test_text_to_image.py -m "advanced_model and example and H100" --run-level "advanced_model"
-    agents:
-      queue: "mithril-h100-pool"
-    plugins:
-      - kubernetes:
-          podSpec:
-            containers:
-              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-                resources:
-                  limits:
-                    nvidia.com/gpu: 2
-                volumeMounts:
-                  - name: devshm
-                    mountPath: /dev/shm
-                  - name: hf-cache
-                    mountPath: /root/.cache/huggingface
-                env:
-                  - name: HF_HOME
-                    value: /root/.cache/huggingface
-                  - name: HF_TOKEN
-                    valueFrom:
-                      secretKeyRef:
-                        name: hf-token-secret
-                        key: token
-            nodeSelector:
-              node.kubernetes.io/instance-type: gpu-h100-sxm
-            volumes:
-              - name: devshm
-                emptyDir:
-                  medium: Memory
-              - name: hf-cache
-                hostPath:
-                  path: /mnt/hf-cache
-                  type: DirectoryOrCreate
-
-  - label: ":full_moon: Omni Model Perf Test & Testcase Statistics with H100"
-    key: nightly-omni-performance
-    timeout_in_minutes: 180
-    depends_on: upload-nightly-pipeline
-    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
-    commands:
-      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - export BENCHMARK_DIR=tests/dfx/perf/results
-      - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
-      - pytest -s -v tests/dfx/perf/scripts/run_benchmark.py
-      - buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
       - python tools/nightly/buildkite_testcase_statistics.py -o tests/dfx/perf/results/buildkite_testcase_statistics.html
       - buildkite-agent artifact upload "tests/dfx/perf/results/*.html"
     agents:
@@ -259,189 +230,13 @@ steps:
                   path: /mnt/hf-cache
                   type: DirectoryOrCreate
 
-  - label: ":full_moon: GEBench Accuracy Test with H100"
-    key: nightly-gebench-accuracy
-    timeout_in_minutes: 60
-    depends_on: upload-nightly-pipeline
-    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
-    commands:
-      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - pytest -s -v tests/e2e/accuracy/test_gebench_h100_smoke.py --run-level advanced_model --gebench-model Qwen/Qwen-Image-2512 --accuracy-judge-model QuantTrio/Qwen3-VL-30B-A3B-Instruct-AWQ --accuracy-gpu 0 --gebench-port 8093 --accuracy-workers 1
-      - buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gebench_qwen-image-2512/summary*.json"
-    agents:
-      queue: "mithril-h100-pool"
-    plugins:
-      - kubernetes:
-          podSpec:
-            containers:
-              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-                resources:
-                  limits:
-                    nvidia.com/gpu: 1
-                volumeMounts:
-                  - name: devshm
-                    mountPath: /dev/shm
-                  - name: hf-cache
-                    mountPath: /root/.cache/huggingface
-                env:
-                  - name: HF_HOME
-                    value: /root/.cache/huggingface
-                  - name: HF_TOKEN
-                    valueFrom:
-                      secretKeyRef:
-                        name: hf-token-secret
-                        key: token
-            nodeSelector:
-              node.kubernetes.io/instance-type: gpu-h100-sxm
-            volumes:
-              - name: devshm
-                emptyDir:
-                  medium: Memory
-              - name: hf-cache
-                hostPath:
-                  path: /mnt/hf-cache
-                  type: DirectoryOrCreate
-
-  - label: ":full_moon: GEdit-Bench Accuracy Test with H100"
-    key: nightly-gedit-bench-accuracy
-    timeout_in_minutes: 60
-    depends_on: upload-nightly-pipeline
-    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
-    commands:
-      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - pytest -s -v tests/e2e/accuracy/test_gedit_bench_h100_smoke.py --run-level advanced_model --gedit-model Qwen/Qwen-Image-Edit --accuracy-judge-model QuantTrio/Qwen3-VL-30B-A3B-Instruct-AWQ --accuracy-gpu 0 --gedit-port 8093 --gedit-samples-per-group 20 --accuracy-workers 1
-      - buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gedit_scores_qwen-image-edit/qwen-image-edit_all_all_vie_score_*.csv"
-      - buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gedit_scores_qwen-image-edit/qwen-image-edit_all_all_summary_*.json"
-    agents:
-      queue: "mithril-h100-pool"
-    plugins:
-      - kubernetes:
-          podSpec:
-            containers:
-              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-                resources:
-                  limits:
-                    nvidia.com/gpu: 1
-                volumeMounts:
-                  - name: devshm
-                    mountPath: /dev/shm
-                  - name: hf-cache
-                    mountPath: /root/.cache/huggingface
-                env:
-                  - name: HF_HOME
-                    value: /root/.cache/huggingface
-                  - name: VLLM_HTTP_TIMEOUT_KEEP_ALIVE
-                    value: "120"
-                  - name: HF_TOKEN
-                    valueFrom:
-                      secretKeyRef:
-                        name: hf-token-secret
-                        key: token
-            nodeSelector:
-              node.kubernetes.io/instance-type: gpu-h100-sxm
-            volumes:
-              - name: devshm
-                emptyDir:
-                  medium: Memory
-              - name: hf-cache
-                hostPath:
-                  path: /mnt/hf-cache
-                  type: DirectoryOrCreate
-
-  - label: ":full_moon: Wan22 I2V Accuracy Test with H100"
-    key: nightly-wan22-i2v-accuracy
-    timeout_in_minutes: 180
-    depends_on: upload-nightly-pipeline
-    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
-    commands:
-      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - pytest -s -v tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py --run-level advanced_model
-    agents:
-      queue: "mithril-h100-pool"
-    plugins:
-      - kubernetes:
-          podSpec:
-            containers:
-              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-                resources:
-                  limits:
-                    nvidia.com/gpu: 2
-                volumeMounts:
-                  - name: devshm
-                    mountPath: /dev/shm
-                  - name: hf-cache
-                    mountPath: /root/.cache/huggingface
-                env:
-                  - name: HF_HOME
-                    value: /root/.cache/huggingface
-                  - name: HF_TOKEN
-                    valueFrom:
-                      secretKeyRef:
-                        name: hf-token-secret
-                        key: token
-            nodeSelector:
-              node.kubernetes.io/instance-type: gpu-h100-sxm
-            volumes:
-              - name: devshm
-                emptyDir:
-                  medium: Memory
-              - name: hf-cache
-                hostPath:
-                  path: /mnt/hf-cache
-                  type: DirectoryOrCreate
-
-  - label: ":full_moon: Diffusion Perf Test with H100"
-    key: nightly-qwen-image-performance
-    timeout_in_minutes: 180
-    depends_on: upload-nightly-pipeline
-    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
-    commands:
-      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
-      - export CACHE_DIT_VERSION=1.3.0
-      - pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json
-      - buildkite-agent artifact upload "tests/dfx/perf/results/benchmark_results_*.json"
-      - buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
-    agents:
-      queue: "mithril-h100-pool"
-    plugins:
-      - kubernetes:
-          podSpec:
-            containers:
-              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-                resources:
-                  limits:
-                    nvidia.com/gpu: 4
-                volumeMounts:
-                  - name: devshm
-                    mountPath: /dev/shm
-                  - name: hf-cache
-                    mountPath: /root/.cache/huggingface
-                env:
-                  - name: HF_HOME
-                    value: /root/.cache/huggingface
-                  - name: HF_TOKEN
-                    valueFrom:
-                      secretKeyRef:
-                        name: hf-token-secret
-                        key: token
-            nodeSelector:
-              node.kubernetes.io/instance-type: gpu-h100-sxm
-            volumes:
-              - name: devshm
-                emptyDir:
-                  medium: Memory
-              - name: hf-cache
-                hostPath:
-                  path: /mnt/hf-cache
-                  type: DirectoryOrCreate
-
   # No need to run this step for PRs with label nightly-test
   - label: ":email: Nightly Collection & Email"
     key: nightly-perf-distribution
     depends_on:
       - nightly-omni-performance
       - nightly-qwen-image-performance
+      - nightly-testcase-statistics
     if: build.env("NIGHTLY") == "1"
     commands:
       - pip install openpyxl

.buildkite/test-ready.yml

@@ -16,11 +16,10 @@ steps:
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
-  - label: "Voxtral TTS CUDA Unit Test"
-    timeout_in_minutes: 10
+  - label: "CUDA Unit Test with single card"
     depends_on: upload-ready-pipeline
     commands:
-      - "timeout 10m pytest -s -v tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py"
+      - timeout 10m pytest -v -s -m 'core_model and cuda and L4 and not distributed_cuda' --ignore=tests/e2e --ignore=tests/engine/test_async_omni_engine_abort.py --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml
     agents:
       queue: "gpu_1_queue"
     plugins:
@@ -33,6 +32,22 @@ steps:
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
+  - label: "CUDA Unit Test with multi cards"
+    depends_on: upload-ready-pipeline
+    commands:
+      - timeout 10m pytest -v -s -m 'core_model and cuda and L4 and distributed_cuda' --ignore=tests/e2e --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml
+    agents:
+      queue: "gpu_4_queue"
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
   - label: "Diffusion Model Test"
     depends_on: upload-ready-pipeline
     commands:
@@ -152,24 +167,6 @@ steps:
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
-  - label: "Diffusion GPU Worker Test"
-    depends_on: upload-ready-pipeline
-    commands:
-      - timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py
-    agents:
-      queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
-    plugins:
-      - docker#v5.2.0:
-          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-          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: "Engine Test"
     depends_on: upload-ready-pipeline

.gitignore

@@ -83,6 +83,9 @@ target/
 profile_default/
 ipython_config.py
 
+# uv
+uv.lock
+
 # pyenv
 #   For a library or package, you might want to ignore these files since the code is
 #   intended to run in multiple environments; otherwise, check them in:

docker/Dockerfile.cuda

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

docs/.nav.yml

@@ -98,6 +98,7 @@ nav:
       - design/feature/ray_based_execution.md
       - design/feature/omni_connectors/
       - design/feature/cfg_parallel.md
+      - design/feature/expert_parallel.md
       - design/feature/sequence_parallel.md
       - design/feature/tensor_parallel.md
       - design/feature/vae_parallel.md

docs/contributing/ci/CI_5levels.md

@@ -271,7 +271,7 @@ Before entering specific testing levels, the project establishes two common spec
 
 L1 and L2 level testing form the foundation of the quality assurance system. L1 level testing focuses on verifying the internal logic correctness of code units (e.g., functions, classes), ensuring each independent component behaves as designed.
 
-L2 level testing builds upon L1 by introducing GPU resources and verifying that the end-to-end (E2E) process of the model in basic deployment scenarios is smooth. For example, it uses dummy models to confirm that core interfaces like the inference pipeline, output format, and streaming response work properly. The common goal of these two levels is to provide developers with rapid feedback, discovering and fixing issues early in the development cycle  .
+L2 level testing builds upon L1 by introducing GPU resources and verifying that the end-to-end (E2E) process of the model in basic deployment scenarios is smooth. For example, it uses dummy models to confirm that core interfaces like the inference pipeline, output format, and streaming response work properly. The common goal of these two levels is to provide developers with rapid feedback, discovering and fixing issues early in the development cycle.
 
 
 

docs/contributing/model/adding_diffusion_model.md

@@ -802,7 +802,7 @@ omni = Omni(model="your-model", enable_layerwise_offload=True)
 
 ```python
 class WanTransformer3DModel(nn.Module):
-    _layerwise_offload_blocks_attr = "blocks"  # Attribute name containing transformer blocks
+    _layerwise_offload_blocks_attrs = ["blocks"]  # Attribute name containing transformer blocks
 
     def __init__(self):
         self.blocks = nn.ModuleList([...])  # Transformer blocks

docs/design/feature/expert_parallel.md

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

docs/design/feature/omni_connectors/mooncake_transfer_engine_connector.md

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

docs/design/feature/vae_parallel.md

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

docs/getting_started/installation/README.md

@@ -6,4 +6,5 @@ vLLM-Omni supports the following hardware platforms:
     - [NVIDIA CUDA](gpu.md)
     - [AMD ROCm](gpu.md)
     - [Intel XPU](gpu.md)
+    - [MThreads MUSA](gpu.md)
 - [NPU](npu.md)

docs/getting_started/installation/gpu.md

@@ -22,6 +22,10 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr
 
     --8<-- "docs/getting_started/installation/gpu/xpu.inc.md:requirements"
 
+=== "MThreads MUSA"
+
+    --8<-- "docs/getting_started/installation/gpu/musa.inc.md:requirements"
+
 ## Set up using Python
 
 ### Create a new Python environment
@@ -44,6 +48,10 @@ Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc
 
     --8<-- "docs/getting_started/installation/gpu/xpu.inc.md:pre-built-wheels"
 
+=== "MThreads MUSA"
+
+    --8<-- "docs/getting_started/installation/gpu/musa.inc.md:pre-built-wheels"
+
 [](){ #build-from-source }
 
 ### Build wheel from source
@@ -60,6 +68,10 @@ Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc
 
     --8<-- "docs/getting_started/installation/gpu/xpu.inc.md:build-wheel-from-source"
 
+=== "MThreads MUSA"
+
+    --8<-- "docs/getting_started/installation/gpu/musa.inc.md:build-wheel-from-source"
+
 ## Set up using Docker
 
 ### Pre-built images
@@ -76,6 +88,10 @@ Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc
 
     --8<-- "docs/getting_started/installation/gpu/xpu.inc.md:pre-built-images"
 
+=== "MThreads MUSA"
+
+    --8<-- "docs/getting_started/installation/gpu/musa.inc.md:pre-built-images"
+
 ### Build your own docker image
 
 === "AMD ROCm"
@@ -85,3 +101,7 @@ Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc
 === "Intel XPU"
 
     --8<-- "docs/getting_started/installation/gpu/xpu.inc.md:build-docker"
+
+=== "MThreads MUSA"
+
+    --8<-- "docs/getting_started/installation/gpu/musa.inc.md:build-docker"

docs/getting_started/installation/gpu/musa.inc.md

@@ -0,0 +1,65 @@
+# --8<-- [start:requirements]
+
+- GPU: Moore Threads GPU with MUSA SDK installed (validated on MTT S5000)
+
+# --8<-- [end:requirements]
+# --8<-- [start:set-up-using-python]
+
+vLLM-Omni for MUSA requires building from source. Pre-built wheels are not currently available.
+
+!!! note
+    MUSA platform requires vLLM-MUSA to be installed first.
+
+# --8<-- [start:pre-built-wheels]
+
+# --8<-- [end:pre-built-wheels]
+
+# --8<-- [start:build-wheel-from-source]
+
+#### Prerequisites
+
+- **MUSA SDK**: Download from [MUSA SDK Download](https://developer.mthreads.com/sdk/download/musa)
+- **torchada**: CUDA→MUSA compatibility layer for PyTorch (`pip install torchada`)
+- **mthreads-ml-py**: MTML Python bindings (`pip install mthreads-ml-py`)
+- **MATE**: MUSA AI Tensor Engine ([GitHub](https://github.com/MooreThreads/mate))
+
+#### Installation of vLLM-MUSA
+
+```bash
+git clone https://github.com/MooreThreads/vllm-musa.git
+cd vllm-musa
+git checkout v0.18.0-dev
+pip install . --no-build-isolation -v
+```
+
+#### Installation of vLLM-Omni
+
+```bash
+git clone https://github.com/vllm-project/vllm-omni.git
+cd vllm-omni
+VLLM_OMNI_TARGET_DEVICE=musa pip install -e . --no-build-isolation
+```
+
+For Gradio demos:
+
+```bash
+pip install -e '.[demo]' --no-build-isolation
+```
+
+#### Environment Variables
+
+```bash
+export MUSA_VISIBLE_DEVICES=0,1
+export VLLM_WORKER_MULTIPROC_METHOD=spawn
+export VLLM_MUSA_CUSTOM_OP_USE_NATIVE=false
+```
+
+# --8<-- [end:build-wheel-from-source]
+
+# --8<-- [start:build-docker]
+
+# --8<-- [end:build-docker]
+
+# --8<-- [start:pre-built-images]
+
+# --8<-- [end:pre-built-images]

docs/models/supported_models.md

@@ -37,6 +37,7 @@ th {
 | `LTX2TwoStagesPipeline` | LTX-2-T2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |
 | `LTX2ImageToVideoTwoStagesPipeline` | LTX-2-I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |
 | `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | |
+| `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | |
 | `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | ✅︎ | ✅︎ | | ✅︎ |
 | `LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
@@ -59,5 +60,6 @@ th {
 | `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | ✅︎ | ✅︎ | | |
 | `HunyuanVideo15ImageToVideoPipeline` | HunyuanVideo-1.5-I2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v` | ✅︎ | ✅︎ | | |
 | `VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✅︎ | ✅︎ | | |
+|`DyninOmniForConditionalGeneration` | Dynin-Omni | `snu-aidas/Dynin-Omni` | ✅︎ | | | |
 
 ✅︎ indicates the model is supported on that backend. Empty cells mean not listed as supported on that backend.

docs/user_guide/diffusion/cpu_offload_diffusion.md

@@ -91,12 +91,19 @@ Models must define the blocks attribute name for layerwise offloading:
 
 ```python
 class WanTransformer3DModel(nn.Module):
-    _layerwise_offload_blocks_attr = "blocks"  # Attribute name containing transformer blocks
+    _layerwise_offload_blocks_attrs = ["blocks"]  # Attribute names containing transformer blocks
 
     def __init__(self):
         self.blocks = nn.ModuleList([...])  # Transformer blocks
 ```
 
+For models with multiple block types:
+
+```python
+class Flux2Transformer2DModel(nn.Module):
+    _layerwise_offload_blocks_attrs = ["transformer_blocks", "single_transformer_blocks"]
+```
+
 ### Limitations
 - Cold start latency increases because of
     1) components are loaded to CPU first at the very first during initialization,
@@ -132,12 +139,16 @@ Factory function `get_offload_backend()` selects the appropriate backend based o
 
 ## Supported Models
 
-| Architecture | Example Models | DiT Class | Model-Level Offload | Layerwise Offload | Blocks Attr (Layerwise specific) |
-|--------------|----------------|-----------|---------------------|-------------------|-------------|
-| Wan22Pipeline | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | `"blocks"` |
-| Wan22I2VPipeline | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | `"blocks"` |
+| Architecture | Example Models | DiT Class | Model-Level Offload | Layerwise Offload | Blocks Attrs (Layerwise specific) |
+|--------------|----------------|-----------|---------------------|-------------------|-----------------------------------|
+| LongCatImagePipeline | `meituan-longcat/LongCat-Image` | `LongCatImageTransformer2DModel` | - | ✓ | `"transformer_blocks"`, `"single_transformer_blocks"` |
+| NextStep11Pipeline | `stepfun-ai/NextStep-1.1` | `NextStepModel` | - | ✓ | `"layers"` |
+| OvisImagePipeline | `AIDC-AI/Ovis-Image-7B` | `OvisImageTransformer2DModel` | - | ✓ | `"transformer"` |
 | QwenImagePipeline | `Qwen/Qwen-Image` | `QwenImageTransformer2DModel` | ✓ | ✓ | `"transformer_blocks"` |
+| StableDiffusion3Pipeline | `stabilityai/stable-diffusion-3.5-medium` | `SD3Transformer2DModel` | - | ✓ | `"transformer_blocks"` |
+| Wan22I2VPipeline | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | `"blocks"` |
+| Wan22Pipeline | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | `"blocks"` |
 
 **Notes:**
 - Model-Level Offloading is expected to be supported by all common diffusion models (DiT and encoders) naturally
-- Layerwise Offloading requires DiT class to define `_layerwise_offload_blocks_attr` pointing to transformer blocks
+- Layerwise Offloading requires DiT class to define `_layerwise_offload_blocks_attrs` pointing to transformer blocks

docs/user_guide/diffusion_features.md

@@ -12,7 +12,7 @@
 
 vLLM-Omni supports various advanced features for diffusion models:
 
-- Acceleration: **cache methods**, **parallelism methods**
+- Acceleration: **cache methods**, **parallelism methods**, **startup optimizations**
 - Memory optimization: **cpu offloading**, **quantization**
 - Extensions: **LoRA inference**
 - Execution modes: **step execution**
@@ -44,6 +44,12 @@ Parallelism methods distribute computation across GPUs without quality loss (mat
 | **[HSDP](diffusion/parallelism/hsdp.md)** | Weight sharding via FSDP2, redistributed on-demand at runtime | Very large models (14B+) on limited VRAM, combinable with SP |
 | **[Expert Parallelism](diffusion/parallelism/expert_parallel.md)** | Shards MoE expert MLP blocks across devices | MoE diffusion models (e.g., HunyuanImage3.0) |
 
+#### Startup Optimization
+
+| Method | Description | Best For |
+|--------|-------------|----------|
+| **[Multi-Thread Weight Loading](#multi-thread-weight-loading)** | Loads safetensors shards in parallel using a thread pool | All diffusion models; reduces startup from minutes to seconds |
+
 **Note:** Some acceleration methods can be combined together for optimized performance. See [Feature Compatibility Table](#feature-compatibility) and [Feature Compatibility Tutorial](feature_compatibility.md) for detailed configuration examples.
 
 ### Memory Optimization
@@ -107,33 +113,34 @@ The following tables show which models support each feature:
 | **FLUX.2-dev** | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
 | **GLM-Image** | ❌ | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
 | **HunyuanImage3** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ |
-| **LongCat-Image** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
-| **LongCat-Image-Edit** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
+| **LongCat-Image** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
+| **LongCat-Image-Edit** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
+| **MagiHuman** | ❌ | ❌ | ❌ | ❓ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
 | **MammothModa2(T2I)** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
-| **Nextstep_1(T2I)** | ❓ | ❓ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
+| **Nextstep_1(T2I)** | ❓ | ❓ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
 | **OmniGen2** | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
-| **Ovis-Image** | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
-| **Qwen-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
-| **Qwen-Image-2512** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
-| **Qwen-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-| **Qwen-Image-Edit-2509** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-| **Qwen-Image-Layered** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-| **Stable-Diffusion3.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ |
-| **Z-Image** | ✅ | ✅ | ✅ | ❓ | ✅ (TP=2 only) | ✅ | ❌ | ✅ | ✅ | ❌ |
+| **Ovis-Image** | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
+| **Qwen-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ✅ |
+| **Qwen-Image-2512** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ✅ |
+| **Qwen-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
+| **Qwen-Image-Edit-2509** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ❌ | ❌ |
+| **Qwen-Image-Layered** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
+| **Stable-Diffusion3.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ (decode) | ❌ | ❌ |
+| **Z-Image** | ✅ | ✅ | ✅ | ❓ | ✅ (TP=2 only) | ✅ | ❌ | ✅ (decode) | ✅ | ❌ |
 
 > Notes:
 > 1. Nextstep_1(T2I) does not support cache acceleration methods such as TeaCache or Cache-DiT.
-> 2. `Tongyi-MAI/Z-Image-Turbo` is a distilled model with minimal NFEs; CFG-Parallel is not necessary.
+> 2. `Tongyi-MAI/Z-Image-Turbo` and `SII-GAIR/daVinci-MagiHuman-Base-1080p` are distilled models with minimal NFEs; CFG-Parallel is not necessary.
 
 ### VideoGen
 
 | Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
 |-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|:----------------:|
-| **Wan2.2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-| **Wan2.1-VACE** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+| **Wan2.2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (encode/decode) | ❌ | ❌ |
+| **Wan2.1-VACE** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
 | **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
 | **Helios** | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
-| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
+| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ❌ |
 | **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
 
 ### AudioGen
@@ -178,6 +185,59 @@ The following tables show which models support each feature:
     6. Step Execution is not compatible with cache backends (TeaCache, Cache-DiT) or LoRA.
 
 
+## Multi-Thread Weight Loading
+
+Large diffusion models can take several minutes to load weights at startup (e.g., ~3 min for Qwen-Image, ~5 min for Wan2.2 I2V 14B). Multi-thread weight loading speeds up this process by loading safetensors shards in parallel using a thread pool instead of sequentially.
+
+This optimization is **enabled by default** with 4 threads. No configuration is needed for the default behavior.
+
+### Configuration
+
+| Parameter | CLI Flag | Default | Description |
+|-----------|----------|---------|-------------|
+| `enable_multithread_weight_load` | `--disable-multithread-weight-load` | `True` (enabled) | Pass the flag to disable multi-thread loading |
+| `num_weight_load_threads` | `--num-weight-load-threads` | `4` | Number of threads for parallel weight loading |
+
+!!! tip
+    The default of 4 threads balances speed and disk I/O contention. On fast NVMe storage you may benefit from more threads (e.g., 8). On HDD or network storage, the default of 4 avoids saturating I/O bandwidth.
+
+### Online Serving
+
+```bash
+# Default (multi-thread enabled, 4 threads)
+vllm serve Qwen/Qwen-Image --omni --port 8091
+
+# Custom thread count
+vllm serve Wan-AI/Wan2.2-I2V-A14B-Diffusers --omni --num-weight-load-threads 8
+
+# Disable multi-thread loading
+vllm serve Qwen/Qwen-Image --omni --disable-multithread-weight-load
+```
+
+### Offline Inference
+
+```python
+from vllm_omni import Omni
+
+# Default (multi-thread enabled, 4 threads)
+omni = Omni(model="Qwen/Qwen-Image")
+
+# Custom thread count
+omni = Omni(
+    model="Wan-AI/Wan2.2-I2V-A14B-Diffusers",
+    num_weight_load_threads=8,
+)
+```
+
+### Benchmarks
+
+Measured on NVIDIA H800:
+
+| Model | Before | After | Speedup |
+|-------|--------|-------|---------|
+| **Qwen/Qwen-Image** (53.7 GiB) | 168s | 27s | **6.2x** |
+| **Wan-AI/Wan2.2-I2V-A14B-Diffusers** (64.5 GiB) | 283s | 56s | **5.1x** |
+
 ## Learn More
 
 **Cache Acceleration:**
@@ -203,6 +263,10 @@ The following tables show which models support each feature:
 
 - **[Step Execution Guide](diffusion/step_execution.md)** - Per-step denoise execution with mid-request abort support
 
+**Startup Optimization:**
+
+- **[Multi-Thread Weight Loading](#multi-thread-weight-loading)** - Speed up model startup by loading safetensors shards in parallel
+
 **Advanced Topics:**
 
 - **[Feature Compatibility](feature_compatibility.md)** - How to combine multiple features for maximum performance

examples/offline_inference/bagel/end2end.py

@@ -2,7 +2,10 @@ import argparse
 import os
 
 from vllm_omni.inputs.data import OmniPromptType
-from vllm_omni.model_executor.stage_input_processors.bagel import GEN_THINK_SYSTEM_PROMPT
+from vllm_omni.model_executor.stage_input_processors.bagel import (
+    GEN_THINK_SYSTEM_PROMPT,
+    VLM_THINK_SYSTEM_PROMPT,
+)
 
 
 def parse_args():
@@ -171,7 +174,10 @@ def main():
         elif args.modality == "img2text":
             if args.image_path:
                 loaded_image = Image.open(args.image_path).convert("RGB")
-                final_prompt_text = f"<|im_start|>user\n<|image_pad|>\n{p}<|im_end|>\n<|im_start|>assistant\n"
+                think_prefix = f"<|im_start|>system\n{VLM_THINK_SYSTEM_PROMPT}<|im_end|>\n" if args.think else ""
+                final_prompt_text = (
+                    f"{think_prefix}<|im_start|>user\n<|image_pad|>\n{p}<|im_end|>\n<|im_start|>assistant\n"
+                )
                 prompt_dict = {
                     "prompt": final_prompt_text,
                     "multi_modal_data": {"image": loaded_image},
@@ -179,7 +185,8 @@ def main():
                 }
                 formatted_prompts.append(prompt_dict)
         elif args.modality == "text2text":
-            final_prompt_text = f"<|im_start|>user\n{p}<|im_end|>\n<|im_start|>assistant\n"
+            think_prefix = f"<|im_start|>{VLM_THINK_SYSTEM_PROMPT}<|im_end|>" if args.think else ""
+            final_prompt_text = f"{think_prefix}<|im_start|>{p}<|im_end|><|im_start|>"
             prompt_dict = {"prompt": final_prompt_text, "modalities": ["text"]}
             formatted_prompts.append(prompt_dict)
         else:
@@ -217,15 +224,11 @@ def main():
     img_idx = 0
     for req_output in omni_outputs:
         if args.think:
-            text_output = getattr(req_output, "text", None) or getattr(req_output, "outputs", None)
-            if text_output:
-                if isinstance(text_output, list) and text_output:
-                    for out in text_output:
-                        txt = getattr(out, "text", str(out))
-                        if txt:
-                            print(f"[Think] {txt}")
-                elif isinstance(text_output, str):
-                    print(f"[Think] {text_output}")
+            ro = getattr(req_output, "request_output", None)
+            if ro and getattr(ro, "outputs", None):
+                txt = "".join(getattr(o, "text", "") or "" for o in ro.outputs)
+                if txt:
+                    print(txt)
 
         images = getattr(req_output, "images", None)
 

examples/offline_inference/dynin_omni/README.md

@@ -0,0 +1,110 @@
+# Dynin-Omni Offline End2End Example
+
+This folder contains a unified offline inference entrypoint:
+
+- `end2end.py`
+
+## 1. Environment Setup
+
+Run from repository root:
+
+```bash
+cd <REPO_ROOT>
+```
+
+If needed, install this repo in editable mode:
+
+```bash
+pip install -e .
+```
+
+## 2. Extra Dependencies (EMOVA)
+
+Install the following packages for EMOVA-related components:
+
+```bash
+pip install \
+  "phonemizer==3.3.0" \
+  "Unidecode==1.4.0" \
+  "hydra-core==1.3.2" \
+  "pytorch-lightning==1.1.0" \
+  "wget==3.2" \
+  "wrapt==2.1.1" \
+  "onnx==1.20.1" \
+  "frozendict==2.4.7" \
+  "inflect==7.5.0" \
+  "braceexpand==0.1.7" \
+  "webdataset==1.0.2" \
+  "torch-stft==0.1.4" \
+  "editdistance==0.8.1"
+```
+
+## 3. Hardware and VRAM Requirements
+
+This example uses a 3-stage pipeline on one GPU by default
+([`dynin_omni.yaml`](../../../vllm_omni/model_executor/stage_configs/dynin_omni.yaml)):
+
+- Stage-0 (`token2text`): `gpu_memory_utilization: 0.5`
+- Stage-1 (`token2image`): `gpu_memory_utilization: 0.1`
+- Stage-2 (`token2audio`): `gpu_memory_utilization: 0.1`
+
+### Requested GPU Memory Budget from `gpu_memory_utilization`
+
+| Stage | Utilization | A100 80GB | H200 141GB |
+| :-- | :-- | :-- | :-- |
+| Stage-0 (token2text) | 0.5 | ~40.0 GB | ~70.5 GB |
+| Stage-1 (token2image) | 0.1 | ~8.0 GB | ~14.1 GB |
+| Stage-2 (token2audio) | 0.1 | ~8.0 GB | ~14.1 GB |
+| Total requested budget | 0.7 | ~56.0 GB | ~98.7 GB |
+
+### Observed Runtime Signal (from your log)
+
+- Stage-0 reported: `Model loading took 15.12 GiB memory` (weights footprint signal).
+- Stages 1/2 can still add runtime memory depending on task path and backend allocations.
+- Keep extra headroom for CUDA/PyTorch overhead and temporary allocations.
+
+### GPU Compatibility
+
+- Confirmed target GPUs for this setup: **NVIDIA H200**, **NVIDIA A100**.
+- CI/e2e coverage in this repo also includes CUDA **L4** markers for Dynin tests.
+
+## 4. End2End Run Examples
+
+```bash
+# t2t
+python <REPO_ROOT>/examples/offline_inference/dynin_omni/end2end.py \
+  --task t2t --model snu-aidas/Dynin-Omni --text <INSTRUCTION_TEXT>
+
+# i2t
+python <REPO_ROOT>/examples/offline_inference/dynin_omni/end2end.py \
+  --task i2t --model snu-aidas/Dynin-Omni --image <IMAGE_PATH> --text "Please describe this image in detail."
+
+# s2t
+python <REPO_ROOT>/examples/offline_inference/dynin_omni/end2end.py \
+  --task s2t --model snu-aidas/Dynin-Omni --audio <AUDIO_PATH> --text "Transcribe the given audio."
+
+# t2i
+python <REPO_ROOT>/examples/offline_inference/dynin_omni/end2end.py \
+  --task t2i --model snu-aidas/Dynin-Omni --text <INSTRUCTION_TEXT>
+
+# v2t
+python <REPO_ROOT>/examples/offline_inference/dynin_omni/end2end.py \
+  --task v2t --model snu-aidas/Dynin-Omni --video <VIDEO_PATH> --text "Describe this video in detail."
+
+# i2i
+python <REPO_ROOT>/examples/offline_inference/dynin_omni/end2end.py \
+  --task i2i --model snu-aidas/Dynin-Omni --image <IMAGE_PATH> --text <INSTRUCTION_TEXT>
+
+# t2s
+python <REPO_ROOT>/examples/offline_inference/dynin_omni/end2end.py \
+  --task t2s --model snu-aidas/Dynin-Omni --text <INSTRUCTION_TEXT>
+```
+
+## 5. Notes
+
+- Outputs are saved under task-specific directories in `/tmp` by default.
+- You can override output path with `--output-dir`.
+- If you want to force local config resolution, pass `--dynin-config-path <PATH_TO_DYNIN_OMNI_YAML>`.
+- If you see the warning
+  `max_num_batched_tokens (32768) exceeds max_num_seqs * max_model_len (4096)`,
+  reduce `max_num_batched_tokens` in stage config (for example, `4096` in CI config).

examples/offline_inference/image_to_image/image_to_image.md

@@ -51,5 +51,6 @@ Key arguments:
 - `--vae-use-tiling`: enable VAE tiling for memory optimization.
 - `--cfg-parallel-size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](../../../docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
 - `--enable-cpu-offload`: enable CPU offloading for diffusion models.
+- `--strength`: **Z-Image only** - controls the denoising start timestep for I2I (default: 0.6). Range: [0.0, 1.0]. Lower values preserve more of the original image; higher values allow more creative changes.
 
 > ℹ️ If you encounter OOM errors, try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage.

examples/offline_inference/image_to_video/image_to_video.py

@@ -146,7 +146,7 @@ def parse_args() -> argparse.Namespace:
         "--audio-sample-rate",
         type=int,
         default=24000,
-        help="Sample rate for audio output when saved (default: 24000 for LTX2).",
+        help="Sample rate for audio output when saved (default: 24000).",
     )
     parser.add_argument(
         "--cache-backend",
@@ -471,15 +471,9 @@ def main():
 
     video_array = _ensure_frame_list(video_array)
 
-    use_ltx2_export = is_ltx2
-    encode_video = None
-    if use_ltx2_export:
-        try:
-            from diffusers.pipelines.ltx2.export_utils import encode_video
-        except ImportError:
-            encode_video = None
+    if audio is not None:
+        from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes
 
-    if use_ltx2_export and encode_video is not None:
         if isinstance(video_array, list):
             frames_np = np.stack(video_array, axis=0)
         elif isinstance(video_array, np.ndarray):
@@ -490,25 +484,24 @@ def main():
         if frames_np.ndim == 4 and frames_np.shape[-1] == 4:
             frames_np = frames_np[..., :3]
 
-        audio_out = None
-        if audio is not None:
-            if isinstance(audio, list):
-                audio = audio[0] if audio else None
-            if isinstance(audio, np.ndarray):
-                audio = torch.from_numpy(audio)
-            if isinstance(audio, torch.Tensor):
-                audio_out = audio
-                if audio_out.dim() > 1:
-                    audio_out = audio_out[0]
-                audio_out = audio_out.float().cpu()
+        frames_u8 = (np.clip(frames_np, 0.0, 1.0) * 255).round().clip(0, 255).astype("uint8")
 
-        encode_video(
-            frames_np,
-            fps=fps,
-            audio=audio_out,
-            audio_sample_rate=args.audio_sample_rate if audio_out is not None else None,
-            output_path=str(output_path),
+        audio_np = audio
+        if isinstance(audio_np, list):
+            audio_np = audio_np[0] if audio_np else None
+        if isinstance(audio_np, torch.Tensor):
+            audio_np = audio_np.detach().cpu().float().numpy()
+        if isinstance(audio_np, np.ndarray):
+            audio_np = np.squeeze(audio_np).astype(np.float32)
+
+        video_bytes = mux_video_audio_bytes(
+            frames_u8,
+            audio_np,
+            fps=float(fps),
+            audio_sample_rate=args.audio_sample_rate,
         )
+        with open(str(output_path), "wb") as f:
+            f.write(video_bytes)
     else:
         export_to_video(video_array, str(output_path), fps=fps)
     print(f"Saved generated video to {output_path}")

examples/offline_inference/magi_human/README.md

@@ -0,0 +1,72 @@
+# MagiHuman Generation
+
+MagiHuman is an advanced, omni-modality model that generates both high-quality video and lip-synced audio from a text prompt.
+
+Because MagiHuman is a very large model featuring a powerful DiT MoE backbone and a ~9B parameter T5Gemma text encoder, it natively supports **Tensor Parallelism (TP)** in vLLM-Omni to run efficiently across multi-GPU setups, reducing device memory bottlenecks.
+
+## Setup
+
+### Install MagiCompiler (recommended)
+
+MagiHuman relies on [MagiCompiler](https://github.com/SandAI-org/MagiCompiler) for custom-op registration used by the DiT attention kernels. While the pipeline can fall back to stub implementations, installing MagiCompiler is **strongly recommended** for correct behaviour.
+
+```bash
+# Clone the repo
+git clone https://github.com/SandAI-org/MagiCompiler.git
+cd MagiCompiler
+
+# System dependencies (optional, for FX graph visualization; Debian/Ubuntu)
+sudo apt update && sudo apt install -y graphviz
+
+# Python dependencies
+pip install -r requirements.txt
+
+# Install MagiCompiler
+pip install .          # end users (recommended)
+# pip install -e .     # developers (editable install)
+```
+
+### Hardware requirements
+
+Ensure your hardware has enough VRAM. For a standard node with 80GB GPUs, running with `--tensor-parallel-size 4` is recommended to shard both the MoE weights and the T5Gemma text encoder across 4 GPUs, reducing the per-GPU peak VRAM overhead significantly (by roughly ~13.5GB per GPU compared to single-device inference).
+
+Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) for further details on allocating memory.
+
+## Run Examples
+
+Get into the example folder:
+```bash
+cd examples/offline_inference/magi_human
+```
+
+### End-to-End Generation (Text to Video+Audio)
+
+Generate a video with synchronized speech natively generated by the model.
+
+```bash
+python end2end.py \
+    --model /proj-tango-pvc/users/zhipeng.wang/workspace/models/daVinci-MagiHuman \
+    --prompt "A young woman with long, wavy golden blonde hair..." \
+    --tensor-parallel-size 4 \
+    --output output_magihuman.mp4
+```
+
+## Common Parameters
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `--model` | *(Required)* | Local model path or HuggingFace ID |
+| `--prompt` | *(built-in demo prompt)* | Highly detailed text prompt dictating visual look and dialogue text |
+| `--tensor-parallel-size` | `4` | Tensor parallelism size (Number of GPUs) |
+| `--height` | `256` | Initial resolution height |
+| `--width` | `448` | Initial resolution width |
+| `--num-inference-steps` | `8` | Denoising steps |
+| `--seed` | `52` | Random seed |
+| `--output` | `output_magihuman.mp4` | Output video with audio path |
+
+## Example materials
+
+??? abstract "end2end.py"
+    ``````py
+    --8<-- "examples/offline_inference/magi_human/end2end.py"
+    ``````

examples/offline_inference/magi_human/end2end.py

@@ -0,0 +1,122 @@
+import argparse
+
+from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes
+from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="End-to-end inference script for MagiHuman.")
+    parser.add_argument("--model", type=str, required=True, help="Path or ID of the MagiHuman model.")
+    parser.add_argument(
+        "--prompt",
+        type=str,
+        default="",
+        help="Text prompt containing visual description, dialogue, and background sound.",
+    )
+    parser.add_argument(
+        "--tensor-parallel-size", "-tp", type=int, default=4, help="Tensor parallel size (number of GPUs)."
+    )
+    parser.add_argument(
+        "--output", type=str, default="output_magihuman.mp4", help="Path to save the generated mp4 file."
+    )
+    parser.add_argument("--height", type=int, default=256, help="Video height.")
+    parser.add_argument("--width", type=int, default=448, help="Video width.")
+    parser.add_argument("--num-inference-steps", type=int, default=8, help="Number of denoising steps.")
+    parser.add_argument("--seed", type=int, default=52, help="Random seed for generation.")
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+
+    print(f"Initializing MagiHuman pipeline with TP={args.tensor_parallel_size}...")
+    omni = Omni(
+        model=args.model,
+        init_timeout=1200,
+        tensor_parallel_size=args.tensor_parallel_size,
+        devices=list(range(args.tensor_parallel_size)),
+    )
+
+    prompt = args.prompt
+    if not prompt:
+        prompt = (
+            "A young woman with long, wavy golden blonde hair and bright blue eyes, "
+            "wearing a fitted ivory silk blouse with a delicate lace collar, sits "
+            "stationary in front of a softly lit, blurred warm-toned interior. Her "
+            "overall disposition is warm, composed, and gently confident. The camera "
+            "holds a static medium close-up, framing her from the shoulders up, "
+            "with shallow depth of field keeping her face in sharp focus. Soft "
+            "directional key light falls from the upper left, casting a gentle "
+            "highlight along her cheekbone and nose bridge. She draws a quiet breath, "
+            "the levator labii superiors relaxing as her lips part. She speaks in "
+            "clear, warm, unhurried American English: "
+            "\"The most beautiful things in life aren't things at all — "
+            "they're moments, feelings, and the people who make you feel truly alive.\" "
+            "Her jaw descends smoothly on each stressed syllable; the orbicularis oris "
+            "shapes each vowel with precision. A faint, genuine smile engages the "
+            "zygomaticus major, lifting her lip corners fractionally. Her brows rest "
+            "in a soft, neutral arch throughout. She maintains steady, forward-facing "
+            "eye contact. Head position remains level; no torso displacement occurs.\n\n"
+            "Dialogue:\n"
+            "<Young blonde woman, American English>: "
+            "\"The most beautiful things in life aren't things at all — "
+            "they're moments, feelings, and the people who make you feel truly alive.\"\n\n"
+            "Background Sound:\n"
+            "<Soft, warm indoor ambience with a faint distant piano melody>"
+        )
+
+    sampling_params = OmniDiffusionSamplingParams(
+        height=args.height,
+        width=args.width,
+        num_inference_steps=args.num_inference_steps,
+        seed=args.seed,
+        extra_args={
+            "seconds": 5,
+            "sr_height": 1080,
+            "sr_width": 1920,
+            "sr_num_inference_steps": 5,
+        },
+    )
+
+    print(f"Generating with prompt: {prompt[:80]}...")
+    outputs = omni.generate(
+        prompts=[prompt],
+        sampling_params_list=[sampling_params],
+    )
+
+    print(f"Generation complete. Output type: {type(outputs)}")
+    if outputs:
+        first = outputs[0]
+
+        if hasattr(first, "images") and first.images:
+            video_frames = first.images[0]
+            print(f"Video frames: shape={video_frames.shape}, dtype={video_frames.dtype}")
+
+            audio_waveform = None
+            mm = first.multimodal_output or {}
+            if mm:
+                audio_waveform = mm.get("audio")
+                if audio_waveform is not None:
+                    print(f"Audio waveform: shape={audio_waveform.shape}, dtype={audio_waveform.dtype}")
+
+            output_fps = float(mm.get("fps", 25))
+            output_sr = int(mm.get("audio_sample_rate", 24000))
+            print(f"Using fps={output_fps}, audio_sample_rate={output_sr} from model output")
+
+            video_bytes = mux_video_audio_bytes(
+                video_frames,
+                audio_waveform,
+                fps=output_fps,
+                audio_sample_rate=output_sr,
+            )
+            with open(args.output, "wb") as f:
+                f.write(video_bytes)
+            print(f"Saved MP4 ({len(video_bytes)} bytes) to {args.output}")
+        print("SUCCESS: MagiHuman pipeline generation completed.")
+    else:
+        print("WARNING: No outputs returned.")
+
+
+if __name__ == "__main__":
+    main()

examples/offline_inference/text_to_image/README.md

@@ -33,6 +33,7 @@ This folder provides several entrypoints for experimenting with text-to-image di
 | `black-forest-labs/FLUX.2-klein-4B` | 1024 x 1024 | 72.7 | 14.9 |
 | `black-forest-labs/FLUX.2-klein-9B` | 1024 x 1024 | 37.1 | 32.3 |
 | `black-forest-labs/FLUX.2-dev` | 1024 x 1024 | 65.7 | >80 (CPU offload required) |
+| `HunyuanImage-3.0` | 1024 x 1024 | 80.0 (TP≥3)  | 160 |
 
 !!! info
 *Peak VRAM:  based on basic single-card usage, batch size =1, without any acceleration/optimization features. FLUX.2-dev requires `--enable-cpu-offload` on a single 80 GiB GPU.
@@ -90,6 +91,8 @@ python text_to_image.py \
 | `--enable-cpu-offload` | flag | off | Enable CPU offloading for diffusion models |
 | `--lora-path` | str | — | Path to PEFT LoRA adapter folder |
 | `--lora-scale` | float | `1.0` | Scale factor for LoRA weights |
+| `--use-system-prompt` | str | `None` | System prompt preset: `en_unified`, `en_vanilla`, `en_recaption`, `en_think_recaption`, `dynamic`, `None`, or custom text. Recommended: `en_unified`. Only for HunyuanImage-3.0.|
+| `--system-prompt` | str | `None` | Custom system prompt text. Only used when `--use-system-prompt` is set to `custom`. Only for HunyuanImage-3.0.|
 
 **NextStep-1.1 specific arguments:**
 

examples/offline_inference/text_to_image/text_to_image.py

@@ -242,6 +242,19 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Enable logging of diffusion pipeline stats.",
     )
+    parser.add_argument(
+        "--use-system-prompt",
+        type=str,
+        default=None,
+        choices=["None", "dynamic", "en_vanilla", "en_recaption", "en_think_recaption", "en_unified", "custom"],
+        help="System prompt preset for generation. Recommended: en_unified.",
+    )
+    parser.add_argument(
+        "--system-prompt",
+        type=str,
+        default=None,
+        help=("Custom system prompt. Used when --use-system-prompt is custom. "),
+    )
     return parser.parse_args()
 
 
@@ -363,7 +376,7 @@ def main():
         f"vae_patch_parallel_size={args.vae_patch_parallel_size}, "
         f"enable_expert_parallel={args.enable_expert_parallel}."
     )
-    print(f"  CPU offload: {args.enable_cpu_offload}")
+    print(f"  CPU offload: {args.enable_cpu_offload}; CPU Layerwise Offload: {args.enable_layerwise_offload}")
     print(f"  Image size: {args.width}x{args.height}")
     if args.lora_path:
         print(f"  LoRA: scale={args.lora_scale}")
@@ -382,13 +395,13 @@ def main():
         )
 
     generation_start = time.perf_counter()
-
     extra_args = {
         "timesteps_shift": args.timesteps_shift,
         "cfg_schedule": args.cfg_schedule,
         "use_norm": args.use_norm,
+        "use_system_prompt": args.use_system_prompt,
+        "system_prompt": args.system_prompt,
     }
-
     if lora_request:
         extra_args["lora_request"] = lora_request
         extra_args["lora_scale"] = args.lora_scale

examples/offline_inference/text_to_video/text_to_video.py

@@ -160,7 +160,7 @@ def parse_args() -> argparse.Namespace:
         "--audio-sample-rate",
         type=int,
         default=24000,
-        help="Sample rate for audio output when saved (default: 24000 for LTX2).",
+        help="Sample rate for audio output when saved (default: 24000).",
     )
     parser.add_argument(
         "--vae-patch-parallel-size",
@@ -430,17 +430,8 @@ def main():
 
     video_array = _ensure_frame_list(video_array)
 
-    use_ltx2_export = False
-    if args.model and "ltx" in str(args.model).lower():
-        use_ltx2_export = True
     if audio is not None:
-        use_ltx2_export = True
-
-    if use_ltx2_export:
-        try:
-            from diffusers.pipelines.ltx2.export_utils import encode_video
-        except ImportError:
-            raise ImportError("diffusers is required for LTX2 encode_video.")
+        from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes
 
         if isinstance(video_array, list):
             frames_np = np.stack(video_array, axis=0)
@@ -449,28 +440,24 @@ def main():
         else:
             frames_np = np.asarray(video_array)
 
-        frames_u8 = (frames_np * 255).round().clip(0, 255).astype("uint8")
-        video_tensor = torch.from_numpy(frames_u8)
+        frames_u8 = (np.clip(frames_np, 0.0, 1.0) * 255).round().clip(0, 255).astype("uint8")
 
-        audio_out = None
-        if audio is not None:
-            if isinstance(audio, list):
-                audio = audio[0] if audio else None
-            if isinstance(audio, np.ndarray):
-                audio = torch.from_numpy(audio)
-            if isinstance(audio, torch.Tensor):
-                audio_out = audio
-                if audio_out.dim() > 1:
-                    audio_out = audio_out[0]
-                audio_out = audio_out.float().cpu()
+        audio_np = audio
+        if isinstance(audio_np, list):
+            audio_np = audio_np[0] if audio_np else None
+        if isinstance(audio_np, torch.Tensor):
+            audio_np = audio_np.detach().cpu().float().numpy()
+        if isinstance(audio_np, np.ndarray):
+            audio_np = np.squeeze(audio_np).astype(np.float32)
 
-        encode_video(
-            video_tensor,
-            fps=args.fps,
-            audio=audio_out,
-            audio_sample_rate=args.audio_sample_rate if audio_out is not None else None,
-            output_path=str(output_path),
+        video_bytes = mux_video_audio_bytes(
+            frames_u8,
+            audio_np,
+            fps=float(args.fps),
+            audio_sample_rate=args.audio_sample_rate,
         )
+        with open(str(output_path), "wb") as f:
+            f.write(video_bytes)
     else:
         export_to_video(video_array, str(output_path), fps=args.fps)
     print(f"Saved generated video to {output_path}")

examples/offline_inference/x_to_video_audio/download_dreamid_omni.py

@@ -82,7 +82,6 @@ def main(output_dir: str):
 
     data = {
         "_class_name": "DreamIDOmniPipeline",
-        "fusion": "DreamID-Omni/dreamid_omni.safetensors",
     }
 
     with open(os.path.join(output_dir, "model_index.json"), "w", encoding="utf-8") as f:
@@ -90,6 +89,12 @@ def main(output_dir: str):
 
     print(f"model_index.json created at {os.path.join(output_dir, 'model_index.json')}")
 
+    transformer_dir = os.path.join(output_dir, "transformer")
+    os.makedirs(transformer_dir, exist_ok=True)
+    with open(os.path.join(transformer_dir, "config.json"), "w", encoding="utf-8") as f:
+        json.dump({"fusion": "DreamID-Omni/dreamid_omni.safetensors"}, f)
+    print(f"transformer/config.json created at {os.path.join(transformer_dir, 'config.json')}")
+
     # now we download the dependency code
     download_dependency()
 

examples/offline_inference/x_to_video_audio/x_to_video_audio.md

@@ -24,7 +24,9 @@ dreamid_omni/
 │   ├── models_t5_umt5-xxl-enc-bf16.pth
 │   ├── Wan2.2_VAE.pth
 │
-├── model_index.json # create by download_dreamid_omni.py
+├── model_index.json
+└── transformer/
+    └── config.json   # create by download_dreamid_omni.py
 ```
 
 ### Run the Inference

examples/offline_inference/x_to_video_audio/x_to_video_audio.py

@@ -132,8 +132,8 @@ def main() -> None:
     if not outputs:
         raise RuntimeError("No output returned from DreamID-Omni.")
     output = outputs[0].request_output
-    generated_video = output[0].images[0][0]
-    generated_audio = output[0].images[0][1]
+    generated_video = output.images[0][0]
+    generated_audio = output.images[0][1]
     try:
         from dreamid_omni.utils.io_utils import save_video
     except Exception as e:

examples/online_serving/dynin_omni/README.md

@@ -0,0 +1,97 @@
+# Dynin-Omni Online Serving Example
+
+## Installation
+
+Please refer to [README.md](../../../README.md).
+
+## Launch the Server
+
+First, find the `transformers_modules` path:
+
+```bash
+python - <<'PY'
+from transformers.utils.hub import HF_MODULES_CACHE
+print(HF_MODULES_CACHE)
+PY
+```
+
+Then export it for both `PYTHONPATH` and `HF_MODULES_CACHE`:
+
+```bash
+export PYTHONPATH=<transformers_modules_path>:$PYTHONPATH
+export HF_MODULES_CACHE=<transformers_modules_path>
+```
+
+Run from repository root:
+
+```bash
+vllm-omni serve snu-aidas/Dynin-Omni \
+  --omni \
+  --port 8091 \
+  --stage-configs-path "$(pwd)/vllm_omni/model_executor/stage_configs/dynin_omni.yaml"
+```
+
+If `vllm-omni` is not in PATH, run:
+
+```bash
+PYTHONPATH="$(pwd)" python -m vllm_omni.entrypoints.cli.main serve snu-aidas/Dynin-Omni \
+  --omni \
+  --port 8091 \
+  --stage-configs-path "$(pwd)/vllm_omni/model_executor/stage_configs/dynin_omni.yaml"
+```
+
+Wait until the server logs show both `All stages initialized successfully` and
+`Application startup complete.` before sending requests.
+
+## Send Requests via Python Client
+
+Move to the example directory:
+
+```bash
+cd examples/online_serving/dynin_omni
+```
+
+### Text -> Image
+
+```bash
+python openai_chat_completion_client_for_multimodal_generation.py \
+  --query-type t2i \
+  --prompt "A realistic indoor living room with natural daylight."
+```
+
+### Image -> Image
+
+```bash
+python openai_chat_completion_client_for_multimodal_generation.py \
+  --query-type i2i \
+  --image-path ../../offline_inference/dynin_omni/data/image/sofa_under_water.jpg \
+  --prompt "Transform this surreal underwater setting into a realistic indoor living room while preserving the sofa layout."
+```
+
+### Text -> Speech
+
+```bash
+python openai_chat_completion_client_for_multimodal_generation.py \
+  --query-type t2s \
+  --prompt "Hello. This is Dynin-omni."
+```
+
+## CLI Arguments
+
+- `--query-type` (`t2i|t2s|i2i`)
+- `--model` (default: `snu-aidas/Dynin-Omni`)
+- `--host` / `--port` (OpenAI-compatible vLLM endpoint)
+- `--prompt` (custom text)
+- `--image-path` (required for `i2i`)
+- `--modalities` (optional output modalities override)
+- `--output-dir` (default: `/tmp/dynin_online_outputs`)
+
+## Notes
+
+- This client currently supports only `t2i`, `t2s`, and `i2i`.
+- `t2t` is intentionally not exposed in this online example.
+- This example intentionally uses the OpenAI-compatible chat completion endpoint.
+- Task routing for non-text outputs relies on Dynin task trigger tokens (`<|t2i|>`, `<|i2i|>`, `<|t2s|>`) injected by the client.
+- Outputs are saved under `/tmp/dynin_online_outputs` by default.
+- Dynin stage-0 warmup can take a while on first startup; do not send requests before startup completes.
+- Dynin itself can execute text-returning tasks such as `t2t`, `s2t`, `i2t`, and `v2t`, but this online serving example currently runs stage-0 in `generation` mode. In that path, the generation worker does not surface the final text as `output.text`, so OpenAI chat responses for those text-output tasks may complete internally but still return empty text.

examples/online_serving/dynin_omni/openai_chat_completion_client_for_multimodal_generation.py

@@ -0,0 +1,342 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import argparse
+import base64
+import json
+import mimetypes
+import os
+import time
+from pathlib import Path
+from typing import Any
+
+DEFAULT_MODEL = "snu-aidas/Dynin-Omni"
+DEFAULT_OUTPUT_DIR = "/tmp/dynin_online_outputs"
+
+QUERY_CHOICES = ("t2i", "t2s", "i2i")
+DEFAULT_PROMPT_BY_QUERY = {
+    "t2i": "A high quality detailed living room interior photo.",
+    "t2s": "Please read this sentence naturally: Hello from Dynin-Omni online serving.",
+    "i2i": "Transform this image into a realistic indoor living room while preserving layout.",
+}
+DEFAULT_MODALITIES_BY_QUERY = {
+    "t2i": ["image"],
+    "t2s": ["audio"],
+    "i2i": ["image"],
+}
+OFFLINE_PARITY_STAGE_COUNT = 3
+OFFLINE_PARITY_STAGE_SAMPLING = {
+    "max_tokens": 1,
+    "temperature": 0.0,
+    "top_p": 1.0,
+    "detokenize": False,
+}
+
+
+def _infer_mime_type(path: Path) -> str:
+    mime_type, _ = mimetypes.guess_type(str(path))
+    return mime_type or "application/octet-stream"
+
+
+def _encode_file_as_data_url(path: Path) -> str:
+    mime_type = _infer_mime_type(path)
+    raw = path.read_bytes()
+    encoded = base64.b64encode(raw).decode("utf-8")
+    return f"data:{mime_type};base64,{encoded}"
+
+
+def _to_image_url(path_or_url: str) -> str:
+    value = str(path_or_url)
+    if value.startswith(("http://", "https://", "data:image/")):
+        return value
+    path = Path(value).expanduser().resolve()
+    if not path.exists():
+        raise FileNotFoundError(f"Image file not found: {path}")
+    return _encode_file_as_data_url(path)
+
+
+def _build_user_content(query_type: str, prompt: str, image_path: str | None) -> list[dict[str, Any]]:
+    if query_type == "t2i":
+        return [{"type": "text", "text": f"<|t2i|> {prompt}"}]
+
+    if query_type == "t2s":
+        return [{"type": "text", "text": f"<|t2s|> {prompt}"}]
+
+    if query_type == "i2i":
+        if not image_path:
+            raise ValueError("--image-path is required for query type i2i")
+        return [
+            {"type": "text", "text": f"<|i2i|> {prompt}"},
+            {"type": "image_url", "image_url": {"url": _to_image_url(image_path)}},
+        ]
+
+    raise ValueError(f"Unsupported query_type: {query_type}")
+
+
+def _collect_text_from_content(content: Any) -> list[str]:
+    texts: list[str] = []
+    if isinstance(content, str):
+        stripped = content.strip()
+        if stripped:
+            texts.append(stripped)
+        return texts
+
+    if isinstance(content, dict):
+        for key in ("text", "content", "value", "output_text"):
+            text_value = content.get(key)
+            if isinstance(text_value, str) and text_value.strip():
+                texts.append(text_value.strip())
+        return texts
+
+    if isinstance(content, list):
+        for item in content:
+            texts.extend(_collect_text_from_content(item))
+        return texts
+
+    content_text = getattr(content, "text", None)
+    if isinstance(content_text, str) and content_text.strip():
+        texts.append(content_text.strip())
+    content_value = getattr(content, "content", None)
+    if isinstance(content_value, str) and content_value.strip():
+        texts.append(content_value.strip())
+    output_text = getattr(content, "output_text", None)
+    if isinstance(output_text, str) and output_text.strip():
+        texts.append(output_text.strip())
+    return texts
+
+
+def _extract_text_outputs(chat_completion: Any) -> list[str]:
+    texts: list[str] = []
+    for choice in getattr(chat_completion, "choices", []) or []:
+        message = getattr(choice, "message", None)
+        if message is None:
+            continue
+        content = getattr(message, "content", None)
+        texts.extend(_collect_text_from_content(content))
+        reasoning_content = getattr(message, "reasoning_content", None)
+        if isinstance(reasoning_content, str) and reasoning_content.strip():
+            texts.append(reasoning_content.strip())
+        choice_text = getattr(choice, "text", None)
+        if isinstance(choice_text, str) and choice_text.strip():
+            texts.append(choice_text.strip())
+    top_level_output_text = getattr(chat_completion, "output_text", None)
+    if isinstance(top_level_output_text, str) and top_level_output_text.strip():
+        texts.append(top_level_output_text.strip())
+    return texts
+
+
+def _extract_image_data_urls(chat_completion: Any) -> list[str]:
+    urls: list[str] = []
+    for choice in getattr(chat_completion, "choices", []) or []:
+        message = getattr(choice, "message", None)
+        if message is None:
+            continue
+        content = getattr(message, "content", None)
+        if not isinstance(content, list):
+            continue
+        for item in content:
+            if not isinstance(item, dict):
+                continue
+            if item.get("type") != "image_url":
+                continue
+            image_url = (item.get("image_url") or {}).get("url")
+            if isinstance(image_url, str) and image_url.startswith("data:image"):
+                urls.append(image_url)
+    return urls
+
+
+def _extract_audio_payloads(chat_completion: Any) -> list[bytes]:
+    payloads: list[bytes] = []
+    for choice in getattr(chat_completion, "choices", []) or []:
+        message = getattr(choice, "message", None)
+        if message is None:
+            continue
+        message_audio = getattr(message, "audio", None)
+        if message_audio is None:
+            continue
+        data_b64 = getattr(message_audio, "data", None)
+        if isinstance(data_b64, str) and data_b64:
+            try:
+                payloads.append(base64.b64decode(data_b64))
+            except Exception:
+                continue
+    return payloads
+
+
+def _decode_data_url(data_url: str) -> tuple[bytes, str]:
+    header, data = data_url.split(",", 1)
+    mime_type = "image/png"
+    if ";" in header and ":" in header:
+        mime_type = header.split(":", 1)[1].split(";", 1)[0]
+    return base64.b64decode(data), mime_type
+
+
+def _image_extension_from_mime(mime_type: str) -> str:
+    if mime_type == "image/jpeg":
+        return ".jpg"
+    if mime_type == "image/webp":
+        return ".webp"
+    if mime_type == "image/gif":
+        return ".gif"
+    return ".png"
+
+
+def _save_outputs(
+    *,
+    query_type: str,
+    chat_completion: Any,
+    output_dir: Path,
+) -> None:
+    output_dir.mkdir(parents=True, exist_ok=True)
+    stamp = time.strftime("%Y%m%d_%H%M%S")
+
+    text_outputs = _extract_text_outputs(chat_completion)
+    image_data_urls = _extract_image_data_urls(chat_completion)
+    audio_payloads = _extract_audio_payloads(chat_completion)
+
+    if text_outputs:
+        text_path = output_dir / f"{query_type}_{stamp}.txt"
+        text_path.write_text("\n\n".join(text_outputs) + "\n", encoding="utf-8")
+        print(f"[dynin-online] text saved: {text_path}")
+        print(text_outputs[0])
+
+    for idx, image_url in enumerate(image_data_urls):
+        image_bytes, mime_type = _decode_data_url(image_url)
+        ext = _image_extension_from_mime(mime_type)
+        image_path = output_dir / f"{query_type}_{stamp}_{idx}{ext}"
+        image_path.write_bytes(image_bytes)
+        print(f"[dynin-online] image saved: {image_path}")
+
+    for idx, audio_bytes in enumerate(audio_payloads):
+        audio_path = output_dir / f"{query_type}_{stamp}_{idx}.wav"
+        audio_path.write_bytes(audio_bytes)
+        print(f"[dynin-online] audio saved: {audio_path}")
+
+    if not text_outputs and not image_data_urls and not audio_payloads:
+        print("[dynin-online] no output extracted from response")
+        raw_path = output_dir / f"{query_type}_{stamp}_raw_response.json"
+        try:
+            if hasattr(chat_completion, "model_dump_json"):
+                serialized = chat_completion.model_dump_json(indent=2)
+            else:
+                if hasattr(chat_completion, "model_dump"):
+                    raw_payload: Any = chat_completion.model_dump(mode="json")
+                else:
+                    raw_payload = chat_completion
+                try:
+                    serialized = json.dumps(raw_payload, ensure_ascii=False, indent=2)
+                except Exception:
+                    serialized = json.dumps({"repr": repr(raw_payload)}, ensure_ascii=False, indent=2)
+            raw_path.write_text(serialized + "\n", encoding="utf-8")
+            print(f"[dynin-online] raw response saved: {raw_path}")
+        except Exception:
+            pass
+
+
+def _build_offline_parity_sampling_params_list() -> list[dict[str, Any]]:
+    return [dict(OFFLINE_PARITY_STAGE_SAMPLING) for _ in range(OFFLINE_PARITY_STAGE_COUNT)]
+
+
+def run_request(args: argparse.Namespace) -> None:
+    from openai import OpenAI
+
+    client = OpenAI(
+        api_key="EMPTY",
+        base_url=f"http://{args.host}:{args.port}/v1",
+    )
+    prompt = args.prompt.strip() if args.prompt else DEFAULT_PROMPT_BY_QUERY[args.query_type]
+    user_content = _build_user_content(
+        query_type=args.query_type,
+        prompt=prompt,
+        image_path=args.image_path,
+    )
+    if args.modalities:
+        modalities = [item.strip() for item in args.modalities.split(",") if item.strip()]
+    else:
+        modalities = DEFAULT_MODALITIES_BY_QUERY[args.query_type]
+
+    extra_body = {
+        "sampling_params_list": _build_offline_parity_sampling_params_list(),
+    }
+    chat_completion = client.chat.completions.create(
+        model=args.model,
+        messages=[{"role": "user", "content": user_content}],
+        modalities=modalities,
+        extra_body=extra_body,
+    )
+    _save_outputs(
+        query_type=args.query_type,
+        chat_completion=chat_completion,
+        output_dir=Path(args.output_dir).expanduser(),
+    )
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Dynin-Omni online chat completion client")
+    parser.add_argument(
+        "--query-type",
+        "-q",
+        type=str,
+        default="t2i",
+        choices=QUERY_CHOICES,
+        help="Dynin query type",
+    )
+    parser.add_argument(
+        "--model",
+        "-m",
+        type=str,
+        default=DEFAULT_MODEL,
+        help="Model name/path",
+    )
+    parser.add_argument(
+        "--host",
+        type=str,
+        default="localhost",
+        help="Host/IP of the vLLM Omni API server",
+    )
+    parser.add_argument(
+        "--port",
+        type=int,
+        default=8091,
+        help="Port of the vLLM Omni API server",
+    )
+    parser.add_argument(
+        "--prompt",
+        "-p",
+        type=str,
+        default="",
+        help="Custom prompt text",
+    )
+    parser.add_argument(
+        "--image-path",
+        "-i",
+        type=str,
+        default=None,
+        help="Image path/URL for i2i",
+    )
+    parser.add_argument(
+        "--modalities",
+        type=str,
+        default="",
+        help="Comma-separated output modalities override (e.g., text,image,audio)",
+    )
+    parser.add_argument(
+        "--output-dir",
+        "-o",
+        type=str,
+        default=DEFAULT_OUTPUT_DIR,
+        help="Directory to save outputs",
+    )
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    os.environ.setdefault("HF_HUB_DISABLE_XET", "1")
+    run_request(args)
+
+
+if __name__ == "__main__":
+    main()

examples/online_serving/image_to_image/README.md

@@ -314,6 +314,7 @@ count, use `size` and `n` rather than `height`, `width`, or
 | `seed`                   | int   | None    | Random seed (reproducible)            |
 | `negative_prompt`        | str   | None    | Negative prompt                       |
 | `num_outputs_per_prompt` | int   | 1       | Number of images to generate          |
+| `strength`               | float | 0.6     | **Z-Image only** - Denoising start timestep for I2I. Range: [0.0, 1.0]. Lower preserves more of original image. |
 | `layers`                 | int   | 4       | Number of layers (Qwen-Image-Layered) |
 | `resolution`             | int   | 640     | Resolution, 640 or 1024 (Qwen-Image-Layered) |
 

examples/online_serving/qwen3_tts/openai_speech_client.py

@@ -71,7 +71,7 @@ def run_tts_generation(args) -> None:
     payload = {
         "model": args.model,
         "input": args.text,
-        "speaker": args.speaker,
+        "voice": args.speaker,
         "response_format": args.response_format,
     }
 

examples/online_serving/text_to_image/README.md

@@ -231,6 +231,8 @@ count, use `size` and `n` rather than `height`, `width`, or
 | `seed`                   | int   | None    | Random seed (reproducible)     |
 | `negative_prompt`        | str   | None    | Negative prompt                |
 | `num_outputs_per_prompt` | int   | 1       | Number of images to generate   |
+| `use_system_prompt` | str | None | System prompt preset: `en_unified`, `en_vanilla`, `en_recaption`, `en_think_recaption`, `dynamic`, `None`, or custom text string. Only for HunyuanImage-3.0. |
+| `system_prompt` | str | None | Custom system prompt text. Only used when `use_system_prompt` is set to `custom`. Only for HunyuanImage-3.0. |
 
 ## Response Format
 

examples/online_serving/text_to_image/openai_chat_client.py

@@ -28,6 +28,8 @@ def generate_image(
     lora_name: str | None = None,
     lora_scale: float | None = None,
     lora_int_id: int | None = None,
+    use_system_prompt: str | None = None,
+    system_prompt: str | None = None,
 ) -> bytes | None:
     """Generate an image using the images generation API.
 
@@ -45,6 +47,8 @@ def generate_image(
         lora_name: LoRA name (optional, defaults to path stem)
         lora_scale: LoRA scale factor (default: 1.0)
         lora_int_id: LoRA integer ID (optional, derived from path if not provided)
+        use_system_prompt: System prompt for generation.
+        system_prompt: Custom system prompt.
 
     Returns:
         Image bytes or None if failed
@@ -70,7 +74,10 @@ def generate_image(
         payload["negative_prompt"] = negative_prompt
     if seed is not None:
         payload["seed"] = seed
-
+    if use_system_prompt is not None:
+        payload["use_system_prompt"] = use_system_prompt
+    if system_prompt is not None:
+        payload["system_prompt"] = system_prompt
     # Add LoRA if provided
     if lora_path:
         lora_body: dict = {
@@ -128,9 +135,21 @@ def main():
         default=None,
         help="LoRA integer id (cache key). If omitted, the server derives a stable id from lora_path.",
     )
-
+    parser.add_argument(
+        "--use-system-prompt",
+        type=str,
+        default=None,
+        help=(
+            "System prompt for generation. Use predefined types: 'en_unified', 'en_vanilla', 'en_recaption', 'en_think_recaption', 'dynamic', or 'None'; Or provide custom text string directly. Recommended en_unified. "
+        ),
+    )
+    parser.add_argument(
+        "--system-prompt",
+        type=str,
+        default=None,
+        help=("Custom system prompt. Used when --use-system-prompt is custom. "),
+    )
     args = parser.parse_args()
-
     print(f"Generating image for: {args.prompt}")
 
     image_bytes = generate_image(
@@ -146,6 +165,8 @@ def main():
         lora_name=args.lora_name,
         lora_scale=args.lora_scale if args.lora_path else None,
         lora_int_id=args.lora_int_id if args.lora_path else None,
+        use_system_prompt=args.use_system_prompt,
+        system_prompt=args.system_prompt,
     )
 
     if image_bytes:

recipes/Qwen/Qwen3-Omni.md

@@ -0,0 +1,90 @@
+# Qwen3-Omni for multimodal chat on 1x A100 80GB
+
+## Summary
+
+- Vendor: Qwen
+- Model: `Qwen/Qwen3-Omni-30B-A3B-Instruct`
+- Task: Multimodal chat with text, image, audio, or video input
+- Mode: Online serving with the OpenAI-compatible API
+- Maintainer: Community
+
+## When to use this recipe
+
+Use this recipe when you want a known-good starting point for serving
+`Qwen/Qwen3-Omni-30B-A3B-Instruct` with vLLM-Omni on a single 80 GB A100 and
+validate the deployment with the existing multimodal client examples in this
+repository.
+
+## References
+
+- Upstream or canonical docs:
+  [`docs/user_guide/examples/online_serving/qwen3_omni.md`](../../docs/user_guide/examples/online_serving/qwen3_omni.md)
+- Related example under `examples/`:
+  [`examples/online_serving/qwen3_omni/README.md`](../../examples/online_serving/qwen3_omni/README.md)
+- Related issue or discussion:
+  [RFC: add recipes folder](https://github.com/vllm-project/vllm-omni/issues/2645)
+
+## Hardware Support
+
+This recipe currently documents one tested-style reference configuration for
+CUDA GPU serving. Add more sections for other hardware as community validation
+lands.
+
+## GPU
+
+### 1x A100 80GB
+
+#### Environment
+
+- OS: Linux
+- Python: 3.10+
+- Driver / runtime: NVIDIA CUDA environment with an A100 80 GB GPU
+- vLLM version: Match the repository requirements for your checkout
+- vLLM-Omni version or commit: Use the commit you are deploying from
+
+#### Command
+
+Start the server from the repository root:
+
+```bash
+vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
+```
+
+To enable async chunking, use the bundled stage config:
+
+```bash
+vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct \
+  --omni \
+  --port 8091 \
+  --stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml
+```
+
+#### Verification
+
+Run one of the existing example clients after the server is ready:
+
+```bash
+python examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py \
+  --model Qwen/Qwen3-Omni-30B-A3B-Instruct \
+  --query-type use_image \
+  --port 8091 \
+  --host localhost
+```
+
+For a quick API smoke test, request text-only output:
+
+```bash
+curl http://localhost:8091/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+    "messages": [{"role": "user", "content": "Describe vLLM in brief."}],
+    "modalities": ["text"]
+  }'
+```
+
+#### Notes
+
+- Memory usage: Size depends on runtime options and output modalities; leave headroom for multimodal workloads.
+- Key flags: `--omni` is required; `--stage-configs-path` is optional for custom or async-chunk stage configs.
+- Known limitations: This starter recipe is intentionally narrow and focuses on the single-GPU online-serving path already documented in the repo examples.

recipes/README.md

@@ -0,0 +1,35 @@
+# Community Recipes
+
+This directory contains community-maintained recipes for answering a
+practical user question:
+
+> How do I run model X on hardware Y for task Z?
+
+Add recipes for this repository under this in-repo `recipes/` directory. To
+keep naming and layout consistent, organize recipes by model vendor in a way
+that is aligned with
+[`vllm-project/recipes`](https://github.com/vllm-project/recipes), but treat
+that external repository as a reference for structure rather than the place to
+add files for this repo. Use one Markdown file per model family by default.
+
+Example layout:
+
+```text
+recipes/
+  Qwen/
+    Qwen3-Omni.md
+    Qwen3-TTS.md
+  Tencent-Hunyuan/
+    HunyuanVideo.md
+```
+
+## Available Recipes
+
+- [`Qwen/Qwen3-Omni.md`](./Qwen/Qwen3-Omni.md): online serving recipe for
+  multimodal chat on `1x A100 80GB`
+
+Within a single recipe file, include different hardware support sections such
+as `GPU`, `ROCm`, and `NPU`, and add concrete tested configurations like
+`1x A100 80GB` or `2x L40S` inside those sections when applicable.
+
+See [TEMPLATE.md](./TEMPLATE.md) for the recommended format.

recipes/TEMPLATE.md

@@ -0,0 +1,82 @@
+# Recipe Title
+
+> Example: Qwen3-Omni for speech chat on 1x A100 80GB
+
+## Summary
+
+- Vendor:
+- Model:
+- Task:
+- Mode:
+- Maintainer:
+
+## When to use this recipe
+
+Briefly describe the concrete scenario this recipe covers.
+
+## References
+
+- Upstream or canonical docs:
+- Related example under `examples/`:
+- Related issue or discussion:
+
+## Hardware Support
+
+Add one section per platform, such as `GPU`, `ROCm`, or `NPU`. Under each
+platform section, document one or more tested hardware configurations.
+
+## GPU
+
+### 1x A100 80GB
+
+#### Environment
+
+- OS:
+- Python:
+- Driver / runtime:
+- vLLM version:
+- vLLM-Omni version or commit:
+
+#### Command
+
+```bash
+# Add the exact command(s) here
+```
+
+#### Verification
+
+```bash
+# Add a quick validation command or expected output here
+```
+
+#### Notes
+
+- Memory usage:
+- Key flags:
+- Known limitations:
+
+### 2x L40S
+
+Repeat the same structure for other hardware setups as needed.
+
+## ROCm
+
+### Example hardware configuration
+
+Repeat the same nested structure for ROCm setups as needed:
+
+- `#### Environment`
+- `#### Command`
+- `#### Verification`
+- `#### Notes`
+
+## NPU
+
+### Example hardware configuration
+
+Repeat the same nested structure for NPU setups as needed:
+
+- `#### Environment`
+- `#### Command`
+- `#### Verification`
+- `#### Notes`

requirements/common.txt

@@ -1,4 +1,5 @@
 # Common dependencies for all platforms
+av>=14.0.0
 omegaconf>=2.3.0
 librosa>=0.11.0
 resampy>=0.4.3

tests/benchmarks/metrics/test_metrics.py

@@ -0,0 +1,67 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+Unit tests for metrics.py
+"""
+
+import pytest
+from vllm.benchmarks.serve import TaskType
+
+from vllm_omni.benchmarks.metrics.metrics import calculate_metrics
+from vllm_omni.benchmarks.patch.patch import MixRequestFuncOutput
+
+pytestmark = [pytest.mark.core_model, pytest.mark.benchmark, pytest.mark.cpu]
+
+
+def _make_output(prompt_len: int, output_tokens: int = 10) -> MixRequestFuncOutput:
+    """Build a minimal successful MixRequestFuncOutput for metrics aggregation."""
+    output = MixRequestFuncOutput()
+    output.success = True
+    output.prompt_len = prompt_len
+    output.output_tokens = output_tokens
+    output.generated_text = "x" * output_tokens
+    output.ttft = 0.1
+    output.text_latency = 1.0
+    output.latency = 1.0
+    output.start_time = 0.0
+    output.itl = [0.1] * max(output_tokens - 1, 0)
+    output.audio_ttfp = 0.0
+    output.audio_rtf = 0.0
+    output.audio_duration = 0.0
+    output.audio_frames = 0
+    output.input_audio_duration = 0.0
+    output.error = ""
+    return output
+
+
+# ============================================================================
+# total_input Tests
+# ============================================================================
+
+
+def test_total_input_aggregated_from_output_prompt_len():
+    """Test that total_input sums outputs[i].prompt_len, not input_requests[i].prompt_len."""
+    outputs = [_make_output(4992), _make_output(3000)]
+
+    metrics, _ = calculate_metrics(
+        input_requests=[],
+        outputs=outputs,
+        dur_s=10.0,
+        tokenizer=None,
+        selected_percentiles=[99.0],
+        goodput_config_dict={},
+        task_type=TaskType.GENERATION,
+        selected_percentile_metrics=[],
+        max_concurrency=None,
+        request_rate=float("inf"),
+        benchmark_duration=10.0,
+    )
+
+    assert metrics.total_input == 7992, (
+        "total_input should aggregate from outputs[i].prompt_len to reflect the true multimodal input token count"
+    )
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v", "-s"])

tests/benchmarks/patch/test_patch.py

@@ -574,5 +574,59 @@ class TestTextLatencyAttribute:
         )
 
 
+# ============================================================================
+# prompt_len Tests
+# ============================================================================
+
+
+@pytest.mark.asyncio
+async def test_prompt_len_assigned_from_usage(mocker: MockerFixture):
+    # Arrange: request claims prompt_len=100, but server reports 4992 (multimodal).
+    request_input = RequestFuncInput(
+        model="test-model",
+        model_name="test-model",
+        prompt="test prompt",
+        api_url="http://test.com/v1/chat/completions",
+        prompt_len=100,
+        output_len=20,
+    )
+
+    chunks = [
+        create_sse_chunk(
+            {
+                "choices": [{"delta": {"content": "Hello"}}],
+                "modality": "text",
+            }
+        ),
+        create_sse_chunk(
+            {
+                "choices": [{"delta": {"content": " world"}}],
+                "modality": "text",
+            }
+        ),
+        # Final usage chunk emitted because stream_options.include_usage=True.
+        create_sse_chunk(
+            {
+                "choices": [],
+                "usage": {"prompt_tokens": 4992, "completion_tokens": 2, "total_tokens": 4994},
+            }
+        ),
+        b"data: [DONE]\n\n",
+    ]
+
+    mock_response = MockResponse(200, chunks)
+    mock_session = mocker.AsyncMock()
+    mock_session.post = mocker.MagicMock(return_value=mock_response)
+
+    # Act
+    output = await async_request_openai_chat_omni_completions(request_input, mock_session)
+
+    # Assert
+    assert output.success is True
+    assert output.prompt_len == 4992, (
+        "prompt_len should be overridden by usage.prompt_tokens to reflect the true multimodal input token count"
+    )
+
+
 if __name__ == "__main__":
     pytest.main([__file__, "-v", "-s"])

tests/conftest.py

@@ -16,6 +16,7 @@ if "VLLM_TARGET_DEVICE" not in os.environ:
     os.environ["VLLM_TARGET_DEVICE"] = "cpu"
 
 import concurrent.futures
+import contextlib
 import gc
 import multiprocessing
 import socket
@@ -52,6 +53,7 @@ from vllm_omni.platforms import current_omni_platform
 
 logger = init_logger(__name__)
 
+
 PromptAudioInput = list[tuple[Any, int]] | tuple[Any, int] | None
 PromptImageInput = list[Any] | Any | None
 PromptVideoInput = list[Any] | Any | None
@@ -337,10 +339,10 @@ def log_test_name_before_test(request):
 
 def _run_pre_test_cleanup(enable_force=False):
     if os.getenv("VLLM_TEST_CLEAN_GPU_MEMORY", "0") != "1" and not enable_force:
-        print("GPU cleanup disabled")
+        print("\nPre-test GPU cleanup skipped(Default off is typical when one worker/instance runs many tests.)\n")
         return
 
-    print("Pre-test GPU status:")
+    print("\nPre-test GPU status:")
 
     num_gpus = torch.cuda.device_count()
     if num_gpus > 0:
@@ -1087,6 +1089,22 @@ def _merge_base64_audio_to_segment(base64_list: list[str]):
     return merged
 
 
+@contextlib.contextmanager
+def _serialize_whisper_small_model_download():
+    """Serialize Whisper ``small`` cache writes across processes (Linux; ``fcntl``)."""
+    import fcntl
+
+    lock_path = Path.home() / ".cache" / "whisper" / ".small_model_download.lock"
+    lock_path.parent.mkdir(parents=True, exist_ok=True)
+    f = open(lock_path, "a+b")
+    try:
+        fcntl.flock(f.fileno(), fcntl.LOCK_EX)
+        yield
+    finally:
+        fcntl.flock(f.fileno(), fcntl.LOCK_UN)
+        f.close()
+
+
 def _whisper_transcribe_in_current_process(output_path: str) -> str:
     import whisper
 
@@ -1107,7 +1125,8 @@ def _whisper_transcribe_in_current_process(output_path: str) -> str:
     else:
         use_accelerator = False
         device = "cpu"
-    model = whisper.load_model("small", device=device)
+    with _serialize_whisper_small_model_download():
+        model = whisper.load_model("small", device=device)
     try:
         text = model.transcribe(
             output_path,
@@ -1126,8 +1145,7 @@ def _whisper_transcribe_in_current_process(output_path: str) -> str:
 
 
 def convert_audio_file_to_text(output_path: str) -> str:
-    """Convert an audio file to text in an isolated subprocess."""
-    # Import locally to avoid impacting test module import time.
+    """Convert an audio file to text in an isolated subprocess (spawn)."""
     ctx = multiprocessing.get_context("spawn")
     with concurrent.futures.ProcessPoolExecutor(max_workers=1, mp_context=ctx) as executor:
         future = executor.submit(_whisper_transcribe_in_current_process, output_path)

tests/core/sched/test_generation_scheduler_restore.py

@@ -9,6 +9,10 @@ those requests are permanently orphaned.
 import unittest
 from collections import deque
 
+import pytest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
 
 class FakeAdapter:
     """Minimal mock of OmniChunkTransferAdapter tracking restore calls."""

tests/dfx/perf/stage_configs/qwen3_tts.yaml

@@ -88,7 +88,7 @@ runtime:
         connector_get_max_wait_first_chunk: 3000
         connector_get_max_wait: 300
         codec_chunk_frames: 25
-        codec_left_context_frames: 25
+        codec_left_context_frames: 72
 
   edges:
     - from: 0

tests/dfx/perf/tests/test_ltx2_vllm_omni.json

@@ -0,0 +1,217 @@
+[
+    {
+        "test_name": "test_ltx2_baseline_eager",
+        "description": "Single-device baseline with enforce-eager (no torch.compile)",
+        "server_type": "vllm-omni",
+        "server_params": {
+            "model": "Lightricks/LTX-2",
+            "serve_args": {
+                "enforce-eager": true,
+                "enable-diffusion-pipeline-profiler": true
+            }
+        },
+        "benchmark_params": [
+            {
+                "name": "256x256_145f_steps6",
+                "dataset": "random",
+                "task": "t2v",
+                "backend": "v1/videos",
+                "width": 256,
+                "height": 256,
+                "num-frames": 145,
+                "fps": 24,
+                "num-inference-steps": 6,
+                "num-prompts": 3,
+                "max-concurrency": 1,
+                "enable-negative-prompt": true
+            },
+            {
+                "name": "480x768_41f_steps20",
+                "dataset": "random",
+                "task": "t2v",
+                "backend": "v1/videos",
+                "width": 768,
+                "height": 480,
+                "num-frames": 41,
+                "fps": 24,
+                "num-inference-steps": 20,
+                "num-prompts": 3,
+                "max-concurrency": 1,
+                "enable-negative-prompt": true
+            }
+        ]
+    },
+
+    {
+        "test_name": "test_ltx2_torch_compile",
+        "description": "Single-device with torch.compile (default, no enforce-eager)",
+        "server_type": "vllm-omni",
+        "server_params": {
+            "model": "Lightricks/LTX-2",
+            "serve_args": {
+                "enable-diffusion-pipeline-profiler": true
+            }
+        },
+        "benchmark_params": [
+            {
+                "name": "256x256_145f_steps6",
+                "dataset": "random",
+                "task": "t2v",
+                "backend": "v1/videos",
+                "width": 256,
+                "height": 256,
+                "num-frames": 145,
+                "fps": 24,
+                "num-inference-steps": 6,
+                "num-prompts": 3,
+                "max-concurrency": 1,
+                "enable-negative-prompt": true
+            },
+            {
+                "name": "480x768_41f_steps20",
+                "dataset": "random",
+                "task": "t2v",
+                "backend": "v1/videos",
+                "width": 768,
+                "height": 480,
+                "num-frames": 41,
+                "fps": 24,
+                "num-inference-steps": 20,
+                "num-prompts": 3,
+                "max-concurrency": 1,
+                "enable-negative-prompt": true
+            }
+        ]
+    },
+
+    {
+        "test_name": "test_ltx2_cfg2_eager",
+        "description": "CFG-parallel=2 with enforce-eager",
+        "server_type": "vllm-omni",
+        "server_params": {
+            "model": "Lightricks/LTX-2",
+            "serve_args": {
+                "cfg-parallel-size": 2,
+                "enforce-eager": true,
+                "enable-diffusion-pipeline-profiler": true
+            }
+        },
+        "benchmark_params": [
+            {
+                "name": "256x256_145f_steps6",
+                "dataset": "random",
+                "task": "t2v",
+                "backend": "v1/videos",
+                "width": 256,
+                "height": 256,
+                "num-frames": 145,
+                "fps": 24,
+                "num-inference-steps": 6,
+                "num-prompts": 3,
+                "max-concurrency": 1,
+                "enable-negative-prompt": true
+            },
+            {
+                "name": "480x768_41f_steps20",
+                "dataset": "random",
+                "task": "t2v",
+                "backend": "v1/videos",
+                "width": 768,
+                "height": 480,
+                "num-frames": 41,
+                "fps": 24,
+                "num-inference-steps": 20,
+                "num-prompts": 3,
+                "max-concurrency": 1,
+                "enable-negative-prompt": true
+            }
+        ]
+    },
+
+    {
+        "test_name": "test_ltx2_cfg2_compile",
+        "description": "CFG-parallel=2 with torch.compile",
+        "server_type": "vllm-omni",
+        "server_params": {
+            "model": "Lightricks/LTX-2",
+            "serve_args": {
+                "cfg-parallel-size": 2,
+                "enable-diffusion-pipeline-profiler": true
+            }
+        },
+        "benchmark_params": [
+            {
+                "name": "256x256_145f_steps6",
+                "dataset": "random",
+                "task": "t2v",
+                "backend": "v1/videos",
+                "width": 256,
+                "height": 256,
+                "num-frames": 145,
+                "fps": 24,
+                "num-inference-steps": 6,
+                "num-prompts": 3,
+                "max-concurrency": 1,
+                "enable-negative-prompt": true
+            },
+            {
+                "name": "480x768_41f_steps20",
+                "dataset": "random",
+                "task": "t2v",
+                "backend": "v1/videos",
+                "width": 768,
+                "height": 480,
+                "num-frames": 41,
+                "fps": 24,
+                "num-inference-steps": 20,
+                "num-prompts": 3,
+                "max-concurrency": 1,
+                "enable-negative-prompt": true
+            }
+        ]
+    },
+
+    {
+        "test_name": "test_ltx2_cache_dit_eager",
+        "description": "CacheDiT with enforce-eager",
+        "server_type": "vllm-omni",
+        "server_params": {
+            "model": "Lightricks/LTX-2",
+            "serve_args": {
+                "cache-backend": "cache_dit",
+                "enforce-eager": true,
+                "enable-diffusion-pipeline-profiler": true
+            }
+        },
+        "benchmark_params": [
+            {
+                "name": "256x256_145f_steps6",
+                "dataset": "random",
+                "task": "t2v",
+                "backend": "v1/videos",
+                "width": 256,
+                "height": 256,
+                "num-frames": 145,
+                "fps": 24,
+                "num-inference-steps": 6,
+                "num-prompts": 3,
+                "max-concurrency": 1,
+                "enable-negative-prompt": true
+            },
+            {
+                "name": "480x768_41f_steps20",
+                "dataset": "random",
+                "task": "t2v",
+                "backend": "v1/videos",
+                "width": 768,
+                "height": 480,
+                "num-frames": 41,
+                "fps": 24,
+                "num-inference-steps": 20,
+                "num-prompts": 3,
+                "max-concurrency": 1,
+                "enable-negative-prompt": true
+            }
+        ]
+    }
+]

tests/diffusion/cache/test_teacache_extractors.py

@@ -21,12 +21,13 @@ from unittest.mock import MagicMock, Mock, patch
 import pytest
 import torch
 
+from tests.utils import hardware_test
 from vllm_omni.diffusion.cache.teacache.extractors import extract_flux2_klein_context
 from vllm_omni.diffusion.models.flux2_klein.flux2_klein_transformer import (
     Flux2Transformer2DModel,
 )
 
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+pytestmark = [pytest.mark.core_model]
 
 
 @pytest.fixture(scope="function", autouse=True)
@@ -113,6 +114,7 @@ class TestFlux2KleinExtractor(BaseExtractorTest):
     def get_sample_inputs(self, sample_inputs):
         return sample_inputs
 
+    @hardware_test(res={"cuda": "L4"}, num_cards=1)
     def test_modulated_input_shape(self, flux2_klein_module, sample_inputs):
         """Test that modulated_input has correct shape matching the model's inner_dim.
 
@@ -126,16 +128,19 @@ class TestFlux2KleinExtractor(BaseExtractorTest):
         inner_dim = flux2_klein_module.inner_dim
         assert context.modulated_input.shape == (batch_size, img_seq_len, inner_dim)
 
+    @hardware_test(res={"cuda": "L4"}, num_cards=1)
     def test_run_transformer_blocks_callable(self, flux2_klein_module, sample_inputs):
         """Test that run_transformer_blocks is callable."""
         context = extract_flux2_klein_context(flux2_klein_module, **sample_inputs)
         assert callable(context.run_transformer_blocks)
 
+    @hardware_test(res={"cuda": "L4"}, num_cards=1)
     def test_postprocess_callable(self, flux2_klein_module, sample_inputs):
         """Test that postprocess is callable."""
         context = extract_flux2_klein_context(flux2_klein_module, **sample_inputs)
         assert callable(context.postprocess)
 
+    @hardware_test(res={"cuda": "L4"}, num_cards=1)
     def test_extra_states_contains_full_transformer(self, flux2_klein_module, sample_inputs):
         """Test that extra_states contains run_flux2_full_transformer_with_single."""
         context = extract_flux2_klein_context(flux2_klein_module, **sample_inputs)
@@ -154,6 +159,7 @@ class TestFlux2KleinExtractor(BaseExtractorTest):
         assert context is not None
         assert context.temb is not None
 
+    @pytest.mark.cpu
     def test_invalid_module_raises_error(self):
         """Test that invalid module without transformer_blocks raises ValueError."""
         invalid_module = Mock()

tests/diffusion/distributed/test_autoencoder_kl_wan.py

@@ -0,0 +1,43 @@
+import pytest
+import torch
+
+from vllm_omni.diffusion.distributed.autoencoders import autoencoder_kl_wan as wan_vae_module
+from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import OmniAutoencoderKLWan
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _DummyOmniAutoencoderKLWan(OmniAutoencoderKLWan):
+    def __init__(self, *, dtype: torch.dtype):
+        torch.nn.Module.__init__(self)
+        self.register_parameter("dummy_weight", torch.nn.Parameter(torch.ones(1, dtype=dtype)))
+
+
+def test_wan_vae_execution_context_handles_fp32():
+    model = _DummyOmniAutoencoderKLWan(dtype=torch.float32)
+    with model._execution_context():
+        output = model.dummy_weight + 1
+    assert output.dtype == torch.float32
+
+
+def test_wan_vae_execution_context_handles_bf16():
+    model = _DummyOmniAutoencoderKLWan(dtype=torch.bfloat16)
+    with model._execution_context():
+        output = model.dummy_weight + 1
+    assert output.dtype == torch.bfloat16
+
+
+def test_wan_vae_execution_context_uses_platform_autocast(mocker):
+    sentinel = object()
+    platform = mocker.Mock()
+    platform.create_autocast_context.return_value = sentinel
+    mocker.patch.object(wan_vae_module, "current_omni_platform", platform)
+
+    model = _DummyOmniAutoencoderKLWan(dtype=torch.bfloat16)
+
+    assert model._execution_context() is sentinel
+    platform.create_autocast_context.assert_called_once_with(
+        device_type=model.dummy_weight.device.type,
+        dtype=torch.bfloat16,
+        enabled=True,
+    )

tests/diffusion/distributed/test_autoencoder_kl_wan_encode.py

@@ -0,0 +1,273 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Unit tests for DistributedAutoencoderKLWan encode parallel (CPU-only)."""
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.cpu, pytest.mark.core_model]
+
+
+class _DummyConfig:
+    def __init__(self, patch_size=None, scale_factor_temporal=4):
+        self.patch_size = patch_size
+        self.scale_factor_temporal = scale_factor_temporal
+
+
+class _DummyWanVae:
+    """Minimal mock of DistributedAutoencoderKLWan for testing encode_tile_split."""
+
+    def __init__(
+        self,
+        config=None,
+        spatial_compression_ratio=8,
+        tile_sample_min_height=256,
+        tile_sample_min_width=256,
+        tile_sample_stride_height=192,
+        tile_sample_stride_width=192,
+    ):
+        self.config = config or _DummyConfig()
+        self.spatial_compression_ratio = spatial_compression_ratio
+        self.tile_sample_min_height = tile_sample_min_height
+        self.tile_sample_min_width = tile_sample_min_width
+        self.tile_sample_stride_height = tile_sample_stride_height
+        self.tile_sample_stride_width = tile_sample_stride_width
+        self.dtype = torch.float32
+
+        # Mock caches
+        self._enc_feat_map = None
+        self._enc_conv_idx = [0]
+
+    def clear_cache(self):
+        self._enc_feat_map = None
+        self._enc_conv_idx = [0]
+
+    def encoder(self, x, feat_cache=None, feat_idx=None):  # noqa: ARG002
+        # Simple mock: just return the input
+        return x
+
+    def quant_conv(self, x):
+        return x
+
+    def blend_v(self, _a, b, _blend_extent):
+        return b
+
+    def blend_h(self, _a, b, _blend_extent):
+        return b
+
+
+def _import_encode_tile_split():
+    """Import the encode_tile_split method from the module."""
+    from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import (
+        DistributedAutoencoderKLWan,
+    )
+
+    return DistributedAutoencoderKLWan.encode_tile_split
+
+
+def _import_encode_tile_exec():
+    from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import (
+        DistributedAutoencoderKLWan,
+    )
+
+    return DistributedAutoencoderKLWan.encode_tile_exec
+
+
+def _import_encode_tile_merge():
+    from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import (
+        DistributedAutoencoderKLWan,
+    )
+
+    return DistributedAutoencoderKLWan.encode_tile_merge
+
+
+class TestEncodeTileSplit:
+    """Tests for encode_tile_split method."""
+
+    def test_basic_split_without_patch_size(self):
+        """Test basic tile splitting without patch_size."""
+        encode_tile_split = _import_encode_tile_split()
+
+        vae = _DummyWanVae(
+            config=_DummyConfig(patch_size=None, scale_factor_temporal=4),
+            spatial_compression_ratio=8,
+            tile_sample_min_height=256,
+            tile_sample_min_width=256,
+            tile_sample_stride_height=192,
+            tile_sample_stride_width=192,
+        )
+
+        # Input: (B, C, T, H, W) = (1, 3, 5, 256, 256)
+        x = torch.randn(1, 3, 5, 256, 256)
+
+        tiletask_list, grid_spec = encode_tile_split(vae, x)
+
+        # With stride 192 and input size 256, we should get:
+        # Height: ceil(256/192) = 2 positions (0, 192) but 192+256 > 256, so only 1
+        # Actually for i in range(0, 256, 192): i = 0, 192 but 192 is out of bounds
+        # So we get 1x1 grid
+        assert len(tiletask_list) >= 1
+        assert grid_spec.grid_shape[0] >= 1
+        assert grid_spec.grid_shape[1] >= 1
+
+        # Check temporal chunking: 5 frames -> 1 + (5-1)//4 = 2 chunks
+        first_task = tiletask_list[0]
+        assert len(first_task.tensor) == 2  # 2 temporal chunks
+
+    def test_split_with_patch_size_scales_coordinates(self):
+        """Test that patch_size properly scales tile coordinates."""
+        encode_tile_split = _import_encode_tile_split()
+
+        # Without patch_size
+        vae_no_patch = _DummyWanVae(
+            config=_DummyConfig(patch_size=None, scale_factor_temporal=4),
+            spatial_compression_ratio=8,
+            tile_sample_min_height=256,
+            tile_sample_min_width=256,
+            tile_sample_stride_height=128,
+            tile_sample_stride_width=128,
+        )
+
+        # With patch_size=2 (simulating patchified input)
+        vae_with_patch = _DummyWanVae(
+            config=_DummyConfig(patch_size=2, scale_factor_temporal=4),
+            spatial_compression_ratio=8,
+            tile_sample_min_height=256,
+            tile_sample_min_width=256,
+            tile_sample_stride_height=128,
+            tile_sample_stride_width=128,
+        )
+
+        # Same patchified input size
+        x = torch.randn(1, 3, 5, 256, 256)
+
+        tasks_no_patch, _ = encode_tile_split(vae_no_patch, x)
+        tasks_with_patch, _ = encode_tile_split(vae_with_patch, x)
+
+        # With patch_size=2, stride becomes 128//2=64, so more tiles
+        assert len(tasks_with_patch) >= len(tasks_no_patch)
+
+    def test_temporal_compression_from_config(self):
+        """Test that temporal compression ratio is read from config."""
+        encode_tile_split = _import_encode_tile_split()
+
+        # temporal_compression=4 (default)
+        vae_4x = _DummyWanVae(
+            config=_DummyConfig(scale_factor_temporal=4),
+            tile_sample_min_height=512,
+            tile_sample_min_width=512,
+            tile_sample_stride_height=512,
+            tile_sample_stride_width=512,
+        )
+
+        # temporal_compression=2
+        vae_2x = _DummyWanVae(
+            config=_DummyConfig(scale_factor_temporal=2),
+            tile_sample_min_height=512,
+            tile_sample_min_width=512,
+            tile_sample_stride_height=512,
+            tile_sample_stride_width=512,
+        )
+
+        # 9 frames input
+        x = torch.randn(1, 3, 9, 512, 512)
+
+        tasks_4x, _ = encode_tile_split(vae_4x, x)
+        tasks_2x, _ = encode_tile_split(vae_2x, x)
+
+        # With 4x compression: 1 + (9-1)//4 = 3 chunks
+        assert len(tasks_4x[0].tensor) == 3
+
+        # With 2x compression: 1 + (9-1)//2 = 5 chunks
+        assert len(tasks_2x[0].tensor) == 5
+
+    def test_grid_spec_latent_dimensions(self):
+        """Test that grid_spec contains correct latent dimensions."""
+        encode_tile_split = _import_encode_tile_split()
+
+        vae = _DummyWanVae(
+            config=_DummyConfig(patch_size=None),
+            spatial_compression_ratio=8,
+            tile_sample_min_height=512,
+            tile_sample_min_width=512,
+            tile_sample_stride_height=512,
+            tile_sample_stride_width=512,
+        )
+
+        # Input: 512x512 with compression 8 -> 64x64 latent
+        x = torch.randn(1, 3, 5, 512, 512)
+
+        _, grid_spec = encode_tile_split(vae, x)
+
+        assert grid_spec.tile_spec["latent_height"] == 64
+        assert grid_spec.tile_spec["latent_width"] == 64
+
+
+class TestEncodeTileExec:
+    """Tests for encode_tile_exec method."""
+
+    def test_basic_exec(self):
+        """Test basic tile execution."""
+        encode_tile_exec = _import_encode_tile_exec()
+
+        vae = _DummyWanVae()
+
+        from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
+            TileTask,
+        )
+
+        # Create a simple task with 2 temporal chunks
+        tile1 = torch.randn(1, 3, 1, 32, 32)
+        tile2 = torch.randn(1, 3, 4, 32, 32)
+        task = TileTask(tile_id=0, grid_coord=(0, 0), tensor=[tile1, tile2])
+
+        result = encode_tile_exec(vae, task)
+
+        # Result should concatenate temporal dimension
+        assert result.shape[2] == 5  # 1 + 4 frames
+
+
+class TestEncodeTileMerge:
+    """Tests for encode_tile_merge method."""
+
+    def test_basic_merge(self):
+        """Test basic tile merging."""
+        encode_tile_merge = _import_encode_tile_merge()
+
+        vae = _DummyWanVae()
+
+        from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
+            GridSpec,
+        )
+
+        # Create 2x2 grid of tiles
+        tile_00 = torch.ones(1, 16, 2, 32, 32) * 0
+        tile_01 = torch.ones(1, 16, 2, 32, 32) * 1
+        tile_10 = torch.ones(1, 16, 2, 32, 32) * 2
+        tile_11 = torch.ones(1, 16, 2, 32, 32) * 3
+
+        coord_tensor_map = {
+            (0, 0): tile_00,
+            (0, 1): tile_01,
+            (1, 0): tile_10,
+            (1, 1): tile_11,
+        }
+
+        grid_spec = GridSpec(
+            split_dims=(3, 4),
+            grid_shape=(2, 2),
+            tile_spec={
+                "latent_height": 48,
+                "latent_width": 48,
+                "blend_height": 8,
+                "blend_width": 8,
+                "tile_latent_stride_height": 24,
+                "tile_latent_stride_width": 24,
+            },
+        )
+
+        result = encode_tile_merge(vae, coord_tensor_map, grid_spec)
+
+        # Output should be (1, 16, 2, 48, 48)
+        assert result.shape == (1, 16, 2, 48, 48)

tests/diffusion/distributed/test_distributed_vae_executor.py

@@ -11,6 +11,8 @@ from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor impor
     TileTask,
 )
 
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
 
 class E2EOperator:
     """tiles with (2, 3) -- (H,W)"""
@@ -59,9 +61,9 @@ class E2EOperator:
 class DummyMixin(DistributedVaeMixin):
     def __init__(self):
         self.use_tiling = True
-        self.distributed_decoder = MagicMock()
-        self.distributed_decoder.parallel_size = 2
-        self.distributed_decoder.group = None
+        self.distributed_executor = MagicMock()
+        self.distributed_executor.parallel_size = 2
+        self.distributed_executor.group = None
 
 
 @pytest.fixture(autouse=True)

tests/diffusion/distributed/test_ulysses_uaa_perf.py

@@ -17,6 +17,7 @@ import pytest
 import torch
 import torch.distributed as dist
 
+from tests.utils import hardware_test
 from vllm_omni.diffusion.attention.parallel.ulysses import (
     _all_gather_int,
     _ulysses_all_to_all_any_o,
@@ -69,6 +70,8 @@ PERF_CASES: list[_PerfCase] = [
 
 
 @pytest.mark.parametrize("case", PERF_CASES)
+@pytest.mark.core_model
+@hardware_test(res={"cuda": "L4"}, num_cards=4)
 def test_ulysses_advanced_uaa_comm_overhead(case: _PerfCase) -> None:
     available_gpus = current_omni_platform.get_device_count()
     if available_gpus < case.world_size:

tests/diffusion/lora/conftest.py

@@ -0,0 +1,56 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Shared test helpers for diffusion LoRA tests."""
+
+from __future__ import annotations
+
+import torch
+from vllm.model_executor.layers.linear import LinearBase
+
+
+class FakeLinearBase(LinearBase):
+    """Minimal LinearBase stub for LoRA layer discovery."""
+
+    def __init__(self):
+        torch.nn.Module.__init__(self)
+
+
+class DummyBaseLayerWithLoRA(torch.nn.Module):
+    """Fake LoRA wrapper that records set/reset/create calls."""
+
+    def __init__(self, base_layer: torch.nn.Module):
+        super().__init__()
+        self.base_layer = base_layer
+
+        self.set_calls: list[
+            tuple[list[torch.Tensor | None] | torch.Tensor, list[torch.Tensor | None] | torch.Tensor]
+        ] = []
+        self.reset_calls: int = 0
+        self.create_calls: int = 0
+
+    def set_lora(self, index: int, lora_a, lora_b):
+        assert index == 0
+        self.set_calls.append((lora_a, lora_b))
+
+    def reset_lora(self, index: int):
+        assert index == 0
+        self.reset_calls += 1
+
+    def create_lora_weights(self, max_loras, lora_config, model_config):
+        self.create_calls += 1
+
+
+def fake_replace_submodule(
+    root: torch.nn.Module,
+    module_name: str,
+    submodule: torch.nn.Module,
+    replace_calls: list[str] | None = None,
+) -> None:
+    """Replace a submodule by traversing dotted paths correctly."""
+    if replace_calls is not None:
+        replace_calls.append(module_name)
+    parts = module_name.split(".")
+    parent = root
+    for attr in parts[:-1]:
+        parent = getattr(parent, attr)
+    setattr(parent, parts[-1], submodule)

tests/diffusion/lora/test_lora_manager.py

@@ -7,8 +7,12 @@ import pytest
 import torch
 from vllm.lora.lora_weights import LoRALayerWeights
 from vllm.lora.utils import get_supported_lora_modules
-from vllm.model_executor.layers.linear import LinearBase
 
+from tests.diffusion.lora.conftest import (
+    DummyBaseLayerWithLoRA,
+    FakeLinearBase,
+    fake_replace_submodule,
+)
 from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager
 from vllm_omni.lora.request import LoRARequest
 
@@ -33,35 +37,9 @@ class _DummyLoRALayer:
         self.reset_calls += 1
 
 
-class _FakeLinearBase(LinearBase):
-    def __init__(self):
-        torch.nn.Module.__init__(self)
-
-
-class _DummyBaseLayerWithLoRA(torch.nn.Module):
-    def __init__(self, base_layer: torch.nn.Module):
-        super().__init__()
-        self.base_layer = base_layer
-
-        self.set_calls: list[
-            tuple[list[torch.Tensor | None] | torch.Tensor, list[torch.Tensor | None] | torch.Tensor]
-        ] = []
-        self.reset_calls: int = 0
-        self.create_calls: int = 0
-
-    def set_lora(self, index: int, lora_a, lora_b):
-        assert index == 0
-        self.set_calls.append((lora_a, lora_b))
-
-    def reset_lora(self, index: int):
-        assert index == 0
-        self.reset_calls += 1
-
-    def create_lora_weights(self, max_loras, lora_config, model_config):
-        # Needs to be callable for scale test when rank changes, but not
-        # actually used since we mock everything and check everything based
-        # on set calls.
-        self.create_calls += 1
+# Aliases for backward compatibility within this file
+_FakeLinearBase = FakeLinearBase
+_DummyBaseLayerWithLoRA = DummyBaseLayerWithLoRA
 
 
 class _DummyPipeline(torch.nn.Module):
@@ -555,3 +533,45 @@ def test_lora_manager_max_rank_validation(monkeypatch, rank):
     req1 = _dummy_lora_request(1)
     with pytest.raises(ValueError):
         manager.add_adapter(req1)
+
+
+def test_lora_manager_discovers_bagel_component(monkeypatch):
+    """Verify that _replace_layers_with_lora finds layers under 'bagel'."""
+    import vllm_omni.diffusion.lora.manager as manager_mod
+
+    monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", _DummyBaseLayerWithLoRA)
+
+    def _fake_from_layer_diffusion(*, layer: torch.nn.Module, **_kwargs):
+        if isinstance(layer, _FakeLinearBase):
+            return _DummyBaseLayerWithLoRA(layer)
+        return layer
+
+    replace_calls: list[str] = []
+
+    monkeypatch.setattr(manager_mod, "from_layer_diffusion", _fake_from_layer_diffusion)
+    monkeypatch.setattr(
+        manager_mod,
+        "replace_submodule",
+        lambda root, name, sub: fake_replace_submodule(root, name, sub, replace_calls),
+    )
+
+    # Pipeline with a 'bagel' component (no 'transformer')
+    pipeline = torch.nn.Module()
+    pipeline.bagel = torch.nn.Module()
+    pipeline.bagel.language_model = torch.nn.Module()
+    pipeline.bagel.language_model.qkv_proj = _FakeLinearBase()
+
+    manager = DiffusionLoRAManager(
+        pipeline=pipeline,
+        device=torch.device("cpu"),
+        dtype=torch.bfloat16,
+        max_cached_adapters=1,
+    )
+
+    peft_helper = type("_PH", (), {"r": 1})()
+    manager._replace_layers_with_lora(peft_helper)
+
+    assert "language_model.qkv_proj" in replace_calls
+    assert "bagel.language_model.qkv_proj" in manager._lora_modules
+    # Verify the module was actually replaced in the tree (not just recorded)
+    assert isinstance(pipeline.bagel.language_model.qkv_proj, _DummyBaseLayerWithLoRA)

tests/diffusion/models/bagel/test_bagel_lora.py

@@ -0,0 +1,248 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for BAGEL LoRA support across Stage 0 (Thinker) and Stage 1 (DiT)."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import pytest
+import torch
+from safetensors.torch import save_file
+
+from tests.diffusion.lora.conftest import (
+    DummyBaseLayerWithLoRA,
+    FakeLinearBase,
+    fake_replace_submodule,
+)
+from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager
+from vllm_omni.lora.request import LoRARequest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_FakeLinearBase = FakeLinearBase
+
+
+# ---------------------------------------------------------------------------
+# Stage 0 (Thinker / AR) -- packed_modules_mapping on the AR model class
+# ---------------------------------------------------------------------------
+
+
+class TestStage0ThinkerLoRA:
+    """Validate that OmniBagelForConditionalGeneration declares correct LoRA metadata."""
+
+    def test_omni_bagel_supports_lora(self):
+        from vllm_omni.model_executor.models.bagel.bagel import (
+            OmniBagelForConditionalGeneration,
+        )
+
+        assert getattr(OmniBagelForConditionalGeneration, "supports_lora", False) is True
+
+    def test_omni_bagel_packed_modules_mapping_complete(self):
+        from vllm_omni.model_executor.models.bagel.bagel import (
+            OmniBagelForConditionalGeneration,
+        )
+
+        mapping = OmniBagelForConditionalGeneration.packed_modules_mapping
+        # Standard Qwen2 projections
+        assert mapping["qkv_proj"] == ["q_proj", "k_proj", "v_proj"]
+        assert mapping["gate_up_proj"] == ["gate_proj", "up_proj"]
+        # MoE generation-mode projections
+        assert mapping["qkv_proj_moe_gen"] == [
+            "q_proj_moe_gen",
+            "k_proj_moe_gen",
+            "v_proj_moe_gen",
+        ]
+        assert mapping["mlp_moe_gen.gate_up_proj"] == [
+            "mlp_moe_gen.gate_proj",
+            "mlp_moe_gen.up_proj",
+        ]
+
+
+# ---------------------------------------------------------------------------
+# Stage 1 (DiT / Diffusion) -- DiffusionLoRAManager with bagel component
+# ---------------------------------------------------------------------------
+
+
+class TestStage1DiTLoRA:
+    """Validate DiffusionLoRAManager discovers BAGEL's packed modules."""
+
+    def test_diffusion_lora_manager_discovers_bagel_packed_modules(self):
+        """Manager should derive packed→sublayer mapping from stacked_params_mapping."""
+        pipeline = torch.nn.Module()
+        pipeline.bagel = torch.nn.Module()
+
+        # Simulate a submodule that exposes stacked_params_mapping
+        # (as Bagel does after load_weights())
+        language_model = torch.nn.Module()
+        language_model.stacked_params_mapping = [
+            (".qkv_proj_moe_gen", ".q_proj_moe_gen", "q"),
+            (".qkv_proj_moe_gen", ".k_proj_moe_gen", "k"),
+            (".qkv_proj_moe_gen", ".v_proj_moe_gen", "v"),
+            (".qkv_proj", ".q_proj", "q"),
+            (".qkv_proj", ".k_proj", "k"),
+            (".qkv_proj", ".v_proj", "v"),
+            (".gate_up_proj", ".gate_proj", 0),
+            (".gate_up_proj", ".up_proj", 1),
+        ]
+        pipeline.bagel.language_model = language_model
+
+        manager = DiffusionLoRAManager(
+            pipeline=pipeline,
+            device=torch.device("cpu"),
+            dtype=torch.bfloat16,
+            max_cached_adapters=1,
+        )
+
+        mapping = manager._packed_modules_mapping
+        assert mapping["qkv_proj"] == ["q_proj", "k_proj", "v_proj"]
+        assert mapping["qkv_proj_moe_gen"] == [
+            "q_proj_moe_gen",
+            "k_proj_moe_gen",
+            "v_proj_moe_gen",
+        ]
+        assert mapping["gate_up_proj"] == ["gate_proj", "up_proj"]
+
+    def test_diffusion_lora_manager_replaces_bagel_packed_layer_via_sublayer_target(self, monkeypatch):
+        """Targeting sublayer 'q_proj' should replace the fused 'qkv_proj' under bagel."""
+        import vllm_omni.diffusion.lora.manager as manager_mod
+
+        monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", DummyBaseLayerWithLoRA)
+
+        def _fake_from_layer_diffusion(*, layer, **_kwargs):
+            return DummyBaseLayerWithLoRA(layer)
+
+        replace_calls: list[str] = []
+
+        monkeypatch.setattr(manager_mod, "from_layer_diffusion", _fake_from_layer_diffusion)
+        monkeypatch.setattr(
+            manager_mod,
+            "replace_submodule",
+            lambda root, name, sub: fake_replace_submodule(root, name, sub, replace_calls),
+        )
+
+        # Build pipeline with bagel component
+        pipeline = torch.nn.Module()
+        pipeline.bagel = torch.nn.Module()
+        lm = torch.nn.Module()
+        lm.stacked_params_mapping = [
+            (".qkv_proj", ".q_proj", "q"),
+            (".qkv_proj", ".k_proj", "k"),
+            (".qkv_proj", ".v_proj", "v"),
+        ]
+        lm.attn = torch.nn.Module()
+        lm.attn.qkv_proj = _FakeLinearBase()
+        pipeline.bagel.language_model = lm
+
+        manager = DiffusionLoRAManager(
+            pipeline=pipeline,
+            device=torch.device("cpu"),
+            dtype=torch.bfloat16,
+            max_cached_adapters=1,
+        )
+
+        # Treat qkv_proj as 3-slice packed layer
+        monkeypatch.setattr(manager, "_get_packed_modules_list", lambda _module: ["q", "k", "v"])
+
+        # Target sublayer "q_proj" -- manager should replace the packed "qkv_proj"
+        peft_helper = type("_PH", (), {"r": 1, "target_modules": ["q_proj"]})()
+        manager._replace_layers_with_lora(peft_helper)
+
+        assert "language_model.attn.qkv_proj" in replace_calls
+        assert "bagel.language_model.attn.qkv_proj" in manager._lora_modules
+        # Verify the module was actually replaced in the tree (not just recorded)
+        assert isinstance(pipeline.bagel.language_model.attn.qkv_proj, DummyBaseLayerWithLoRA)
+
+
+# ---------------------------------------------------------------------------
+# Round-trip: synthetic checkpoint → set_active_adapter → verify weights
+# ---------------------------------------------------------------------------
+
+
+def _write_synthetic_lora(
+    adapter_dir: Path,
+    module_name: str,
+    rank: int,
+    in_dim: int,
+    out_dim: int,
+) -> str:
+    """Write a minimal LoRA adapter (safetensors + config) to *adapter_dir*."""
+    adapter_dir.mkdir(parents=True, exist_ok=True)
+    lora_a = torch.ones((rank, in_dim), dtype=torch.float32)
+    lora_b = torch.ones((out_dim, rank), dtype=torch.float32) * 2.0
+    save_file(
+        {
+            f"base_model.model.{module_name}.lora_A.weight": lora_a,
+            f"base_model.model.{module_name}.lora_B.weight": lora_b,
+        },
+        str(adapter_dir / "adapter_model.safetensors"),
+    )
+    (adapter_dir / "adapter_config.json").write_text(
+        json.dumps({"r": rank, "lora_alpha": rank, "target_modules": [module_name]}),
+        encoding="utf-8",
+    )
+    return str(adapter_dir)
+
+
+class TestBagelLoRARoundTrip:
+    """End-to-end: synthetic checkpoint → load → activate → verify weights in fused layer."""
+
+    def test_set_active_adapter_loads_and_activates_bagel_lora(self, tmp_path, monkeypatch):
+        """Full round-trip through set_active_adapter for a bagel component module."""
+        import vllm_omni.diffusion.lora.manager as manager_mod
+
+        monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", DummyBaseLayerWithLoRA)
+
+        # Build pipeline with bagel.language_model.foo (simple non-packed layer)
+        pipeline = torch.nn.Module()
+        pipeline.bagel = torch.nn.Module()
+        lm = torch.nn.Module()
+        lm.foo = _FakeLinearBase()
+        pipeline.bagel.language_model = lm
+
+        def _fake_from_layer(*, layer, **_kwargs):
+            if isinstance(layer, FakeLinearBase):
+                return DummyBaseLayerWithLoRA(layer)
+            return layer
+
+        monkeypatch.setattr(manager_mod, "from_layer_diffusion", _fake_from_layer)
+        monkeypatch.setattr(
+            manager_mod,
+            "replace_submodule",
+            lambda root, name, sub: fake_replace_submodule(root, name, sub),
+        )
+
+        manager = DiffusionLoRAManager(
+            pipeline=pipeline,
+            device=torch.device("cpu"),
+            dtype=torch.bfloat16,
+            max_cached_adapters=1,
+        )
+
+        # Write synthetic adapter targeting bagel.language_model.foo
+        module_name = "bagel.language_model.foo"
+        rank = 2
+        in_dim = 4
+        out_dim = 4
+        lora_dir = _write_synthetic_lora(tmp_path / "lora", module_name, rank, in_dim, out_dim)
+
+        lora_request = LoRARequest(
+            lora_name="test_bagel",
+            lora_int_id=42,
+            lora_path=lora_dir,
+        )
+
+        # Full round-trip: load from disk → replace layer → activate weights
+        manager.set_active_adapter(lora_request, lora_scale=0.5)
+
+        # Verify the layer was replaced and weights were set
+        replaced_layer = pipeline.bagel.language_model.foo
+        assert isinstance(replaced_layer, DummyBaseLayerWithLoRA), "Layer should be wrapped with LoRA"
+        assert len(replaced_layer.set_calls) == 1, "set_lora should have been called once"
+
+        lora_a, lora_b = replaced_layer.set_calls[0]
+        # A weights should be ones (as written)
+        assert torch.all(lora_a == 1.0), f"lora_a should be all ones, got {lora_a}"
+        # B weights should be 2.0 * scale(0.5) = 1.0
+        assert torch.allclose(lora_b, torch.ones_like(lora_b)), f"lora_b should be 2.0 * 0.5 = 1.0, got {lora_b}"

tests/diffusion/models/bagel/test_trajectory_recording.py

@@ -0,0 +1,236 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for BAGEL trajectory recording in the denoising loop."""
+
+import types
+from dataclasses import dataclass
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.bagel.bagel_transformer import (
+    Bagel,
+    NaiveCache,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+NUM_TOKENS = 8
+HIDDEN_DIM = 16
+NUM_TIMESTEPS = 5
+# generate_image uses timesteps[:-1], so actual steps = NUM_TIMESTEPS - 1
+EXPECTED_STEPS = NUM_TIMESTEPS - 1
+
+
+def _make_mock_bagel():
+    """Create a mock Bagel with forward returning constant velocity."""
+    mock = MagicMock(spec=Bagel)
+    mock._sp_size = 1
+
+    # forward returns a small constant velocity so x_t changes each step
+    def fake_forward(self, x_t, **kwargs):
+        return torch.ones_like(x_t) * 0.1
+
+    mock.forward = types.MethodType(fake_forward, mock)
+    # _merge_naive_caches is called in the batched CFG path
+    mock._merge_naive_caches = types.MethodType(lambda self, caches: NaiveCache(1), mock)
+
+    # Bind the real generate_image to our mock
+    mock.generate_image = types.MethodType(Bagel.generate_image, mock)
+    return mock
+
+
+def _make_generate_args(num_tokens=NUM_TOKENS, hidden_dim=HIDDEN_DIM, cfg=False):
+    """Tensor arguments for generate_image.
+
+    Args:
+        cfg: If True, enable batched CFG path (cfg_text_scale > 1.0).
+    """
+    seq_len = num_tokens + 2  # packed_seqlens includes 2 extra tokens
+    base = dict(
+        packed_text_ids=torch.zeros(2, dtype=torch.long),
+        packed_text_indexes=torch.tensor([0, 1], dtype=torch.long),
+        packed_init_noises=torch.randn(num_tokens, hidden_dim),
+        packed_vae_position_ids=torch.arange(num_tokens, dtype=torch.long),
+        packed_vae_token_indexes=torch.arange(2, seq_len, dtype=torch.long),
+        packed_seqlens=torch.tensor([seq_len], dtype=torch.int),
+        packed_position_ids=torch.arange(seq_len, dtype=torch.long),
+        packed_indexes=torch.arange(seq_len, dtype=torch.long),
+        past_key_values=NaiveCache(1),
+        key_values_lens=torch.tensor([0], dtype=torch.int),
+        packed_key_value_indexes=torch.zeros(0, dtype=torch.long),
+        num_timesteps=NUM_TIMESTEPS,
+        timestep_shift=1.0,
+        cfg_text_scale=1.0,
+        cfg_img_scale=1.0,
+    )
+    if cfg:
+        base |= dict(
+            cfg_text_scale=4.0,
+            cfg_text_packed_query_indexes=torch.arange(seq_len, dtype=torch.long),
+            cfg_text_packed_position_ids=torch.arange(seq_len, dtype=torch.long),
+            cfg_text_past_key_values=NaiveCache(1),
+            cfg_text_key_values_lens=torch.tensor([0], dtype=torch.int),
+            cfg_text_packed_key_value_indexes=torch.zeros(0, dtype=torch.long),
+        )
+    return base
+
+
+@pytest.fixture(params=[False, True], ids=["no_cfg", "batched_cfg"])
+def bagel_and_args(request):
+    """Mock Bagel instance and generate_image arguments.
+
+    Parametrized over CFG mode so every test runs on both the no-CFG
+    and batched-CFG code paths.
+    """
+    cfg = request.param
+    with patch(
+        "vllm_omni.diffusion.models.bagel.bagel_transformer.get_classifier_free_guidance_world_size",
+        return_value=1,
+    ):
+        yield _make_mock_bagel(), _make_generate_args(cfg=cfg)
+
+
+class TestTrajectoryRecording:
+    """Tests for trajectory latent/timestep recording in generate_image."""
+
+    def test_trajectory_disabled_returns_none(self, bagel_and_args):
+        bagel, args = bagel_and_args
+
+        unpacked, trajectory_latents, trajectory_timesteps, trajectory_log_probs = bagel.generate_image(
+            **args, return_trajectory_latents=False
+        )
+
+        assert isinstance(unpacked, (list, tuple))
+        assert len(unpacked) == 1  # one sequence
+        assert trajectory_latents is None
+        assert trajectory_timesteps is None
+        assert trajectory_log_probs is None
+
+    def test_trajectory_enabled_returns_correct_count(self, bagel_and_args):
+        bagel, args = bagel_and_args
+
+        _, trajectory_latents, trajectory_timesteps, trajectory_log_probs = bagel.generate_image(
+            **args, return_trajectory_latents=True
+        )
+
+        assert trajectory_latents is not None
+        assert trajectory_timesteps is not None
+        assert len(trajectory_latents) == EXPECTED_STEPS
+        assert len(trajectory_timesteps) == EXPECTED_STEPS
+        # log_probs is None without a scheduler (default ODE path)
+        assert trajectory_log_probs is None
+
+    def test_trajectory_latents_shape_matches_input(self, bagel_and_args):
+        bagel, args = bagel_and_args
+        expected_shape = args["packed_init_noises"].shape
+
+        _, trajectory_latents, *_ = bagel.generate_image(**args, return_trajectory_latents=True)
+
+        for i, lat in enumerate(trajectory_latents):
+            assert lat.shape == expected_shape, f"Step {i}: expected {expected_shape}, got {lat.shape}"
+
+    def test_trajectory_latents_are_distinct(self, bagel_and_args):
+        bagel, args = bagel_and_args
+
+        _, trajectory_latents, *_ = bagel.generate_image(**args, return_trajectory_latents=True)
+
+        for i in range(1, len(trajectory_latents)):
+            assert not torch.equal(trajectory_latents[i], trajectory_latents[i - 1]), (
+                f"Steps {i - 1} and {i} should differ"
+            )
+
+    def test_trajectory_timesteps_are_decreasing(self, bagel_and_args):
+        bagel, args = bagel_and_args
+
+        _, _, trajectory_timesteps, _ = bagel.generate_image(**args, return_trajectory_latents=True)
+
+        for i in range(1, len(trajectory_timesteps)):
+            assert trajectory_timesteps[i] < trajectory_timesteps[i - 1], (
+                f"Timestep {i} ({trajectory_timesteps[i]:.4f}) should be less than "
+                f"timestep {i - 1} ({trajectory_timesteps[i - 1]:.4f})"
+            )
+
+    def test_trajectory_final_latent_matches_output(self, bagel_and_args):
+        bagel, args = bagel_and_args
+
+        unpacked, trajectory_latents, *_ = bagel.generate_image(**args, return_trajectory_latents=True)
+
+        # Reconstruct the full final latent from unpacked pieces
+        final_latent = torch.cat(unpacked, dim=0)
+        assert torch.allclose(trajectory_latents[-1], final_latent, atol=1e-6), (
+            "Last trajectory latent should match the final output"
+        )
+
+
+# ---------------------------------------------------------------------------
+# Mock scheduler for log-prob tests
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class _MockStepOutput:
+    prev_sample: torch.Tensor
+    log_prob: torch.Tensor
+
+
+class _MockScheduler:
+    """Minimal scheduler: Euler step + constant log-prob per step."""
+
+    def step(self, model_output, sigma, sample, dt, **kwargs):
+        prev_sample = sample - model_output * dt
+        log_prob = torch.tensor(-1.0)
+        return _MockStepOutput(prev_sample=prev_sample, log_prob=log_prob)
+
+
+class TestTrajectoryLogProbs:
+    """Tests for log-prob recording when a scheduler is provided."""
+
+    @pytest.fixture()
+    def bagel_scheduler_args(self):
+        with patch(
+            "vllm_omni.diffusion.models.bagel.bagel_transformer.get_classifier_free_guidance_world_size",
+            return_value=1,
+        ):
+            yield _make_mock_bagel(), _make_generate_args(), _MockScheduler()
+
+    def test_log_probs_recorded_with_scheduler(self, bagel_scheduler_args):
+        bagel, args, scheduler = bagel_scheduler_args
+
+        _, _, _, trajectory_log_probs = bagel.generate_image(
+            **args, return_trajectory_latents=True, scheduler=scheduler
+        )
+
+        assert trajectory_log_probs is not None
+        assert len(trajectory_log_probs) == EXPECTED_STEPS
+
+    def test_log_probs_are_finite(self, bagel_scheduler_args):
+        bagel, args, scheduler = bagel_scheduler_args
+
+        _, _, _, trajectory_log_probs = bagel.generate_image(
+            **args, return_trajectory_latents=True, scheduler=scheduler
+        )
+
+        for i, lp in enumerate(trajectory_log_probs):
+            assert torch.isfinite(lp).all(), f"Step {i}: log_prob is not finite"
+
+    def test_log_probs_none_without_scheduler(self, bagel_scheduler_args):
+        bagel, args, _ = bagel_scheduler_args
+
+        _, _, _, trajectory_log_probs = bagel.generate_image(**args, return_trajectory_latents=True, scheduler=None)
+
+        assert trajectory_log_probs is None
+
+    def test_scheduler_updates_latents(self, bagel_scheduler_args):
+        """Verify the scheduler's prev_sample is used (not the raw Euler step)."""
+        bagel, args, scheduler = bagel_scheduler_args
+
+        _, traj_with_sched, *_ = bagel.generate_image(**args, return_trajectory_latents=True, scheduler=scheduler)
+        _, traj_without, *_ = bagel.generate_image(**args, return_trajectory_latents=True, scheduler=None)
+
+        # Mock scheduler does the same Euler step, so latents should match
+        for i in range(len(traj_with_sched)):
+            assert torch.allclose(traj_with_sched[i], traj_without[i], atol=1e-5), (
+                f"Step {i}: scheduler and ODE paths should produce same latents"
+            )

tests/diffusion/models/flux2/test_flux2_transformer_tp.py

@@ -3,6 +3,7 @@ from unittest.mock import MagicMock, patch
 import pytest
 import torch
 
+from tests.utils import hardware_test
 from vllm_omni.diffusion.models.flux2.flux2_transformer import (
     Flux2PosEmbed,
     Flux2Transformer2DModel,
@@ -24,6 +25,8 @@ def setup_tp_group():
 class TestFlux2TransformerWeightLoading:
     """Test Flux2Transformer weight loading functionality"""
 
+    @pytest.mark.core_model
+    @hardware_test(res={"cuda": "L4"}, num_cards=1)
     def test_weight_loading_tp2(self, setup_tp_group):
         """Verify weights load correctly with TP=2"""
         # Prepare test data
@@ -78,6 +81,8 @@ class TestFlux2TransformerWeightLoading:
 class TestFlux2RopePositionEmbedding:
     """Test Flux2 RoPE position embedding functionality"""
 
+    @pytest.mark.core_model
+    @pytest.mark.cpu
     def test_rope_position_embedding(self):
         """Verify RoPE produces correct embeddings for 4D coordinates"""
         # Prepare test data - use model default configuration
@@ -132,6 +137,8 @@ class TestFlux2RopePositionEmbedding:
 class TestFlux2PackedModuleMapping:
     """Test Flux2 packed module mapping functionality"""
 
+    @pytest.mark.core_model
+    @hardware_test(res={"cuda": "L4"}, num_cards=1)
     def test_packed_module_mapping(self, setup_tp_group):
         """Verify to_qkv packing matches HF checkpoint"""
         model = Flux2Transformer2DModel(
@@ -208,6 +215,8 @@ class TestFlux2PackedModuleMapping:
             f"add_kv_proj weight dimension should be {expected_add_kv_shape}, got {attn_block.add_kv_proj.weight.shape}"
         )
 
+    @pytest.mark.core_model
+    @hardware_test(res={"cuda": "L4"}, num_cards=1)
     def test_packed_mapping_edge_cases(self, setup_tp_group):
         """Test edge cases for packed mapping"""
         model = Flux2Transformer2DModel(

tests/diffusion/models/ltx2/test_ltx2_cfg_parallel_adaptation.py

@@ -0,0 +1,58 @@
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_pipeline(sequence_parallel_size: int = 1) -> LTX2Pipeline:
+    pipeline = object.__new__(LTX2Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.audio_vae_temporal_compression_ratio = 4
+    pipeline.audio_vae_mel_compression_ratio = 4
+    pipeline.od_config = SimpleNamespace(parallel_config=SimpleNamespace(sequence_parallel_size=sequence_parallel_size))
+    # Mock audio_vae with identity normalization (mean=0, std=1) so
+    # _normalize_audio_latents is a no-op and test values are preserved.
+    pipeline.audio_vae = SimpleNamespace(
+        latents_mean=torch.tensor(0.0),
+        latents_std=torch.tensor(1.0),
+    )
+    return pipeline
+
+
+def test_prepare_audio_latents_pads_packed_sequence_dim_for_provided_latents():
+    pipeline = _make_pipeline(sequence_parallel_size=4)
+    latents = torch.arange(40, dtype=torch.float32).view(1, 10, 4)
+
+    padded, original_num_frames, padded_num_frames = pipeline.prepare_audio_latents(
+        batch_size=1,
+        num_channels_latents=2,
+        num_mel_bins=8,
+        audio_latent_length=10,
+        dtype=torch.float32,
+        device=torch.device("cpu"),
+        latents=latents,
+    )
+
+    assert original_num_frames == 10
+    assert padded_num_frames == 12
+    assert padded.shape == (1, 12, 4)
+    torch.testing.assert_close(padded[:, :10], latents)
+    torch.testing.assert_close(padded[:, 10:], torch.zeros(1, 2, 4))
+
+
+def test_unpad_audio_latents_restores_original_frames_before_unpack():
+    pipeline = _make_pipeline()
+    original = torch.arange(40, dtype=torch.float32).view(1, 10, 4)
+    padded = torch.cat([original, torch.full((1, 2, 4), 999.0)], dim=1)
+
+    unpadded = pipeline._unpad_audio_latents(padded, 10)
+    unpacked = pipeline._unpack_audio_latents(unpadded, latent_length=10, num_mel_bins=2)
+    expected = pipeline._unpack_audio_latents(original, latent_length=10, num_mel_bins=2)
+
+    assert unpacked.shape == (1, 2, 10, 2)
+    assert not (unpacked == 999.0).any()
+    torch.testing.assert_close(unpacked, expected)

tests/diffusion/offloader/test_layerwise_backend.py

@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-"""Unit tests for LayerwiseOffloadHook."""
+"""Unit tests for LayerwiseOffloadHook and LayerWiseOffloadBackend utilities."""
 
 import gc
 import os
@@ -15,7 +15,7 @@ from torch import nn
 from torch.distributed.tensor import DeviceMesh, DTensor, Replicate
 
 import vllm_omni.diffusion.offloader.layerwise_backend as layerwise_backend_module
-from vllm_omni.diffusion.offloader.layerwise_backend import LayerwiseOffloadHook
+from vllm_omni.diffusion.offloader.layerwise_backend import LayerWiseOffloadBackend, LayerwiseOffloadHook
 from vllm_omni.platforms import current_omni_platform
 
 pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
@@ -127,3 +127,116 @@ class TestLayerwiseOffloadHook:
         assert current_block.weight.to_local().is_meta
         assert current_block.weight.to_local().shape == torch.Size([4])
         assert not hook.is_materialized
+
+
+class _DummyBlock(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.weight = nn.Parameter(torch.randn(10, 10))
+
+
+class _SingleBlockModel(nn.Module):
+    _layerwise_offload_blocks_attrs = ["blocks"]
+
+    def __init__(self, num_blocks: int = 3):
+        super().__init__()
+        self.blocks = nn.ModuleList([_DummyBlock() for _ in range(num_blocks)])
+
+
+class _MultiBlockModel(nn.Module):
+    _layerwise_offload_blocks_attrs = ["transformer_blocks", "single_transformer_blocks"]
+
+    def __init__(self, num_transformer: int = 2, num_single: int = 2):
+        super().__init__()
+        self.transformer_blocks = nn.ModuleList([_DummyBlock() for _ in range(num_transformer)])
+        self.single_transformer_blocks = nn.ModuleList([_DummyBlock() for _ in range(num_single)])
+
+
+class _EmptyBlocksModel(nn.Module):
+    _layerwise_offload_blocks_attrs = ["blocks"]
+
+    def __init__(self):
+        super().__init__()
+        self.blocks = nn.ModuleList([])
+
+
+class _InvalidAttrModel(nn.Module):
+    _layerwise_offload_blocks_attrs = ["nonexistent_blocks", "blocks"]
+
+    def __init__(self, num_blocks: int = 2):
+        super().__init__()
+        self.blocks = nn.ModuleList([_DummyBlock() for _ in range(num_blocks)])
+
+
+class _DeprecatedSingleAttrModel(nn.Module):
+    _layerwise_offload_blocks_attr = "blocks"
+
+    def __init__(self, num_blocks: int = 2):
+        super().__init__()
+        self.blocks = nn.ModuleList([_DummyBlock() for _ in range(num_blocks)])
+
+
+class _NoAttrsModel(nn.Module):
+    def __init__(self, num_blocks: int = 2):
+        super().__init__()
+        self.blocks = nn.ModuleList([_DummyBlock() for _ in range(num_blocks)])
+
+
+class TestGetBlocksFromDit:
+    def test_get_blocks_from_dit_single_block_attr(self):
+        model = _SingleBlockModel(num_blocks=3)
+        attr_names, blocks = LayerWiseOffloadBackend.get_blocks_from_dit(model)
+        assert attr_names == ["blocks"]
+        assert len(blocks) == 3
+        assert all(isinstance(b, _DummyBlock) for b in blocks)
+
+    def test_get_blocks_from_dit_multi_block_attrs(self):
+        model = _MultiBlockModel(num_transformer=2, num_single=3)
+        attr_names, blocks = LayerWiseOffloadBackend.get_blocks_from_dit(model)
+        assert set(attr_names) == {"transformer_blocks", "single_transformer_blocks"}
+        assert len(blocks) == 5
+        assert all(isinstance(b, _DummyBlock) for b in blocks)
+
+    def test_get_blocks_from_dit_empty_blocks(self):
+        model = _EmptyBlocksModel()
+        attr_names, blocks = LayerWiseOffloadBackend.get_blocks_from_dit(model)
+        assert attr_names == []
+        assert blocks == []
+
+    def test_get_blocks_from_dit_invalid_attr_name(self):
+        model = _InvalidAttrModel(num_blocks=2)
+        with pytest.raises(
+            AttributeError,
+            match="Attribute 'nonexistent_blocks' declared in _layerwise_offload_blocks_attrs does not exist",
+        ):
+            LayerWiseOffloadBackend.get_blocks_from_dit(model)
+
+    def test_get_blocks_from_dit_no_attrs_defined(self):
+        model = _NoAttrsModel(num_blocks=3)
+        attr_names, blocks = LayerWiseOffloadBackend.get_blocks_from_dit(model)
+        assert attr_names == []
+        assert blocks == []
+
+    def test_get_blocks_from_dit_deprecated_single_attr(self):
+        model = _DeprecatedSingleAttrModel(num_blocks=2)
+        attr_names, blocks = LayerWiseOffloadBackend.get_blocks_from_dit(model)
+        assert attr_names == ["blocks"]
+        assert len(blocks) == 2
+
+
+class TestGetBlocksAttrNames:
+    def test_get_blocks_attr_names_new_format(self):
+        model = _MultiBlockModel()
+        attrs = LayerWiseOffloadBackend.get_blocks_attr_names(model)
+        assert attrs == ["transformer_blocks", "single_transformer_blocks"]
+
+    def test_get_blocks_attr_names_no_attrs(self):
+        model = _NoAttrsModel()
+        attrs = LayerWiseOffloadBackend.get_blocks_attr_names(model)
+        assert attrs == []
+
+    def test_set_blocks_attr_names(self):
+        model = _NoAttrsModel()
+        LayerWiseOffloadBackend.set_blocks_attr_names(model, ["new_blocks"])
+        assert hasattr(model.__class__, "_layerwise_offload_blocks_attrs")
+        assert model.__class__._layerwise_offload_blocks_attrs == ["new_blocks"]

tests/diffusion/quantization/test_fp8_config.py

@@ -5,7 +5,7 @@
 import pytest
 from torch import nn
 
-pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
 
 
 def test_build_quant_config_fp8():

tests/diffusion/test_diffusion_model_runner.py

@@ -8,9 +8,10 @@ import pytest
 import torch
 
 import vllm_omni.diffusion.worker.diffusion_model_runner as model_runner_module
+from tests.utils import hardware_test
 from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
 
-pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+pytestmark = [pytest.mark.diffusion]
 
 
 @contextmanager
@@ -64,6 +65,8 @@ def _make_runner(cache_backend, cache_backend_name: str, enable_cache_dit_summar
     return runner
 
 
+@pytest.mark.core_model
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
 def test_execute_model_skips_cache_summary_without_active_cache_backend(monkeypatch):
     """Guard cache diagnostics with runtime backend state to avoid stale-config crashes."""
     runner = _make_runner(cache_backend=None, cache_backend_name="cache_dit")
@@ -84,6 +87,8 @@ def test_execute_model_skips_cache_summary_without_active_cache_backend(monkeypa
     assert cache_summary_calls == []
 
 
+@pytest.mark.core_model
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
 def test_execute_model_emits_cache_summary_with_active_cache_dit_backend(monkeypatch):
     class _EnabledCacheBackend:
         def is_enabled(self):
@@ -107,6 +112,8 @@ def test_execute_model_emits_cache_summary_with_active_cache_dit_backend(monkeyp
     assert cache_summary_calls == [(runner.pipeline, True)]
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
 def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch):
     class _DummyLoader:
         def __init__(self, load_config, od_config=None):

tests/diffusion/test_stage_diffusion_proc.py

@@ -0,0 +1,75 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import asyncio
+from concurrent.futures import ThreadPoolExecutor
+from dataclasses import asdict
+from types import SimpleNamespace
+
+import pytest
+
+from vllm_omni.diffusion.stage_diffusion_proc import StageDiffusionProc
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def test_process_batch_request_preserves_parent_request_id_and_kv_sender_info():
+    async def run_test():
+        captured = {}
+
+        def step(request):
+            captured["request"] = request
+            return [
+                SimpleNamespace(
+                    images=["img-1"],
+                    _multimodal_output={},
+                    _custom_output={},
+                    metrics={},
+                    stage_durations={},
+                    peak_memory_mb=0.0,
+                    latents=None,
+                    trajectory_latents=None,
+                    trajectory_timesteps=None,
+                    trajectory_log_probs=None,
+                    trajectory_decoded=None,
+                    final_output_type="image",
+                ),
+                SimpleNamespace(
+                    images=["img-2"],
+                    _multimodal_output={},
+                    _custom_output={},
+                    metrics={},
+                    stage_durations={},
+                    peak_memory_mb=0.0,
+                    latents=None,
+                    trajectory_latents=None,
+                    trajectory_timesteps=None,
+                    trajectory_log_probs=None,
+                    trajectory_decoded=None,
+                    final_output_type="image",
+                ),
+            ]
+
+        proc = object.__new__(StageDiffusionProc)
+        proc._engine = SimpleNamespace(step=step)
+        proc._executor = ThreadPoolExecutor(max_workers=1)
+
+        try:
+            result = await proc._process_batch_request(
+                request_id="req-parent",
+                prompts=["hello", "world"],
+                sampling_params_dict=asdict(OmniDiffusionSamplingParams()),
+                kv_sender_info={0: {"host": "10.0.0.2", "zmq_port": 50151}},
+            )
+        finally:
+            proc._executor.shutdown(wait=True)
+
+        request = captured["request"]
+        assert request.request_id == "req-parent"
+        assert request.request_ids == ["req-parent-0", "req-parent-1"]
+        assert request.kv_sender_info == {0: {"host": "10.0.0.2", "zmq_port": 50151}}
+        assert result.request_id == "req-parent"
+        assert result.images == ["img-1", "img-2"]
+
+    asyncio.run(run_test())

tests/distributed/omni_connectors/test_basic_connectors.py

@@ -9,7 +9,7 @@ from vllm_omni.distributed.omni_connectors.factory import OmniConnectorFactory
 from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec
 from vllm_omni.distributed.omni_connectors.utils.serialization import OmniSerializer
 
-# pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
 def test_basic_serialization():
@@ -120,3 +120,61 @@ def test_get_invalid_metadata(shm_connector):
 
     result = shm_connector.get("stage_0", "stage_1", "req_3", {"unknown": "format"})
     assert result is None
+
+
+def test_mooncake_connector_defaults_missing_host_to_detected_ip(monkeypatch: pytest.MonkeyPatch):
+    import vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector as mooncake_module
+
+    class _FakePool:
+        is_cuda = False
+
+        def pin_memory(self):
+            return self
+
+        def data_ptr(self):
+            return 1234
+
+    class _FakeTransferEngine:
+        def initialize(self, host, mode, protocol, device_name):
+            self.host = host
+            self.mode = mode
+            self.protocol = protocol
+            self.device_name = device_name
+            return 0
+
+        def get_rpc_port(self):
+            return 23456
+
+        def register_memory(self, base_ptr, pool_size):
+            del base_ptr, pool_size
+            return 0
+
+        def unregister_memory(self, base_ptr):
+            del base_ptr
+            return 0
+
+    monkeypatch.setattr(mooncake_module, "TransferEngine", _FakeTransferEngine)
+    monkeypatch.setattr(mooncake_module.torch, "empty", lambda *args, **kwargs: _FakePool())
+    monkeypatch.setattr(
+        mooncake_module.MooncakeTransferEngineConnector,
+        "_get_local_ip",
+        lambda self: "10.20.30.40",
+    )
+    monkeypatch.setattr(
+        mooncake_module.MooncakeTransferEngineConnector,
+        "_zmq_listener_loop",
+        lambda self: self._listener_ready.set(),
+    )
+
+    connector = mooncake_module.MooncakeTransferEngineConnector(
+        {
+            "zmq_port": 50051,
+            "memory_pool_size": 4096,
+        }
+    )
+    try:
+        assert connector.host == "10.20.30.40"
+        assert connector.engine.host == "10.20.30.40"
+        assert connector.get_connection_info()["host"] == "10.20.30.40"
+    finally:
+        connector.close()

tests/distributed/omni_connectors/test_chunk_transfer_adapter.py

@@ -133,6 +133,22 @@ def test_save_async(build_adapter):
     assert task["is_finished"] is False
 
 
+def test_send_single_request_cleans_up_after_finished_payload(build_adapter, monkeypatch):
+    adapter, _ = build_adapter(stage_id=1)
+    request = _req("req-finished", RequestStatus.FINISHED_STOPPED, external_req_id="ext-finished")
+
+    adapter.custom_process_next_stage_input_func = lambda **kwargs: {"x": [1], "finished": True}
+    cleanup_calls = []
+    monkeypatch.setattr(adapter, "cleanup", lambda *a, **kw: cleanup_calls.append((a, kw)))
+
+    adapter._send_single_request({"pooling_output": None, "request": request, "is_finished": True})
+
+    assert len(cleanup_calls) == 1
+    args, _ = cleanup_calls[0]
+    assert args[0] == "req-finished"
+    assert args[1] == "ext-finished"
+
+
 def test_update_request_payload(build_adapter):
     adapter, _ = build_adapter()
 
@@ -409,3 +425,86 @@ def test_generation_scheduler_calls_cleanup_on_finished(monkeypatch, mocker: Moc
     args, _ = cleanup_calls[0]
     assert args[0] == "req-s1"
     assert args[1] == "ext-s1"
+
+
+def test_ar_scheduler_defers_cleanup_and_queues_save_on_finished(mocker: MockerFixture):
+    """OmniARScheduler should enqueue save; adapter cleanup is handled in save thread."""
+    cleanup_calls = []
+    save_calls = []
+
+    adapter_mock = mocker.MagicMock()
+    adapter_mock.cleanup = lambda *a, **kw: cleanup_calls.append((a, kw))
+    adapter_mock.save_async = lambda *a, **kw: save_calls.append((a, kw))
+
+    from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
+
+    scheduler = mocker.MagicMock()
+    scheduler.chunk_transfer_adapter = adapter_mock
+    scheduler.connector = None
+    scheduler.perf_metrics = None
+    scheduler.log_stats = False
+    scheduler.recompute_kv_load_failures = False
+    scheduler.structured_output_manager = mocker.MagicMock()
+    scheduler.structured_output_manager.should_advance.return_value = False
+    scheduler.finished_req_ids_dict = {}
+    scheduler.kv_cache_manager = mocker.MagicMock()
+    scheduler.kv_cache_manager.take_events.return_value = None
+    scheduler.kv_event_publisher = mocker.MagicMock()
+    scheduler.waiting_for_transfer_free = set()
+    scheduler.transfer_triggered_requests = set()
+    scheduler.active_kv_transfers = set()
+
+    request = _HashableRequest(
+        request_id="req-ar",
+        external_req_id="ext-ar",
+        status=RequestStatus.RUNNING,
+        is_finished=lambda: False,
+        num_computed_tokens=1,
+        num_prompt_tokens=1,
+        prompt_token_ids=[1],
+        num_output_placeholders=0,
+        sampling_params=None,
+        pooling_params=None,
+        stop_reason=None,
+        client_index=0,
+        take_events=lambda: [],
+        trace_headers=None,
+        num_cached_tokens=0,
+        num_external_computed_tokens=0,
+        num_nans_in_logits=0,
+        get_finished_reason=lambda: "stop",
+    )
+    scheduler.requests = {"req-ar": request}
+
+    scheduler._update_request_with_output = mocker.MagicMock(return_value=([], True))
+    scheduler._process_kv_transfer_trigger = mocker.MagicMock(return_value=False)
+    scheduler._handle_stopped_request = mocker.MagicMock(return_value=True)
+    scheduler._free_request = mocker.MagicMock(return_value=None)
+    scheduler._get_routed_experts = mocker.MagicMock(return_value=None)
+    scheduler.running = [request]
+    scheduler.waiting = mocker.MagicMock()
+    scheduler.waiting.remove_requests = mocker.MagicMock()
+    scheduler.make_spec_decoding_stats = mocker.MagicMock(return_value=None)
+    scheduler.make_stats = mocker.MagicMock(return_value=None)
+
+    scheduler_output = SimpleNamespace(
+        num_scheduled_tokens={"req-ar": 1},
+        scheduled_spec_decode_tokens={},
+        num_invalid_spec_tokens=0,
+    )
+    model_runner_output = SimpleNamespace(
+        sampled_token_ids=[[123]],
+        logprobs=None,
+        prompt_logprobs_dict={},
+        pooler_output=None,
+        num_nans_in_logits=None,
+        kv_connector_output=None,
+        cudagraph_stats=None,
+        req_id_to_index={"req-ar": 0},
+        kv_extracted_req_ids=None,
+    )
+
+    OmniARScheduler.update_from_output(scheduler, scheduler_output, model_runner_output)
+
+    assert len(cleanup_calls) == 0
+    assert len(save_calls) == 1

tests/distributed/omni_connectors/test_kv_flow.py

@@ -1,8 +1,14 @@
+import json
+import struct
+
+import numpy as np
 import pytest
 import torch
 
+import vllm_omni.distributed.omni_connectors.kv_transfer_manager as kv_transfer_manager_module
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.distributed.omni_connectors.kv_transfer_manager import (
+    KVCacheTransferData,
     OmniKVCacheConfig,
     OmniKVTransferManager,
 )
@@ -60,6 +66,35 @@ def common_constants():
     }
 
 
+def _decode_stored_payload(data):
+    if isinstance(data, torch.Tensor) and data.dtype == torch.uint8 and data.dim() == 1:
+        return KVCacheTransferData.from_bytes(data.cpu().numpy().tobytes())
+
+    if isinstance(data, (bytes, bytearray, memoryview)):
+        return KVCacheTransferData.from_bytes(data)
+
+    return data
+
+
+def _make_serialized_payload() -> tuple[bytes, torch.Tensor]:
+    key_tensor = torch.arange(12, dtype=torch.float32).reshape(3, 4)
+    payload = KVCacheTransferData(
+        request_id="req-payload",
+        layer_blocks={"key_cache": [key_tensor], "value_cache": [None]},
+        block_ids=[1],
+        metadata={"seq_len": 3},
+    ).to_bytes()
+    return payload, key_tensor
+
+
+def _rewrite_serialized_header(payload: bytes, mutate_header) -> bytes:
+    header_len = struct.unpack(">I", payload[:4])[0]
+    header = json.loads(payload[4 : 4 + header_len])
+    mutate_header(header)
+    new_header = json.dumps(header, separators=(",", ":")).encode("utf-8")
+    return struct.pack(">I", len(new_header)) + new_header + payload[4 + header_len :]
+
+
 def test_manager_extraction(kv_config, mock_connector, common_constants):
     """Test extraction and sending logic in OmniKVTransferManager."""
     num_layers = common_constants["num_layers"]
@@ -95,7 +130,7 @@ def test_manager_extraction(kv_config, mock_connector, common_constants):
     expected_key = f"stage1->stage2:{full_request_id}"
     assert expected_key in mock_connector.store
 
-    data = mock_connector.store[expected_key]
+    data = _decode_stored_payload(mock_connector.store[expected_key])
     assert data["request_id"] == req_id
     assert "layer_blocks" in data
     assert len(data["layer_blocks"]["key_cache"]) == num_layers
@@ -106,6 +141,116 @@ def test_manager_extraction(kv_config, mock_connector, common_constants):
     assert data["layer_blocks"]["key_cache"][0].shape == expected_shape
 
 
+def test_from_bytes_rejects_out_of_bounds_header_len():
+    payload, _ = _make_serialized_payload()
+    bad_payload = struct.pack(">I", len(payload)) + payload[4:]
+
+    with pytest.raises(ValueError, match="header_len"):
+        KVCacheTransferData.from_bytes(bad_payload)
+
+    with pytest.raises(ValueError, match="header_len"):
+        KVCacheTransferData.from_bytes_gpu(torch.tensor(list(bad_payload), dtype=torch.uint8))
+
+
+def test_from_bytes_rejects_out_of_bounds_tensor_span():
+    payload, _ = _make_serialized_payload()
+    bad_payload = _rewrite_serialized_header(payload, lambda header: header["td"][0].update({"o": 4096}))
+
+    with pytest.raises(ValueError, match="tensor span"):
+        KVCacheTransferData.from_bytes(bad_payload)
+
+    with pytest.raises(ValueError, match="tensor span"):
+        KVCacheTransferData.from_bytes_gpu(torch.tensor(list(bad_payload), dtype=torch.uint8))
+
+
+def test_from_bytes_rejects_unsupported_dtype():
+    payload, _ = _make_serialized_payload()
+    bad_payload = _rewrite_serialized_header(payload, lambda header: header["td"][0].update({"d": "cuda"}))
+
+    with pytest.raises(ValueError, match="Unsupported dtype"):
+        KVCacheTransferData.from_bytes(bad_payload)
+
+    with pytest.raises(ValueError, match="Unsupported dtype"):
+        KVCacheTransferData.from_bytes_gpu(torch.tensor(list(bad_payload), dtype=torch.uint8))
+
+
+def test_from_bytes_uses_explicit_layer_index_descriptor():
+    payload, key_tensor = _make_serialized_payload()
+    payload_with_explicit_index = _rewrite_serialized_header(
+        payload,
+        lambda header: header["td"][0].update({"n": "key_cache_extra_suffix", "i": 0}),
+    )
+
+    data = KVCacheTransferData.from_bytes(payload_with_explicit_index)
+
+    assert torch.equal(data["layer_blocks"]["key_cache"][0], key_tensor)
+
+
+def test_update_sender_info_uses_configured_source_stage():
+    config = OmniKVCacheConfig(
+        connector_config={"type": "mock"},
+        stage_id=2,
+        engine_input_source=[1],
+        need_recv_cache=True,
+    )
+    manager = OmniKVTransferManager(config)
+
+    manager.update_sender_info(
+        {
+            0: {"host": "10.0.0.1", "zmq_port": 50151},
+            1: {"host": "10.0.0.2", "zmq_port": 50152},
+        }
+    )
+
+    assert manager.config.connector_config["sender_host"] == "10.0.0.2"
+    assert manager.config.connector_config["sender_zmq_port"] == 50152
+
+
+def test_clone_received_payload_tensors_breaks_buffer_alias():
+    payload, key_tensor = _make_serialized_payload()
+    raw = np.frombuffer(bytearray(payload), dtype=np.uint8)
+    data = KVCacheTransferData.from_bytes(memoryview(raw))
+
+    OmniKVTransferManager._clone_received_payload_tensors(data)
+    raw[:] = 0
+
+    assert torch.equal(data["layer_blocks"]["key_cache"][0], key_tensor)
+
+
+def test_receive_kv_cache_uses_exponential_backoff(monkeypatch):
+    config = OmniKVCacheConfig(
+        connector_config={"type": "mock"},
+        from_stage="sender",
+        stage_id="receiver",
+        need_recv_cache=True,
+        recv_timeout=0.3,
+    )
+    manager = OmniKVTransferManager(config)
+
+    class _NeverReadyConnector:
+        def get(self, **kwargs):
+            del kwargs
+            return None
+
+    manager._connector = _NeverReadyConnector()
+
+    now = {"value": 0.0}
+    sleep_intervals = []
+
+    monkeypatch.setattr(kv_transfer_manager_module.time, "time", lambda: now["value"])
+
+    def _fake_sleep(interval: float) -> None:
+        sleep_intervals.append(interval)
+        now["value"] += interval
+
+    monkeypatch.setattr(kv_transfer_manager_module.time, "sleep", _fake_sleep)
+
+    data, size = manager.receive_kv_cache_for_request("req-backoff")
+
+    assert (data, size) == (None, 0)
+    assert sleep_intervals == pytest.approx([0.01, 0.02, 0.04, 0.08, 0.16])
+
+
 def test_manager_extraction_tuple_layout(kv_config, mock_connector, common_constants):
     """Test extraction with tuple layout."""
     num_layers = common_constants["num_layers"]
@@ -135,7 +280,7 @@ def test_manager_extraction_tuple_layout(kv_config, mock_connector, common_const
     expected_key = f"stage1->stage2:{full_request_id}"
     assert expected_key in mock_connector.store
 
-    data = mock_connector.store[expected_key]
+    data = _decode_stored_payload(mock_connector.store[expected_key])
     expected_shape = (seq_len, num_heads, head_dim)
     for idx in range(len(kv_caches)):
         assert data["layer_blocks"]["key_cache"][idx].shape == expected_shape
@@ -165,7 +310,7 @@ def test_manager_extraction_mismatched_kv_block_counts(kv_config, mock_connector
     expected_key = f"stage1->stage2:{full_request_id}"
     assert expected_key in mock_connector.store
 
-    data = mock_connector.store[expected_key]
+    data = _decode_stored_payload(mock_connector.store[expected_key])
     expected_shape = (2 * block_size, num_heads, head_dim)
     assert data["layer_blocks"]["key_cache"][0].shape == expected_shape
     assert data["layer_blocks"]["value_cache"][0].shape == expected_shape
@@ -254,6 +399,82 @@ def test_manager_reception(kv_config, mock_connector, common_constants):
     assert req.kv_metadata["seq_len"] == seq_len
 
 
+def test_manager_reception_prefers_parent_request_id_for_batched_request(kv_config, mock_connector, common_constants):
+    """Batched diffusion requests must fetch KV using the parent/global request ID."""
+    num_layers = common_constants["num_layers"]
+    num_heads = common_constants["num_heads"]
+    head_dim = common_constants["head_dim"]
+    seq_len = common_constants["seq_len"]
+    parent_req_id = common_constants["req_id"]
+
+    expected_shape = (seq_len, num_heads, head_dim)
+    key_cache = [torch.randn(expected_shape) for _ in range(num_layers)]
+    value_cache = [torch.randn(expected_shape) for _ in range(num_layers)]
+
+    data_to_receive = {
+        "request_id": parent_req_id,
+        "layer_blocks": {"key_cache": key_cache, "value_cache": value_cache},
+        "metadata": {"seq_len": seq_len},
+        "block_ids": [],
+    }
+
+    manager = OmniKVTransferManager(kv_config)
+    manager._connector = mock_connector
+
+    full_request_id = f"omni_stage1_to_stage2_kv_cache_{parent_req_id}"
+    store_key = f"stage1->stage2:{full_request_id}"
+    mock_connector.store[store_key] = data_to_receive
+
+    req = OmniDiffusionRequest(
+        prompts=["prompt-a", "prompt-b"],
+        sampling_params=OmniDiffusionSamplingParams(),
+        request_ids=[f"{parent_req_id}-0", f"{parent_req_id}-1"],
+        request_id=parent_req_id,
+    )
+
+    success = manager.receive_kv_cache(req, target_device=torch.device("cpu"))
+
+    assert success
+    assert req.kv_metadata["seq_len"] == seq_len
+    assert torch.allclose(req.past_key_values.key_cache[0], key_cache[0])
+
+
+def test_receive_multi_kv_cache_uses_parent_request_id_for_cfg_collection(kv_config):
+    manager = OmniKVTransferManager(kv_config)
+
+    seen = {}
+
+    def collect_cfg(request_id, cfg_request_ids, kv_transfer_manager, target_device):
+        seen["request_id"] = request_id
+        seen["cfg_request_ids"] = cfg_request_ids
+        seen["kv_transfer_manager"] = kv_transfer_manager
+        seen["target_device"] = target_device
+        return {"cfg_text_kv_metadata": {"ok": True}}
+
+    req = OmniDiffusionRequest(
+        prompts=["prompt-a", "prompt-b"],
+        sampling_params=OmniDiffusionSamplingParams(),
+        request_ids=["req-parent-0", "req-parent-1"],
+        request_id="req-parent",
+    )
+    req.sampling_params.cfg_kv_request_ids = {"cfg_text": "req-parent__cfg_text"}
+
+    manager.receive_kv_cache = lambda request, target_device=None: request is req
+
+    success = manager.receive_multi_kv_cache(
+        req,
+        cfg_kv_collect_func=collect_cfg,
+        target_device=torch.device("cpu"),
+    )
+
+    assert success
+    assert seen["request_id"] == "req-parent"
+    assert seen["cfg_request_ids"] == {"cfg_text": "req-parent__cfg_text"}
+    assert seen["kv_transfer_manager"] is manager
+    assert seen["target_device"] == torch.device("cpu")
+    assert req.sampling_params.cfg_text_kv_metadata == {"ok": True}
+
+
 def test_integration_flow(common_constants):
     """Simulate extraction -> connector -> reception."""
     num_layers = common_constants["num_layers"]

tests/distributed/omni_coordinator/test_load_balancer.py

@@ -3,12 +3,16 @@
 
 from time import time
 
+import pytest
+
 from vllm_omni.distributed.omni_coordinator import (
     InstanceInfo,
     RandomBalancer,
     StageStatus,
 )
 
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
 
 def test_load_balancer_select_returns_valid_index():
     """Verify RandomBalancer.select() returns a valid index for instances."""

tests/distributed/omni_coordinator/test_omni_coord_client_for_hub.py

@@ -12,6 +12,8 @@ from vllm_omni.distributed.omni_coordinator import (
     OmniCoordClientForHub,
 )
 
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
 
 def _bind_pub() -> tuple[zmq.Context, zmq.Socket, str]:
     ctx = zmq.Context.instance()

tests/distributed/omni_coordinator/test_omni_coord_client_for_stage.py

@@ -2,13 +2,22 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import json
+import threading
 
+import pytest
 import zmq
 
 from vllm_omni.distributed.omni_coordinator import (
     OmniCoordClientForStage,
     StageStatus,
 )
+from vllm_omni.distributed.omni_coordinator import (
+    omni_coord_client_for_stage as stage_client_module,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
 def _bind_router() -> tuple[zmq.Context, zmq.Socket, str]:
@@ -19,7 +28,8 @@ def _bind_router() -> tuple[zmq.Context, zmq.Socket, str]:
     return ctx, router, endpoint
 
 
-def _recv_event(router: zmq.Socket) -> dict:
+def _recv_event(router: zmq.Socket, timeout_ms: int = 2000) -> dict:
+    assert router.poll(timeout=timeout_ms) != 0, "Timed out waiting for coordinator event"
     frames = router.recv_multipart()
     # ROUTER adds identity frame; the last frame is the payload.
     payload = frames[-1]
@@ -108,3 +118,197 @@ def test_stage_client_close_sends_down_status():
 
     router.close(0)
     ctx.term()
+
+
+def test_stage_client_reconnects_after_send_failure(mocker):
+    """Verify send failure path invokes reconnect before retrying send."""
+    ctx, router, endpoint = _bind_router()
+
+    client = OmniCoordClientForStage(
+        endpoint,
+        "tcp://stage:reconnect-in",
+        "tcp://stage:reconnect-out",
+        0,
+    )
+
+    # Discard initial registration event from the real socket.
+    _recv_event(router)
+
+    class _FlakySocket:
+        def __init__(self):
+            self.send_calls = 0
+            self.closed = False
+
+        def send(self, *_args, **_kwargs):
+            self.send_calls += 1
+            if self.send_calls == 1:
+                raise RuntimeError("simulated send failure")
+
+        def close(self, *_args, **_kwargs):
+            self.closed = True
+
+    flaky_socket = _FlakySocket()
+    client._socket = flaky_socket
+    client._reconnect = mocker.Mock(return_value=True)
+
+    client.update_info(queue_length=1)
+
+    client._reconnect.assert_called_once_with(max_retries=3)
+    assert flaky_socket.send_calls == 2
+
+    client.close()
+    router.close(0)
+    ctx.term()
+
+
+def test_stage_client_raises_when_reconnect_fails(mocker):
+    """Verify send failure is propagated when reconnect cannot recover."""
+    ctx, router, endpoint = _bind_router()
+
+    client = OmniCoordClientForStage(
+        endpoint,
+        "tcp://stage:reconnect-fail-in",
+        "tcp://stage:reconnect-fail-out",
+        0,
+    )
+
+    # Discard initial registration event from the real socket.
+    _recv_event(router)
+
+    class _AlwaysFailSocket:
+        def send(self, *_args, **_kwargs):
+            raise RuntimeError("simulated send failure")
+
+        def close(self, *_args, **_kwargs):
+            pass
+
+    client._socket = _AlwaysFailSocket()
+    client._reconnect = mocker.Mock(return_value=False)
+
+    with pytest.raises(RuntimeError, match="simulated send failure"):
+        client.update_info(queue_length=2)
+
+    client._reconnect.assert_called_once_with(max_retries=3)
+    client.close()
+    router.close(0)
+    ctx.term()
+
+
+def test_stage_client_close_handles_runtime_error_in_final_update(mocker):
+    """Verify close() still releases resources when final update raises RuntimeError."""
+    ctx, router, endpoint = _bind_router()
+
+    client = OmniCoordClientForStage(
+        endpoint,
+        "tcp://stage:close-runtime-in",
+        "tcp://stage:close-runtime-out",
+        0,
+    )
+
+    # Discard initial registration event from the real socket.
+    _recv_event(router)
+
+    client._send_event = mocker.Mock(side_effect=RuntimeError("simulated close-time failure"))
+    client.close()
+
+    assert client._closed
+    assert client._socket.closed
+
+    router.close(0)
+    ctx.term()
+
+
+def test_reconnect_respects_retry_limit(monkeypatch):
+    """Verify _reconnect stops after max_retries on repeated failures."""
+    attempts = {"connect": 0}
+
+    class _FailSocket:
+        def close(self, *_args, **_kwargs):
+            pass
+
+        def connect(self, *_args, **_kwargs):
+            attempts["connect"] += 1
+            raise zmq.ZMQError("simulated reconnect failure")
+
+    class _FailContext:
+        def socket(self, *_args, **_kwargs):
+            return _FailSocket()
+
+        def term(self):
+            pass
+
+    client = OmniCoordClientForStage.__new__(OmniCoordClientForStage)
+    client._closed = False
+    client._coord_zmq_addr = "tcp://127.0.0.1:9999"
+    client._stop_event = threading.Event()
+    client._send_lock = threading.RLock()
+    client._socket = _FailSocket()
+    client._ctx = _FailContext()
+
+    monkeypatch.setattr(stage_client_module.zmq, "Context", lambda: _FailContext())
+    monkeypatch.setattr(stage_client_module.time, "sleep", lambda *_args, **_kwargs: None)
+
+    assert client._reconnect(max_retries=3, retry_interval=5.0) is False
+    assert attempts["connect"] == 3
+
+
+def test_heartbeat_loop_retries_after_transient_send_failure():
+    """Verify heartbeat loop continues after one transient send failure."""
+
+    class _FakeStopEvent:
+        def __init__(self):
+            self.wait_calls = 0
+            self._set = False
+
+        def wait(self, timeout=None):
+            _ = timeout
+            self.wait_calls += 1
+            # Run two loop iterations, then stop.
+            return self._set or self.wait_calls >= 3
+
+        def is_set(self):
+            return self._set
+
+        def set(self):
+            self._set = True
+
+    client = OmniCoordClientForStage.__new__(OmniCoordClientForStage)
+    client._closed = False
+    client._heartbeat_interval = 0.0
+    client._stop_event = _FakeStopEvent()
+
+    calls = {"count": 0}
+
+    def _fake_send(event_type):
+        assert event_type == "heartbeat"
+        calls["count"] += 1
+        if calls["count"] == 1:
+            raise RuntimeError("transient heartbeat failure")
+
+    client._send_event = _fake_send
+
+    client._heartbeat_loop()
+
+    assert calls["count"] == 2
+
+
+def test_update_info_rejected_while_closing():
+    """Verify update_info is rejected once client enters closing state."""
+    ctx, router, endpoint = _bind_router()
+
+    client = OmniCoordClientForStage(
+        endpoint,
+        "tcp://stage:closing-in",
+        "tcp://stage:closing-out",
+        0,
+    )
+    _recv_event(router)
+
+    client._closing = True
+    with pytest.raises(RuntimeError, match="closing"):
+        client.update_info(queue_length=3)
+
+    client._closing = False
+    client.close()
+    router.close(0)
+    ctx.term()

tests/distributed/omni_coordinator/test_omni_coordinator.py

@@ -4,6 +4,7 @@
 import json
 import time
 
+import pytest
 import zmq
 from vllm.v1.utils import get_engine_client_zmq_addr
 
@@ -13,6 +14,8 @@ from vllm_omni.distributed.omni_coordinator import (
     StageStatus,
 )
 
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
 
 def _recv_instance_list(sub: zmq.Socket, timeout_ms: int = 2000) -> dict | None:
     """Receive InstanceList JSON from SUB socket. Returns None on timeout."""

tests/e2e/offline_inference/custom_pipeline/qwen_image_pipeline_with_logprob.py

@@ -6,7 +6,8 @@
 This pipeline follows the structure of the user's reference implementation:
 - supports pre-tokenized prompt IDs via OmniCustomPrompt-style dict input
 - uses an SDE scheduler that can return step logprobs
-- returns rich custom_output fields for testing
+- returns structured trajectory_* fields (latents, timesteps, log_probs)
+  consistent with the BAGEL trajectory recording design
 """
 
 from __future__ import annotations
@@ -393,10 +394,10 @@ class QwenImagePipelineWithLogProbForTest(QwenImagePipeline):
 
         return DiffusionOutput(
             output=_maybe_to_cpu(image),
+            trajectory_latents=_maybe_to_cpu(all_latents),
+            trajectory_log_probs=_maybe_to_cpu(all_log_probs),
+            trajectory_timesteps=_maybe_to_cpu(all_timesteps),
             custom_output={
-                "all_latents": _maybe_to_cpu(all_latents),
-                "all_log_probs": _maybe_to_cpu(all_log_probs),
-                "all_timesteps": _maybe_to_cpu(all_timesteps),
                 "prompt_embeds": _maybe_to_cpu(prompt_embeds),
                 "prompt_embeds_mask": _maybe_to_cpu(prompt_embeds_mask),
                 "negative_prompt_embeds": _maybe_to_cpu(negative_prompt_embeds),

tests/e2e/offline_inference/custom_pipeline/test_async_omni_qwen_image_generate.py

@@ -1,7 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-"""E2E tests for AsyncOmni Qwen-Image generation flow (no Ray, no HTTP server)."""
+"""E2E tests for AsyncOmni Qwen-Image generation with trajectory_* fields.
+
+Validates that the custom Qwen-Image pipeline returns structured trajectory
+outputs (latents, timesteps, log_probs) via OmniRequestOutput's trajectory_*
+fields instead of the legacy custom_output dict.
+"""
 
 from __future__ import annotations
 
@@ -191,10 +196,17 @@ async def test_async_omni_generate_with_logprobs():
 
         _assert_valid_image_output(output)
 
-        all_log_probs = output.custom_output.get("all_log_probs")
-        assert all_log_probs is not None, "all_log_probs should be present when logprobs=True"
-        assert hasattr(all_log_probs, "shape")
-        assert all_log_probs.numel() > 0
+        assert output.trajectory_latents is not None, "trajectory_latents should be present"
+        assert hasattr(output.trajectory_latents, "shape")
+        assert output.trajectory_latents.numel() > 0
+
+        assert output.trajectory_timesteps is not None, "trajectory_timesteps should be present"
+        assert hasattr(output.trajectory_timesteps, "shape")
+        assert output.trajectory_timesteps.numel() > 0
+
+        assert output.trajectory_log_probs is not None, "trajectory_log_probs should be present when logprobs=True"
+        assert hasattr(output.trajectory_log_probs, "shape")
+        assert output.trajectory_log_probs.numel() > 0
 
 
 @pytest.mark.core_model

tests/e2e/offline_inference/test_bagel_lora.py

@@ -0,0 +1,198 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+End-to-end test for BAGEL LoRA support (Stage 1 / DiT).
+
+Validates that LoRA adapters are correctly loaded, applied with controllable
+scale, and cleanly deactivated.  Uses a synthetic rank-1 adapter targeting the
+first decoder layer's QKV projection.
+
+Assertions:
+  (a) LoRA at scale=1.0 visibly changes the output  (diff > 0.5)
+  (b) scale=2.0 produces a larger delta than scale=1.0  (linearity)
+  (c) The delta is bounded  (diff < 80, not corrupted)
+  (d) Deactivating LoRA exactly restores the baseline  (diff == 0)
+"""
+
+import json
+import os
+
+from vllm_omni.inputs.data import OmniSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
+
+from pathlib import Path
+
+import numpy as np
+import pytest
+import torch
+from PIL import Image
+from safetensors.torch import save_file
+
+from tests.conftest import modify_stage_config
+from tests.utils import hardware_test
+from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.lora.request import LoRARequest
+from vllm_omni.lora.utils import stable_lora_int_id
+
+MODEL = "ByteDance-Seed/BAGEL-7B-MoT"
+BAGEL_STAGE_CONFIG = str(Path(__file__).parent / "stage_configs" / "bagel_sharedmemory_ci.yaml")
+DEFAULT_PROMPT = "<|im_start|>A cute cat<|im_end|>"
+
+
+# ---------------------------------------------------------------------------
+# Helpers (reused from test_bagel_text2img.py patterns)
+# ---------------------------------------------------------------------------
+
+
+def _resolve_stage_config(config_path: str, run_level: str) -> str:
+    if run_level == "advanced_model":
+        return modify_stage_config(
+            config_path,
+            deletes={
+                "stage_args": {
+                    0: ["engine_args.load_format"],
+                    1: ["engine_args.load_format"],
+                }
+            },
+        )
+    return config_path
+
+
+def _configure_sampling_params(omni: Omni, num_inference_steps: int = 10) -> list[OmniSamplingParams]:
+    params_list = omni.default_sampling_params_list
+    if len(params_list) > 1:
+        params_list[1].num_inference_steps = num_inference_steps
+        params_list[1].extra_args = {
+            "cfg_text_scale": 4.0,
+            "cfg_img_scale": 1.5,
+        }
+    return params_list
+
+
+def _extract_generated_image(omni_outputs: list[OmniRequestOutput]) -> Image.Image | None:
+    for req_output in omni_outputs:
+        if req_output.images:
+            return req_output.images[0]
+    return None
+
+
+def _generate_bagel_image(omni: Omni) -> Image.Image:
+    params_list = _configure_sampling_params(omni)
+    params_list[1].lora_request = None
+    outputs = list(
+        omni.generate(
+            prompts=[{"prompt": DEFAULT_PROMPT, "modalities": ["image"]}],
+            sampling_params_list=params_list,
+        )
+    )
+    img = _extract_generated_image(outputs)
+    assert img is not None, "No image generated"
+    return img
+
+
+def _generate_bagel_image_with_lora(
+    omni: Omni,
+    lora_request: LoRARequest,
+    lora_scale: float = 1.0,
+) -> Image.Image:
+    params_list = _configure_sampling_params(omni)
+    params_list[1].lora_request = lora_request
+    params_list[1].lora_scale = lora_scale
+    outputs = list(
+        omni.generate(
+            prompts=[{"prompt": DEFAULT_PROMPT, "modalities": ["image"]}],
+            sampling_params_list=params_list,
+        )
+    )
+    img = _extract_generated_image(outputs)
+    assert img is not None, "No image generated with LoRA"
+    return img
+
+
+# BAGEL uses GQA: hidden_size=3584, 28 Q heads, 4 KV heads, head_dim=128
+# QKV packed dim = 28*128 + 4*128 + 4*128 = 3584 + 512 + 512 = 4608
+_LORA_DIM = 3584
+_LORA_QKV_DIM = 4608
+_LORA_MODULE = "bagel.language_model.model.layers.0.self_attn.qkv_proj"
+_LORA_RANK = 4
+
+
+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_QKV_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.core_model
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"})
+def test_bagel_lora_scale_and_deactivation(run_level, tmp_path):
+    """Validate LoRA effect, bounded perturbation, and clean deactivation."""
+    config_path = _resolve_stage_config(BAGEL_STAGE_CONFIG, run_level)
+    omni = Omni(model=MODEL, stage_configs_path=config_path, stage_init_timeout=300)
+    try:
+        lora_request = _make_file_lora_request(tmp_path / "bagel_lora")
+
+        # 1) Baseline (no LoRA)
+        baseline = _generate_bagel_image(omni)
+
+        # 2) LoRA with scale=1.0
+        img_1x = _generate_bagel_image_with_lora(omni, lora_request, lora_scale=1.0)
+
+        # 3) LoRA with scale=2.0
+        img_2x = _generate_bagel_image_with_lora(omni, lora_request, lora_scale=2.0)
+
+        # 4) No LoRA again (deactivation)
+        restored = _generate_bagel_image(omni)
+
+        baseline_arr = np.array(baseline, dtype=np.int16)
+        img_1x_arr = np.array(img_1x, dtype=np.int16)
+        img_2x_arr = np.array(img_2x, dtype=np.int16)
+        restored_arr = np.array(restored, dtype=np.int16)
+
+        diff_1x = np.abs(baseline_arr - img_1x_arr).mean()
+        diff_2x = np.abs(baseline_arr - img_2x_arr).mean()
+        diff_restored = np.abs(baseline_arr - restored_arr).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 < 80, 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}"
+    finally:
+        omni.close()

tests/e2e/offline_inference/test_bagel_understanding.py

@@ -0,0 +1,144 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+End-to-end tests for Bagel text2text and img2text (understanding) tasks.
+
+These tests validate that the Bagel multistage pipeline correctly generates
+text output for understanding tasks, matching reference results.
+
+Equivalent to running:
+    python3 examples/offline_inference/bagel/end2end.py \
+        --modality text2text \
+        --prompts "Where is the capital of France?"
+
+    python3 examples/offline_inference/bagel/end2end.py \
+        --modality img2text \
+        --prompts "Please describe this image" \
+        --image-path 2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg
+"""
+
+import os
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
+from pathlib import Path
+
+import pytest
+from vllm.assets.image import ImageAsset
+
+from tests.conftest import modify_stage_config
+from tests.utils import hardware_test
+from vllm_omni.entrypoints.omni import Omni
+
+MODEL_NAME = "ByteDance-Seed/BAGEL-7B-MoT"
+STAGE_CONFIG = str(Path(__file__).parent / "stage_configs" / "bagel_sharedmemory_ci.yaml")
+
+REFERENCE_TEXT_TEXT2TEXT = "The capital of France is Paris."
+
+REFERENCE_TEXT_IMG2TEXT = (
+    "This is a photo of a wooden boardwalk or pathway that leads through "
+    "tall green grass. The path appears to be in a natural setting, possibly "
+    "a wetland or marsh area. The sky above is blue with some scattered "
+    "clouds, suggesting it might be a sunny day. The overall scene looks "
+    "peaceful and serene."
+)
+
+
+def _resolve_stage_config(config_path: str, run_level: str) -> str:
+    """Strip load_format: dummy for advanced_model (real weights)."""
+    if run_level == "advanced_model":
+        return modify_stage_config(
+            config_path,
+            deletes={
+                "stage_args": {
+                    0: ["engine_args.load_format"],
+                    1: ["engine_args.load_format"],
+                }
+            },
+        )
+    return config_path
+
+
+def _extract_text(omni_outputs: list) -> str:
+    """Extract generated text from OmniRequestOutput list."""
+    for req_output in omni_outputs:
+        ro = getattr(req_output, "request_output", None)
+        if ro and getattr(ro, "outputs", None):
+            return "".join(getattr(o, "text", "") or "" for o in ro.outputs)
+    return ""
+
+
+@pytest.mark.core_model
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"})
+def test_bagel_text2text(run_level):
+    """Test Bagel text2text produces correct text output."""
+    config_path = _resolve_stage_config(STAGE_CONFIG, run_level)
+    omni = Omni(
+        model=MODEL_NAME,
+        stage_configs_path=config_path,
+        stage_init_timeout=300,
+    )
+
+    try:
+        prompt = "<|im_start|>user\nWhere is the capital of France?<|im_end|>\n<|im_start|>assistant\n"
+        params_list = omni.default_sampling_params_list
+        omni_outputs = list(
+            omni.generate(
+                prompts=[{"prompt": prompt, "modalities": ["text"]}],
+                sampling_params_list=params_list,
+            )
+        )
+
+        assert len(omni_outputs) > 0, "No outputs returned"
+        text = _extract_text(omni_outputs)
+        assert len(text) > 0, "Generated text is empty"
+
+        if run_level == "advanced_model":
+            assert text == REFERENCE_TEXT_TEXT2TEXT, (
+                f"Text mismatch: expected {REFERENCE_TEXT_TEXT2TEXT!r}, got {text!r}"
+            )
+    finally:
+        omni.close()
+
+
+@pytest.mark.core_model
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"})
+def test_bagel_img2text(run_level):
+    """Test Bagel img2text produces correct text output."""
+    input_image = ImageAsset("2560px-Gfp-wisconsin-madison-the-nature-boardwalk").pil_image.convert("RGB")
+    config_path = _resolve_stage_config(STAGE_CONFIG, run_level)
+    omni = Omni(
+        model=MODEL_NAME,
+        stage_configs_path=config_path,
+        stage_init_timeout=300,
+    )
+
+    try:
+        prompt = "<|im_start|>user\n<|image_pad|>\nPlease describe this image<|im_end|>\n<|im_start|>assistant\n"
+        params_list = omni.default_sampling_params_list
+        omni_outputs = list(
+            omni.generate(
+                prompts=[
+                    {
+                        "prompt": prompt,
+                        "multi_modal_data": {"image": input_image},
+                        "modalities": ["text"],
+                    }
+                ],
+                sampling_params_list=params_list,
+            )
+        )
+
+        assert len(omni_outputs) > 0, "No outputs returned"
+        text = _extract_text(omni_outputs)
+        assert len(text) > 0, "Generated text is empty"
+
+        if run_level == "advanced_model":
+            assert text == REFERENCE_TEXT_IMG2TEXT, f"Text mismatch: expected {REFERENCE_TEXT_IMG2TEXT!r}, got {text!r}"
+    finally:
+        omni.close()

tests/e2e/offline_inference/test_cosyvoice3.py

@@ -0,0 +1,218 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Offline E2E smoke test for CosyVoice3 zero-shot reference inference.
+
+This test uses the official upstream zero-shot prompt text/audio pair and
+verifies a stable reference recipe:
+- config-derived top_p/top_k and token-length ratios
+- model EOS token as the stop token
+- a conservative repetition penalty to avoid degenerate loops
+"""
+
+from __future__ import annotations
+
+import functools
+import io
+import os
+import tempfile
+from pathlib import Path
+from urllib.request import urlopen
+
+import numpy as np
+import pytest
+import soundfile as sf
+import yaml
+from huggingface_hub import snapshot_download
+from vllm.sampling_params import SamplingParams
+
+from tests.conftest import OmniRunner
+from tests.utils import hardware_test
+from vllm_omni.model_executor.models.cosyvoice3.config import CosyVoice3Config
+from vllm_omni.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokenizer
+
+MODEL = "FunAudioLLM/Fun-CosyVoice3-0.5B-2512"
+MODEL_DIR_ENV = "VLLM_OMNI_COSYVOICE3_MODEL_DIR"
+
+REFERENCE_PROMPT_WAV_URL = "https://raw.githubusercontent.com/FunAudioLLM/CosyVoice/main/asset/zero_shot_prompt.wav"
+REFERENCE_PROMPT_TEXT = "You are a helpful assistant.<|endofprompt|>希望你以后能够做的比我还好呦。"
+REFERENCE_SYNTH_TEXT = (
+    "CosyVoice is undergoing a comprehensive upgrade, providing more accurate, "
+    "stable, faster, and better voice generation capabilities."
+)
+REFERENCE_STAGE0_TEMPERATURE = 1.0
+REFERENCE_STAGE0_REPETITION_PENALTY = 2.0
+
+
+def _stage_config(name: str) -> str:
+    return str(Path(__file__).parent.parent.parent.parent / "vllm_omni" / "model_executor" / "stage_configs" / name)
+
+
+STAGE_CONFIGS = [
+    _stage_config("cosyvoice3.yaml"),
+    _stage_config("cosyvoice3_async_chunk.yaml"),
+]
+
+
+@functools.lru_cache(maxsize=1)
+def _load_reference_prompt_wav() -> tuple[np.ndarray, int]:
+    with urlopen(REFERENCE_PROMPT_WAV_URL, timeout=30) as resp:
+        data = resp.read()
+    audio, sr = sf.read(io.BytesIO(data), dtype="float32", always_2d=False)
+    if isinstance(audio, np.ndarray) and audio.ndim > 1:
+        audio = np.mean(audio, axis=-1)
+    return np.asarray(audio, dtype=np.float32), int(sr)
+
+
+@functools.lru_cache(maxsize=1)
+def _resolve_model_dir() -> Path:
+    override = os.environ.get(MODEL_DIR_ENV)
+    if override:
+        return Path(override).expanduser().resolve()
+    return Path(snapshot_download(MODEL, allow_patterns=["*"]))
+
+
+def _reference_zero_shot_stage0_sampling(*, text: str) -> SamplingParams:
+    config = CosyVoice3Config()
+    sampling_cfg = config.llm.get("sampling", {})
+    eos_token_id = int(config.llm["eos_token_id"])
+    model_dir = _resolve_model_dir()
+    tokenizer = get_qwen_tokenizer(
+        token_path=str(model_dir / config.qwen_pretrain_path),
+        skip_special_tokens=config.skip_special_tokens,
+        version=config.version,
+    )
+    text_len = max(1, len(tokenizer.encode(text, allowed_special=config.allowed_special)))
+    return SamplingParams(
+        temperature=REFERENCE_STAGE0_TEMPERATURE,
+        top_p=float(sampling_cfg.get("top_p", 0.8)),
+        top_k=int(sampling_cfg.get("top_k", 25)),
+        repetition_penalty=REFERENCE_STAGE0_REPETITION_PENALTY,
+        stop_token_ids=[eos_token_id],
+        min_tokens=int(text_len * config.min_token_text_ratio),
+        max_tokens=int(text_len * config.max_token_text_ratio),
+    )
+
+
+def _concat_audio(audio_val) -> np.ndarray:
+    import torch
+
+    if isinstance(audio_val, list):
+        tensors = []
+        for t in audio_val:
+            if t is None:
+                continue
+            if hasattr(t, "detach"):
+                t = t.detach()
+            if hasattr(t, "cpu"):
+                t = t.cpu()
+            if hasattr(t, "float"):
+                t = t.float()
+            if isinstance(t, torch.Tensor):
+                tensors.append(t.reshape(-1))
+        if not tensors:
+            return np.zeros((0,), dtype=np.float32)
+        return torch.cat(tensors, dim=-1).numpy().astype(np.float32, copy=False)
+
+    if hasattr(audio_val, "detach"):
+        audio_val = audio_val.detach()
+    if hasattr(audio_val, "cpu"):
+        audio_val = audio_val.cpu()
+    if hasattr(audio_val, "float"):
+        audio_val = audio_val.float()
+    if hasattr(audio_val, "numpy"):
+        audio_val = audio_val.numpy()
+    audio_np = np.asarray(audio_val, dtype=np.float32)
+    return audio_np.reshape(-1)
+
+
+def _get_stage_engine_outputs(omni_runner: OmniRunner, stage_id: int):
+    stage_list = getattr(omni_runner.omni, "stage_list", None)
+    if stage_list is not None:
+        return getattr(stage_list[stage_id], "engine_outputs", None) or []
+
+    stage_clients = getattr(getattr(omni_runner.omni, "engine", None), "stage_clients", None)
+    if stage_clients is not None:
+        return getattr(stage_clients[stage_id], "engine_outputs", None) or []
+
+    raise AttributeError("Unable to locate stage outputs on Omni runner")
+
+
+def _patched_stage_config(base_stage_config: str, model_dir: Path, tmp_dir: Path) -> str:
+    cfg = yaml.safe_load(Path(base_stage_config).read_text(encoding="utf-8"))
+    tokenizer_path = str(model_dir / "CosyVoice-BlankEN")
+    for stage in cfg.get("stage_args", []):
+        engine_args = stage.setdefault("engine_args", {})
+        engine_args["tokenizer"] = tokenizer_path
+        engine_args["enforce_eager"] = True
+        engine_args["hf_overrides"] = {"architectures": ["CosyVoice3Model"]}
+    out_path = tmp_dir / Path(base_stage_config).name
+    out_path.write_text(yaml.safe_dump(cfg, sort_keys=False), encoding="utf-8")
+    return str(out_path)
+
+
+def _build_reference_inputs(prompt_audio: tuple[np.ndarray, int]) -> list[dict[str, object]]:
+    return [
+        {
+            "prompt": REFERENCE_SYNTH_TEXT,
+            "multi_modal_data": {"audio": prompt_audio},
+            "modalities": ["audio"],
+            "mm_processor_kwargs": {"prompt_text": REFERENCE_PROMPT_TEXT},
+        }
+    ]
+
+
+@pytest.mark.core_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("base_stage_config", STAGE_CONFIGS)
+def test_cosyvoice3_offline_reference_zero_shot(base_stage_config: str) -> None:
+    """CosyVoice3 zero-shot reference inference should stop cleanly and produce sane audio."""
+    prompt_audio, prompt_sr = _load_reference_prompt_wav()
+    model_dir = _resolve_model_dir()
+    expected_stop_token = int(CosyVoice3Config().llm["eos_token_id"])
+
+    with tempfile.TemporaryDirectory(prefix="cv3-e2e-") as tmp:
+        stage_config = _patched_stage_config(base_stage_config, model_dir, Path(tmp))
+        with OmniRunner(
+            str(model_dir), seed=42, stage_configs_path=stage_config, stage_init_timeout=300
+        ) as omni_runner:
+            sampling_params_list = omni_runner.get_default_sampling_params_list()
+            sampling_params_list[0] = _reference_zero_shot_stage0_sampling(text=REFERENCE_SYNTH_TEXT)
+
+            outputs = omni_runner.omni.generate(
+                _build_reference_inputs((prompt_audio, prompt_sr)), sampling_params_list
+            )
+
+            assert outputs, "No outputs returned"
+            audio_mm = outputs[0].multimodal_output
+            assert "audio" in audio_mm, "No audio output found"
+
+            audio = _concat_audio(audio_mm["audio"])
+            assert audio.size > 0, "Generated audio is empty"
+
+            sr_val = audio_mm.get("sr", 24000)
+            if isinstance(sr_val, list) and sr_val:
+                sr_val = sr_val[-1]
+            if hasattr(sr_val, "item"):
+                sr_val = sr_val.item()
+            sr = int(sr_val)
+            assert sr == 24000, f"Unexpected sample_rate={sr}"
+
+            duration_s = audio.size / sr
+            assert 2.8 <= duration_s <= 8.8, f"Unexpected duration={duration_s:.3f}s (samples={audio.size}, sr={sr})"
+
+            stage0_outputs = _get_stage_engine_outputs(omni_runner, 0)
+            if stage0_outputs:
+                completion = stage0_outputs[0].outputs[0]
+                finish_reason = getattr(completion, "finish_reason", None)
+                stop_reason = getattr(completion, "stop_reason", None)
+                num_tokens = len(getattr(completion, "token_ids", []) or [])
+
+                assert finish_reason == "stop", f"Stage-0 finish_reason={finish_reason}, expected 'stop'"
+                assert int(stop_reason) == expected_stop_token, (
+                    f"Stage-0 stop_reason={stop_reason}, expected {expected_stop_token}"
+                )
+                assert 80 <= num_tokens <= 220, f"Stage-0 num_tokens={num_tokens}, expected sane stop-bound range"
+            else:
+                assert "async_chunk" in Path(base_stage_config).name, "Stage-0 produced no engine outputs"

tests/e2e/offline_inference/test_dynin_omni.py

@@ -0,0 +1,419 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+E2E offline smoke tests for Dynin-Omni.
+
+- model: "snu-aidas/Dynin-Omni"
+- stage config: tests/e2e/stage_configs/dynin_omni_ci.yaml
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import pytest
+import torch
+from transformers import AutoTokenizer
+
+from tests.conftest import OmniRunner
+from tests.utils import hardware_test
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
+
+_REPO_ROOT = Path(__file__).resolve().parents[3]
+_DEFAULT_DYNIN_CONFIG_PATH: Path | None = None
+_DEFAULT_STAGE_CONFIG_PATH = _REPO_ROOT / "tests" / "e2e" / "stage_configs" / "dynin_omni_ci.yaml"
+
+models = ["snu-aidas/Dynin-Omni"]
+stage_configs = [str(_DEFAULT_STAGE_CONFIG_PATH)]
+test_params = [(model, stage_config) for model in models for stage_config in stage_configs]
+
+DYNIN_CONFIG_PATH = str(_DEFAULT_DYNIN_CONFIG_PATH) if _DEFAULT_DYNIN_CONFIG_PATH is not None else None
+
+pytestmark = [
+    pytest.mark.core_model,
+    pytest.mark.omni,
+]
+
+
+# prompting util
+def _build_mmu_prompt(tokenizer: Any, question: str, dynin_config_path: str | None) -> dict[str, Any]:
+    encoded = tokenizer(question, return_tensors="pt", add_special_tokens=True)
+    token_ids = [int(v) for v in encoded["input_ids"][0].tolist()]
+    attention_mask = [int(v) for v in encoded["attention_mask"][0].tolist()]
+    additional_information: dict[str, Any] = {
+        "task": ["mmu"],
+        "detok_id": [0],
+        "prompt_length": [len(token_ids)],
+        "attention_mask": [attention_mask],
+        "max_new_tokens": [64],
+        "steps": [64],
+        "block_length": [16],
+        "temperature": [0.0],
+    }
+    if dynin_config_path:
+        additional_information["dynin_config_path"] = [str(dynin_config_path)]
+    return {
+        "prompt_token_ids": token_ids,
+        "additional_information": additional_information,
+        "modalities": ["text"],
+    }
+
+
+def _build_mmu_multimodal_prompt(
+    tokenizer: Any,
+    question: str,
+    dynin_config_path: str | None,
+    *,
+    image: Any | None = None,
+    audio: tuple[np.ndarray, int] | None = None,
+) -> dict[str, Any]:
+    if image is None and audio is None:
+        raise ValueError("At least one multimodal input (image or audio) must be provided.")
+
+    prefix_chunks: list[str] = []
+    mm_data: dict[str, Any] = {}
+    if image is not None:
+        prefix_chunks.append("<|soi|><|image|><|eoi|>")
+        mm_data["image"] = image
+    if audio is not None:
+        prefix_chunks.append("<|soa|><|audio|><|eoa|>")
+        mm_data["audio"] = audio
+
+    prefixed_question = " ".join(prefix_chunks + [question]).strip()
+    prompt = _build_mmu_prompt(
+        tokenizer=tokenizer,
+        question=prefixed_question,
+        dynin_config_path=dynin_config_path,
+    )
+    prompt["multi_modal_data"] = mm_data
+    prompt["modalities"] = ["text"]
+    return prompt
+
+
+def _generate_synthetic_image(width: int = 224, height: int = 224) -> np.ndarray:
+    x = np.linspace(0, 255, width, dtype=np.uint8)
+    y = np.linspace(0, 255, height, dtype=np.uint8)[:, None]
+    red = np.tile(x, (height, 1))
+    green = np.tile(y, (1, width))
+    blue = ((red.astype(np.uint16) + green.astype(np.uint16)) // 2).astype(np.uint8)
+    return np.stack([red, green, blue], axis=-1)
+
+
+def _generate_synthetic_audio(duration_s: int = 5, sample_rate: int = 48_000) -> tuple[np.ndarray, int]:
+    t = np.linspace(0, duration_s, int(sample_rate * duration_s), endpoint=False, dtype=np.float32)
+    waveform = 0.1 * np.sin(2.0 * np.pi * 440.0 * t)
+    return waveform.astype(np.float32), sample_rate
+
+
+# prompting util
+def _build_t2s_decode_prompt(dynin_config_path: str | None) -> dict[str, Any]:
+    # Bypass stage-0 generation and directly validate token->audio decode path.
+    generated_audio_token_ids = [int(v) for v in ([10, 11, 12, 13, 14] * 32)]
+    additional_information: dict[str, Any] = {
+        "task": ["t2s"],
+        "detok_id": [1],
+        "generated_token_ids": [generated_audio_token_ids],
+        "audio_codebook_size": [4096],
+    }
+    if dynin_config_path:
+        additional_information["dynin_config_path"] = [str(dynin_config_path)]
+    return {
+        "prompt_token_ids": [0],
+        "additional_information": additional_information,
+        "modalities": ["audio"],
+    }
+
+
+# prompting util
+def _build_t2i_decode_prompt(dynin_config_path: str | None) -> dict[str, Any]:
+    # Bypass stage-0 generation and directly validate token->image decode path.
+    # MAGVIT decode path expects a square token grid; 1024 tokens -> 32x32.
+    generated_image_token_ids = [int(v) for v in ([10, 11, 12, 13, 14, 15, 16, 17] * 128)]
+    additional_information: dict[str, Any] = {
+        "task": ["t2i"],
+        "detok_id": [2],
+        "generated_token_ids": [generated_image_token_ids],
+        "codebook_size": [8192],
+    }
+    if dynin_config_path:
+        additional_information["dynin_config_path"] = [str(dynin_config_path)]
+    return {
+        "prompt_token_ids": [0],
+        "additional_information": additional_information,
+        "modalities": ["image"],
+    }
+
+
+def _configure_dynin_config_env() -> None:
+    if DYNIN_CONFIG_PATH:
+        os.environ["DYNIN_CONFIG_PATH"] = str(DYNIN_CONFIG_PATH)
+    else:
+        os.environ.pop("DYNIN_CONFIG_PATH", None)
+
+
+def _is_finished_request_output(request_output: Any) -> bool:
+    if request_output is None:
+        return False
+    req_list = request_output if isinstance(request_output, list) else [request_output]
+    for req in req_list:
+        if req is not None and bool(getattr(req, "finished", False)):
+            return True
+    return False
+
+
+def _find_stage_output(outputs: list[Any], output_type: str) -> Any | None:
+    matched = [
+        stage_output for stage_output in outputs if getattr(stage_output, "final_output_type", None) == output_type
+    ]
+    if not matched:
+        return None
+
+    # Prefer the latest finished chunk to avoid picking an intermediate stream output.
+    for stage_output in reversed(matched):
+        if _is_finished_request_output(getattr(stage_output, "request_output", None)):
+            return stage_output
+    return matched[-1]
+
+
+def _to_token_list(value: Any) -> list[int]:
+    if value is None:
+        return []
+    if hasattr(value, "detach"):
+        value = value.detach()
+    if hasattr(value, "cpu"):
+        value = value.cpu()
+    if hasattr(value, "flatten"):
+        value = value.flatten().tolist()
+    if isinstance(value, tuple):
+        value = list(value)
+    if not isinstance(value, list):
+        return []
+    out: list[int] = []
+    for token in value:
+        if isinstance(token, bool):
+            continue
+        try:
+            out.append(int(token))
+        except Exception:
+            continue
+    return out
+
+
+def _extract_text(stage_output: Any, tokenizer: Any | None = None) -> str:
+    request_output = getattr(stage_output, "request_output", None)
+    if request_output is None:
+        return ""
+    req_list = request_output if isinstance(request_output, list) else [request_output]
+    for req in req_list:
+        completions = getattr(req, "outputs", None) or []
+        if not completions:
+            continue
+        completion = completions[0]
+        mm_out = (
+            getattr(completion, "multimodal_output", None)
+            or getattr(req, "multimodal_output", None)
+            or getattr(stage_output, "multimodal_output", None)
+            or {}
+        )
+        text = mm_out.get("text")
+        if isinstance(text, list) and text:
+            text = text[-1]
+        if isinstance(text, str) and text.strip():
+            return text.strip()
+        if tokenizer is not None:
+            for key in ("text_tokens", "token_ids"):
+                token_ids = _to_token_list(mm_out.get(key))
+                if not token_ids:
+                    continue
+                decoded = tokenizer.decode(token_ids, skip_special_tokens=True)
+                if isinstance(decoded, str) and decoded.strip():
+                    return decoded.strip()
+        fallback = getattr(completion, "text", None)
+        if isinstance(fallback, str) and fallback.strip():
+            return fallback.strip()
+    return ""
+
+
+def _extract_audio(stage_output: Any) -> Any | None:
+    request_output = getattr(stage_output, "request_output", None)
+    if request_output is None:
+        return None
+    req_list = request_output if isinstance(request_output, list) else [request_output]
+    for req in req_list:
+        completions = getattr(req, "outputs", None) or []
+        if not completions:
+            continue
+        completion = completions[0]
+        mm_out = getattr(completion, "multimodal_output", None) or {}
+        if "audio" in mm_out:
+            return mm_out["audio"]
+    return None
+
+
+def _extract_image(stage_output: Any) -> Any | None:
+    request_output = getattr(stage_output, "request_output", None)
+    if request_output is None:
+        return None
+    req_list = request_output if isinstance(request_output, list) else [request_output]
+    for req in req_list:
+        completions = getattr(req, "outputs", None) or []
+        if not completions:
+            continue
+        completion = completions[0]
+        mm_out = getattr(completion, "multimodal_output", None) or {}
+        if "image" in mm_out:
+            return mm_out["image"]
+    return None
+
+
+def _numel(value: Any) -> int:
+    if value is None:
+        return 0
+    if isinstance(value, torch.Tensor):
+        return int(value.numel())
+    shape = getattr(value, "shape", None)
+    if shape is not None:
+        try:
+            total = 1
+            for dim in shape:
+                total *= int(dim)
+            return int(total)
+        except Exception:
+            pass
+    if isinstance(value, (list, tuple)):
+        return len(value)
+    return 0
+
+
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"})
+@pytest.mark.parametrize("test_config", test_params)
+def test_dynin_t2i_decode_to_image(test_config: tuple[str, str]) -> None:
+    model, stage_config_path = test_config
+    _configure_dynin_config_env()
+    prompt = _build_t2i_decode_prompt(dynin_config_path=DYNIN_CONFIG_PATH)
+
+    with OmniRunner(
+        model,
+        seed=42,
+        stage_configs_path=stage_config_path,
+        stage_init_timeout=600,
+        init_timeout=600,
+    ) as runner:
+        outputs = runner.generate([prompt])
+
+    image_output = _find_stage_output(outputs, "image")
+    assert image_output is not None
+    image_value = _extract_image(image_output)
+    assert image_value is not None
+    assert _numel(image_value) > 0
+
+
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"})
+@pytest.mark.parametrize("test_config", test_params)
+def test_dynin_mmu_to_text(test_config: tuple[str, str]) -> None:
+    model, stage_config_path = test_config
+    _configure_dynin_config_env()
+    tokenizer = AutoTokenizer.from_pretrained(model, trust_remote_code=True)
+    prompt = _build_mmu_prompt(
+        tokenizer=tokenizer,
+        question="What is 2 + 2? Answer in one short sentence.",
+        dynin_config_path=DYNIN_CONFIG_PATH,
+    )
+
+    with OmniRunner(
+        model,
+        seed=42,
+        stage_configs_path=stage_config_path,
+        stage_init_timeout=600,
+        init_timeout=600,
+    ) as runner:
+        outputs = runner.generate([prompt])
+
+    text_output = _find_stage_output(outputs, "text")
+    assert text_output is not None
+    text_content = _extract_text(text_output, tokenizer=tokenizer)
+    assert text_content
+
+
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"})
+@pytest.mark.parametrize("test_config", test_params)
+def test_dynin_image_to_text(test_config: tuple[str, str]) -> None:
+    model, stage_config_path = test_config
+    _configure_dynin_config_env()
+    tokenizer = AutoTokenizer.from_pretrained(model, trust_remote_code=True)
+    prompt = _build_mmu_multimodal_prompt(
+        tokenizer=tokenizer,
+        question="Describe the image briefly in one sentence.",
+        dynin_config_path=DYNIN_CONFIG_PATH,
+        image=_generate_synthetic_image(),
+    )
+
+    with OmniRunner(
+        model,
+        seed=42,
+        stage_configs_path=stage_config_path,
+        stage_init_timeout=600,
+        init_timeout=600,
+    ) as runner:
+        outputs = runner.generate([prompt])
+
+    text_output = _find_stage_output(outputs, "text")
+    assert text_output is not None
+    text_content = _extract_text(text_output, tokenizer=tokenizer)
+    assert text_content
+
+
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"})
+@pytest.mark.parametrize("test_config", test_params)
+def test_dynin_speech_to_text(test_config: tuple[str, str]) -> None:
+    model, stage_config_path = test_config
+    _configure_dynin_config_env()
+    tokenizer = AutoTokenizer.from_pretrained(model, trust_remote_code=True)
+    prompt = _build_mmu_multimodal_prompt(
+        tokenizer=tokenizer,
+        question="Transcribe the audio briefly in one sentence.",
+        dynin_config_path=DYNIN_CONFIG_PATH,
+        audio=_generate_synthetic_audio(),
+    )
+
+    with OmniRunner(
+        model,
+        seed=42,
+        stage_configs_path=stage_config_path,
+        stage_init_timeout=600,
+        init_timeout=600,
+    ) as runner:
+        outputs = runner.generate([prompt])
+
+    text_output = _find_stage_output(outputs, "text")
+    assert text_output is not None
+    text_content = _extract_text(text_output, tokenizer=tokenizer)
+    assert text_content
+
+
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"})
+@pytest.mark.parametrize("test_config", test_params)
+def test_dynin_t2s_decode_to_audio(test_config: tuple[str, str]) -> None:
+    model, stage_config_path = test_config
+    _configure_dynin_config_env()
+    prompt = _build_t2s_decode_prompt(dynin_config_path=DYNIN_CONFIG_PATH)
+
+    with OmniRunner(
+        model,
+        seed=42,
+        stage_configs_path=stage_config_path,
+        stage_init_timeout=600,
+        init_timeout=600,
+    ) as runner:
+        outputs = runner.generate([prompt])
+
+    audio_output = _find_stage_output(outputs, "audio")
+    assert audio_output is not None
+    audio_value = _extract_audio(audio_output)
+    assert audio_value is not None
+    assert _numel(audio_value) > 0

tests/e2e/offline_inference/test_hunyuanimage3_text2img.py

@@ -0,0 +1,347 @@
+# ruff: noqa: E501
+from collections.abc import Generator
+from pathlib import Path
+
+import pytest
+import torch
+import torch.nn.functional as F
+from PIL import Image
+from transformers import CLIPModel, CLIPProcessor
+
+from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.platforms import current_omni_platform
+
+PROMPT = "A brown and white dog is running on the grass"
+MODEL_NAME = "tencent/HunyuanImage-3.0"
+LOCAL_CLIP_PATH = "openai/clip-vit-base-patch32"
+REPO_ROOT = Path(__file__).resolve().parents[3]
+STAGE_CONFIG_PATH = REPO_ROOT / "vllm_omni" / "model_executor" / "stage_configs" / "hunyuan_image_3_moe.yaml"
+
+pytestmark = [pytest.mark.advanced_model, pytest.mark.diffusion]
+
+# System prompt type. Options: None, dynamic, en_vanilla, en_recaption, en_think_recaption, en_unified
+# Below are the CLIP embedding tensors from the official HunyuanImage model (seed=1234, prompt: "A brown and white dog is running on the grass").
+# SEED_1234 denotes the output without system prompt, while the remaining entries correspond to outputs generated with different system prompts.
+# fmt: off
+SEED_1234 = torch.tensor(
+    [
+        0.027797, 0.028964, -0.005051, 0.001059, 0.017021, -0.034029, 0.021989, 0.033318, -0.000308, 0.016179, 0.010504, -0.034201, 0.050230, -0.021170, 0.083530, -0.003621,
+        0.040758, 0.039913, 0.044305, -0.019285, -0.058387, -0.001099, 0.042782, -0.036136, -0.014955, 0.002147, 0.009439, 0.012943, -0.028732, -0.018349, 0.002861, 0.013019,
+        0.014362, -0.038833, 0.029413, 0.020724, 0.002714, 0.010416, -0.020527, 0.050266, -0.081026, -0.006814, -0.007457, -0.032333, 0.008417, -0.122455, -0.006085, -0.025610,
+        0.012614, 0.025817, -0.005419, 0.038657, 0.000789, 0.067111, 0.002818, 0.028696, 0.047305, -0.009993, -0.019508, 0.038604, 0.099657, 0.026728, 0.012361, 0.013626,
+        0.023164, -0.037186, 0.007535, 0.054645, -0.009012, -0.019383, -0.005234, -0.018715, -0.000346, 0.051317, -0.028744, 0.029933, -0.006382, -0.018414, -0.033906, -0.028892,
+        -0.015301, -0.004276, 0.014626, -0.008505, 0.013717, -0.027323, -0.001332, -0.040227, 0.047021, -0.019082, -0.037260, -0.029780, -0.594026, 0.016573, -0.010523, 0.042616,
+        -0.013136, 0.030540, -0.151685, -0.005367, 0.016209, -0.034183, 0.009852, 0.038452, 0.005494, -0.017887, -0.007167, 0.017262, -0.038980, 0.011995, 0.021952, -0.031660,
+        0.020507, -0.035880, 0.035183, -0.026975, -0.050788, -0.002553, 0.037774, -0.020082, -0.015403, 0.045022, 0.072167, -0.029237, 0.003895, -0.051250, 0.008581, 0.023545,
+        -0.026827, 0.020895, 0.041780, -0.040766, -0.008146, 0.080630, 0.000404, 0.032003, -0.005279, -0.090707, -0.013813, 0.010204, -0.001513, 0.016394, -0.001321, 0.020535,
+        -0.038645, 0.024858, 0.024378, 0.018717, -0.056314, 0.024402, 0.018694, 0.029009, -0.008502, -0.014694, -0.028345, 0.005202, 0.046116, -0.032166, -0.030706, -0.038738,
+        -0.031356, -0.009683, 0.040069, 0.001596, -0.012621, 0.018590, -0.024138, 0.035330, 0.011546, 0.015791, -0.026932, 0.004531, 0.022455, -0.012871, 0.013915, -0.009567,
+        -0.010976, 0.013497, 0.042590, 0.002072, -0.052718, -0.045494, 0.013036, -0.005403, -0.005947, -0.003437, 0.016653, -0.016805, -0.040291, 0.007927, 0.001296, -0.008319,
+        0.021514, -0.001452, -0.121998, 0.015396, -0.022594, -0.006977, -0.040108, -0.035550, -0.021872, -0.014721, 0.019799, 0.036556, 0.015072, -0.057988, -0.011684, -0.045220,
+        -0.026295, 0.052647, 0.013741, -0.013428, 0.061794, 0.021431, -0.011316, -0.009963, 0.008198, 0.027746, 0.074219, -0.019499, 0.042673, 0.016028, 0.007214, -0.010650,
+        -0.019682, 0.001902, 0.038867, -0.007333, 0.031749, 0.004391, 0.018688, 0.044654, 0.030615, -0.027816, 0.031711, -0.056952, -0.033499, -0.039368, 0.025801, -0.027610,
+        -0.009329, -0.001799, 0.024061, -0.012593, -0.050266, -0.012512, 0.019528, -0.083434, 0.018238, 0.034138, -0.020120, -0.009910, -0.002280, 0.035325, 0.034440, -0.055205,
+        -0.017698, -0.000439, -0.034703, 0.013356, -0.037287, 0.048494, -0.018570, 0.028069, 0.019269, -0.007263, -0.008521, 0.000426, -0.016677, 0.056162, -0.011944, 0.017322,
+        0.022219, -0.014266, -0.009292, -0.009979, 0.014973, 0.011623, -0.017799, 0.032925, -0.024668, 0.007312, -0.025035, -0.008967, -0.026827, 0.011889, -0.138517, -0.009608,
+        -0.020592, -0.001272, 0.015676, -0.025706, 0.031775, -0.004195, 0.026876, -0.014748, -0.025966, -0.008741, 0.035437, 0.017139, -0.005140, -0.007101, -0.012510, -0.023600,
+        0.032969, -0.005510, 0.020010, 0.032567, 0.015558, 0.004265, -0.036300, 0.048210, 0.080424, -0.052820, -0.002063, -0.020875, 0.052530, -0.001638, -0.020299, -0.035202,
+        0.087818, 0.034614, -0.032735, 0.033201, -0.001751, 0.029574, 0.009926, 0.011619, -0.001267, -0.020149, -0.003826, -0.029860, 0.011437, -0.051276, 0.024344, 0.003096,
+        -0.011573, 0.038228, -0.005730, -0.052328, 0.001909, -0.025877, 0.019976, -0.010160, 0.023892, 0.049161, -0.028978, 0.018700, -0.026460, 0.001090, -0.072128, -0.008406,
+        0.010828, 0.020621, -0.005706, 0.023797, 0.036231, -0.112069, 0.017601, 0.007496, 0.045999, 0.016771, 0.021977, 0.022305, 0.018377, 0.002036, -0.029815, -0.082922,
+        -0.012710, -0.026355, 0.003790, 0.017472, -0.023148, -0.002901, -0.057854, 0.028393, 0.230866, -0.023486, 0.051094, 0.047508, 0.018957, -0.037130, 0.001054, -0.026126,
+        0.021970, -0.046915, -0.019419, -0.014077, 0.002502, -0.079454, -0.057149, -0.081701, 0.041979, -0.043074, -0.009425, -0.035776, -0.021794, -0.004826, -0.057263, -0.072940,
+        0.037651, -0.013991, -0.043863, -0.020581, 0.034319, -0.052566, -0.010355, -0.022963, 0.027144, -0.017339, 0.088930, -0.000670, -0.026547, -0.026586, -0.032531, 0.040314,
+        0.010148, 0.021104, 0.009228, -0.073227, 0.036650, -0.019337, 0.010211, -0.089620, -0.024676, -0.020729, -0.004070, 0.000784, -0.110561, 0.015390, 0.027151, -0.003228,
+        -0.066704, -0.004797, -0.026117, -0.018131, -0.090114, 0.020659, -0.007157, 0.013608, -0.022324, 0.027487, 0.018873, 0.027854, 0.045085, -0.039992, -0.017829, 0.011071,
+        -0.011393, -0.004454, -0.037189, -0.030299, 0.059668, 0.005064, 0.024655, -0.037239, 0.046882, -0.010356, -0.009690, 0.061909, -0.024736, 0.016849, 0.000784, 0.000201,
+        0.066165, 0.010234, -0.012134, -0.002823, -0.060847, 0.008953, 0.010348, 0.022292, -0.044602, -0.020981, 0.038839, 0.006616, -0.016836, -0.043995, -0.005463, -0.036413,
+        0.034895, -0.018008, -0.009543, -0.025080, -0.035243, 0.042696, -0.028911, -0.030676, -0.038542, -0.027798, -0.026607, 0.019467, 0.070629, -0.037356, -0.042648, -0.000284,
+        0.033095, 0.077781, -0.052930, 0.022515, -0.029926, -0.033821, -0.003277, -0.000038, -0.026871, 0.018223, -0.004221, 0.023454, -0.030611, -0.006396, -0.009873, -0.008402,
+    ],
+    dtype=torch.float32,
+)
+SYSTEM_PROMPT_DYNAMIC = torch.tensor(
+    [
+        0.010809, 0.021177, -0.017600, -0.016814, 0.012351, -0.024554, 0.018299, 0.039305, 0.003331, 0.030473, 0.005557, -0.040898, 0.047294, -0.016136, 0.076989, -0.002723,
+        0.017622, 0.042330, 0.058266, -0.016232, -0.029502, 0.004529, 0.033543, -0.041481, -0.017631, 0.002727, 0.018874, 0.019932, -0.030052, -0.009997, 0.004582, 0.002135,
+        -0.003720, -0.030923, 0.021174, 0.034033, -0.007096, 0.011522, -0.009518, 0.055688, -0.092351, -0.003914, 0.004589, -0.032635, 0.012479, -0.140607, -0.014141, -0.031821,
+        0.001396, 0.026780, -0.007623, 0.039957, 0.006434, 0.047516, 0.014377, 0.015237, 0.034212, 0.003576, -0.027357, 0.038888, 0.087272, 0.020248, 0.015165, 0.016002,
+        0.020781, -0.040509, -0.008929, 0.080857, -0.002642, -0.009738, -0.005683, -0.000615, -0.012801, 0.046457, -0.045004, 0.024689, 0.002498, -0.017333, -0.027366, -0.023231,
+        -0.006064, -0.021505, 0.007405, -0.021249, 0.026252, -0.018690, 0.020093, -0.036954, 0.037510, -0.032027, -0.030871, -0.011173, -0.618627, 0.021213, -0.004366, 0.029555,
+        -0.004324, 0.020221, -0.143832, -0.021386, 0.010482, -0.042113, 0.016164, 0.040350, 0.014627, -0.011778, -0.018102, 0.035380, -0.020305, 0.010590, 0.009227, -0.011415,
+        0.018623, -0.036384, 0.031003, -0.017073, -0.056456, -0.010423, 0.033029, -0.023511, -0.008717, 0.045716, 0.068273, -0.027886, 0.009665, -0.039801, 0.001465, 0.024361,
+        -0.015039, 0.022903, 0.033362, -0.022804, 0.008631, 0.076518, 0.000619, 0.022786, -0.015435, -0.095242, -0.006092, 0.015496, -0.009081, 0.015740, 0.004280, 0.013103,
+        -0.031836, 0.034241, 0.031836, 0.032636, -0.053721, 0.034370, 0.019172, 0.018383, 0.006907, -0.036039, -0.027927, 0.008646, 0.040496, -0.060314, -0.039116, -0.021488,
+        -0.031682, -0.005077, 0.034920, 0.002148, -0.008087, 0.002024, -0.008480, 0.041096, 0.011401, 0.020380, -0.025078, 0.005002, 0.022252, -0.014577, 0.008051, -0.014476,
+        -0.007078, 0.021075, 0.036965, 0.005343, -0.038671, -0.037222, 0.014052, -0.009952, -0.003958, -0.001878, 0.017848, -0.016608, -0.030813, 0.010921, 0.001068, 0.003095,
+        0.007076, -0.001936, -0.102996, 0.006838, -0.005243, -0.009140, -0.043796, -0.027227, -0.008426, -0.013177, 0.015602, 0.021036, 0.025484, -0.064836, -0.003593, -0.038036,
+        -0.023102, 0.064053, 0.007850, 0.000771, 0.039297, 0.011903, -0.015866, -0.017612, 0.006308, 0.024342, 0.086761, -0.016705, 0.039239, 0.025079, -0.006452, 0.003174,
+        -0.010146, 0.010787, 0.035932, -0.015346, 0.037191, 0.010990, 0.011573, 0.044958, 0.035560, -0.017339, 0.018878, -0.025394, -0.044339, -0.029852, 0.015951, -0.032248,
+        -0.012019, 0.013497, 0.012224, -0.001284, -0.034041, -0.015768, 0.000230, -0.086076, 0.024878, 0.031929, -0.016668, -0.019815, -0.001325, 0.007944, 0.017674, -0.036097,
+        -0.019651, -0.001272, -0.032842, 0.002056, -0.037140, 0.043191, -0.003710, 0.011767, 0.020313, -0.018396, -0.015935, 0.010228, -0.017349, 0.049363, -0.010007, 0.019533,
+        0.018076, 0.016608, -0.005523, -0.007793, 0.016868, 0.019341, -0.008236, 0.026765, -0.025324, -0.007849, -0.023648, -0.007791, -0.018508, 0.015357, -0.166499, -0.003718,
+        -0.035447, -0.005229, 0.019327, -0.014207, 0.028433, -0.002619, 0.013888, -0.033146, -0.017015, 0.004677, 0.039554, 0.003803, -0.014592, -0.018886, -0.023868, -0.022708,
+        0.033661, 0.008626, 0.015687, 0.046395, 0.014173, 0.015083, -0.025994, 0.039120, 0.076334, -0.061165, 0.001791, -0.017579, 0.067567, -0.002415, -0.032495, -0.025576,
+        0.079027, 0.036370, -0.013303, 0.030510, -0.009061, 0.019135, 0.015627, 0.024864, 0.015093, -0.017066, -0.014075, -0.021907, 0.017388, -0.033492, 0.013317, -0.000040,
+        0.003396, 0.044030, -0.009194, -0.049524, -0.005015, -0.040007, 0.009104, 0.000580, 0.005603, 0.035891, -0.038913, 0.023239, -0.017022, -0.002695, -0.095759, 0.018503,
+        0.017365, 0.011104, -0.003433, 0.024113, 0.052609, -0.085274, 0.027565, -0.005833, 0.020700, 0.015842, 0.019148, 0.020203, -0.000698, -0.005337, -0.037400, -0.060144,
+        -0.031893, -0.038396, -0.001949, 0.018901, -0.014268, -0.004721, -0.055913, 0.013814, 0.215024, -0.011357, 0.057530, 0.050092, 0.016513, -0.059254, 0.001494, -0.031472,
+        0.032190, -0.047512, -0.020501, -0.002571, 0.007844, -0.063630, -0.043938, -0.079595, 0.032820, -0.021659, -0.003738, -0.035267, -0.013794, -0.021172, -0.046356, -0.077079,
+        0.021526, -0.007447, -0.050276, -0.029743, 0.022208, -0.039137, -0.021426, -0.029825, 0.029390, -0.002943, 0.073158, -0.000435, -0.032029, -0.038524, -0.029886, 0.017473,
+        0.013513, 0.022738, 0.000632, -0.073718, 0.029219, -0.018896, 0.007302, -0.116122, -0.013324, -0.012214, -0.005960, -0.003720, -0.155869, 0.019896, 0.016919, -0.021133,
+        -0.066911, -0.000926, -0.020871, -0.015295, -0.086108, 0.014918, -0.009284, 0.001689, -0.038155, 0.039163, 0.015988, 0.014413, 0.034205, -0.053273, 0.001687, 0.012227,
+        -0.007341, -0.006123, -0.005731, -0.026863, 0.060196, 0.028929, 0.019328, -0.033709, 0.038789, -0.015624, 0.013323, 0.053821, -0.015538, -0.001610, 0.012959, -0.013897,
+        0.082010, 0.012866, -0.017269, 0.000017, -0.059458, 0.015870, 0.028455, 0.025234, -0.051163, -0.022976, 0.011866, -0.005613, -0.008738, -0.047658, -0.002155, -0.029432,
+        0.039242, -0.013491, -0.001641, -0.024210, -0.019187, 0.026716, -0.025698, -0.027591, -0.034678, -0.002473, -0.019391, 0.017597, 0.064385, -0.029104, -0.034501, -0.004955,
+        0.015008, 0.060749, -0.051693, 0.020279, -0.027170, -0.027003, 0.000254, 0.011352, -0.028116, 0.028938, -0.007224, 0.019978, -0.025379, -0.004874, -0.019361, -0.020278,
+    ],
+    dtype=torch.float32,
+)
+SYSTEM_EN_RECAPTION = torch.tensor(
+    [
+        0.007721, 0.015421, -0.019305, -0.000920, 0.016031, -0.019730, 0.029683, 0.026810, -0.010510, 0.021463, 0.008833, -0.040851, 0.043260, -0.007042, 0.057224, 0.011995,
+        0.007818, 0.046369, 0.059838, -0.028548, -0.047399, -0.000983, 0.024343, -0.052259, -0.013638, 0.006856, 0.009186, 0.014235, -0.031497, -0.008644, -0.009349, 0.018900,
+        0.002913, -0.022475, 0.039518, 0.019052, -0.007600, 0.010634, -0.011830, 0.075675, -0.071738, -0.014947, 0.004995, -0.025804, -0.002553, -0.093262, 0.002881, -0.033744,
+        -0.007234, 0.013659, 0.009897, 0.039185, -0.005366, 0.041534, -0.005924, 0.019786, 0.048566, -0.009356, -0.027360, 0.042557, 0.091286, 0.009286, 0.015410, 0.028166,
+        0.022476, -0.025162, 0.012144, 0.084603, -0.003150, -0.008549, -0.002099, -0.014987, -0.019480, 0.046843, -0.030613, 0.015557, -0.008965, -0.008798, -0.027032, -0.014112,
+        0.018703, -0.014749, -0.000928, -0.024660, 0.024004, 0.004560, 0.028156, -0.028467, 0.025444, -0.038699, -0.014927, -0.031593, -0.648498, 0.018529, 0.003378, 0.030188,
+        -0.002314, 0.014950, -0.146615, -0.009005, 0.016579, -0.037867, 0.020907, 0.033160, 0.007877, -0.026345, -0.056428, 0.031255, -0.018404, 0.013334, 0.009988, -0.022790,
+        0.020803, -0.036862, 0.036222, -0.006646, -0.058084, -0.012036, 0.044199, -0.027665, -0.015779, 0.051554, 0.059970, -0.025977, 0.003967, -0.035247, -0.000488, 0.023182,
+        0.000468, 0.019190, 0.047268, -0.032279, -0.005302, 0.078669, -0.001915, 0.024918, -0.014952, -0.078905, -0.018333, 0.001362, -0.015115, 0.005435, 0.002313, 0.018766,
+        -0.032773, 0.037344, 0.024061, 0.012143, -0.057106, 0.029490, 0.019537, 0.009099, 0.026064, -0.015927, -0.037047, 0.006002, 0.025191, -0.035318, -0.032245, -0.047822,
+        -0.023568, -0.004533, 0.025100, 0.002758, -0.002649, -0.012287, -0.012139, 0.043080, 0.003295, 0.024667, -0.021050, 0.006752, 0.025315, -0.011127, 0.009800, -0.021343,
+        -0.024866, 0.010098, 0.026954, 0.012467, -0.035866, -0.031780, 0.007479, -0.003388, -0.012619, -0.012099, 0.014974, -0.001908, -0.032700, 0.004703, 0.003238, -0.007498,
+        0.023241, 0.002715, -0.111739, 0.003317, 0.006475, -0.019792, -0.046558, -0.032593, -0.020762, -0.005059, 0.016934, 0.029195, 0.028744, -0.050633, 0.001907, -0.028791,
+        -0.016695, 0.052143, 0.010439, 0.007204, 0.028502, 0.012607, -0.012414, -0.031238, 0.007305, 0.032309, 0.087924, -0.010530, 0.029925, 0.032666, -0.002202, 0.017539,
+        -0.009091, -0.001631, 0.024906, -0.013102, 0.031772, 0.018465, 0.012035, 0.031460, 0.030193, 0.005289, 0.025859, -0.038971, -0.046577, -0.025852, 0.035235, -0.038514,
+        0.001042, 0.013012, 0.023701, -0.014630, -0.029269, -0.011981, 0.008219, -0.067347, -0.003456, 0.028198, -0.008657, -0.017773, 0.010540, 0.023964, 0.021012, -0.034465,
+        -0.023748, 0.004065, -0.021598, 0.008440, -0.031533, 0.038390, -0.007680, -0.003852, 0.016136, -0.017906, -0.008927, 0.006300, -0.001251, 0.029337, -0.008632, 0.020568,
+        0.021560, -0.007222, 0.005313, -0.013089, 0.012299, 0.031303, -0.013951, 0.016547, -0.024771, -0.008753, -0.030908, -0.014421, -0.017656, 0.014044, -0.114986, 0.000956,
+        -0.035588, 0.003756, 0.015383, -0.013358, 0.009385, -0.001359, 0.012623, -0.028724, 0.001607, 0.012809, 0.032668, 0.011834, -0.015587, -0.007170, -0.021344, -0.019664,
+        0.017690, -0.014538, 0.016511, 0.038037, 0.029919, 0.020907, -0.018565, 0.032964, 0.078548, -0.050386, -0.003012, -0.016965, 0.064131, 0.008077, -0.025879, -0.035820,
+        0.095075, 0.019901, -0.019114, 0.022832, 0.003741, 0.027148, 0.018231, 0.027741, 0.020328, 0.001700, -0.006939, -0.024154, 0.018523, -0.029819, 0.008050, -0.004477,
+        0.006087, 0.056878, -0.009083, -0.061537, -0.011531, -0.037551, 0.000434, -0.005843, 0.024739, 0.032020, -0.053119, 0.020704, -0.012385, -0.002726, -0.082489, 0.009072,
+        0.013341, 0.000316, 0.001899, 0.022868, 0.034407, -0.066857, 0.020589, 0.012195, 0.023211, -0.001520, 0.000897, 0.029670, -0.015930, 0.006509, -0.035172, -0.061215,
+        -0.014099, -0.038584, -0.012213, 0.018613, -0.012365, -0.002777, -0.055184, 0.017146, 0.214358, -0.015750, 0.052488, 0.045205, 0.025334, -0.054615, 0.002117, -0.038122,
+        0.012402, -0.053418, -0.025405, 0.007235, 0.013208, -0.092481, -0.048700, -0.085186, 0.029039, -0.036767, -0.000777, -0.017625, -0.012556, -0.004887, -0.033660, -0.082310,
+        0.013387, -0.003256, -0.062981, -0.019886, 0.017624, -0.037421, -0.020743, -0.020894, 0.041974, -0.008502, 0.088413, -0.018697, -0.029398, -0.029389, -0.043721, 0.013872,
+        0.003944, 0.030361, 0.005355, -0.081355, 0.041843, -0.016395, 0.011954, -0.060440, -0.000966, -0.019101, 0.006803, -0.011310, -0.148581, 0.020342, 0.012795, -0.016473,
+        -0.053300, -0.012340, -0.016640, -0.029834, -0.082405, 0.011859, -0.004255, -0.004396, -0.012515, 0.031962, 0.030438, 0.013792, 0.031557, -0.047200, 0.006485, 0.024815,
+        -0.019376, -0.011454, -0.034184, -0.021329, 0.050115, 0.021720, 0.002874, -0.047163, 0.044031, -0.014663, 0.020534, 0.056017, 0.007017, 0.003323, 0.005734, -0.002777,
+        0.082836, 0.012048, -0.023236, -0.007401, -0.071598, 0.016760, 0.017282, 0.028306, -0.026220, -0.008016, -0.000202, -0.020271, -0.019828, -0.046986, -0.005805, -0.039647,
+        0.042879, -0.004463, 0.007753, -0.028916, -0.020612, 0.028833, -0.039839, -0.052447, -0.013275, -0.002407, -0.018937, 0.033216, 0.075535, -0.045026, -0.009901, 0.016637,
+        -0.000322, 0.073925, -0.055701, 0.014912, -0.045671, -0.021189, 0.006761, -0.002015, -0.027410, 0.018250, -0.015916, 0.016254, -0.044964, 0.029261, -0.029319, -0.005222,
+    ],
+    dtype=torch.float32,
+)
+SYSTEM_EN_THINK_RECAPTION = torch.tensor(
+    [
+        0.011004, 0.017341, -0.019959, -0.018314, 0.016520, -0.027395, 0.017946, 0.039665, 0.000645, 0.035903, 0.002499, -0.045664, 0.039472, -0.013479, 0.081302, 0.000182,
+        0.006947, 0.042845, 0.059741, -0.010796, -0.035240, 0.004176, 0.029557, -0.043467, -0.017271, 0.006896, 0.010997, 0.022498, -0.023308, -0.013046, -0.000742, 0.016209,
+        -0.007152, -0.029868, 0.028747, 0.033743, -0.000227, 0.018419, -0.015023, 0.050376, -0.098475, -0.002375, 0.007897, -0.023936, 0.007843, -0.122463, -0.011680, -0.027267,
+        -0.007270, 0.021869, -0.011415, 0.043770, 0.000551, 0.048573, 0.003132, 0.014233, 0.037080, -0.004818, -0.028738, 0.044468, 0.073843, 0.016947, 0.014484, 0.021931,
+        0.020110, -0.032309, -0.003811, 0.095704, -0.006950, -0.007237, -0.005529, -0.020573, -0.016259, 0.041909, -0.038748, 0.018029, 0.005066, -0.021186, -0.020102, -0.019719,
+        0.006239, -0.021284, 0.004213, -0.024963, 0.032345, -0.012557, 0.037268, -0.038075, 0.040998, -0.032766, -0.023509, -0.016426, -0.627412, 0.022675, 0.000101, 0.023162,
+        -0.002081, 0.015922, -0.138671, -0.027995, 0.011579, -0.042859, 0.019935, 0.038077, 0.012640, -0.017377, -0.027456, 0.035151, -0.015756, 0.018530, 0.004646, -0.002589,
+        0.019645, -0.043736, 0.034947, -0.010166, -0.061165, -0.019195, 0.028909, -0.019415, -0.009485, 0.049566, 0.068621, -0.038644, 0.011278, -0.036133, 0.000564, 0.022611,
+        -0.013612, 0.020854, 0.030614, -0.025578, 0.005673, 0.076526, -0.004887, 0.027769, -0.022605, -0.092657, -0.013218, 0.008081, -0.015227, 0.018031, -0.005145, 0.015028,
+        -0.027193, 0.034767, 0.028710, 0.032007, -0.053175, 0.033528, 0.019437, 0.011517, 0.012107, -0.027679, -0.026937, 0.008612, 0.036909, -0.051484, -0.039971, -0.034372,
+        -0.023825, -0.003025, 0.033648, -0.001852, 0.007309, 0.000714, -0.001075, 0.038534, 0.007586, 0.016213, -0.025223, -0.001099, 0.015852, -0.011477, 0.020635, -0.010696,
+        -0.019634, 0.025613, 0.034374, 0.007169, -0.035000, -0.032268, 0.015114, -0.014217, -0.005229, -0.005495, 0.018189, -0.011360, -0.026755, 0.007036, -0.002333, -0.001174,
+        0.014729, 0.001739, -0.108591, 0.004699, 0.002048, -0.014801, -0.042855, -0.028846, -0.009609, -0.004500, 0.019466, 0.021848, 0.022140, -0.063035, -0.004272, -0.030798,
+        -0.018452, 0.055169, 0.012240, -0.003555, 0.038293, 0.008503, -0.016608, -0.021309, 0.000690, 0.027093, 0.088054, -0.008881, 0.034087, 0.030647, 0.003284, 0.005038,
+        -0.008359, 0.006311, 0.032462, -0.009699, 0.035283, 0.015261, 0.012827, 0.038169, 0.033959, -0.018048, 0.018122, -0.025259, -0.040084, -0.030879, 0.019853, -0.042558,
+        -0.011938, 0.019602, 0.016537, -0.003378, -0.027890, -0.014909, -0.005464, -0.071862, 0.012335, 0.021899, -0.017008, -0.023228, 0.003263, 0.004571, 0.016447, -0.029446,
+        -0.022645, -0.001261, -0.018573, 0.007431, -0.027587, 0.035362, -0.006785, -0.000614, 0.026044, -0.009056, -0.009843, 0.010467, -0.011929, 0.042025, -0.014068, 0.023113,
+        0.023880, 0.014948, 0.004370, -0.005262, 0.012587, 0.021608, -0.001783, 0.023697, -0.024945, -0.011533, -0.020953, -0.007205, -0.024693, 0.012961, -0.168760, 0.001767,
+        -0.041265, -0.007044, 0.015021, -0.008407, 0.029642, -0.000956, 0.008607, -0.035365, -0.012187, 0.011744, 0.032612, 0.006226, -0.015891, -0.017747, -0.022565, -0.024505,
+        0.031279, 0.004188, 0.011939, 0.038032, 0.008798, 0.012314, -0.024830, 0.034484, 0.076395, -0.060108, 0.001019, -0.016138, 0.067729, 0.003899, -0.029845, -0.019960,
+        0.086663, 0.040965, -0.010458, 0.027808, -0.006394, 0.017343, 0.014788, 0.024756, 0.016446, -0.012537, -0.008406, -0.028109, 0.013369, -0.033571, 0.012170, -0.002199,
+        0.005263, 0.052280, -0.018171, -0.047898, -0.010087, -0.038632, 0.006773, -0.000838, 0.011197, 0.038187, -0.049525, 0.021689, -0.007385, -0.005987, -0.094551, 0.019019,
+        0.012760, 0.009617, -0.002262, 0.030228, 0.047823, -0.079764, 0.023391, -0.005561, 0.018866, 0.012817, 0.020878, 0.027037, -0.013905, -0.002874, -0.035522, -0.046266,
+        -0.032448, -0.036010, -0.007776, 0.016512, -0.012279, -0.005665, -0.057974, 0.016967, 0.202836, -0.009066, 0.066093, 0.045689, 0.018319, -0.048465, 0.000242, -0.040874,
+        0.027824, -0.049045, -0.015616, -0.000307, 0.009163, -0.072975, -0.042979, -0.082254, 0.040549, -0.027049, 0.000725, -0.034118, -0.019604, -0.019097, -0.042483, -0.075446,
+        0.019387, -0.005218, -0.053573, -0.029975, 0.008195, -0.036608, -0.018920, -0.025610, 0.028426, -0.002688, 0.074996, -0.003423, -0.032505, -0.030565, -0.028142, 0.014437,
+        0.013359, 0.019376, 0.008356, -0.069731, 0.031824, -0.011103, 0.019327, -0.117090, -0.009352, -0.010290, -0.002129, -0.009198, -0.172915, 0.021232, 0.017274, -0.030060,
+        -0.061449, -0.006598, -0.013069, -0.012857, -0.081220, 0.019058, -0.004841, 0.003066, -0.037741, 0.041806, 0.018281, 0.009458, 0.036761, -0.044987, 0.003557, 0.008890,
+        -0.008011, -0.004063, -0.013474, -0.022090, 0.055398, 0.037475, 0.006991, -0.035962, 0.045503, -0.017162, 0.022391, 0.052754, -0.005924, -0.005936, 0.012673, -0.017922,
+        0.084548, 0.014695, -0.013817, 0.000421, -0.065167, 0.018269, 0.023317, 0.023523, -0.034229, -0.019588, 0.007911, -0.002426, -0.017109, -0.050870, 0.002848, -0.033077,
+        0.043451, -0.010609, -0.000375, -0.023206, -0.018155, 0.027102, -0.036006, -0.035115, -0.023922, 0.005989, -0.015372, 0.027123, 0.075210, -0.035302, -0.029799, 0.003642,
+        0.007714, 0.063498, -0.053234, 0.015699, -0.040459, -0.027354, -0.002433, 0.010923, -0.020134, 0.029292, -0.010176, 0.013508, -0.032403, 0.004323, -0.017504, -0.015237,
+    ],
+    dtype=torch.float32,
+)
+SYSTEM_EN_VANILLA = torch.tensor(
+    [
+        0.010809, 0.021177, -0.017600, -0.016814, 0.012351, -0.024554, 0.018299, 0.039305, 0.003331, 0.030473, 0.005557, -0.040898, 0.047294, -0.016136, 0.076989, -0.002723,
+        0.017622, 0.042330, 0.058266, -0.016232, -0.029502, 0.004529, 0.033543, -0.041481, -0.017631, 0.002727, 0.018874, 0.019932, -0.030052, -0.009997, 0.004582, 0.002135,
+        -0.003720, -0.030923, 0.021174, 0.034033, -0.007096, 0.011522, -0.009518, 0.055688, -0.092351, -0.003914, 0.004589, -0.032635, 0.012479, -0.140607, -0.014141, -0.031821,
+        0.001396, 0.026780, -0.007623, 0.039957, 0.006434, 0.047516, 0.014377, 0.015237, 0.034212, 0.003576, -0.027357, 0.038888, 0.087272, 0.020248, 0.015165, 0.016002,
+        0.020781, -0.040509, -0.008929, 0.080857, -0.002642, -0.009738, -0.005683, -0.000615, -0.012801, 0.046457, -0.045004, 0.024689, 0.002498, -0.017333, -0.027366, -0.023231,
+        -0.006064, -0.021505, 0.007405, -0.021249, 0.026252, -0.018690, 0.020093, -0.036954, 0.037510, -0.032027, -0.030871, -0.011173, -0.618627, 0.021213, -0.004366, 0.029555,
+        -0.004324, 0.020221, -0.143832, -0.021386, 0.010482, -0.042113, 0.016164, 0.040350, 0.014627, -0.011778, -0.018102, 0.035380, -0.020305, 0.010590, 0.009227, -0.011415,
+        0.018623, -0.036384, 0.031003, -0.017073, -0.056456, -0.010423, 0.033029, -0.023511, -0.008717, 0.045716, 0.068273, -0.027886, 0.009665, -0.039801, 0.001465, 0.024361,
+        -0.015039, 0.022903, 0.033362, -0.022804, 0.008631, 0.076518, 0.000619, 0.022786, -0.015435, -0.095242, -0.006092, 0.015496, -0.009081, 0.015740, 0.004280, 0.013103,
+        -0.031836, 0.034241, 0.031836, 0.032636, -0.053721, 0.034370, 0.019172, 0.018383, 0.006907, -0.036039, -0.027927, 0.008646, 0.040496, -0.060314, -0.039116, -0.021488,
+        -0.031682, -0.005077, 0.034920, 0.002148, -0.008087, 0.002024, -0.008480, 0.041096, 0.011401, 0.020380, -0.025078, 0.005002, 0.022252, -0.014577, 0.008051, -0.014476,
+        -0.007078, 0.021075, 0.036965, 0.005343, -0.038671, -0.037222, 0.014052, -0.009952, -0.003958, -0.001878, 0.017848, -0.016608, -0.030813, 0.010921, 0.001068, 0.003095,
+        0.007076, -0.001936, -0.102996, 0.006838, -0.005243, -0.009140, -0.043796, -0.027227, -0.008426, -0.013177, 0.015602, 0.021036, 0.025484, -0.064836, -0.003593, -0.038036,
+        -0.023102, 0.064053, 0.007850, 0.000771, 0.039297, 0.011903, -0.015866, -0.017612, 0.006308, 0.024342, 0.086761, -0.016705, 0.039239, 0.025079, -0.006452, 0.003174,
+        -0.010146, 0.010787, 0.035932, -0.015346, 0.037191, 0.010990, 0.011573, 0.044958, 0.035560, -0.017339, 0.018878, -0.025394, -0.044339, -0.029852, 0.015951, -0.032248,
+        -0.012019, 0.013497, 0.012224, -0.001284, -0.034041, -0.015768, 0.000230, -0.086076, 0.024878, 0.031929, -0.016668, -0.019815, -0.001325, 0.007944, 0.017674, -0.036097,
+        -0.019651, -0.001272, -0.032842, 0.002056, -0.037140, 0.043191, -0.003710, 0.011767, 0.020313, -0.018396, -0.015935, 0.010228, -0.017349, 0.049363, -0.010007, 0.019533,
+        0.018076, 0.016608, -0.005523, -0.007793, 0.016868, 0.019341, -0.008236, 0.026765, -0.025324, -0.007849, -0.023648, -0.007791, -0.018508, 0.015357, -0.166499, -0.003718,
+        -0.035447, -0.005229, 0.019327, -0.014207, 0.028433, -0.002619, 0.013888, -0.033146, -0.017015, 0.004677, 0.039554, 0.003803, -0.014592, -0.018886, -0.023868, -0.022708,
+        0.033661, 0.008626, 0.015687, 0.046395, 0.014173, 0.015083, -0.025994, 0.039120, 0.076334, -0.061165, 0.001791, -0.017579, 0.067567, -0.002415, -0.032495, -0.025576,
+        0.079027, 0.036370, -0.013303, 0.030510, -0.009061, 0.019135, 0.015627, 0.024864, 0.015093, -0.017066, -0.014075, -0.021907, 0.017388, -0.033492, 0.013317, -0.000040,
+        0.003396, 0.044030, -0.009194, -0.049524, -0.005015, -0.040007, 0.009104, 0.000580, 0.005603, 0.035891, -0.038913, 0.023239, -0.017022, -0.002695, -0.095759, 0.018503,
+        0.017365, 0.011104, -0.003433, 0.024113, 0.052609, -0.085274, 0.027565, -0.005833, 0.020700, 0.015842, 0.019148, 0.020203, -0.000698, -0.005337, -0.037400, -0.060144,
+        -0.031893, -0.038396, -0.001949, 0.018901, -0.014268, -0.004721, -0.055913, 0.013814, 0.215024, -0.011357, 0.057530, 0.050092, 0.016513, -0.059254, 0.001494, -0.031472,
+        0.032190, -0.047512, -0.020501, -0.002571, 0.007844, -0.063630, -0.043938, -0.079595, 0.032820, -0.021659, -0.003738, -0.035267, -0.013794, -0.021172, -0.046356, -0.077079,
+        0.021526, -0.007447, -0.050276, -0.029743, 0.022208, -0.039137, -0.021426, -0.029825, 0.029390, -0.002943, 0.073158, -0.000435, -0.032029, -0.038524, -0.029886, 0.017473,
+        0.013513, 0.022738, 0.000632, -0.073718, 0.029219, -0.018896, 0.007302, -0.116122, -0.013324, -0.012214, -0.005960, -0.003720, -0.155869, 0.019896, 0.016919, -0.021133,
+        -0.066911, -0.000926, -0.020871, -0.015295, -0.086108, 0.014918, -0.009284, 0.001689, -0.038155, 0.039163, 0.015988, 0.014413, 0.034205, -0.053273, 0.001687, 0.012227,
+        -0.007341, -0.006123, -0.005731, -0.026863, 0.060196, 0.028929, 0.019328, -0.033709, 0.038789, -0.015624, 0.013323, 0.053821, -0.015538, -0.001610, 0.012959, -0.013897,
+        0.082010, 0.012866, -0.017269, 0.000017, -0.059458, 0.015870, 0.028455, 0.025234, -0.051163, -0.022976, 0.011866, -0.005613, -0.008738, -0.047658, -0.002155, -0.029432,
+        0.039242, -0.013491, -0.001641, -0.024210, -0.019187, 0.026716, -0.025698, -0.027591, -0.034678, -0.002473, -0.019391, 0.017597, 0.064385, -0.029104, -0.034501, -0.004955,
+        0.015008, 0.060749, -0.051693, 0.020279, -0.027170, -0.027003, 0.000254, 0.011352, -0.028116, 0.028938, -0.007224, 0.019978, -0.025379, -0.004874, -0.019361, -0.020278,
+    ],
+    dtype=torch.float32,
+)
+SYSTEM_EN_UNIFIED = torch.tensor(
+    [
+        0.011409, 0.014191, -0.023163, -0.020119, 0.019190, -0.029559, 0.019616, 0.035872, 0.010434, 0.028709, 0.011616, -0.039422, 0.038369, -0.004631, 0.081177, 0.007400,
+        0.008903, 0.040408, 0.055323, -0.011950, -0.026940, 0.004916, 0.028101, -0.046200, -0.016732, 0.005115, 0.012100, 0.016136, -0.026057, -0.013827, -0.004914, 0.015261,
+        -0.010824, -0.028188, 0.022934, 0.026204, -0.003855, 0.013797, -0.014518, 0.050289, -0.100077, -0.002962, 0.009050, -0.028205, 0.016294, -0.128956, -0.012730, -0.023647,
+        -0.009306, 0.020066, 0.000033, 0.043619, 0.003250, 0.053425, 0.005889, 0.021529, 0.036032, -0.003254, -0.029715, 0.048345, 0.077978, 0.010674, 0.019296, 0.018721,
+        0.019244, -0.040115, -0.004245, 0.085214, -0.005280, -0.010746, -0.000164, -0.023405, -0.015641, 0.040193, -0.038735, 0.018966, -0.004031, -0.017879, -0.023017, -0.030379,
+        0.006468, -0.015959, 0.000532, -0.026530, 0.042640, -0.006095, 0.037899, -0.043658, 0.040965, -0.034682, -0.023729, -0.019291, -0.630840, 0.029658, 0.005462, 0.026650,
+        -0.000292, 0.013954, -0.149594, -0.019405, 0.015321, -0.045104, 0.030332, 0.031727, 0.012349, -0.009553, -0.022371, 0.034043, -0.014838, 0.015398, -0.003657, 0.000477,
+        0.021084, -0.041406, 0.029946, -0.013832, -0.057358, -0.018086, 0.031598, -0.031835, -0.006697, 0.040866, 0.068602, -0.042203, 0.007362, -0.036959, 0.003794, 0.026533,
+        -0.011873, 0.017343, 0.028333, -0.021804, 0.004007, 0.075133, 0.003340, 0.025326, -0.015068, -0.092280, -0.011514, 0.006827, -0.008254, 0.021181, -0.005035, 0.022263,
+        -0.022443, 0.043919, 0.026637, 0.028568, -0.056881, 0.036740, 0.024430, 0.015891, 0.012257, -0.031126, -0.030108, 0.007229, 0.026998, -0.051685, -0.033003, -0.031170,
+        -0.024021, 0.004235, 0.030164, 0.002674, 0.008018, 0.005532, 0.001621, 0.044790, 0.006413, 0.027160, -0.015022, 0.000911, 0.019723, -0.016244, 0.020077, -0.006847,
+        -0.014110, 0.022461, 0.031656, 0.002760, -0.039078, -0.026893, 0.006628, -0.011775, -0.000240, -0.005908, 0.014943, -0.012131, -0.021755, 0.004732, -0.005297, -0.002922,
+        0.014631, -0.002010, -0.112400, 0.000842, -0.002732, -0.014861, -0.052099, -0.034167, -0.011613, -0.006101, 0.013278, 0.018867, 0.026530, -0.068150, -0.003306, -0.032801,
+        -0.018523, 0.050875, 0.005488, -0.007241, 0.045707, 0.023119, -0.021519, -0.022683, 0.004806, 0.024827, 0.091371, -0.014424, 0.043836, 0.033094, 0.002390, 0.005450,
+        -0.004893, 0.013608, 0.031272, -0.002449, 0.031607, 0.014646, 0.014146, 0.043995, 0.028826, -0.012219, 0.021008, -0.020911, -0.036967, -0.036256, 0.013328, -0.038382,
+        -0.012084, 0.018183, 0.018782, -0.004697, -0.024284, -0.015474, -0.001463, -0.076015, 0.013923, 0.022125, -0.018765, -0.010793, 0.008409, 0.002067, 0.017961, -0.029716,
+        -0.020915, -0.001779, -0.009217, -0.001933, -0.036081, 0.042577, 0.000118, -0.013920, 0.014901, -0.016486, -0.010278, -0.000449, -0.017234, 0.042453, -0.009893, 0.021087,
+        0.017671, 0.009861, -0.004210, 0.004944, 0.015627, 0.014370, -0.001128, 0.030247, -0.019552, -0.014017, -0.020859, -0.002614, -0.024405, 0.016532, -0.173204, -0.001196,
+        -0.037415, -0.010990, 0.010449, -0.006124, 0.019211, 0.003695, 0.011679, -0.031852, -0.009764, 0.005773, 0.035793, 0.003455, -0.011772, -0.020532, -0.027434, -0.024761,
+        0.027483, -0.001554, 0.010411, 0.037888, 0.015619, 0.019186, -0.021204, 0.038158, 0.074991, -0.064521, -0.002503, -0.014499, 0.068165, 0.006145, -0.032891, -0.021540,
+        0.091385, 0.047584, -0.009590, 0.028004, -0.002962, 0.021061, 0.014854, 0.025840, 0.016068, -0.014364, -0.016418, -0.033454, 0.011734, -0.036518, 0.013015, -0.003966,
+        0.000855, 0.051373, -0.010960, -0.047078, -0.011048, -0.042015, 0.006818, 0.005483, 0.010251, 0.034951, -0.046162, 0.021258, -0.013397, -0.005259, -0.093775, 0.019974,
+        0.014992, 0.004043, -0.005931, 0.035662, 0.050723, -0.083293, 0.028047, -0.008042, 0.020763, 0.016763, 0.022913, 0.027129, -0.014314, -0.009854, -0.039019, -0.044870,
+        -0.028101, -0.038026, -0.006294, 0.018265, -0.015425, -0.007866, -0.052784, 0.010470, 0.200260, -0.007798, 0.064482, 0.046612, 0.025353, -0.059695, -0.001831, -0.039643,
+        0.025148, -0.042752, -0.014928, -0.010216, 0.014195, -0.069149, -0.041424, -0.078360, 0.036999, -0.021357, 0.011032, -0.026564, -0.016214, -0.023440, -0.044723, -0.064498,
+        0.018283, -0.007165, -0.051802, -0.026299, 0.005867, -0.034691, -0.020621, -0.030512, 0.024458, -0.011330, 0.066558, -0.004069, -0.031624, -0.030639, -0.037451, 0.013079,
+        0.015152, 0.008058, 0.009223, -0.069514, 0.030702, -0.009681, 0.014826, -0.115441, -0.005514, -0.011925, 0.001046, -0.007148, -0.164128, 0.018043, 0.017001, -0.026352,
+        -0.049691, -0.011637, -0.013045, -0.014851, -0.079469, 0.017692, -0.006575, 0.001063, -0.028299, 0.038777, 0.019930, 0.010641, 0.036955, -0.039004, -0.006477, 0.004278,
+        -0.001006, -0.002514, -0.017242, -0.023927, 0.049113, 0.038393, 0.011633, -0.031537, 0.041725, -0.012146, 0.023445, 0.049999, -0.008538, 0.001319, 0.012732, -0.021170,
+        0.082096, 0.009610, -0.025717, 0.002566, -0.060849, 0.017403, 0.032650, 0.018658, -0.030629, -0.025032, 0.005555, 0.000522, -0.009667, -0.043099, 0.005939, -0.027156,
+        0.045634, -0.011986, 0.002713, -0.032225, -0.015494, 0.028734, -0.036528, -0.033101, -0.027174, 0.009490, -0.016537, 0.029435, 0.065709, -0.037711, -0.020497, -0.005578,
+        0.011768, 0.061035, -0.044676, 0.016113, -0.042945, -0.022579, 0.002430, 0.012474, -0.018198, 0.030468, -0.016646, 0.019020, -0.035804, 0.001175, -0.018312, -0.010760,
+    ],
+    dtype=torch.float32,
+)
+# fmt: on
+SYSTEM_PROMPT_CASES = [
+    pytest.param("none", None, SEED_1234, id="none"),
+    pytest.param("dynamic", "dynamic", SYSTEM_PROMPT_DYNAMIC, id="dynamic"),
+    pytest.param("en_vanilla", "en_vanilla", SYSTEM_EN_VANILLA, id="en_vanilla"),
+    pytest.param("en_recaption", "en_recaption", SYSTEM_EN_RECAPTION, id="en_recaption"),
+    pytest.param("en_think_recaption", "en_think_recaption", SYSTEM_EN_THINK_RECAPTION, id="en_think_recaption"),
+    pytest.param("en_unified", "en_unified", SYSTEM_EN_UNIFIED, id="en_unified"),
+]
+
+
+@pytest.fixture(scope="session")
+def clip_bundle() -> tuple[CLIPModel, CLIPProcessor]:
+    try:
+        model = CLIPModel.from_pretrained(LOCAL_CLIP_PATH, local_files_only=True)
+        processor = CLIPProcessor.from_pretrained(LOCAL_CLIP_PATH, local_files_only=True)
+    except OSError as exc:
+        pytest.skip(f"Could not load CLIP model from local cache ({LOCAL_CLIP_PATH}): {exc}")
+
+    model.eval()
+    return model, processor
+
+
+@pytest.fixture(scope="module")
+def omni() -> Generator[Omni, None, None]:
+    engine = Omni(
+        model=MODEL_NAME,
+        stage_configs_path=str(STAGE_CONFIG_PATH),
+        stage_init_timeout=600,
+        init_timeout=900,
+    )
+    try:
+        yield engine
+    finally:
+        engine.close()
+
+
+def _extract_generated_image(outputs: list[object]) -> Image.Image:
+    if not outputs:
+        raise AssertionError("No outputs were returned from Omni.generate()")
+
+    first_output = outputs[0]
+    if images := getattr(first_output, "images", None):
+        return images[0]
+
+    request_output = getattr(first_output, "request_output", None)
+    if request_output is not None and (images := getattr(request_output, "images", None)):
+        return images[0]
+
+    raise AssertionError("No generated image found in Omni output")
+
+
+def extract_embedding(image: Image.Image, clip_model: CLIPModel, clip_processor: CLIPProcessor) -> torch.Tensor:
+    inputs = clip_processor(images=image.convert("RGB"), return_tensors="pt")
+    with torch.inference_mode():
+        features = clip_model.get_image_features(**inputs)
+        features = F.normalize(features, p=2, dim=-1)
+    return features.squeeze(0)
+
+
+def compare_semantic(
+    expected_embedding: torch.Tensor,
+    image: Image.Image,
+    clip_model: CLIPModel,
+    clip_processor: CLIPProcessor,
+) -> float:
+    features = extract_embedding(image, clip_model, clip_processor)
+    expected = F.normalize(expected_embedding, p=2, dim=-1)
+    return torch.dot(expected, features).item()
+
+
+def _generate_image(omni: Omni, use_system_prompt: str | None) -> Image.Image:
+    generator_device = current_omni_platform.device_type or "cuda"
+    sampling_params = OmniDiffusionSamplingParams(
+        seed=1234,
+        generator=torch.Generator(device=generator_device).manual_seed(1234),
+        num_outputs_per_prompt=1,
+    )
+    if use_system_prompt is not None:
+        sampling_params.extra_args = {"use_system_prompt": use_system_prompt}
+
+    outputs = omni.generate({"prompt": PROMPT}, sampling_params)
+    return _extract_generated_image(outputs)
+
+
+@pytest.mark.skipif(torch.cuda.device_count() < 8, reason="Need at least 8 CUDA GPUs for this test.")
+@pytest.mark.parametrize("system_prompt_name,use_system_prompt,expected_embedding", SYSTEM_PROMPT_CASES)
+def test_system_prompt_scores(
+    omni: Omni,
+    clip_bundle: tuple[CLIPModel, CLIPProcessor],
+    system_prompt_name: str,
+    use_system_prompt: str | None,
+    expected_embedding: torch.Tensor,
+) -> None:
+    clip_model, clip_processor = clip_bundle
+    generated_image = _generate_image(omni, use_system_prompt)
+    score = compare_semantic(expected_embedding, generated_image, clip_model, clip_processor)
+
+    print(f"{system_prompt_name}: CLIP cosine similarity = {score:.6f}")

tests/e2e/offline_inference/test_ltx2_cfg_parallel_parity.py

@@ -0,0 +1,243 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import hashlib
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+import numpy as np
+import pytest
+from PIL import Image
+
+from tests.utils import hardware_test
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+T2V_EXAMPLE = REPO_ROOT / "examples" / "offline_inference" / "text_to_video" / "text_to_video.py"
+I2V_EXAMPLE = REPO_ROOT / "examples" / "offline_inference" / "image_to_video" / "image_to_video.py"
+
+T2V_PROMPT = (
+    "At sunrise, a glowing paper lantern boat drifts through a narrow canal between mossy stone walls, "
+    "soft fog above the water, the camera slowly gliding forward as golden reflections shimmer across "
+    "the ripples, cinematic, realistic, highly detailed."
+)
+T2V_NEGATIVE_PROMPT = "worst quality, blurry, jittery motion, distorted, oversaturated, artifacts"
+I2V_PROMPT = "A cinematic dolly shot of a boat drifting on calm water at sunset"
+I2V_NEGATIVE_PROMPT = "worst quality, blurry, jittery motion"
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+
+def _get_ltx2_model() -> str:
+    return os.environ.get("VLLM_TEST_LTX2_MODEL", "Lightricks/LTX-2")
+
+
+def _md5(path: Path) -> str:
+    digest = hashlib.md5(usedforsecurity=False)
+    with path.open("rb") as f:
+        for chunk in iter(lambda: f.read(1024 * 1024), b""):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def _make_deterministic_test_image(path: Path) -> None:
+    """Create a deterministic 256x256 test image for I2V tests."""
+    rng = np.random.RandomState(42)
+    img = Image.fromarray(rng.randint(0, 255, (256, 256, 3), dtype=np.uint8))
+    img.save(path)
+
+
+def _run_and_check(cmd: list[str], env: dict, output_path: Path, expected_md5: str) -> None:
+    result = subprocess.run(cmd, cwd=REPO_ROOT, env=env, capture_output=True, text=True, check=False)
+    assert result.returncode == 0, (
+        f"Command failed (exit {result.returncode}).\nstdout:\n{result.stdout}\nstderr:\n{result.stderr}"
+    )
+    generated_md5 = _md5(output_path)
+    assert generated_md5 == expected_md5, (
+        f"Unexpected output md5: {generated_md5} != {expected_md5}.\nstdout:\n{result.stdout}\nstderr:\n{result.stderr}"
+    )
+
+
+# ── T2V tests ──
+
+
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+@pytest.mark.parallel
+@pytest.mark.slow
+@hardware_test(res={"cuda": "L4"}, num_cards=2)
+def test_ltx2_t2v_cfg_parallel(tmp_path: Path):
+    """T2V with CFG=4.0, cfg-parallel-size=2."""
+    output = tmp_path / "t2v_cfg4.mp4"
+    env = os.environ.copy()
+    env.setdefault("CUDA_VISIBLE_DEVICES", "0,1")
+    cmd = [
+        sys.executable,
+        str(T2V_EXAMPLE),
+        "--model",
+        _get_ltx2_model(),
+        "--prompt",
+        T2V_PROMPT,
+        "--negative-prompt",
+        T2V_NEGATIVE_PROMPT,
+        "--height",
+        "256",
+        "--width",
+        "256",
+        "--num-frames",
+        "145",
+        "--num-inference-steps",
+        "6",
+        "--guidance-scale",
+        "4.0",
+        "--frame-rate",
+        "24",
+        "--fps",
+        "24",
+        "--seed",
+        "42",
+        "--cfg-parallel-size",
+        "2",
+        "--enforce-eager",
+        "--output",
+        str(output),
+    ]
+    _run_and_check(cmd, env, output, expected_md5="08e606b9c522fee4b6f30cee8b77db40")
+
+
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+@pytest.mark.slow
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_ltx2_t2v_no_cfg(tmp_path: Path):
+    """T2V with CFG=1.0 (no classifier-free guidance)."""
+    output = tmp_path / "t2v_nocfg.mp4"
+    env = os.environ.copy()
+    env.setdefault("CUDA_VISIBLE_DEVICES", "0")
+    cmd = [
+        sys.executable,
+        str(T2V_EXAMPLE),
+        "--model",
+        _get_ltx2_model(),
+        "--prompt",
+        T2V_PROMPT,
+        "--height",
+        "256",
+        "--width",
+        "256",
+        "--num-frames",
+        "145",
+        "--num-inference-steps",
+        "6",
+        "--guidance-scale",
+        "1.0",
+        "--frame-rate",
+        "24",
+        "--fps",
+        "24",
+        "--seed",
+        "42",
+        "--enforce-eager",
+        "--output",
+        str(output),
+    ]
+    _run_and_check(cmd, env, output, expected_md5="a83994b94b6e67c54a524e0383c45ce8")
+
+
+# ── I2V tests ──
+
+
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+@pytest.mark.parallel
+@pytest.mark.slow
+@hardware_test(res={"cuda": "L4"}, num_cards=2)
+def test_ltx2_i2v_cfg_parallel(tmp_path: Path):
+    """I2V with CFG=4.0, cfg-parallel-size=2."""
+    test_image = tmp_path / "test_input.png"
+    _make_deterministic_test_image(test_image)
+    output = tmp_path / "i2v_cfg4.mp4"
+    env = os.environ.copy()
+    env.setdefault("CUDA_VISIBLE_DEVICES", "0,1")
+    cmd = [
+        sys.executable,
+        str(I2V_EXAMPLE),
+        "--model",
+        _get_ltx2_model(),
+        "--model-class-name",
+        "LTX2ImageToVideoPipeline",
+        "--image",
+        str(test_image),
+        "--prompt",
+        I2V_PROMPT,
+        "--negative-prompt",
+        I2V_NEGATIVE_PROMPT,
+        "--height",
+        "256",
+        "--width",
+        "256",
+        "--num-frames",
+        "73",
+        "--num-inference-steps",
+        "6",
+        "--guidance-scale",
+        "4.0",
+        "--frame-rate",
+        "24",
+        "--fps",
+        "24",
+        "--seed",
+        "42",
+        "--cfg-parallel-size",
+        "2",
+        "--enforce-eager",
+        "--output",
+        str(output),
+    ]
+    _run_and_check(cmd, env, output, expected_md5="aed7e56084b36373244d8f839b16d115")
+
+
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+@pytest.mark.slow
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_ltx2_i2v_no_cfg(tmp_path: Path):
+    """I2V with CFG=1.0 (no classifier-free guidance)."""
+    test_image = tmp_path / "test_input.png"
+    _make_deterministic_test_image(test_image)
+    output = tmp_path / "i2v_nocfg.mp4"
+    env = os.environ.copy()
+    env.setdefault("CUDA_VISIBLE_DEVICES", "0")
+    cmd = [
+        sys.executable,
+        str(I2V_EXAMPLE),
+        "--model",
+        _get_ltx2_model(),
+        "--model-class-name",
+        "LTX2ImageToVideoPipeline",
+        "--image",
+        str(test_image),
+        "--prompt",
+        I2V_PROMPT,
+        "--height",
+        "256",
+        "--width",
+        "256",
+        "--num-frames",
+        "73",
+        "--num-inference-steps",
+        "6",
+        "--guidance-scale",
+        "1.0",
+        "--frame-rate",
+        "24",
+        "--fps",
+        "24",
+        "--seed",
+        "42",
+        "--enforce-eager",
+        "--output",
+        str(output),
+    ]
+    _run_and_check(cmd, env, output, expected_md5="81b21ede12753e9e14a357a6c548b666")

tests/e2e/offline_inference/test_magi_human.py

@@ -0,0 +1,144 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""End-to-end tests for MagiHuman pipeline via vLLM-Omni."""
+
+import io
+
+import av
+import numpy as np
+import pytest
+
+from tests.utils import hardware_test
+from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes
+from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+
+def _validate_mp4(video_bytes: bytes, min_frames: int = 10) -> None:
+    """Validate that the MP4 contains meaningful video and audio tracks."""
+    container = av.open(io.BytesIO(video_bytes))
+
+    v_streams = [s for s in container.streams if s.type == "video"]
+    assert len(v_streams) >= 1, "No video stream found in MP4"
+
+    a_streams = [s for s in container.streams if s.type == "audio"]
+    assert len(a_streams) >= 1, "No audio stream found in MP4"
+
+    v_stream = v_streams[0]
+    assert v_stream.width >= 1080, f"Unexpected video width: {v_stream.width}"
+    assert v_stream.height >= 1056, f"Unexpected video height: {v_stream.height}"
+
+    frame_count = 0
+    for frame in container.decode(video=0):
+        frame_count += 1
+        if frame_count >= min_frames:
+            break
+    assert frame_count >= min_frames, f"Video has only {frame_count} frames (expected >= {min_frames})"
+
+    container.close()
+
+
+@pytest.mark.core_model
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "H100"}, num_cards=2)
+def test_magi_human_e2e(run_level):
+    """End-to-end test for MagiHuman generating video and audio."""
+    if run_level != "advanced_model":
+        pytest.skip("MagiHuman e2e test requires advanced_model run level with real weights.")
+
+    model_path = "SII-GAIR/daVinci-MagiHuman-Base-1080p"
+
+    omni = Omni(
+        model=model_path,
+        init_timeout=1200,
+        tensor_parallel_size=2,
+    )
+
+    prompt = (
+        "A young woman with long, wavy golden blonde hair and bright blue eyes, "
+        "wearing a fitted ivory silk blouse with a delicate lace collar, sits "
+        "stationary in front of a softly lit, blurred warm-toned interior. Her "
+        "overall disposition is warm, composed, and gently confident. The camera "
+        "holds a static medium close-up, framing her from the shoulders up, "
+        "with shallow depth of field keeping her face in sharp focus. Soft "
+        "directional key light falls from the upper left, casting a gentle "
+        "highlight along her cheekbone and nose bridge. She draws a quiet breath, "
+        "the levator labii superiors relaxing as her lips part. She speaks in "
+        "clear, warm, unhurried American English: "
+        "\"The most beautiful things in life aren't things at all — "
+        "they're moments, feelings, and the people who make you feel truly alive.\" "
+        "Her jaw descends smoothly on each stressed syllable; the orbicularis oris "
+        "shapes each vowel with precision. A faint, genuine smile engages the "
+        "zygomaticus major, lifting her lip corners fractionally. Her brows rest "
+        "in a soft, neutral arch throughout. She maintains steady, forward-facing "
+        "eye contact. Head position remains level; no torso displacement occurs.\n\n"
+        "Dialogue:\n"
+        "<Young blonde woman, American English>: "
+        "\"The most beautiful things in life aren't things at all — "
+        "they're moments, feelings, and the people who make you feel truly alive.\"\n\n"
+        "Background Sound:\n"
+        "<Soft, warm indoor ambience with a faint distant piano melody>"
+    )
+
+    sampling_params = OmniDiffusionSamplingParams(
+        height=256,
+        width=448,
+        num_inference_steps=8,
+        seed=52,
+        extra_args={
+            "seconds": 5,
+            "sr_height": 1080,
+            "sr_width": 1920,
+            "sr_num_inference_steps": 5,
+        },
+    )
+
+    try:
+        outputs = list(
+            omni.generate(
+                prompts=[prompt],
+                sampling_params_list=[sampling_params],
+            )
+        )
+
+        assert len(outputs) > 0, "No outputs returned"
+        first = outputs[0]
+
+        assert hasattr(first, "images") and first.images, "No video frames in output"
+        video_frames = first.images[0]
+        assert isinstance(video_frames, np.ndarray), f"Expected numpy array, got {type(video_frames)}"
+        assert video_frames.ndim == 4, f"Expected 4D array (T,H,W,3), got shape {video_frames.shape}"
+
+        mm = first.multimodal_output
+        assert mm, "multimodal_output is empty or missing"
+
+        audio_waveform = mm.get("audio")
+        assert audio_waveform is not None, "No audio waveform in multimodal_output"
+
+        audio_sample_rate = mm.get("audio_sample_rate")
+        assert audio_sample_rate is not None, (
+            "audio_sample_rate not found in multimodal_output; model post-process must propagate it"
+        )
+        assert isinstance(audio_sample_rate, (int, float)), (
+            f"audio_sample_rate should be numeric, got {type(audio_sample_rate)}"
+        )
+        assert int(audio_sample_rate) > 0, f"audio_sample_rate must be positive, got {audio_sample_rate}"
+
+        fps = mm.get("fps")
+        assert fps is not None, "fps not found in multimodal_output; model post-process must propagate it"
+        assert isinstance(fps, (int, float)), f"fps should be numeric, got {type(fps)}"
+        assert int(fps) > 0, f"fps must be positive, got {fps}"
+
+        video_bytes = mux_video_audio_bytes(
+            video_frames,
+            audio_waveform,
+            fps=float(fps),
+            audio_sample_rate=int(audio_sample_rate),
+        )
+        assert isinstance(video_bytes, bytes), f"Expected MP4 bytes, got {type(video_bytes)}"
+        assert len(video_bytes) > 1000, f"MP4 too small ({len(video_bytes)} bytes)"
+
+        _validate_mp4(video_bytes)
+    finally:
+        omni.close()

tests/e2e/online_serving/test_cosyvoice3_tts.py

@@ -50,8 +50,18 @@ tts_server_params = [
     )
 ]
 
+tts_async_chunk_server_params = [
+    pytest.param(
+        OmniServerParams(
+            model=MODEL,
+            stage_config_path=get_stage_config("cosyvoice3_async_chunk.yaml"),
+            server_args=["--trust-remote-code", "--disable-log-stats"],
+        ),
+        id="cosyvoice3_async_chunk",
+    )
+]
+
 
-@pytest.mark.advanced_model
 @pytest.mark.core_model
 @pytest.mark.omni
 @hardware_test(res={"cuda": "H100"}, num_cards=1)
@@ -76,16 +86,16 @@ def test_voice_clone_zh_001(omni_server, openai_client) -> None:
     openai_client.send_audio_speech_request(request_config)
 
 
-@pytest.mark.advanced_model
+@pytest.mark.core_model
 @pytest.mark.omni
 @hardware_test(res={"cuda": "H100"}, num_cards=1)
-@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
+@pytest.mark.parametrize("omni_server", tts_async_chunk_server_params, indirect=True)
 def test_voice_clone_zh_002(omni_server, openai_client) -> None:
     """
-    Test voice cloning TTS with Chinese text via OpenAI API.
-    Deploy Setting: default yaml
+    Test voice cloning TTS with Chinese text via async_chunk streaming.
+    Deploy Setting: cosyvoice3_async_chunk.yaml
     Input Modal: text + ref_audio + ref_text
-    Output Modal: audio
+    Output Modal: audio (streamed)
     Input Setting: stream=True
     Datasets: single request
     """
@@ -100,7 +110,7 @@ def test_voice_clone_zh_002(omni_server, openai_client) -> None:
     openai_client.send_audio_speech_request(request_config)
 
 
-@pytest.mark.advanced_model
+@pytest.mark.core_model
 @pytest.mark.omni
 @hardware_test(res={"cuda": "H100"}, num_cards=1)
 @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)

tests/e2e/online_serving/test_dynin_omni_expansion.py

@@ -0,0 +1,160 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Example online tests for Dynin-Omni model.
+"""
+
+import base64
+import gc
+import os
+from io import BytesIO
+from pathlib import Path
+
+import numpy as np
+import pytest
+import soundfile as sf
+from vllm.assets.image import ImageAsset
+
+from tests import conftest as tests_conftest
+from tests.conftest import OmniServerParams
+from tests.utils import hardware_test
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
+
+MODEL = "snu-aidas/Dynin-Omni"
+STAGE_CONFIG = str(Path(__file__).parent.parent / "stage_configs" / "dynin_omni_ci.yaml")
+_WHISPER_SAMPLE_RATE_HZ = 16_000
+
+T2I_PROMPT = "A high quality detailed living room interior photo."
+T2S_PROMPT = "Please read this sentence naturally: Hello from Dynin-Omni online serving."
+I2I_PROMPT = "Transform this outdoor nature boardwalk scene into a painting style with vivid colors."
+
+TEST_PARAMS = [OmniServerParams(model=MODEL, stage_config_path=STAGE_CONFIG)]
+_STAGE_COUNT = 3
+_I2I_STAGE_SAMPLING = {"max_tokens": 1, "temperature": 0.0, "top_p": 1.0, "detokenize": False}
+
+
+def _prepare_audio_waveform_for_whisper(audio_data: np.ndarray, samplerate: int) -> np.ndarray:
+    """Normalize decoded audio into a mono 16 kHz float32 waveform for Whisper."""
+    if samplerate <= 0:
+        raise ValueError(f"Invalid audio sample rate: {samplerate}")
+
+    waveform = np.asarray(audio_data, dtype=np.float32)
+    if waveform.ndim == 0:
+        raise ValueError("Audio waveform must have at least one dimension")
+    if waveform.ndim > 1:
+        waveform = np.mean(waveform, axis=1)
+    if waveform.size == 0:
+        raise ValueError("Empty audio waveform")
+
+    if samplerate != _WHISPER_SAMPLE_RATE_HZ:
+        target_num_samples = max(int(round(waveform.shape[0] * _WHISPER_SAMPLE_RATE_HZ / samplerate)), 1)
+        source_positions = np.arange(waveform.shape[0], dtype=np.float64)
+        target_positions = np.linspace(
+            0.0,
+            max(waveform.shape[0] - 1, 0),
+            num=target_num_samples,
+            dtype=np.float64,
+        )
+        waveform = np.interp(target_positions, source_positions, waveform).astype(np.float32)
+
+    return np.ascontiguousarray(np.clip(waveform, -1.0, 1.0), dtype=np.float32)
+
+
+def _convert_audio_bytes_to_text_without_ffmpeg(raw_bytes: bytes) -> str:
+    """Dynin t2s keeps Whisper transcription local to this test module and avoids ffmpeg."""
+    import whisper
+
+    data, samplerate = sf.read(BytesIO(raw_bytes), dtype="float32", always_2d=True)
+    audio_waveform = _prepare_audio_waveform_for_whisper(data, samplerate)
+
+    model = whisper.load_model("small", device="cpu")
+    try:
+        transcript = model.transcribe(
+            audio_waveform,
+            temperature=0.0,
+            word_timestamps=True,
+            condition_on_previous_text=False,
+        )["text"]
+    finally:
+        del model
+        gc.collect()
+
+    return transcript or ""
+
+
+@pytest.fixture
+def dynin_t2s_openai_client(openai_client, monkeypatch):
+    monkeypatch.setattr(
+        tests_conftest,
+        "convert_audio_bytes_to_text",
+        _convert_audio_bytes_to_text_without_ffmpeg,
+    )
+    return openai_client
+
+
+def _build_t2i_messages(prompt: str) -> list[dict]:
+    return [{"role": "user", "content": [{"type": "text", "text": f"<|t2i|> {prompt}"}]}]
+
+
+def _build_t2s_messages(prompt: str) -> list[dict]:
+    return [{"role": "user", "content": [{"type": "text", "text": f"<|t2s|> {prompt}"}]}]
+
+
+def _build_i2i_messages(prompt: str) -> list[dict]:
+    input_image = ImageAsset("2560px-Gfp-wisconsin-madison-the-nature-boardwalk").pil_image.convert("RGB")
+    buffer = BytesIO()
+    input_image.save(buffer, format="JPEG")
+    image_b64 = base64.b64encode(buffer.getvalue()).decode("utf-8")
+    return [
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": f"<|i2i|> {prompt}"},
+                {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{image_b64}"}},
+            ],
+        }
+    ]
+
+
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"})
+@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
+def test_send_i2i_request_001(omni_server, openai_client) -> None:
+    request_config = {
+        "model": omni_server.model,
+        "messages": _build_i2i_messages(I2I_PROMPT),
+        "modalities": ["image"],
+        "extra_body": {
+            "sampling_params_list": [dict(_I2I_STAGE_SAMPLING) for _ in range(_STAGE_COUNT)],
+        },
+    }
+    openai_client.send_diffusion_request(request_config)
+
+
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"})
+@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
+def test_send_t2i_request_001(omni_server, openai_client) -> None:
+    request_config = {
+        "model": omni_server.model,
+        "messages": _build_t2i_messages(T2I_PROMPT),
+        "modalities": ["image"],
+    }
+    openai_client.send_diffusion_request(request_config)
+
+
+@pytest.mark.core_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"})
+@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
+def test_send_t2s_request_001(omni_server, dynin_t2s_openai_client) -> None:
+    request_config = {
+        "model": omni_server.model,
+        "messages": _build_t2s_messages(T2S_PROMPT),
+        "modalities": ["audio"],
+    }
+    dynin_t2s_openai_client.send_omni_request(request_config)

tests/e2e/online_serving/test_flux2_expansion.py

@@ -1,6 +1,12 @@
 """
 Tests for Flux2 Klein; currently Dev is implemented separately,
 but ideally these models will fold together in the future.
+
+Coverage:
+- FP8 + CacheDiT + Ulysses=2 + TP=2
+- Layerwise CPU offload + Ulysses=2 + Ring=2
+- Layerwise CPU offload + TP=2
+- Layerwise CPU offload + HSDP
 """
 
 import pytest
@@ -42,6 +48,45 @@ def _get_diffusion_feature_cases(model: str):
             ),
             marks=FOUR_CARD_FEATURE_MARKS,
         ),
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=[
+                    "--enable-layerwise-offload",
+                    "--ulysses-degree",
+                    "2",
+                    "--ring",
+                    "2",
+                ],
+            ),
+            id="layerwise_ulysses2_ring2",
+            marks=FOUR_CARD_FEATURE_MARKS,
+        ),
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=[
+                    "--enable-layerwise-offload",
+                    "--tensor-parallel-size",
+                    "2",
+                ],
+            ),
+            id="layerwise_tp2",
+            marks=FOUR_CARD_FEATURE_MARKS,
+        ),
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=[
+                    "--enable-layerwise-offload",
+                    "--use-hsdp",
+                    "--hsdp-shard-size",
+                    "2",
+                ],
+            ),
+            id="layerwise_hsdp",
+            marks=FOUR_CARD_FEATURE_MARKS,
+        ),
     ]
 
 

tests/e2e/online_serving/test_zimage_expansion.py

@@ -3,9 +3,12 @@ Tests of common diffusion feature combinations in online serving mode
 for Z-Image.
 
 Coverage is intentionally limited to the minimal 4xL4 cases that
-exercise Z-Image's supported parallel feature combinations:
+exercise Z-Image's supported feature combinations:
 - CacheDiT + FP8 + Ring=2 + TP=2
 - TeaCache + FP8 + Ulysses=2 + Ring=2
+- Layerwise CPU offload + Ulysses=2 + Ring=2
+- Layerwise CPU offload + TP=2
+- Layerwise CPU offload + HSDP
 """
 
 import pytest
@@ -64,12 +67,39 @@ def _get_diffusion_feature_cases():
             OmniServerParams(
                 model=MODEL,
                 server_args=[
+                    "--enable-layerwise-offload",
+                    "--ulysses-degree",
+                    "2",
+                    "--ring",
+                    "2",
+                ],
+            ),
+            id="layerwise_ulysses2_ring2",
+            marks=FOUR_CARD_MARKS,
+        ),
+        pytest.param(
+            OmniServerParams(
+                model=MODEL,
+                server_args=[
+                    "--enable-layerwise-offload",
+                    "--tensor-parallel-size",
+                    "2",
+                ],
+            ),
+            id="layerwise_tp2",
+            marks=FOUR_CARD_MARKS,
+        ),
+        pytest.param(
+            OmniServerParams(
+                model=MODEL,
+                server_args=[
+                    "--enable-layerwise-offload",
                     "--use-hsdp",
                     "--hsdp-shard-size",
                     "2",
                 ],
             ),
-            id="parallel_hsdp",
+            id="layerwise_hsdp",
             marks=[*FOUR_CARD_MARKS, pytest.mark.skip(reason="issue #2435")],
         ),
     ]

tests/e2e/stage_configs/dynin_omni_ci.yaml

@@ -0,0 +1,84 @@
+# stage config for running dynin_omni with a 3-stage architecture.
+# this config is intended for e2e smoke tests.
+
+stage_args:
+  - stage_id: 0
+    stage_type: llm
+    runtime:
+      process: true
+      devices: "0"
+      max_batch_size: 1
+    engine_args:
+      model_stage: token2text
+      model_arch: DyninOmniForConditionalGeneration
+      worker_type: generation
+      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
+      engine_output_type: latent
+      trust_remote_code: true
+      gpu_memory_utilization: 0.5
+      enforce_eager: true
+      enable_prefix_caching: false
+      async_scheduling: false
+      max_num_batched_tokens: 4096
+    is_comprehension: true
+    final_output: true
+    final_output_type: text
+
+  - stage_id: 1
+    stage_type: llm
+    runtime:
+      process: true
+      devices: "0"
+      max_batch_size: 1
+    engine_args:
+      model_stage: token2image
+      model_arch: DyninOmniForConditionalGeneration
+      worker_type: generation
+      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
+      engine_output_type: latent
+      trust_remote_code: true
+      gpu_memory_utilization: 0.2
+      enforce_eager: true
+      enable_prefix_caching: false
+      async_scheduling: false
+      max_num_batched_tokens: 4096
+    engine_input_source: [0]
+    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.dynin_omni.token2text_to_token2image
+    final_output: true
+    final_output_type: image
+
+  - stage_id: 2
+    stage_type: llm
+    runtime:
+      process: true
+      devices: "0"
+      max_batch_size: 1
+    engine_args:
+      model_stage: token2audio
+      model_arch: DyninOmniForConditionalGeneration
+      worker_type: generation
+      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
+      engine_output_type: latent
+      trust_remote_code: true
+      gpu_memory_utilization: 0.2
+      enforce_eager: true
+      enable_prefix_caching: false
+      async_scheduling: false
+      max_num_batched_tokens: 4096
+    engine_input_source: [1]
+    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.dynin_omni.token2image_to_token2audio
+    final_output: true
+    final_output_type: audio
+
+runtime:
+  enabled: true
+  defaults:
+    window_size: -1
+    max_inflight: 1
+  edges:
+    - from: 0
+      to: 1
+      window_size: -1
+    - from: 1
+      to: 2
+      window_size: -1

tests/engine/test_cross_stage_lora.py

@@ -0,0 +1,44 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for cross-stage LoRA routing in the orchestrator."""
+
+from __future__ import annotations
+
+import pytest
+from vllm.lora.request import LoRARequest
+from vllm.sampling_params import SamplingParams
+
+from vllm_omni.engine.orchestrator import build_engine_core_request_from_tokens
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class TestBuildEngineCoreRequestLoRA:
+    """Verify build_engine_core_request_from_tokens passes LoRA from params."""
+
+    def test_lora_extracted_from_diffusion_params(self):
+        lr = LoRARequest(lora_name="test", lora_int_id=1, lora_path="/tmp/fake")
+        params = OmniDiffusionSamplingParams(lora_request=lr)
+
+        # OmniDiffusionSamplingParams is not a SamplingParams, so
+        # build_engine_core_request_from_tokens takes the pooling path.
+        # We only care that lora_request is extracted via getattr.
+        request = build_engine_core_request_from_tokens(
+            request_id="req-1",
+            prompt={"prompt_token_ids": [1, 2, 3]},
+            params=params,
+            model_config=None,
+        )
+        assert request.lora_request is lr
+
+    def test_no_lora_on_sampling_params(self):
+        params = SamplingParams(max_tokens=10)
+
+        request = build_engine_core_request_from_tokens(
+            request_id="req-2",
+            prompt={"prompt_token_ids": [1, 2, 3]},
+            params=params,
+            model_config=None,
+        )
+        assert request.lora_request is None

tests/engine/test_orchestrator.py

@@ -0,0 +1,510 @@
+from __future__ import annotations
+
+import asyncio
+import concurrent.futures
+import queue
+import threading
+import time
+from dataclasses import dataclass
+from types import SimpleNamespace
+from typing import Any
+
+import janus
+import pytest
+from vllm.outputs import CompletionOutput, RequestOutput
+from vllm.sampling_params import SamplingParams
+
+from vllm_omni.engine.orchestrator import Orchestrator
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@dataclass
+class OrchestratorFixture:
+    orchestrator: Orchestrator
+    request_sync_q: Any
+    output_sync_q: Any
+    queues: tuple[janus.Queue, ...]
+    thread: threading.Thread
+    result_future: concurrent.futures.Future[None]
+
+
+class FakeStageClient:
+    def __init__(
+        self,
+        *,
+        stage_type: str = "llm",
+        final_output: bool = False,
+        final_output_type: str = "text",
+        next_inputs: list[dict] | None = None,
+    ) -> None:
+        self.stage_type = stage_type
+        self.final_output = final_output
+        self.final_output_type = final_output_type
+        self.next_inputs = list(next_inputs or [])
+        self.custom_process_input_func = None
+        self.add_request_calls: list[tuple] = []
+        self.abort_calls: list[list[str]] = []
+        self.shutdown_calls = 0
+        self._engine_core_outputs = queue.Queue()
+        self._diffusion_outputs = queue.Queue()
+
+    # Orchestrator-facing interface.
+    async def add_request_async(self, *args, **_kwargs) -> None:
+        self.add_request_calls.append(args)
+
+    async def get_output_async(self):
+        try:
+            return self._engine_core_outputs.get_nowait()
+        except queue.Empty:
+            return SimpleNamespace(outputs=[])
+
+    def get_diffusion_output_nowait(self):
+        try:
+            return self._diffusion_outputs.get_nowait()
+        except queue.Empty:
+            return None
+
+    def set_engine_outputs(self, outputs) -> None:
+        return None
+
+    def process_engine_inputs(self, stage_list, prompt=None):
+        return list(self.next_inputs)
+
+    async def abort_requests_async(self, request_ids: list[str]) -> None:
+        self.abort_calls.append(list(request_ids))
+
+    def shutdown(self) -> None:
+        self.shutdown_calls += 1
+
+    # Test helpers for seeding fake stage outputs.
+    def push_engine_core_outputs(self, outputs) -> None:
+        self._engine_core_outputs.put_nowait(outputs)
+
+    def push_diffusion_output(self, output) -> None:
+        self._diffusion_outputs.put_nowait(output)
+
+
+class FakeOutputProcessor:
+    def __init__(self, *, request_outputs: list[object] | None = None) -> None:
+        self.request_outputs = list(request_outputs or [])
+
+    def add_request(self, *_args, **_kwargs) -> None:
+        return None
+
+    def process_outputs(self, *_args, **_kwargs):
+        return SimpleNamespace(
+            request_outputs=list(self.request_outputs),
+            reqs_to_abort=[],
+        )
+
+    def update_scheduler_stats(self, _scheduler_stats) -> None:
+        return None
+
+
+def _sampling_params(max_tokens: int = 4) -> SamplingParams:
+    return SamplingParams(max_tokens=max_tokens)
+
+
+def _engine_core_outputs(tag: str, timestamp: float) -> SimpleNamespace:
+    return SimpleNamespace(outputs=[tag], timestamp=timestamp, scheduler_stats=None)
+
+
+def _build_request_output(
+    request_id: str,
+    *,
+    token_ids: list[int] | None = None,
+    prompt_token_ids: list[int] | None = None,
+    finished: bool = True,
+    text: str = "test",
+) -> RequestOutput:
+    completion = CompletionOutput(
+        index=0,
+        text=text,
+        token_ids=list(token_ids or [1, 2]),
+        cumulative_logprob=0.0,
+        logprobs=None,
+        finish_reason="stop" if finished else None,
+        stop_reason=None,
+    )
+    return RequestOutput(
+        request_id=request_id,
+        prompt="prompt",
+        prompt_token_ids=list(prompt_token_ids or [10, 11]),
+        prompt_logprobs=None,
+        outputs=[completion],
+        finished=finished,
+        metrics=None,
+        lora_request=None,
+    )
+
+
+def _build_harness(
+    stage_clients: list[object],
+    *,
+    output_processors: list[object] | None = None,
+    stage_vllm_configs: list[object] | None = None,
+    async_chunk: bool = False,
+) -> OrchestratorFixture:
+    if output_processors is None:
+        output_processors = [FakeOutputProcessor() for _ in stage_clients]
+    if stage_vllm_configs is None:
+        stage_vllm_configs = [SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)) for _ in stage_clients]
+
+    ready_future: concurrent.futures.Future[tuple[Orchestrator, janus.Queue, janus.Queue, janus.Queue]] = (
+        concurrent.futures.Future()
+    )
+    result_future: concurrent.futures.Future[None] = concurrent.futures.Future()
+
+    def _runner() -> None:
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+
+        async def _run() -> None:
+            request_queue = janus.Queue()
+            output_queue = janus.Queue()
+            rpc_queue = janus.Queue()
+            orchestrator = Orchestrator(
+                request_async_queue=request_queue.async_q,
+                output_async_queue=output_queue.async_q,
+                rpc_async_queue=rpc_queue.async_q,
+                stage_clients=stage_clients,
+                output_processors=output_processors,
+                stage_vllm_configs=stage_vllm_configs,
+                async_chunk=async_chunk,
+            )
+            ready_future.set_result((orchestrator, request_queue, output_queue, rpc_queue))
+            await orchestrator.run()
+
+        try:
+            loop.run_until_complete(_run())
+            result_future.set_result(None)
+        except Exception as exc:
+            result_future.set_exception(exc)
+        finally:
+            try:
+                pending = [task for task in asyncio.all_tasks(loop) if not task.done()]
+                for task in pending:
+                    task.cancel()
+                if pending:
+                    loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True))
+                loop.run_until_complete(loop.shutdown_asyncgens())
+            finally:
+                asyncio.set_event_loop(None)
+                loop.close()
+
+    thread = threading.Thread(target=_runner, daemon=True, name="test-orchestrator")
+    thread.start()
+
+    orchestrator, request_queue, output_queue, rpc_queue = ready_future.result(timeout=5)
+    return OrchestratorFixture(
+        orchestrator=orchestrator,
+        request_sync_q=request_queue.sync_q,
+        output_sync_q=output_queue.sync_q,
+        queues=(request_queue, output_queue, rpc_queue),
+        thread=thread,
+        result_future=result_future,
+    )
+
+
+async def _shutdown_orchestrator(orchestrator_fixture: OrchestratorFixture) -> None:
+    orchestrator_fixture.request_sync_q.put_nowait({"type": "shutdown"})
+    await asyncio.to_thread(orchestrator_fixture.thread.join, 5)
+    if orchestrator_fixture.thread.is_alive():
+        raise AssertionError("Timed out waiting for orchestrator thread shutdown")
+    orchestrator_fixture.result_future.result(timeout=0)
+
+
+async def _wait_for(predicate, *, timeout: float = 2.0) -> None:
+    deadline = time.monotonic() + timeout
+    while not predicate():
+        if time.monotonic() >= deadline:
+            raise AssertionError("Timed out waiting for predicate")
+        await asyncio.sleep(0.01)
+
+
+async def _get_output_message(orchestrator_fixture: OrchestratorFixture, *, timeout: float = 2.0) -> dict:
+    deadline = time.monotonic() + timeout
+    while True:
+        if time.monotonic() >= deadline:
+            raise AssertionError("Timed out waiting for orchestrator output")
+        try:
+            msg = orchestrator_fixture.output_sync_q.get_nowait()
+        except queue.Empty:
+            await asyncio.sleep(0.01)
+            continue
+        if msg.get("type") == "output":
+            return msg
+
+
+async def _enqueue_add_request(
+    orchestrator_fixture: OrchestratorFixture,
+    *,
+    request_id: str,
+    prompt,
+    original_prompt,
+    sampling_params_list,
+    final_stage_id: int,
+) -> None:
+    orchestrator_fixture.request_sync_q.put_nowait(
+        {
+            "type": "add_request",
+            "request_id": request_id,
+            "prompt": prompt,
+            "original_prompt": original_prompt,
+            "sampling_params_list": sampling_params_list,
+            "final_stage_id": final_stage_id,
+        }
+    )
+
+
+async def _enqueue_abort_request(orchestrator_fixture: OrchestratorFixture, request_ids: list[str]) -> None:
+    orchestrator_fixture.request_sync_q.put_nowait(
+        {
+            "type": "abort",
+            "request_ids": request_ids,
+        }
+    )
+
+
+@pytest.fixture
+def orchestrator_factory():
+    fixtures: list[OrchestratorFixture] = []
+
+    def _factory(*args, **kwargs) -> OrchestratorFixture:
+        fixture = _build_harness(*args, **kwargs)
+        fixtures.append(fixture)
+        return fixture
+
+    yield _factory
+
+    for fixture in fixtures:
+        if fixture.thread.is_alive():
+            fixture.request_sync_q.put_nowait({"type": "shutdown"})
+            fixture.thread.join(timeout=5)
+        for q in fixture.queues:
+            q.close()
+
+
+@pytest.mark.asyncio
+async def test_run_two_stage_llm(orchestrator_factory) -> None:
+    stage0 = FakeStageClient(stage_type="llm", final_output=False)
+    stage1 = FakeStageClient(
+        stage_type="llm",
+        final_output=True,
+        next_inputs=[{"prompt_token_ids": [7, 8, 9]}],
+    )
+    processors = [
+        FakeOutputProcessor(request_outputs=[_build_request_output("req-llm", token_ids=[3, 4], finished=True)]),
+        FakeOutputProcessor(request_outputs=[_build_request_output("req-llm", token_ids=[10, 11], finished=True)]),
+    ]
+    orchestrator_fixture = orchestrator_factory([stage0, stage1], output_processors=processors)
+    request = SimpleNamespace(request_id="req-llm", prompt_token_ids=[1, 2, 3])
+
+    try:
+        await _enqueue_add_request(
+            orchestrator_fixture,
+            request_id="req-llm",
+            prompt=request,
+            original_prompt={"prompt": "hello"},
+            sampling_params_list=[_sampling_params(), _sampling_params()],
+            final_stage_id=1,
+        )
+
+        await _wait_for(lambda: len(stage0.add_request_calls) == 1)
+        stage0.push_engine_core_outputs(_engine_core_outputs("stage0-raw", 1.0))
+
+        await _wait_for(lambda: len(stage1.add_request_calls) == 1)
+        stage1_request = stage1.add_request_calls[0][0]
+        assert stage1_request.request_id == "req-llm"
+        assert stage1_request.prompt_token_ids == [7, 8, 9]
+
+        stage1.push_engine_core_outputs(_engine_core_outputs("stage1-raw", 2.0))
+
+        output_msg = await _get_output_message(orchestrator_fixture)
+
+        assert output_msg["request_id"] == "req-llm"
+        assert output_msg["stage_id"] == 1
+        assert output_msg["finished"] is True
+        assert output_msg["engine_outputs"].request_id == "req-llm"
+        assert "req-llm" not in orchestrator_fixture.orchestrator.request_states
+    finally:
+        await _shutdown_orchestrator(orchestrator_fixture)
+
+
+@pytest.mark.asyncio
+async def test_run_single_stage_diffusion(orchestrator_factory) -> None:
+    stage0 = FakeStageClient(stage_type="diffusion", final_output=True, final_output_type="image")
+    orchestrator_fixture = orchestrator_factory([stage0])
+    params = OmniDiffusionSamplingParams()
+
+    try:
+        await _enqueue_add_request(
+            orchestrator_fixture,
+            request_id="req-diff",
+            prompt={"prompt": "draw a cat"},
+            original_prompt={"prompt": "draw a cat"},
+            sampling_params_list=[params],
+            final_stage_id=0,
+        )
+
+        await _wait_for(lambda: len(stage0.add_request_calls) == 1)
+        stage0.push_diffusion_output(
+            OmniRequestOutput.from_diffusion(
+                request_id="req-diff",
+                images=[],
+                final_output_type="image",
+            )
+        )
+
+        output_msg = await _get_output_message(orchestrator_fixture)
+
+        assert output_msg["request_id"] == "req-diff"
+        assert output_msg["stage_id"] == 0
+        assert output_msg["finished"] is True
+        assert output_msg["engine_outputs"].request_id == "req-diff"
+        assert "req-diff" not in orchestrator_fixture.orchestrator.request_states
+    finally:
+        await _shutdown_orchestrator(orchestrator_fixture)
+
+
+@pytest.mark.asyncio
+async def test_run_llm_to_diffusion(orchestrator_factory) -> None:
+    stage0 = FakeStageClient(stage_type="llm", final_output=False)
+    stage1 = FakeStageClient(stage_type="diffusion", final_output=True, final_output_type="image")
+    processors = [
+        FakeOutputProcessor(request_outputs=[_build_request_output("req-img", token_ids=[3, 4], finished=True)]),
+        FakeOutputProcessor(),
+    ]
+    orchestrator_fixture = orchestrator_factory([stage0, stage1], output_processors=processors)
+    request = SimpleNamespace(request_id="req-img", prompt_token_ids=[1, 2, 3])
+    params = OmniDiffusionSamplingParams()
+    original_prompt = {"prompt": "draw a fox"}
+
+    try:
+        await _enqueue_add_request(
+            orchestrator_fixture,
+            request_id="req-img",
+            prompt=request,
+            original_prompt=original_prompt,
+            sampling_params_list=[_sampling_params(), params],
+            final_stage_id=1,
+        )
+
+        await _wait_for(lambda: len(stage0.add_request_calls) == 1)
+        stage0.push_engine_core_outputs(_engine_core_outputs("stage0-raw", 1.0))
+
+        await _wait_for(lambda: len(stage1.add_request_calls) == 1)
+        assert stage1.add_request_calls[0] == ("req-img", original_prompt, params)
+
+        stage1.push_diffusion_output(
+            OmniRequestOutput.from_diffusion(
+                request_id="req-img",
+                images=[],
+                final_output_type="image",
+            )
+        )
+
+        output_msg = await _get_output_message(orchestrator_fixture)
+
+        assert output_msg["request_id"] == "req-img"
+        assert output_msg["stage_id"] == 1
+        assert output_msg["finished"] is True
+        assert output_msg["engine_outputs"].request_id == "req-img"
+        assert "req-img" not in orchestrator_fixture.orchestrator.request_states
+    finally:
+        await _shutdown_orchestrator(orchestrator_fixture)
+
+
+@pytest.mark.asyncio
+async def test_run_async_chunk(orchestrator_factory) -> None:
+    stage0 = FakeStageClient(stage_type="llm", final_output=False)
+    stage1 = FakeStageClient(stage_type="llm", final_output=True)
+    processors = [
+        FakeOutputProcessor(request_outputs=[_build_request_output("req-async", token_ids=[1], finished=True)]),
+        FakeOutputProcessor(request_outputs=[_build_request_output("req-async", token_ids=[20, 21], finished=True)]),
+    ]
+    orchestrator_fixture = orchestrator_factory(
+        [stage0, stage1],
+        output_processors=processors,
+        async_chunk=True,
+    )
+    request = SimpleNamespace(request_id="req-async", prompt_token_ids=[1, 2, 3, 4])
+
+    try:
+        await _enqueue_add_request(
+            orchestrator_fixture,
+            request_id="req-async",
+            prompt=request,
+            original_prompt={"prompt": "hello async"},
+            sampling_params_list=[_sampling_params(), _sampling_params()],
+            final_stage_id=1,
+        )
+
+        await _wait_for(lambda: len(stage1.add_request_calls) == 1)
+        prewarmed_request = stage1.add_request_calls[0][0]
+        assert prewarmed_request.request_id == "req-async"
+        assert prewarmed_request.prompt_token_ids
+        assert all(token_id == 0 for token_id in prewarmed_request.prompt_token_ids)
+
+        stage1.push_engine_core_outputs(_engine_core_outputs("stage1-final", 3.0))
+
+        output_msg = await _get_output_message(orchestrator_fixture)
+
+        assert output_msg["request_id"] == "req-async"
+        assert output_msg["stage_id"] == 1
+        assert output_msg["finished"] is True
+        assert "req-async" not in orchestrator_fixture.orchestrator.request_states
+    finally:
+        await _shutdown_orchestrator(orchestrator_fixture)
+
+
+@pytest.mark.asyncio
+async def test_run_shutdown(orchestrator_factory) -> None:
+    stages = [
+        FakeStageClient(stage_type="llm", final_output=False),
+        FakeStageClient(stage_type="diffusion", final_output=True, final_output_type="image"),
+    ]
+    orchestrator_fixture = orchestrator_factory(stages)
+
+    await _shutdown_orchestrator(orchestrator_fixture)
+
+    assert not orchestrator_fixture.thread.is_alive()
+    for stage in stages:
+        assert stage.shutdown_calls == 1
+
+
+@pytest.mark.asyncio
+async def test_run_abort(orchestrator_factory) -> None:
+    stages = [
+        FakeStageClient(stage_type="llm", final_output=False),
+        FakeStageClient(stage_type="llm", final_output=True),
+    ]
+    processors = [
+        FakeOutputProcessor(request_outputs=[_build_request_output("req-abort", token_ids=[1], finished=True)]),
+        FakeOutputProcessor(request_outputs=[_build_request_output("req-abort", token_ids=[2], finished=True)]),
+    ]
+    orchestrator_fixture = orchestrator_factory(stages, output_processors=processors)
+    request = SimpleNamespace(request_id="req-abort", prompt_token_ids=[1, 2, 3])
+
+    try:
+        await _enqueue_add_request(
+            orchestrator_fixture,
+            request_id="req-abort",
+            prompt=request,
+            original_prompt={"prompt": "cancel me"},
+            sampling_params_list=[_sampling_params(), _sampling_params()],
+            final_stage_id=1,
+        )
+        await _wait_for(lambda: len(stages[0].add_request_calls) == 1)
+
+        await _enqueue_abort_request(orchestrator_fixture, ["req-abort"])
+        await _wait_for(lambda: all(stage.abort_calls for stage in stages))
+
+        for stage in stages:
+            assert stage.abort_calls == [["req-abort"]]
+        assert "req-abort" not in orchestrator_fixture.orchestrator.request_states
+    finally:
+        await _shutdown_orchestrator(orchestrator_fixture)

tests/engine/test_orchestrator_kv_sender_info.py

@@ -0,0 +1,207 @@
+import asyncio
+from types import SimpleNamespace
+
+import pytest
+from vllm import SamplingParams
+
+from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState
+from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClient
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _DummySenderStage:
+    def __init__(self, sender_info):
+        self._sender_info = sender_info
+        self.engine_outputs = None
+
+    def set_engine_outputs(self, outputs):
+        self.engine_outputs = outputs
+
+    def get_kv_sender_info(self):
+        return self._sender_info
+
+
+class _DummyDiffusionStage:
+    stage_type = "diffusion"
+    custom_process_input_func = None
+
+    def __init__(self, engine_input_source=None):
+        self.engine_input_source = engine_input_source or [0]
+        self.calls = []
+
+    async def add_request_async(self, request_id, prompt, sampling_params, kv_sender_info=None):
+        self.calls.append(
+            {
+                "request_id": request_id,
+                "prompt": prompt,
+                "sampling_params": sampling_params,
+                "kv_sender_info": kv_sender_info,
+            }
+        )
+
+
+def test_stage_engine_core_client_builds_kv_sender_info_from_tcp_address():
+    client = object.__new__(StageEngineCoreClient)
+    client.stage_id = 0
+    client.client_addresses = {"input_address": "tcp://10.20.30.40:1234"}
+    client._omni_kv_config = None
+    client._kv_sender_info = None
+    client._kv_sender_initialized = False
+    client._kv_sender_host = client._resolve_contact_host()
+    client._initialize_kv_sender_endpoint()
+
+    assert client.get_kv_sender_info() == {
+        "host": "10.20.30.40",
+        "zmq_port": 50151,
+    }
+
+
+def test_stage_engine_core_client_falls_back_to_detected_ip_for_loopback(monkeypatch):
+    client = object.__new__(StageEngineCoreClient)
+    client.stage_id = 1
+    client.client_addresses = {"input_address": "tcp://127.0.0.1:1234"}
+    client._omni_kv_config = None
+    client._kv_sender_info = None
+    client._kv_sender_initialized = False
+    monkeypatch.setattr(client, "_detect_local_ip", lambda: "192.168.0.12")
+    client._kv_sender_host = client._resolve_contact_host()
+    client._initialize_kv_sender_endpoint()
+
+    assert client.get_kv_sender_info() == {
+        "host": "192.168.0.12",
+        "zmq_port": 50152,
+    }
+
+
+def test_stage_engine_core_client_uses_connector_config_for_sender_port():
+    client = object.__new__(StageEngineCoreClient)
+    client.stage_id = 3
+    client.client_addresses = {"input_address": "tcp://10.20.30.40:1234"}
+    client._kv_sender_info = None
+    client._kv_sender_initialized = False
+    client._omni_kv_config = {
+        "omni_from_stage": "3",
+        "connector_config": {
+            "type": "MooncakeTransferEngineConnector",
+            "role": "sender",
+            "host": "10.20.30.99",
+            "zmq_port": 51000,
+        },
+    }
+    client._kv_sender_host = client._resolve_contact_host()
+    client._initialize_kv_sender_endpoint()
+
+    assert client.get_kv_sender_info() == {
+        "host": "10.20.30.99",
+        "zmq_port": 51103,
+    }
+
+
+def test_stage_engine_core_client_preserves_explicit_loopback_sender_host():
+    client = object.__new__(StageEngineCoreClient)
+    client.stage_id = 2
+    client.client_addresses = {"input_address": "tcp://10.20.30.40:1234"}
+    client._kv_sender_info = None
+    client._kv_sender_initialized = False
+    client._omni_kv_config = {
+        "omni_from_stage": "2",
+        "connector_config": {
+            "type": "MooncakeTransferEngineConnector",
+            "role": "sender",
+            "host": "127.0.0.1",
+            "zmq_port": 51000,
+        },
+    }
+    client._kv_sender_host = client._resolve_contact_host()
+    client._initialize_kv_sender_endpoint()
+
+    assert client.get_kv_sender_info() == {
+        "host": "127.0.0.1",
+        "zmq_port": 51102,
+    }
+
+
+def test_forward_to_diffusion_attaches_kv_sender_info():
+    orchestrator = object.__new__(Orchestrator)
+    sender_stage = _DummySenderStage({"host": "10.0.0.2", "zmq_port": 50151})
+    diffusion_stage = _DummyDiffusionStage(engine_input_source=[0])
+
+    orchestrator.num_stages = 2
+    orchestrator.stage_clients = [sender_stage, diffusion_stage]
+    orchestrator._companion_map = {}
+    orchestrator.stage_vllm_configs = [None, None]
+    orchestrator.output_processors = [None, None]
+
+    params = OmniDiffusionSamplingParams()
+    req_state = OrchestratorRequestState(
+        request_id="req-1",
+        prompt={"prompt": "hello"},
+        sampling_params_list=[SamplingParams(max_tokens=4), params],
+        final_stage_id=1,
+    )
+
+    output = SimpleNamespace(request_id="req-1", finished=True)
+    asyncio.run(Orchestrator._forward_to_next_stage(orchestrator, "req-1", 0, output, req_state))
+
+    assert sender_stage.engine_outputs == [output]
+    assert diffusion_stage.calls[0]["request_id"] == "req-1"
+    assert diffusion_stage.calls[0]["kv_sender_info"] == {
+        0: {"host": "10.0.0.2", "zmq_port": 50151},
+    }
+    assert req_state.stage_submit_ts[1] > 0
+
+
+def test_forward_to_diffusion_uses_engine_input_source_for_kv_sender_info():
+    orchestrator = object.__new__(Orchestrator)
+    source_stage = _DummySenderStage({"host": "10.0.0.2", "zmq_port": 50151})
+    previous_stage = _DummySenderStage({"host": "10.0.0.9", "zmq_port": 59999})
+    diffusion_stage = _DummyDiffusionStage(engine_input_source=[0])
+
+    orchestrator.num_stages = 3
+    orchestrator.stage_clients = [source_stage, previous_stage, diffusion_stage]
+    orchestrator._companion_map = {}
+    orchestrator.stage_vllm_configs = [None, None, None]
+    orchestrator.output_processors = [None, None, None]
+
+    params = OmniDiffusionSamplingParams()
+    req_state = OrchestratorRequestState(
+        request_id="req-3",
+        prompt={"prompt": "hello"},
+        sampling_params_list=[SamplingParams(max_tokens=4), SamplingParams(max_tokens=4), params],
+        final_stage_id=2,
+    )
+
+    output = SimpleNamespace(request_id="req-3", finished=True)
+    asyncio.run(Orchestrator._forward_to_next_stage(orchestrator, "req-3", 1, output, req_state))
+
+    assert previous_stage.engine_outputs == [output]
+    assert diffusion_stage.calls[0]["kv_sender_info"] == {
+        0: {"host": "10.0.0.2", "zmq_port": 50151},
+    }
+
+
+def test_prewarm_diffusion_attaches_kv_sender_info():
+    orchestrator = object.__new__(Orchestrator)
+    sender_stage = _DummySenderStage({"host": "10.0.0.3", "zmq_port": 50151})
+    diffusion_stage = _DummyDiffusionStage(engine_input_source=[0])
+
+    orchestrator.stage_clients = [sender_stage, diffusion_stage]
+    orchestrator.num_stages = 2
+
+    req_state = OrchestratorRequestState(
+        request_id="req-2",
+        prompt={"prompt": "hello"},
+        sampling_params_list=[SamplingParams(max_tokens=4), OmniDiffusionSamplingParams()],
+        final_stage_id=1,
+    )
+
+    stage0_request = SimpleNamespace(prompt_token_ids=[1, 2, 3])
+    asyncio.run(Orchestrator._prewarm_async_chunk_stages(orchestrator, "req-2", stage0_request, req_state))
+
+    assert diffusion_stage.calls[0]["request_id"] == "req-2"
+    assert diffusion_stage.calls[0]["kv_sender_info"] == {
+        0: {"host": "10.0.0.3", "zmq_port": 50151},
+    }
+    assert req_state.stage_submit_ts[1] > 0

tests/engine/test_output_modality.py

@@ -12,6 +12,7 @@ import pytest
 import torch
 
 # ── Load modules without triggering vllm_omni.__init__ ─────────────
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 _ENGINE_DIR = Path(__file__).resolve().parents[2] / "vllm_omni" / "engine"
 

tests/entrypoints/openai_api/test_image_server.py

@@ -178,7 +178,7 @@ def test_client(mock_async_diffusion):
     )
     app.state.args = Namespace(
         default_sampling_params='{"0": {"num_inference_steps":4, "guidance_scale":7.5}}',
-        max_generated_image_size=4096,  # 64*64
+        max_generated_image_size=1024 * 1792,
     )
 
     return TestClient(app)
@@ -245,7 +245,7 @@ def async_omni_stage_configs_only_client():
     # AsyncOmni exposes stage_configs on the engine instance.
     app.state.args = Namespace(
         default_sampling_params='{"1": {"num_inference_steps":4, "guidance_scale":7.5}}',
-        max_generated_image_size=4096,  # 64*64
+        max_generated_image_size=1024 * 1792,
     )
     return TestClient(app)
 
@@ -392,6 +392,18 @@ def test_image_edits_async_omni_stage_configs_only(async_omni_stage_configs_only
     assert len(captured) == 2
 
 
+def test_generate_images_max_size_rejected(async_omni_test_client):
+    """Test that a size exceeding max_generated_image_size returns 400."""
+    response = async_omni_test_client.post(
+        "/v1/images/generations",
+        json={
+            "prompt": "a cat",
+            "size": "2048x2048",  # 4,194,304 pixels > max_generated_image_size (1,048,576)
+        },
+    )
+    assert response.status_code == 400
+
+
 def test_generate_multiple_images(test_client):
     """Test generating multiple images"""
     response = test_client.post(
@@ -982,12 +994,13 @@ def test_image_edit_parameter_default_single_stage(test_client):
     assert captured_sampling_params.num_inference_steps == 4
     assert captured_sampling_params.guidance_scale == 7.5
 
+    # Size exceeding max_generated_image_size (1024*1792) returns 400
     response = test_client.post(
         "/v1/images/edits",
         files=[("image", img_bytes_1)],
         data={
             "prompt": "hello world.",
-            "size": "96x96",
+            "size": "2048x2048",
         },
     )
     assert response.status_code == 400

tests/entrypoints/openai_api/test_serving_speech.py

@@ -658,11 +658,13 @@ class TestTTSMethods:
         mock_engine_client.tts_max_instructions_length = None
         mock_models = mocker.MagicMock()
         mock_models.is_base_model.return_value = True
-        return OmniOpenAIServingSpeech(
+        server = OmniOpenAIServingSpeech(
             engine_client=mock_engine_client,
             models=mock_models,
             request_logger=mocker.MagicMock(),
         )
+        yield server
+        server.shutdown()
 
     def test_is_tts_detection_no_stage(self, speech_server):
         """Test TTS model detection when no TTS stage exists."""
@@ -1026,6 +1028,181 @@ class TestTTSMethods:
 
         assert result is None
 
+    # ── speaker field alias ──
+
+    def test_speaker_alias_accepted_as_voice(self):
+        """The 'speaker' JSON key should be accepted as an alias for 'voice'."""
+        req = OpenAICreateSpeechRequest.model_validate({"input": "Hello", "speaker": "custom_voice"})
+        assert req.voice == "custom_voice"
+
+    def test_voice_field_still_accepted(self):
+        """The canonical 'voice' JSON key should still work."""
+        req = OpenAICreateSpeechRequest.model_validate({"input": "Hello", "voice": "custom_voice"})
+        assert req.voice == "custom_voice"
+
+    def test_speaker_alias_in_base_task_with_uploaded_voice(self, speech_server):
+        """Using 'speaker' key with an uploaded voice should work for Base task."""
+        speech_server.uploaded_speakers = {
+            "utesf": {
+                "name": "UTESF",
+                "file_path": "/tmp/voice_samples/utesf.wav",
+                "mime_type": "audio/wav",
+                "ref_text": None,
+            }
+        }
+        req = OpenAICreateSpeechRequest.model_validate({"input": "Hello", "speaker": "UTESF", "task_type": "Base"})
+        assert req.voice == "UTESF"
+        with patch("pathlib.Path.exists", return_value=True):
+            result = speech_server._validate_qwen_tts_request(req)
+        assert result is None
+
+    # ── uploaded voice with embedding ──
+
+    def test_build_tts_params_with_uploaded_voice_embedding(self, speech_server):
+        """Test _build_tts_params loads embedding for embedding-uploaded voices."""
+        speech_server.uploaded_speakers = {
+            "emb_voice": {
+                "name": "emb_voice",
+                "file_path": "/tmp/voice_samples/emb_voice.safetensors",
+                "mime_type": "application/x-safetensors",
+                "embedding_source": "direct",
+                "embedding_dim": 1024,
+                "cache_status": "ready",
+                "cache_file": "/tmp/voice_samples/emb_voice.safetensors",
+            }
+        }
+        speech_server.supported_speakers = {"ryan", "vivian", "emb_voice"}
+
+        fake_embedding = [0.1] * 1024
+        with patch.object(speech_server, "_get_uploaded_speaker_embedding") as mock_get_emb:
+            mock_get_emb.return_value = fake_embedding
+            req = OpenAICreateSpeechRequest(input="Hello", voice="emb_voice")
+            params = speech_server._build_tts_params(req)
+
+            assert "voice_clone_prompt" in params
+            assert params["voice_clone_prompt"][0]["ref_spk_embedding"] == fake_embedding
+            assert params["task_type"] == ["Base"]
+            assert params["x_vector_only_mode"] == [True]
+            assert "ref_audio" not in params
+
+    # ── regression: full flow from issue #1603 ──
+
+    def test_regression_1603_speaker_key_with_uploaded_audio_voice(self, speech_server):
+        """Regression test for #1603: upload audio voice, then invoke TTS with 'speaker' key.
+
+        Verifies the full validate → build_params pipeline works end-to-end.
+        """
+        speech_server.uploaded_speakers = {
+            "utesf": {
+                "name": "UTESF",
+                "file_path": "/tmp/voice_samples/utesf.wav",
+                "mime_type": "audio/wav",
+                "ref_text": "Hola, esta es una prueba.",
+            }
+        }
+        # Parse with 'speaker' alias (the key users actually send)
+        req = OpenAICreateSpeechRequest.model_validate(
+            {"input": "Hello world", "speaker": "UTESF", "task_type": "Base"}
+        )
+        assert req.voice == "UTESF"
+
+        # Validation should pass (file exists)
+        with patch("pathlib.Path.exists", return_value=True):
+            err = speech_server._validate_qwen_tts_request(req)
+        assert err is None, f"Validation failed: {err}"
+
+        # Build params should auto-set ref_audio from stored file
+        with patch.object(speech_server, "_get_uploaded_audio_data") as mock_audio:
+            mock_audio.return_value = "data:audio/wav;base64,ZmFrZQ=="
+            params = speech_server._build_tts_params(req)
+
+        assert params["task_type"] == ["Base"]
+        assert params["ref_audio"] == ["data:audio/wav;base64,ZmFrZQ=="]
+        assert params["ref_text"] == ["Hola, esta es una prueba."]
+        assert params["x_vector_only_mode"] == [False]
+        assert params["speaker"] == ["utesf"]
+
+    def test_regression_1603_speaker_key_with_uploaded_embedding_voice(self, speech_server):
+        """Regression test for #1603: upload embedding voice, then invoke TTS with 'speaker' key.
+
+        Verifies embedding-uploaded voices are loaded as voice_clone_prompt, not as audio.
+        """
+        speech_server.uploaded_speakers = {
+            "myvoice": {
+                "name": "myvoice",
+                "file_path": "/tmp/voice_samples/myvoice.safetensors",
+                "mime_type": "application/x-safetensors",
+                "embedding_source": "direct",
+                "embedding_dim": 1024,
+                "cache_status": "ready",
+                "cache_file": "/tmp/voice_samples/myvoice.safetensors",
+            }
+        }
+        # Parse with 'speaker' alias
+        req = OpenAICreateSpeechRequest.model_validate(
+            {"input": "Hello world", "speaker": "myvoice", "task_type": "Base"}
+        )
+        assert req.voice == "myvoice"
+
+        # Validation should pass
+        with patch("pathlib.Path.exists", return_value=True):
+            err = speech_server._validate_qwen_tts_request(req)
+        assert err is None, f"Validation failed: {err}"
+
+        # Build params should use embedding, NOT audio
+        fake_emb = [0.1] * 1024
+        with patch.object(speech_server, "_get_uploaded_speaker_embedding") as mock_emb:
+            mock_emb.return_value = fake_emb
+            params = speech_server._build_tts_params(req)
+
+        assert params["task_type"] == ["Base"]
+        assert params["x_vector_only_mode"] == [True]
+        assert "voice_clone_prompt" in params
+        assert params["voice_clone_prompt"][0]["ref_spk_embedding"] == fake_emb
+        # Must NOT have ref_audio — that would fail for safetensors files
+        assert "ref_audio" not in params
+
+    def test_validate_rejects_embedding_voice_with_pending_cache(self, speech_server):
+        """Validation should reject embedding voices whose cache is not yet ready."""
+        speech_server.uploaded_speakers = {
+            "myvoice": {
+                "name": "myvoice",
+                "file_path": "/tmp/myvoice.safetensors",
+                "mime_type": "application/x-safetensors",
+                "embedding_source": "direct",
+                "cache_status": "pending",
+                "cache_file": None,
+            }
+        }
+        req = OpenAICreateSpeechRequest.model_validate({"input": "Hello", "speaker": "myvoice", "task_type": "Base"})
+        with patch("pathlib.Path.exists", return_value=True):
+            err = speech_server._validate_qwen_tts_request(req)
+        assert err is not None
+        assert "not yet ready" in err
+
+    def test_x_vector_only_mode_not_overwritten_for_uploaded_embedding(self, speech_server):
+        """x_vector_only_mode set by uploaded embedding must not be overwritten by request field."""
+        speech_server.uploaded_speakers = {
+            "emb_voice": {
+                "name": "emb_voice",
+                "file_path": "/tmp/emb_voice.safetensors",
+                "mime_type": "application/x-safetensors",
+                "embedding_source": "direct",
+                "embedding_dim": 1024,
+                "cache_status": "ready",
+                "cache_file": "/tmp/emb_voice.safetensors",
+            }
+        }
+        fake_emb = [0.1] * 1024
+        with patch.object(speech_server, "_get_uploaded_speaker_embedding") as mock_emb:
+            mock_emb.return_value = fake_emb
+            # Client explicitly sends x_vector_only_mode=False, but embedding requires True
+            req = OpenAICreateSpeechRequest(input="Hello", voice="emb_voice", x_vector_only_mode=False)
+            params = speech_server._build_tts_params(req)
+
+        assert params["x_vector_only_mode"] == [True]
+        assert "voice_clone_prompt" in params
+
     def test_max_instructions_length_default(self, speech_server):
         """Test default max instructions length (500) when no config provided."""
         # Fixture creates server with no CLI override and no TTS stage
@@ -1639,11 +1816,13 @@ def fish_speech_server(mocker: MockerFixture):
     mock_models = mocker.MagicMock()
     mock_models.is_base_model.return_value = True
 
-    return OmniOpenAIServingSpeech(
+    server = OmniOpenAIServingSpeech(
         engine_client=mock_engine_client,
         models=mock_models,
         request_logger=mocker.MagicMock(),
     )
+    yield server
+    server.shutdown()
 
 
 class TestFishSpeechServing:
@@ -1717,7 +1896,7 @@ class TestFishSpeechServing:
             fish_speech_server._build_fish_speech_prompt(request)
 
     def test_prepare_speech_generation_overrides_fish_default_max_tokens(self, fish_speech_server):
-        fish_speech_server._build_fish_speech_prompt = MagicMock(
+        fish_speech_server._build_fish_speech_prompt_async = AsyncMock(
             return_value={
                 "prompt_token_ids": [1, 2, 3],
                 "additional_information": {},
@@ -1730,13 +1909,14 @@ class TestFishSpeechServing:
 
         assert request_id.startswith("speech-")
         assert generator == "generator"
+        fish_speech_server._build_fish_speech_prompt_async.assert_awaited_once()
         fish_speech_server.engine_client.generate.assert_called_once()
         sampling_params_list = fish_speech_server.engine_client.generate.call_args.kwargs["sampling_params_list"]
         assert sampling_params_list[0].max_tokens == 4096
         assert fish_speech_server.engine_client.default_sampling_params_list[0].max_tokens == 2048
 
     def test_prepare_speech_generation_uses_stage_default_max_tokens(self, fish_speech_server):
-        fish_speech_server._build_fish_speech_prompt = MagicMock(
+        fish_speech_server._build_fish_speech_prompt_async = AsyncMock(
             return_value={
                 "prompt_token_ids": [1, 2, 3],
                 "additional_information": {},
@@ -1985,3 +2165,115 @@ class TestCosyVoice3Serving:
         assert generator == "generator"
         assert tts_params == {}
         cosyvoice3_server._build_cosyvoice3_prompt.assert_awaited_once()
+
+
+class TestTTSAsyncOffloading:
+    """Tests for event-loop-safe offloading of blocking TTS operations."""
+
+    def test_build_voxtral_prompt_is_sync(self):
+        """_build_voxtral_prompt should be a regular function, not a coroutine."""
+        assert not asyncio.iscoroutinefunction(OmniOpenAIServingSpeech._build_voxtral_prompt)
+
+    @pytest.fixture
+    def voxtral_server(self, mocker: MockerFixture):
+        mocker.patch.object(OmniOpenAIServingSpeech, "_load_supported_speakers", return_value=set())
+        mocker.patch.object(OmniOpenAIServingSpeech, "_load_codec_frame_rate", return_value=None)
+        mock_engine_client = mocker.MagicMock()
+        mock_engine_client.errored = False
+        mock_engine_client.model_config = mocker.MagicMock(model="mistralai/Voxtral")
+        mock_engine_client.default_sampling_params_list = [SimpleNamespace(max_tokens=2048)]
+        mock_engine_client.tts_batch_max_items = 32
+        mock_engine_client.generate = mocker.MagicMock(return_value="generator")
+        mock_engine_client.stage_configs = [
+            SimpleNamespace(
+                engine_args=SimpleNamespace(model_stage="audio_generation"),
+                tts_args={},
+            )
+        ]
+        mock_models = mocker.MagicMock()
+        mock_models.is_base_model.return_value = True
+        server = OmniOpenAIServingSpeech(
+            engine_client=mock_engine_client,
+            models=mock_models,
+            request_logger=mocker.MagicMock(),
+        )
+        yield server
+        server.shutdown()
+
+    @pytest.fixture
+    def qwen3_tts_server(self, mocker: MockerFixture):
+        mocker.patch.object(OmniOpenAIServingSpeech, "_load_supported_speakers", return_value=set())
+        mocker.patch.object(OmniOpenAIServingSpeech, "_load_codec_frame_rate", return_value=None)
+        mock_engine_client = mocker.MagicMock()
+        mock_engine_client.errored = False
+        mock_engine_client.model_config = mocker.MagicMock(model="Qwen/Qwen3-TTS", hf_config=mocker.MagicMock())
+        mock_engine_client.default_sampling_params_list = [SimpleNamespace(max_tokens=2048)]
+        mock_engine_client.tts_batch_max_items = 32
+        mock_engine_client.generate = mocker.MagicMock(return_value="generator")
+        mock_engine_client.tts_max_instructions_length = None
+        mock_engine_client.stage_configs = [
+            SimpleNamespace(
+                engine_args=SimpleNamespace(model_stage="qwen3_tts"),
+                tts_args={},
+            )
+        ]
+        mock_models = mocker.MagicMock()
+        mock_models.is_base_model.return_value = True
+        server = OmniOpenAIServingSpeech(
+            engine_client=mock_engine_client,
+            models=mock_models,
+            request_logger=mocker.MagicMock(),
+        )
+        yield server
+        server.shutdown()
+
+    def test_prepare_speech_generation_awaits_voxtral_async(self, voxtral_server):
+        """Voxtral path in _prepare_speech_generation should call the async wrapper."""
+        voxtral_server._build_voxtral_prompt_async = AsyncMock(
+            return_value={
+                "prompt_token_ids": [1, 2, 3],
+                "additional_information": {"voice": ["test"]},
+            }
+        )
+        request = OpenAICreateSpeechRequest(input="hello", voice="test")
+        asyncio.run(voxtral_server._prepare_speech_generation(request))
+        voxtral_server._build_voxtral_prompt_async.assert_awaited_once()
+
+    def test_prepare_speech_generation_awaits_qwen3_tts_async(self, qwen3_tts_server):
+        """Qwen3 TTS path should call _estimate_prompt_len_async."""
+        qwen3_tts_server._validate_tts_request = MagicMock(return_value=None)
+        qwen3_tts_server._build_tts_params = MagicMock(
+            return_value={"text": ["hello"], "task_type": ["CustomVoice"], "speaker": ["Vivian"]}
+        )
+        qwen3_tts_server._estimate_prompt_len_async = AsyncMock(return_value=512)
+        request = OpenAICreateSpeechRequest(input="hello")
+        asyncio.run(qwen3_tts_server._prepare_speech_generation(request))
+        qwen3_tts_server._build_tts_params.assert_called_once()
+        qwen3_tts_server._estimate_prompt_len_async.assert_awaited_once()
+
+    def test_shutdown_is_idempotent(self, mocker: MockerFixture):
+        """Calling shutdown() twice should not raise."""
+        mocker.patch.object(OmniOpenAIServingSpeech, "_load_supported_speakers", return_value=set())
+        mocker.patch.object(OmniOpenAIServingSpeech, "_load_codec_frame_rate", return_value=None)
+        mock_engine_client = mocker.MagicMock()
+        mock_engine_client.errored = False
+        mock_engine_client.stage_configs = []
+        mock_engine_client.tts_max_instructions_length = None
+        mock_models = mocker.MagicMock()
+        mock_models.is_base_model.return_value = True
+        server = OmniOpenAIServingSpeech(
+            engine_client=mock_engine_client,
+            models=mock_models,
+            request_logger=mocker.MagicMock(),
+        )
+        assert server._tts_executor is not None
+        server.shutdown()
+        assert server._tts_executor is None
+        server.shutdown()  # Should not raise
+        assert server._tts_executor is None
+
+    def test_diffusion_instance_shutdown_safe(self):
+        """Diffusion instances (created via for_diffusion) should have safe shutdown."""
+        server = OmniOpenAIServingSpeech.for_diffusion(diffusion_engine=MagicMock(), model_name="test-model")
+        assert server._tts_executor is None
+        server.shutdown()  # Should not raise

tests/entrypoints/openai_api/test_serving_speech_stream.py

@@ -385,3 +385,18 @@ class TestStreamingSpeechWebSocket:
 
         speech_service.engine_client.abort.assert_awaited_once_with("req-abort")
         assert websocket.send_json.await_count == 2
+
+
+class TestGeneratePcmChunksContract:
+    """Guard: _generate_pcm_chunks must exist on OmniOpenAIServingSpeech.
+
+    The WebSocket handler calls speech_service._generate_pcm_chunks()
+    at runtime. If the method is removed, all WS TTS streaming breaks
+    with an AttributeError. This test catches that at CI time.
+    """
+
+    def test_generate_pcm_chunks_defined(self):
+        assert hasattr(OmniOpenAIServingSpeech, "_generate_pcm_chunks")
+        assert asyncio.iscoroutinefunction(OmniOpenAIServingSpeech._generate_pcm_chunks) or callable(
+            OmniOpenAIServingSpeech._generate_pcm_chunks
+        )

tests/entrypoints/openai_api/test_text_splitter.py

@@ -4,7 +4,7 @@ import pytest
 
 from vllm_omni.entrypoints.openai.text_splitter import SentenceSplitter
 
-pytestmark = [pytest.mark.openai, pytest.mark.speech]
+pytestmark = [pytest.mark.openai, pytest.mark.speech, pytest.mark.core_model, pytest.mark.cpu]
 
 
 class TestSentenceSplitterEnglish:

tests/entrypoints/test_stage_utils.py

@@ -6,8 +6,6 @@ from pytest_mock import MockerFixture
 
 from vllm_omni.entrypoints.stage_utils import set_stage_devices
 
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
 
 def _make_dummy_torch(call_log):
     class _Props:
@@ -55,6 +53,8 @@ def _make_mock_platform(mocker, device_type: str = "cuda", env_var: str = "CUDA_
     return mock_platform
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
 @pytest.mark.usefixtures("clean_gpu_memory_between_tests")
 def test_set_stage_devices_respects_logical_ids(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch):
     # Preserve an existing logical mapping and ensure devices "0,1" map through it.
@@ -75,6 +75,8 @@ def test_set_stage_devices_respects_logical_ids(mocker: MockerFixture, monkeypat
     assert os.environ["CUDA_VISIBLE_DEVICES"] == "6,7"
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
 @pytest.mark.usefixtures("clean_gpu_memory_between_tests")
 def test_set_stage_devices_handles_not_enough_devices(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch):
     # Preserve an existing logical mapping and ensure devices "0,1" map through it.
@@ -90,9 +92,10 @@ def test_set_stage_devices_handles_not_enough_devices(mocker: MockerFixture, mon
         mock_platform,
     )
 
-    # Raise since we need 4 GPUs, but we only have 2 visible
-    with pytest.raises(ValueError):
-        set_stage_devices(stage_id=0, devices="0,1,2,3")
+    # Keep the logical mapping and resolve to the visible subset.
+    set_stage_devices(stage_id=0, devices="0,1,2,3")
+
+    assert os.environ["CUDA_VISIBLE_DEVICES"] == "6,7"
 
 
 @pytest.mark.usefixtures("clean_gpu_memory_between_tests")

tests/entrypoints/test_utils.py

@@ -5,14 +5,17 @@ from collections import Counter
 from dataclasses import dataclass
 
 import pytest
+import torch
 from pytest_mock import MockerFixture
 
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.engine.arg_utils import OmniEngineArgs
+from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
 from vllm_omni.entrypoints.utils import (
     _convert_dataclasses_to_dict,
     _filter_dict_like_object,
     filter_dataclass_kwargs,
+    load_and_resolve_stage_configs,
     resolve_model_config_path,
 )
 
@@ -304,3 +307,18 @@ class TestResolveModelConfigPath:
 
         assert result is not None
         assert "glm_image.yaml" in result
+
+
+class TestLoadAndResolveStageConfigs:
+    def test_load_and_resolve_with_kwargs(self):
+        """Ensure that dtype survives default stage creation."""
+        kwargs = {"dtype": torch.float32}
+        config_path, stage_configs = load_and_resolve_stage_configs(
+            model="black-forest-labs/FLUX.2-klein-4B",
+            stage_configs_path=None,
+            kwargs=kwargs,
+            default_stage_cfg_factory=lambda: AsyncOmniEngine._create_default_diffusion_stage_cfg(kwargs),
+        )
+        assert config_path is None
+        assert len(stage_configs) == 1
+        assert "dtype" in stage_configs[0]["engine_args"]

tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py

@@ -2,10 +2,14 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Unit tests for CosyVoice3 components."""
 
+from types import SimpleNamespace
+
 import pytest
 import torch
 import torch.nn as nn
 
+from tests.utils import hardware_test
+
 
 class TestPreLookaheadLayer:
     """Tests for PreLookaheadLayer."""
@@ -16,6 +20,8 @@ class TestPreLookaheadLayer:
 
         return PreLookaheadLayer(in_channels=512, channels=512, pre_lookahead_len=3)
 
+    @pytest.mark.core_model
+    @pytest.mark.cpu
     def test_forward_shape(self, layer):
         """Test that output shape matches input shape."""
         batch, seq_len, channels = 2, 10, 512
@@ -25,6 +31,8 @@ class TestPreLookaheadLayer:
 
         assert out.shape == x.shape
 
+    @pytest.mark.core_model
+    @pytest.mark.cpu
     def test_forward_with_context(self, layer):
         """Test forward with context for streaming."""
         batch, seq_len, channels = 1, 10, 512
@@ -36,6 +44,8 @@ class TestPreLookaheadLayer:
 
         assert out.shape == x.shape
 
+    @pytest.mark.core_model
+    @pytest.mark.cpu
     def test_residual_connection(self, layer):
         """Test that residual connection is applied."""
         batch, seq_len, channels = 1, 5, 512
@@ -57,6 +67,8 @@ class TestDiTAttention:
 
         return DiTAttention(dim=512, heads=8, dim_head=64, dropout=0.0)
 
+    @pytest.mark.core_model
+    @hardware_test(res={"cuda": "L4"}, num_cards=1)
     def test_forward_shape(self, attention):
         """Test attention output shape."""
         batch, seq_len, dim = 2, 16, 512
@@ -66,6 +78,8 @@ class TestDiTAttention:
 
         assert out.shape == x.shape
 
+    @pytest.mark.core_model
+    @hardware_test(res={"cuda": "L4"}, num_cards=1)
     def test_forward_with_mask(self, attention):
         """Test attention with mask."""
         batch, seq_len, dim = 2, 16, 512
@@ -79,6 +93,8 @@ class TestDiTAttention:
         # Masked positions should be zero
         assert torch.allclose(out[:, -3:], torch.zeros_like(out[:, -3:]))
 
+    @pytest.mark.core_model
+    @hardware_test(res={"cuda": "L4"}, num_cards=1)
     def test_qkv_projections(self, attention):
         """Test that Q/K/V projections exist and have correct dimensions."""
         assert hasattr(attention, "to_q")
@@ -98,6 +114,8 @@ class TestDiTBlock:
 
         return DiTBlock(dim=512, heads=8, dim_head=64, ff_mult=4, dropout=0.0)
 
+    @pytest.mark.core_model
+    @hardware_test(res={"cuda": "L4"}, num_cards=1)
     def test_forward_shape(self, block):
         """Test block output shape."""
         batch, seq_len, dim = 2, 16, 512
@@ -108,6 +126,8 @@ class TestDiTBlock:
 
         assert out.shape == x.shape
 
+    @pytest.mark.core_model
+    @hardware_test(res={"cuda": "L4"}, num_cards=1)
     def test_adalayernorm_modulation(self, block):
         """Test that AdaLayerNorm modulates based on timestep."""
         batch, seq_len, dim = 1, 8, 512
@@ -142,6 +162,8 @@ class TestDiT:
             long_skip_connection=True,
         )
 
+    @pytest.mark.core_model
+    @hardware_test(res={"cuda": "L4"}, num_cards=1)
     def test_forward_shape(self, dit):
         """Test DiT forward output shape."""
         batch, mel_dim, seq_len = 1, 80, 32
@@ -156,6 +178,8 @@ class TestDiT:
 
         assert out.shape == (batch, mel_dim, seq_len)
 
+    @pytest.mark.core_model
+    @hardware_test(res={"cuda": "L4"}, num_cards=1)
     def test_timestep_embedding(self, dit):
         """Test that different timesteps produce different outputs."""
         batch, mel_dim, seq_len = 1, 80, 16
@@ -188,6 +212,8 @@ class TestCFM:
 
         return DummyEstimator()
 
+    @pytest.mark.core_model
+    @pytest.mark.cpu
     def test_causal_conditional_cfm_forward(self, dummy_estimator):
         """Test CausalConditionalCFM forward pass."""
         from omegaconf import DictConfig
@@ -226,6 +252,8 @@ class TestCFM:
 class TestSDPAFallback:
     """Test SDPA fallback for float32 inputs."""
 
+    @pytest.mark.core_model
+    @hardware_test(res={"cuda": "L4"}, num_cards=1)
     def test_float32_uses_sdpa(self):
         """Test that float32 inputs use SDPA fallback."""
         from vllm_omni.diffusion.attention.layer import Attention
@@ -247,3 +275,32 @@ class TestSDPAFallback:
 
         assert out.shape == (batch, seq_len, heads, dim)
         assert out.dtype == torch.float32
+
+
+def test_code2wav_forward_finalizes_hift_tail():
+    from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3_code2wav import CosyVoice3Code2Wav
+
+    class DummyHiFT(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.m_source = SimpleNamespace(l_linear=SimpleNamespace(weight=torch.ones(1, dtype=torch.float32)))
+            self.finalize_calls: list[bool] = []
+
+        def inference(self, speech_feat, finalize=True):
+            self.finalize_calls.append(bool(finalize))
+            return torch.zeros((speech_feat.shape[0], 1, speech_feat.shape[-1]), dtype=speech_feat.dtype), None
+
+    model = object.__new__(CosyVoice3Code2Wav)
+    nn.Module.__init__(model)
+    model.hift = DummyHiFT()
+    model._forward_mel = lambda **_: torch.ones((1, 80, 8), dtype=torch.float32)
+
+    out = model.forward(
+        token=torch.tensor([[1, 2, 3]], dtype=torch.int32),
+        prompt_token=torch.tensor([[4, 5]], dtype=torch.int32),
+        prompt_feat=torch.ones((1, 4, 80), dtype=torch.float32),
+        embedding=torch.ones((1, 192), dtype=torch.float32),
+    )
+
+    assert out.shape == (1, 1, 8)
+    assert model.hift.finalize_calls == [True]

tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py

@@ -0,0 +1,463 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from threading import Lock
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.nn as nn
+from vllm.v1.outputs import SamplerOutput
+from vllm.v1.sample.logits_processor.state import LogitsProcessors
+from vllm.v1.sample.metadata import SamplingMetadata
+
+from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3 import CosyVoice3Model
+from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _DummyCode2Wav:
+    def __init__(
+        self,
+        vocab_size: int,
+        num_samples: int = 32,
+        outputs: list[tuple[torch.Tensor, dict[str, object] | None]] | None = None,
+    ):
+        self.input_embedding = SimpleNamespace(num_embeddings=vocab_size)
+        self.num_samples = num_samples
+        self.outputs = list(outputs or [])
+        self.forward_calls: list[dict[str, object]] = []
+        self.forward_streaming_calls: list[dict[str, object]] = []
+
+    def forward(self, **kwargs):
+        self.forward_calls.append(kwargs)
+        token = kwargs["token"]
+        num_samples = int(token.shape[-1])
+        return torch.linspace(-1.0, 1.0, max(num_samples, 1), dtype=torch.float32).reshape(1, 1, -1)
+
+    def forward_streaming(self, **kwargs):
+        self.forward_streaming_calls.append(kwargs)
+        if self.outputs:
+            return self.outputs.pop(0)
+
+        token = kwargs["token"]
+        num_samples = int(token.shape[-1])
+        audio = torch.linspace(-1.0, 1.0, max(num_samples, 1), dtype=torch.float32).reshape(1, 1, -1)
+        new_state = None
+        if not kwargs.get("finalize", False):
+            new_state = {
+                "mel": torch.ones((1, 80, max(num_samples, 1)), dtype=torch.float32),
+                "speech_offset": audio.shape[-1],
+            }
+        return audio, new_state
+
+
+def _make_code2wav_model(
+    *,
+    with_stride_cfg: bool = False,
+    num_samples: int = 32,
+    outputs: list[tuple[torch.Tensor, dict[str, object] | None]] | None = None,
+) -> CosyVoice3Model:
+    model = object.__new__(CosyVoice3Model)
+    nn.Module.__init__(model)
+    model.model_stage = "cosyvoice3_code2wav"
+    hift_cfg = {} if not with_stride_cfg else {"upsample_rates": [8, 5, 3], "istft_params": {"hop_len": 4}}
+    model.config = SimpleNamespace(
+        sample_rate=24000,
+        hift=hift_cfg,
+        token_frame_rate=25 if with_stride_cfg else 0,
+        token_mel_ratio=2 if with_stride_cfg else 0,
+    )
+    model.code2wav = _DummyCode2Wav(vocab_size=4, num_samples=num_samples, outputs=outputs)
+    model.source_cache_len = 4
+    model.speech_window = torch.hamming_window(8, periodic=False)
+    model._stream_audio_cache_by_req = {}
+    model._stream_audio_cache_lock = Lock()
+    model._stream_vocoder_cache_by_req = {}
+    return model
+
+
+def _make_talker_model() -> CosyVoice3Model:
+    model = object.__new__(CosyVoice3Model)
+    nn.Module.__init__(model)
+    model.model_stage = "cosyvoice3_talker"
+    model.config = SimpleNamespace(
+        llm={
+            "speech_token_size": 6561,
+            "eos_token_id": 6562,
+            "sampling": {
+                "top_p": 0.8,
+                "top_k": 25,
+                "win_size": 10,
+                "tau_r": 0.1,
+            },
+        },
+        vocab_size=151923,
+    )
+    return model
+
+
+def _make_sampling_metadata(
+    *,
+    output_token_ids: list[list[int]],
+    repetition_penalty: float = 2.0,
+) -> SamplingMetadata:
+    return SamplingMetadata(
+        temperature=torch.tensor([1.0], dtype=torch.float32),
+        all_greedy=False,
+        all_random=True,
+        top_p=torch.tensor([0.8], dtype=torch.float32),
+        top_k=torch.tensor([25], dtype=torch.int32),
+        generators={},
+        max_num_logprobs=None,
+        no_penalties=False,
+        prompt_token_ids=None,
+        frequency_penalties=torch.zeros(1, dtype=torch.float32),
+        presence_penalties=torch.zeros(1, dtype=torch.float32),
+        repetition_penalties=torch.tensor([repetition_penalty], dtype=torch.float32),
+        output_token_ids=output_token_ids,
+        allowed_token_ids_mask=None,
+        bad_words_token_ids={},
+        logitsprocs=LogitsProcessors(),
+    )
+
+
+def test_split_request_ids_uses_seq_token_counts():
+    ids = torch.tensor([10, 11, 12, 13, 14], dtype=torch.long)
+    chunks = CosyVoice3Model._split_request_ids(ids, [2, 2, 2])
+    assert [c.tolist() for c in chunks] == [[10, 11], [12, 13], [14]]
+
+
+def test_split_request_ids_honors_single_request_seq_token_counts():
+    ids = torch.tensor([10, 11, 12, 13, 14], dtype=torch.long)
+    chunks = CosyVoice3Model._split_request_ids(ids, [3])
+    assert [c.tolist() for c in chunks] == [[10, 11, 12]]
+
+
+def test_sanitize_codec_tokens_filters_out_of_range():
+    model = _make_code2wav_model()
+    raw = torch.tensor([-1, 0, 3, 4, 99], dtype=torch.long)
+    clean = model._sanitize_codec_tokens(raw)
+    assert clean.tolist() == [0, 3]
+
+
+def test_forward_prefers_token_offset_when_present():
+    model = _make_code2wav_model()
+
+    runtime_info = [
+        {
+            "speech_token": torch.tensor([[1, 2, 3]], dtype=torch.long),
+            "speech_feat": torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.float32),
+            "embedding": torch.tensor([[0.5, 0.6]], dtype=torch.float32),
+            "token_offset": 2,
+            "left_context_size": 1,
+        }
+    ]
+
+    out = model.forward(
+        input_ids=torch.tensor([0, 1, 2], dtype=torch.long),
+        positions=torch.tensor([0, 1, 2], dtype=torch.long),
+        model_intermediate_buffer=runtime_info,
+        seq_token_counts=[3],
+    )
+
+    assert len(out.multimodal_outputs["audio"]) == 1
+    assert out.multimodal_outputs["audio"][0].numel() > 0
+    assert len(model.code2wav.forward_streaming_calls) == 1
+    call = model.code2wav.forward_streaming_calls[0]
+    assert call["token"].shape == (1, 3)
+    assert call["token_offset_tokens"] == 2
+    assert call["finalize"] is False
+
+
+def test_forward_falls_back_to_left_context_size_for_backward_compat():
+    model = _make_code2wav_model()
+
+    runtime_info = [
+        {
+            "speech_token": torch.tensor([[1, 2, 3]], dtype=torch.long),
+            "speech_feat": torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.float32),
+            "embedding": torch.tensor([[0.5, 0.6]], dtype=torch.float32),
+            "left_context_size": 2,
+        }
+    ]
+
+    model.forward(
+        input_ids=torch.tensor([0, 1, 2], dtype=torch.long),
+        positions=torch.tensor([0, 1, 2], dtype=torch.long),
+        model_intermediate_buffer=runtime_info,
+        seq_token_counts=[3],
+    )
+
+    assert model.code2wav.forward_streaming_calls[0]["token_offset_tokens"] == 2
+
+
+def test_forward_ignores_single_request_padded_tail_tokens():
+    model = _make_code2wav_model(with_stride_cfg=True)
+    runtime_info = [
+        {
+            "speech_token": torch.tensor([[1, 2, 3]], dtype=torch.long),
+            "speech_feat": torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.float32),
+            "embedding": torch.tensor([[0.5, 0.6]], dtype=torch.float32),
+            "token_offset": 0,
+        }
+    ]
+
+    out = model.forward(
+        input_ids=torch.tensor([0, 1, 2, 3, 3], dtype=torch.long),
+        positions=torch.tensor([0, 1, 2, 3, 4], dtype=torch.long),
+        model_intermediate_buffer=runtime_info,
+        seq_token_counts=[3],
+    )
+
+    # The padded tail must not contribute to code2wav length.
+    assert out.multimodal_outputs["audio"][0].numel() == 3
+    assert model.code2wav.forward_streaming_calls[0]["token"].tolist() == [[0, 1, 2]]
+
+
+def test_forward_uses_non_stream_decode_without_chunk_metadata():
+    model = _make_code2wav_model()
+
+    runtime_info = [
+        {
+            "speech_token": torch.tensor([[1, 2, 3]], dtype=torch.long),
+            "speech_feat": torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.float32),
+            "embedding": torch.tensor([[0.5, 0.6]], dtype=torch.float32),
+            "prefix_ids": [101, 102],
+            "generated_len": 3,
+        }
+    ]
+
+    out = model.forward(
+        input_ids=torch.tensor([0, 1, 2], dtype=torch.long),
+        positions=torch.tensor([0, 1, 2], dtype=torch.long),
+        model_intermediate_buffer=runtime_info,
+        seq_token_counts=[3],
+    )
+
+    assert out.multimodal_outputs["audio"][0].numel() == 3
+    assert len(model.code2wav.forward_calls) == 1
+    assert len(model.code2wav.forward_streaming_calls) == 0
+    call = model.code2wav.forward_calls[0]
+    assert call["token"].tolist() == [[0, 1, 2]]
+
+
+def test_forward_reuses_streaming_cache_state_between_chunks():
+    model = _make_code2wav_model(
+        outputs=[
+            (
+                torch.arange(4, dtype=torch.float32).reshape(1, 1, -1),
+                {"mel": torch.ones((1, 80, 3), dtype=torch.float32), "speech_offset": 4},
+            ),
+            (
+                torch.full((1, 1, 2), 9.0, dtype=torch.float32),
+                {"mel": torch.ones((1, 80, 5), dtype=torch.float32), "speech_offset": 6},
+            ),
+        ]
+    )
+    runtime_info = [
+        {
+            "req_id": ["rid-stream"],
+            "speech_token": torch.tensor([[1, 2, 3]], dtype=torch.long),
+            "speech_feat": torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.float32),
+            "embedding": torch.tensor([[0.5, 0.6]], dtype=torch.float32),
+            "token_offset": 0,
+            "stream_finished": torch.tensor(False),
+        }
+    ]
+
+    out1 = model.forward(
+        input_ids=torch.tensor([0, 1, 2], dtype=torch.long),
+        positions=torch.tensor([0, 1, 2], dtype=torch.long),
+        model_intermediate_buffer=runtime_info,
+        seq_token_counts=[3],
+    )
+    assert out1.multimodal_outputs["audio"][0].tolist() == [0.0, 1.0, 2.0, 3.0]
+    assert model.code2wav.forward_streaming_calls[0]["cache_state"] is None
+
+    out2 = model.forward(
+        input_ids=torch.tensor([0, 1, 2], dtype=torch.long),
+        positions=torch.tensor([0, 1, 2], dtype=torch.long),
+        model_intermediate_buffer=runtime_info,
+        seq_token_counts=[3],
+    )
+    assert out2.multimodal_outputs["audio"][0].tolist() == [9.0, 9.0]
+    cache_state = model.code2wav.forward_streaming_calls[1]["cache_state"]
+    assert cache_state is not None
+    assert cache_state["speech_offset"] == 4
+    assert "rid-stream" in model._stream_vocoder_cache_by_req
+
+
+def test_forward_clears_streaming_cache_on_terminal_chunk():
+    model = _make_code2wav_model(
+        outputs=[
+            (
+                torch.arange(4, dtype=torch.float32).reshape(1, 1, -1),
+                {"mel": torch.ones((1, 80, 3), dtype=torch.float32), "speech_offset": 4},
+            ),
+            (
+                torch.full((1, 1, 1), 7.0, dtype=torch.float32),
+                None,
+            ),
+        ]
+    )
+    runtime_info = [
+        {
+            "req_id": ["rid-stream"],
+            "speech_token": torch.tensor([[1, 2, 3]], dtype=torch.long),
+            "speech_feat": torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.float32),
+            "embedding": torch.tensor([[0.5, 0.6]], dtype=torch.float32),
+            "token_offset": 0,
+            "stream_finished": torch.tensor(False),
+        }
+    ]
+
+    model.forward(
+        input_ids=torch.tensor([0, 1, 2], dtype=torch.long),
+        positions=torch.tensor([0, 1, 2], dtype=torch.long),
+        model_intermediate_buffer=runtime_info,
+        seq_token_counts=[3],
+    )
+    assert "rid-stream" in model._stream_vocoder_cache_by_req
+
+    runtime_info[0]["stream_finished"] = torch.tensor(True)
+    out = model.forward(
+        input_ids=torch.tensor([0, 1, 2], dtype=torch.long),
+        positions=torch.tensor([0, 1, 2], dtype=torch.long),
+        model_intermediate_buffer=runtime_info,
+        seq_token_counts=[3],
+    )
+    assert out.multimodal_outputs["audio"][0].tolist() == [7.0]
+    assert "rid-stream" not in model._stream_vocoder_cache_by_req
+
+
+def test_sample_uses_ras_rejection_for_recent_repetition():
+    model = _make_talker_model()
+    metadata = _make_sampling_metadata(output_token_ids=[[1] * 10])
+    logits = torch.tensor([[-1e9, 10.0, 0.0]], dtype=torch.float32)
+
+    out = model.sample(logits, metadata)
+
+    assert out is not None
+    assert out.sampled_token_ids.tolist() == [[2]]
+
+
+def test_sample_tolerates_padded_rows_without_history():
+    model = _make_talker_model()
+    metadata = _make_sampling_metadata(output_token_ids=[[1] * 10])
+    logits = torch.tensor(
+        [
+            [-1e9, 10.0, 0.0],
+            [-1e9, 0.0, 10.0],
+        ],
+        dtype=torch.float32,
+    )
+
+    out = model.sample(logits, metadata)
+
+    assert out is not None
+    assert out.sampled_token_ids.shape == (2, 1)
+
+
+def test_gpu_ar_model_runner_prefers_model_sampler_when_opted_in():
+    metadata = _make_sampling_metadata(output_token_ids=[[1, 2, 3]])
+    expected = SamplerOutput(
+        sampled_token_ids=torch.tensor([[7]], dtype=torch.int32),
+        logprobs_tensors=None,
+    )
+    calls: list[torch.Tensor] = []
+
+    class _DummyInputBatch:
+        def __init__(self):
+            self.sampling_metadata = metadata
+            self.updated = False
+
+        def update_async_output_token_ids(self):
+            self.updated = True
+
+    runner = object.__new__(GPUARModelRunner)
+    runner.input_batch = _DummyInputBatch()
+    runner.model = SimpleNamespace(
+        prefer_model_sampler=True,
+        sample=lambda logits, sampling_metadata: calls.append(logits.clone()) or expected,
+    )
+    runner.sampler = lambda **_: (_ for _ in ()).throw(AssertionError("fallback sampler should not be used"))
+
+    out = runner._sample(torch.tensor([[0.1, 0.2]], dtype=torch.float32), spec_decode_metadata=None)
+
+    assert out is expected
+    assert runner.input_batch.updated is False
+    assert len(calls) == 1
+
+
+def test_gpu_ar_model_runner_supplies_req_output_history_to_model_sampler():
+    metadata = _make_sampling_metadata(output_token_ids=[])
+    seen_histories: list[list[list[int]]] = []
+
+    class _DummyInputBatch:
+        def __init__(self):
+            self.sampling_metadata = metadata
+            self.req_output_token_ids = [[1, 2, 3]]
+            self.req_ids = ["rid-1"]
+            self.sampled_token_ids_cpu = None
+            self.async_copy_ready_event = None
+            self.prev_req_id_to_index = None
+
+        def update_async_output_token_ids(self):
+            raise AssertionError("fallback async repair should not run for model sampler path")
+
+    runner = object.__new__(GPUARModelRunner)
+    runner.input_batch = _DummyInputBatch()
+    runner.model = SimpleNamespace(
+        prefer_model_sampler=True,
+        sample=lambda logits, sampling_metadata: seen_histories.append(
+            [list(x) for x in sampling_metadata.output_token_ids]
+        )
+        or SamplerOutput(sampled_token_ids=torch.tensor([[7]], dtype=torch.int32), logprobs_tensors=None),
+    )
+    runner.sampler = lambda **_: (_ for _ in ()).throw(AssertionError("fallback sampler should not be used"))
+
+    runner._sample(torch.tensor([[0.1, 0.2]], dtype=torch.float32), spec_decode_metadata=None)
+
+    assert seen_histories == [[[1, 2, 3]]]
+
+
+def test_gpu_ar_model_runner_repairs_async_placeholders_for_model_sampler():
+    metadata = _make_sampling_metadata(output_token_ids=[])
+    seen_histories: list[list[list[int]]] = []
+
+    class _ReadyEvent:
+        def __init__(self):
+            self.synced = False
+
+        def synchronize(self):
+            self.synced = True
+
+    class _DummyInputBatch:
+        def __init__(self):
+            self.sampling_metadata = metadata
+            self.req_output_token_ids = [[11, -1]]
+            self.req_ids = ["rid-1"]
+            self.sampled_token_ids_cpu = torch.tensor([[29]], dtype=torch.int32)
+            self.async_copy_ready_event = _ReadyEvent()
+            self.prev_req_id_to_index = {"rid-1": 0}
+
+        def update_async_output_token_ids(self):
+            raise AssertionError("fallback async repair should not run for model sampler path")
+
+    runner = object.__new__(GPUARModelRunner)
+    runner.input_batch = _DummyInputBatch()
+    runner.model = SimpleNamespace(
+        prefer_model_sampler=True,
+        sample=lambda logits, sampling_metadata: seen_histories.append(
+            [list(x) for x in sampling_metadata.output_token_ids]
+        )
+        or SamplerOutput(sampled_token_ids=torch.tensor([[7]], dtype=torch.int32), logprobs_tensors=None),
+    )
+    runner.sampler = lambda **_: (_ for _ in ()).throw(AssertionError("fallback sampler should not be used"))
+
+    runner._sample(torch.tensor([[0.1, 0.2]], dtype=torch.float32), spec_decode_metadata=None)
+
+    assert runner.input_batch.async_copy_ready_event.synced is True
+    assert seen_histories == [[[11, 29]]]