Enable twine upload to PyPI in release pipeline

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

.buildkite/nightly-release-pipeline.yaml

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

.buildkite/pipeline-npu.yaml

@@ -78,15 +78,19 @@ steps:
       echo "--- Image pushed successfully"
       echo "$${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG}"
 
-  # - label: "Upload Ready Pipeline"
+  - label: "Upload Ready Pipeline"
-  #   depends_on: image-build
+    depends_on: image-build
-  #   key: upload-ready-pipeline
+    key: upload-ready-pipeline
-  #   if: build.branch != "main" && build.pull_request.labels includes "npu-test"
+    if: >-
-  #   agents:
+      build.branch != "main" && (
-  #     queue: "ascend-a2b3"
+        build.pull_request.labels includes "npu-test" ||
-  #     resource_class: "npu-1"
+        build.pull_request.labels includes "ready"
-  #   command: |
+      )
-  #     buildkite-agent pipeline upload .buildkite/test-npu-ready.yaml
+    agents:
+      queue: "ascend-a2b3"
+      resource_class: "npu-1"
+    command: |
+      buildkite-agent pipeline upload .buildkite/test-npu-ready.yaml
 
   # L4 Test — main+NIGHTLY=1 (scheduled), or PR with specific label (e.g. add label then Rebuild)
   - label: "Upload Nightly Pipeline"

.buildkite/rebase-pipeline.yaml

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

.buildkite/release-pipeline.yaml

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

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

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

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

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

.buildkite/test-merge.yml

@@ -1,5 +1,6 @@
 env:
   VLLM_WORKER_MULTIPROC_METHOD: spawn
+  VLLM_USE_FLASHINFER_MOE_FP16: 0
   HF_HUB_DOWNLOAD_TIMEOUT: 300
   HF_HUB_ETAG_TIMEOUT: 60
 
@@ -172,7 +173,7 @@ steps:
               pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py tests/e2e/offline_inference/test_qwen3_tts_customvoice.py -m 'advanced_model and cuda' --run-level 'advanced_model'
             "
         agents:
-          queue: "gpu_1_queue"
+          queue: "gpu_4_queue"
         plugins:
           - docker#v5.2.0:
               image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
@@ -193,7 +194,7 @@ steps:
               pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py tests/e2e/offline_inference/test_qwen3_tts_base.py -m 'advanced_model and cuda' --run-level 'advanced_model'
             "
         agents:
-          queue: "gpu_1_queue"
+          queue: "gpu_4_queue"
         plugins:
           - docker#v5.2.0:
               image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT

.buildkite/test-nightly.yml

@@ -1,5 +1,6 @@
 env:
   VLLM_WORKER_MULTIPROC_METHOD: spawn
+  VLLM_USE_FLASHINFER_MOE_FP16: 0
   HF_HUB_DOWNLOAD_TIMEOUT: 300
   HF_HUB_ETAG_TIMEOUT: 60
 
@@ -388,7 +389,7 @@ steps:
                   - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
                     resources:
                       limits:
-                        nvidia.com/gpu: 1
+                        nvidia.com/gpu: 2
                     volumeMounts:
                       - name: devshm
                         mountPath: /dev/shm
@@ -659,87 +660,6 @@ steps:
                       path: /mnt/hf-cache
                       type: DirectoryOrCreate
 
-      - label: ":full_moon: Diffusion X2I(&A&T) · GEBench Accuracy Test"
-        timeout_in_minutes: 60
-        commands:
-          - pytest -s -v tests/e2e/accuracy/test_gebench_h100_smoke.py --run-level full_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 X2I(&A&T) · GEdit-Bench Accuracy Test"
-        timeout_in_minutes: 60
-        commands:
-          - pytest -s -v tests/e2e/accuracy/test_gedit_bench_h100_smoke.py --run-level full_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 X2I(&A&T) · Accuracy Test"
         timeout_in_minutes: 180
         commands:
@@ -1039,7 +959,7 @@ steps:
       - label: ":full_moon: Diffusion X2V · Function Test"
         timeout_in_minutes: 90
         commands:
-          - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py tests/e2e/online_serving/test_wan_2_1_vace_expansion.py tests/e2e/online_serving/test_hunyuan_video_15_expansion.py  --run-level "full_model"
+          - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py tests/e2e/online_serving/test_wan_2_1_vace_expansion.py tests/e2e/online_serving/test_hunyuan_video_15_expansion.py  -m "full_model and cuda" --run-level "full_model"
         agents:
           queue: "mithril-h100-pool"
         plugins:

.buildkite/test-npu-ready.yaml

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

.buildkite/test-ready.yml

@@ -1,5 +1,6 @@
 env:
   VLLM_WORKER_MULTIPROC_METHOD: spawn
+  VLLM_USE_FLASHINFER_MOE_FP16: 0
   HF_HUB_DOWNLOAD_TIMEOUT: 300
   HF_HUB_ETAG_TIMEOUT: 60
 
@@ -303,7 +304,7 @@ steps:
               pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m 'core_model' --run-level 'core_model'
             "
         agents:
-          queue: "gpu_1_queue"
+          queue: "gpu_4_queue"
         plugins:
           - docker#v5.2.0:
               image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
@@ -326,7 +327,7 @@ steps:
               pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py -m 'core_model' --run-level 'core_model'
             "
         agents:
-          queue: "gpu_1_queue"
+          queue: "gpu_4_queue"
         plugins:
           - docker#v5.2.0:
               image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT

.buildkite/test-weekly.yml

@@ -1,5 +1,6 @@
 env:
   VLLM_WORKER_MULTIPROC_METHOD: spawn
+  VLLM_USE_FLASHINFER_MOE_FP16: 0
   HF_HUB_DOWNLOAD_TIMEOUT: 300
   HF_HUB_ETAG_TIMEOUT: 60
 

README.md

@@ -9,19 +9,19 @@ Easy, fast, and cheap omni-modality model serving for everyone
 </h3>
 
 <p align="center">
-| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> | <a href="docs/assets/WeChat.jpg"><b>WeChat</b></a> | <a href="https://arxiv.org/abs/2602.02204"><b>Paper</b></a> | <a href="https://docs.google.com/presentation/d/1XJWgv79lORl8rbaVvp2d5Sqs6ZEBgAgj/edit?slide=id.p1#slide=id.p1"><b>Slides</b></a> |
+| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> | <a href="docs/assets/WeChat.jpg"><b>WeChat</b></a> | <a href="https://arxiv.org/abs/2602.02204"><b>Paper</b></a> | <a href="https://docs.google.com/presentation/d/111-L8zF7A1j_YI_cR8JsblofdScdRr2f/edit?usp=sharing&ouid=110473603432222024453&rtpof=true&sd=true"><b>Slides</b></a> |
 </p>
 
 
 ---
 
 *Latest News* 🔥
+- [2026/05] We released [0.20.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.20.0) - refreshes the serving/runtime stack for large-scale omni workloads, and improves diffusion model performance, quantization, and hardware readiness across CUDA, ROCm, MUSA, NPU, and XPU backends.
 - [2026/03] We released [0.18.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.18.0) - strengthens the core runtime through a large entrypoint refactor and scheduler/runtime cleanups, expands unified quantization and diffusion execution, broadens multimodal model coverage, and improves production readiness across audio, omni, image, video, RL, and multi-platform deployments.
 - [2026/03] Check out our first public [project deepdive](https://youtu.be/sgwNfsNnR9I) at the vLLM Hong Kong Meetup!
 - [2026/03] **[vllm-omni-skills](https://github.com/hsliuustc0106/vllm-omni-skills)** is a community-driven collection of AI assistant skills that help developers work with vLLM-Omni more effectively. These skills can be used with popular agentic AI coding assistants like **Cursor IDE**, **Claude**, **Codex**, and more.
 - [2026/02] We released [0.16.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.16.0) - A major alignment + capability release that rebases onto **upstream vLLM v0.16.0** and significantly expands performance, distributed execution, and production readiness across **Qwen3-Omni / Qwen3-TTS**, **Bagel**, **MiMo-Audio**, **GLM-Image** and the **Diffusion (DiT) image/video stack**—while also improving platform coverage (CUDA / ROCm / NPU / XPU), CI quality, and documentation.
 - [2026/02] We released [0.14.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.14.0) - This is the first **stable release** of vLLM-Omni that expands Omni’s diffusion / image-video generation and audio / TTS stack, improves distributed execution and memory efficiency, and broadens platform/backend coverage (GPU/ROCm/NPU/XPU). It also brings meaningful upgrades to serving APIs, profiling & benchmarking, and overall stability. Please check our latest [paper](https://arxiv.org/abs/2602.02204) for architecture design and performance results.
-- [2026/01] We released [0.12.0rc1](https://github.com/vllm-project/vllm-omni/releases/tag/v0.12.0rc1) - a major RC milestone focused on maturing the diffusion stack, strengthening OpenAI-compatible serving, expanding omni-model coverage, and improving stability across platforms (GPU/NPU/ROCm).
 - [2025/11] vLLM community officially released [vllm-project/vllm-omni](https://github.com/vllm-project/vllm-omni) in order to support omni-modality models serving.
 
 ---

benchmarks/build_dataset/seed_tts_design/en/meta.lst

@@ -18,3 +18,13 @@ vd017|||The quarterly results exceed expectations across all major metrics.|A sh
 vd018|||Chapter one. The morning sun filtered gently through the forest canopy.|A smooth, rich male audiobook narrator voice, expressive and engaging.
 vd019|||To be or not to be, that is the question.|A theatrical female voice, dramatic and expressive, stage projection quality.
 vd020|||And that is all for tonight. Stay well out there, everyone.|A warm, velvety male late-night radio DJ voice, smooth and intimate.
+vd021|||Please fasten your seatbelts; we will be experiencing some turbulence shortly.|A calm, professional female flight attendant voice, reassuring and clear.
+vd022|||The fossil record suggests these creatures lived more than sixty million years ago.|A measured male documentary narrator voice, scientific and curious.
+vd023|||Hey kiddo, did you finish your homework before dinner today?|A gentle, caring father voice, warm and slightly tired after a long day.
+vd024|||Live from courtside, the home team has just taken a commanding lead.|A high-energy male sports broadcaster, urgent and engaged.
+vd025|||Please remain on the line while I transfer you to the next available agent.|A polite, neutral female call-center voice, professional and patient.
+vd026|||I have a dream that one day this nation will rise up and live out the true meaning of its creed.|A powerful, resonant male orator voice, slow and impassioned.
+vd027|||Mama always said life was like a box of chocolates; you never know what you're gonna get.|A folksy male voice with a Southern American accent, warm and reflective.
+vd028|||Ahoy there, matey! Hoist the colors and prepare to set sail.|A theatrical pirate male voice, gruff and boisterous.
+vd029|||Ladies and gentlemen, please welcome to the stage our keynote speaker.|A polished female event MC voice, gracious and articulate.
+vd030|||Once you've completed the form, click submit and we'll review it within two business days.|A friendly female product-onboarding voice, encouraging and clear.

benchmarks/build_dataset/seed_tts_smoke/en/meta.lst

@@ -18,3 +18,13 @@ smoke017|||In the distance, the mountains were shrouded in thick morning fog.
 smoke018|||Our company reported record revenue for the fourth quarter of the fiscal year.
 smoke019|||She explained the new policy in detail during the staff meeting this morning.
 smoke020|||The children laughed and played in the garden until the sun began to set.
+smoke021|||A long line of cars stretched along the coastal highway after the lighthouse opened.
+smoke022|||The chef recommended the seared scallops with a lemon butter sauce as the special.
+smoke023|||Our quarterly engineering review begins promptly at two in the afternoon on Thursday.
+smoke024|||Despite the heavy rain, the marathon runners kept a steady pace through downtown.
+smoke025|||She finished the crossword puzzle in under fifteen minutes while sipping her tea.
+smoke026|||The new release adds streaming output, faster startup, and better error reporting.
+smoke027|||His grandfather taught him how to whittle small wooden birds from cedar offcuts.
+smoke028|||Travelers were advised to check the timetable since trains were running on a holiday schedule.
+smoke029|||The startup launched a public beta after eighteen months of closed development.
+smoke030|||At night the desert sky filled with stars so bright the Milky Way was clearly visible.

docker/Dockerfile.ci

@@ -1,5 +1,5 @@
 ARG VLLM_BASE_IMAGE=vllm/vllm-openai
-ARG VLLM_BASE_TAG=v0.20.0
+ARG VLLM_BASE_TAG=v0.21.0
 FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG}
 ARG APP_DIR=/workspace/vllm-omni
 WORKDIR ${APP_DIR}

docker/Dockerfile.npu.ci

@@ -15,7 +15,8 @@ RUN apt-get update && \
         libsm6 \
         libxext6 \
         libxrender1 \
-        libxcb1 && \
+        libxcb1 \
+        espeak-ng && \
     rm -rf /var/cache/apt/* && \
     rm -rf /var/lib/apt/lists/*
 
@@ -29,6 +30,11 @@ RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
       "pytest-cov>=4.0.0" \
       "pytest-mock>=3.10.0" \
       "openpyxl>=3.0.0" \
+      "psutil>=7.2.0" \
+      "soundfile>=0.13.1" \
+      "pyttsx3>=2.99" \
+      "opencv-python>=4.12.0.88" \
+      "imageio[ffmpeg]>=0.6.0" \
       --index-url http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple \
       --trusted-host cache-service.nginx-pypi-cache.svc.cluster.local \
       --default-timeout=60 && \

docker/Dockerfile.npu.ci.a3

@@ -15,7 +15,8 @@ RUN apt-get update && \
         libsm6 \
         libxext6 \
         libxrender1 \
-        libxcb1 && \
+        libxcb1 \
+        espeak-ng && \
     rm -rf /var/cache/apt/* && \
     rm -rf /var/lib/apt/lists/*
 
@@ -29,6 +30,11 @@ RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
       "pytest-cov>=4.0.0" \
       "pytest-mock>=3.10.0" \
       "openpyxl>=3.0.0" \
+      "psutil>=7.2.0" \
+      "soundfile>=0.13.1" \
+      "pyttsx3>=2.99" \
+      "opencv-python>=4.12.0.88" \
+      "imageio[ffmpeg]>=0.6.0" \
       --index-url http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple \
       --trusted-host cache-service.nginx-pypi-cache.svc.cluster.local \
       --default-timeout=60 && \

docker/Dockerfile.rocm

@@ -1,4 +1,4 @@
-ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.20.0
+ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.21.0
 FROM ${BASE_IMAGE} AS base
 
 # Declare a variable to know if we want to use the nightly build or the stable build.

docs/api/README.md

@@ -99,7 +99,6 @@ Configuration classes.
 - [vllm_omni.diffusion.cache.teacache.config.TeaCacheConfig][]
 - [vllm_omni.distributed.omni_connectors.utils.config.ConnectorSpec][]
 - [vllm_omni.distributed.omni_connectors.utils.config.OmniTransferConfig][]
-- [vllm_omni.model_executor.models.cosyvoice3.config.CosyVoice3Config][]
 - [vllm_omni.model_executor.models.fish_speech.configuration_fish_speech.FishSpeechConfig][]
 - [vllm_omni.model_executor.models.fish_speech.configuration_fish_speech.FishSpeechFastARConfig][]
 - [vllm_omni.model_executor.models.fish_speech.configuration_fish_speech.FishSpeechSlowARConfig][]
@@ -116,10 +115,14 @@ Configuration classes.
 - [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1DecoderConfig][]
 - [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1DecoderDiTConfig][]
 - [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1EncoderConfig][]
+- [vllm_omni.transformers_utils.configs.cosyvoice3.CosyVoice3Config][]
 - [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Config][]
 - [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Qwen2_5_VLConfig][]
 - [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Qwen2_5_VLTextConfig][]
 - [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Qwen2_5_VLVisionConfig][]
+- [vllm_omni.transformers_utils.configs.omnivoice.OmniVoiceConfig][]
+- [vllm_omni.transformers_utils.configs.voxcpm.VoxCPMConfig][]
+- [vllm_omni.transformers_utils.configs.voxcpm2.VoxCPM2Config][]
 
 ## Workers
 

docs/getting_started/installation/gpu.md

@@ -34,7 +34,7 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr
 
 ### Pre-built wheels
 
-Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, and 0.18.0. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
+Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, and 0.21.0. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
 
 === "NVIDIA CUDA"
 

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

@@ -20,7 +20,7 @@ Therefore, it is recommended to install vLLM and vLLM-Omni with a **fresh new**
 
 vLLM-Omni is built based on vLLM. Please install it with command below.
 ```bash
-uv pip install vllm --torch-backend=auto
+uv pip install vllm==0.21.0 --torch-backend=auto
 ```
 
 #### Installation of vLLM-Omni
@@ -39,13 +39,13 @@ uv pip install 'vllm-omni[demo]'
 # --8<-- [start:build-wheel-from-source]
 
 #### Installation of vLLM
-If you do not need to modify source code of vLLM, you can directly install the stable 0.20.0 release version of the library
+If you do not need to modify source code of vLLM, you can directly install the stable 0.21.0 release version of the library
 
 ```bash
-uv pip install vllm==0.20.0 --torch-backend=auto
+uv pip install vllm==0.21.0 --torch-backend=auto
 ```
 
-The 0.20.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
+The 0.21.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
 
 #### Installation of vLLM-Omni
 Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -66,12 +66,12 @@ If you want to check, modify or debug with source code of vLLM, install the libr
 ```bash
 git clone https://github.com/vllm-project/vllm.git
 cd vllm
-git checkout v0.20.0
+git checkout v0.21.0
 ```
 Set up environment variables to get pre-built wheels. If there are internet problems, just download the whl file manually. And set `VLLM_PRECOMPILED_WHEEL_LOCATION` as your local absolute path of whl file.
 ```bash
-#For CUDA 13.0 (the default for v0.20.0; the wheel filename has no `+cu130` suffix)
+#For CUDA 13.0 (the default for v0.21.0; the wheel filename has no `+cu130` suffix)
-export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.20.0/vllm-0.20.0-cp38-abi3-manylinux_2_35_x86_64.whl
+export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.21.0/vllm-0.21.0-cp38-abi3-manylinux_2_35_x86_64.whl
 ```
 Install vllm with command below (If you have no existing PyTorch).
 ```bash
@@ -102,7 +102,7 @@ docker run --runtime nvidia --gpus 2 \
     --env "HF_TOKEN=$HF_TOKEN" \
     -p 8091:8091 \
     --ipc=host \
-    vllm/vllm-omni:v0.18.0 \
+    vllm/vllm-omni:v0.21.0 \
     --model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 8091
 ```
 

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

@@ -13,7 +13,7 @@ vLLM-Omni current recommends the steps in under setup through Docker Images.
 
 vLLM-Omni is built based on vLLM. Please install it with command below.
 ```bash
-uv pip install vllm==0.20.0+rocm721 --extra-index-url https://wheels.vllm.ai/rocm/0.20.0/rocm721
+uv pip install vllm==0.21.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.21.0/rocm722
 ```
 
 #### Installation of vLLM-Omni
@@ -34,13 +34,13 @@ uv pip install onnxruntime-rocm
 # --8<-- [start:build-wheel-from-source]
 
 #### Installation of vLLM
-If you do not need to modify source code of vLLM, you can directly install the stable 0.20.0 release version of the library
+If you do not need to modify source code of vLLM, you can directly install the stable 0.21.0 release version of the library
 
 ```bash
-uv pip install vllm==0.20.0+rocm721 --extra-index-url https://wheels.vllm.ai/rocm/0.20.0/rocm721
+uv pip install vllm==0.21.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.21.0/rocm722
 ```
 
-The pre-built 0.20.0 vLLM wheel targets ROCm 7.2.1. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
+The pre-built 0.21.0 vLLM wheel targets ROCm 7.2.2. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
 
 #### Installation of vLLM-Omni
 Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -58,7 +58,7 @@ If you want to check, modify or debug with source code of vLLM, install the libr
 ```bash
 git clone https://github.com/vllm-project/vllm.git
 cd vllm
-git checkout v0.20.0
+git checkout v0.21.0
 python3 -m pip install -r requirements/rocm.txt
 python3 setup.py develop
 ```
@@ -130,7 +130,7 @@ docker run --rm \
   -v ~/.cache/huggingface:/root/.cache/huggingface \
   --env "HF_TOKEN=$HF_TOKEN" \
   -p 8091:8091 \
-  vllm/vllm-omni-rocm:v0.20.0 \
+  vllm/vllm-omni-rocm:v0.21.0 \
   --model Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 ```
 
@@ -149,7 +149,7 @@ docker run --rm -it \
   -v ~/.cache/huggingface:/root/.cache/huggingface \
   --env "HF_TOKEN=$HF_TOKEN" \
   --entrypoint bash \
-  vllm/vllm-omni-rocm:v0.20.0
+  vllm/vllm-omni-rocm:v0.21.0
 ```
 
 # --8<-- [end:pre-built-images]

docs/getting_started/quickstart.md

@@ -19,10 +19,10 @@ uv venv --python 3.12 --seed
 source .venv/bin/activate
 
 # On CUDA
-uv pip install vllm==0.20.0 --torch-backend=auto
+uv pip install vllm==0.21.0 --torch-backend=auto
 
 # On ROCm
-uv pip install vllm==0.20.0+rocm721 --extra-index-url https://wheels.vllm.ai/rocm/0.20.0/rocm721
+uv pip install vllm==0.21.0+rocm721 --extra-index-url https://wheels.vllm.ai/rocm/0.21.0/rocm721
 
 git clone https://github.com/vllm-project/vllm-omni.git
 cd vllm-omni
@@ -35,7 +35,7 @@ For additional installation methods — please see the [installation guide](inst
 !!! note
     It is important to install the same major & minor version of vLLM and vLLM Omni, otherwise things may not work as expected. If the versions are misaligned, you will see a warning when you import vLLM Omni.
 
-    If you are seeing strange behavior with the `vllm` command not handling the `--omni` flag correctly, you most likely have a version mismatch with vLLM < `0.20.0` and vLLM Omni `0.20.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue.
+    If you are seeing strange behavior with the `vllm` command not handling the `--omni` flag correctly, you most likely have a version mismatch with vLLM < `0.21.0` and vLLM Omni `0.21.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue.
 
 ## Offline Inference
 

docs/user_guide/diffusion_features.md

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

docs/user_guide/quantization/quantized_kvcache.md

@@ -8,15 +8,20 @@ vLLM-Omni supports online FP8 quantization for eligible diffusion Flash
 Attention (FA) to reduce FA latency while keeping model weights in their
 original dtype.
 
-This feature is configured through `kv_cache_dtype`, matching the option name
+This feature is configured through `diffusion_kv_cache_dtype` on
-used by vLLM's language-model KV-cache quantization. In vLLM-Omni diffusion
+`OmniDiffusionConfig` (CLI: `--diffusion-kv-cache-dtype`). It is intentionally
-pipelines, however, it is a runtime FA path: Q/K/V tensors are dynamically
+**not** the same as vLLM's `--kv-cache-dtype`, which controls autoregressive
-quantized before the attention operator. It does not quantize model weights and
+language-model KV cache storage and defaults to `"auto"`. Diffusion FA
-is separate from [FP8 W8A8](fp8.md), [Int8 W8A8](int8.md), or pre-quantized
+quantization uses the dedicated diffusion flags so omni serve does not inherit
-checkpoint formats.
+that default.
 
-If `kv_cache_dtype` is not set, behavior is unchanged and attention runs in the
+In vLLM-Omni diffusion pipelines, this is a runtime FA path: Q/K/V tensors are
-native dtype.
+dynamically quantized before the attention operator. It does not quantize model
+weights and is separate from [FP8 W8A8](fp8.md), [Int8 W8A8](int8.md), or
+pre-quantized checkpoint formats.
+
+If `diffusion_kv_cache_dtype` is not set, behavior is unchanged and attention
+runs in the native dtype.
 
 ## Hardware Support
 
@@ -30,8 +35,8 @@ native dtype.
 Legend: `✅` supported, `❌` unsupported.
 
 FP8 FA is currently implemented only for the NPU Flash Attention backend. Other
-backends do not support `kv_cache_dtype="fp8"` for diffusion attention and fall
+backends do not support `diffusion_kv_cache_dtype="fp8"` for diffusion attention
-back to native dtype execution.
+and fall back to native dtype execution.
 
 ## Model Type Support
 
@@ -40,7 +45,7 @@ back to native dtype execution.
 | Model | Scope | Status | Notes |
 |-------|-------|--------|-------|
 | Wan2.2 | Eligible DiT full-attention FA on Ascend NPU | Tested | Compare quality and latency against a BF16 baseline before production use |
-| Other diffusion models | Eligible DiT full-attention FA on Ascend NPU | Not tested | You can try `kv_cache_dtype="fp8"`; tune `kv_cache_skip_steps` and `kv_cache_skip_layers` when higher precision is needed |
+| Other diffusion models | Eligible DiT full-attention FA on Ascend NPU | Not tested | You can try `diffusion_kv_cache_dtype="fp8"`; tune `diffusion_kv_cache_skip_steps` and `diffusion_kv_cache_skip_layers` when higher precision is needed |
 
 ### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
 
@@ -50,8 +55,8 @@ guide documents support.
 ### Multi-Stage Diffusion Model (BAGEL, GLM-Image)
 
 Not tested. If the diffusion stage uses the same NPU Flash Attention backend,
-`kv_cache_dtype` may apply in theory; validate quality and latency for each
+`diffusion_kv_cache_dtype` may apply in theory; validate quality and latency for
-stage and model.
+each stage and model.
 
 ## Configuration
 
@@ -67,15 +72,15 @@ python examples/offline_inference/image_to_video/image_to_video.py \
     --num-inference-steps 4 \
     --ulysses-degree 4 \
     --vae-patch-parallel-size 4 \
-    --kv-cache-dtype fp8 \
+    --diffusion-kv-cache-dtype fp8 \
-    --kv-cache-skip-steps "0,1" \
+    --diffusion-kv-cache-skip-steps "0,1" \
-    --kv-cache-skip-layers "0-2"
+    --diffusion-kv-cache-skip-layers "0-2"
 ```
 
 Online serving:
 
 ```bash
-vllm serve <your-model> --omni --kv-cache-dtype fp8
+vllm serve <your-model> --omni --diffusion-kv-cache-dtype fp8
 ```
 
 Stage config:
@@ -86,18 +91,23 @@ stage_args:
     stage_type: diffusion
     engine_args:
       model_stage: dit
-      kv_cache_dtype: "fp8"
+      diffusion_kv_cache_dtype: "fp8"
-      kv_cache_skip_steps: "0,1"
+      diffusion_kv_cache_skip_steps: "0,1"
-      kv_cache_skip_layers: "0-2"
+      diffusion_kv_cache_skip_layers: "0-2"
 ```
 
+Legacy YAML keys `kv_cache_dtype`, `kv_cache_skip_steps`, and
+`kv_cache_skip_layers` are still accepted when constructing
+`OmniDiffusionConfig` (for example via `from_kwargs`); prefer the `diffusion_*`
+names for new configs.
+
 ## Parameters
 
 | Parameter | Type | Default | Description |
 |-----------|------|---------|-------------|
-| `kv_cache_dtype` | str \| None | `None` | Set to `"fp8"` to enable dynamic FP8 FA on supported attention backends |
+| `diffusion_kv_cache_dtype` | str \| None | `None` | Set to `"fp8"` to enable dynamic FP8 FA on supported attention backends |
-| `kv_cache_skip_steps` | str \| None | `None` | Denoising step selector to keep in native dtype, for example `"0,1,4-6"` |
+| `diffusion_kv_cache_skip_steps` | str \| None | `None` | Denoising step selector to keep in native dtype, for example `"0,1,4-6"` |
-| `kv_cache_skip_layers` | str \| None | `None` | Transformer layer selector to keep in native dtype, for example `"0-2,10"` |
+| `diffusion_kv_cache_skip_layers` | str \| None | `None` | Transformer layer selector to keep in native dtype, for example `"0-2,10"` |
 
 Selectors use comma-separated integers and inclusive ranges. Listed steps or
 layers skip FP8 FA; all other eligible full-attention forwards use the FP8 path.
@@ -106,9 +116,9 @@ layers skip FP8 FA; all other eligible full-attention forwards use the FP8 path.
 
 1. Compare generated images or videos against a BF16 baseline with the same
    seed, prompt, resolution, frame count, and denoising steps.
-2. Use `kv_cache_skip_steps` for denoising steps where quality is more
+2. Use `diffusion_kv_cache_skip_steps` for denoising steps where quality is more
    sensitive.
-3. Use `kv_cache_skip_layers` for transformer layers that show visible quality
+3. Use `diffusion_kv_cache_skip_layers` for transformer layers that show visible quality
    regressions.
 4. Report both latency and quality results when enabling this option for a new
    model. For image or video models, include visual comparison and quantitative

examples/offline_inference/hunyuan_image3/README.md

@@ -112,6 +112,7 @@ python end2end.py --modality text2img \
                   --additional-config '{"torchair_graph_config":{"enabled":true}}'
 ```
 
+
 ## Key Arguments
 
 | Argument | Description |
@@ -123,16 +124,15 @@ python end2end.py --modality text2img \
 | `--steps` | Number of diffusion inference steps for image generation. |
 | `--guidance-scale` | Classifier-free guidance scale for image generation. |
 | `--height`, `--width` | Output image size for `text2img`. |
-| `--bot-task` | Prompt behavior. `auto` selects the default from `--modality`; `think` adds `<think>`; `recaption` adds `<recaption>`; `vanilla` uses the text-to-image pretrain template. |
+| `--bot-task` | Override prompt mode. `none`, `think`, `recaption`, `think_recaption`, or `vanilla`. |
 | `--sys-type` | Override the system prompt type, for example `en_unified` or `en_vanilla`. |
 | `--vae-use-tiling` | Enable VAE tiling for memory reduction. |
 
 ## Notes
 
-- `hunyuan_image3_ar.yaml` is a 4-card AR-only text/comprehension deploy. It sets `engine_output_type: text`, `final_output_type: text`, and text sampling defaults.
+- `hunyuan_image3_ar.yaml` is a 4-card AR-only text/comprehension deploy.
-- `hunyuan_image3_dit.yaml` is a single-stage DiT deploy with `stage_id: 0`; it does not require stage 1 or a running AR stage.
+- `hunyuan_image3_dit.yaml` is a single-stage DiT deploy with `stage_id: 0`.
 - The old HunyuanImage3 YAMLs under `model_executor/stage_configs/` and `platforms/*/stage_configs/` have been folded into the deploy YAMLs.
-- This PR does not keep the HunyuanImage3 AR-to-DiT KV reuse wiring. The deploy YAMLs describe the topology and platform settings only.
 
 ## Prompt Format
 
@@ -148,22 +148,8 @@ Assistant: {trigger_tag?}
 
 - `<img>`: Placeholder for each input image (single token; expanded by the multimodal pipeline).
 - Trigger tags: `<think>` for CoT and `<recaption>` for recaptioning, placed after `Assistant: `.
-- System prompt: Auto-selected based on task.
+- System prompt: Auto-selected from `task` and `bot_task`.
-- `t2i_vanilla` is the only task that uses the bare pretrain template without chat structure.
+- `bot_task='vanilla'` with `task='t2i'` uses the bare pretrain template.
-- The example composes the internal prompt task from `--modality` and `--bot-task`
-  before calling `prompt_utils`; for example, `img2text + think` becomes
-  `i2t_think` for prompt and stop-token lookup.
 
 The shared `vllm_omni.diffusion.models.hunyuan_image3.prompt_utils.build_prompt_tokens()`
 helper handles segment-by-segment tokenization and matches HF `apply_chat_template`.
-
-## FAQ
-
-- **OOM errors**: Decrease `gpu_memory_utilization` in the deploy YAML, use a smaller `max_num_batched_tokens`, or enable VAE tiling with `--vae-use-tiling`.
-- **Custom image sizes**: Use `--height` and `--width` flags (multiples of 16 recommended).
-
-| Stage | VRAM (approx) |
-| :--- | :--- |
-| Stage 0 (AR) | ~15 GiB + KV Cache |
-| Stage 1 (DiT) | ~30 GiB |
-| Total (8-GPU) | ~45 GiB + KV Cache |

examples/offline_inference/hunyuan_image3/end2end.py

@@ -1,16 +1,5 @@
 """
 HunyuanImage-3.0-Instruct unified end-to-end inference script.
-
-Supports all modalities through a single entry point:
-  - text2img:  Text → AR → DiT → Image
-  - img2img:   Text+Image → AR → DiT → Edited Image (IT2I)
-  - img2text:  Image+Text → AR → Text description (I2T)
-  - text2text: Text → AR → Text (comprehension, no image)
-
-Usage:
-    python end2end.py --modality text2img --prompts "A cute cat"
-    python end2end.py --modality img2img --image-path input.png --prompts "Make it snowy"
-    python end2end.py --modality img2text --image-path input.png --prompts "Describe this image"
 """
 
 import argparse
@@ -19,18 +8,25 @@ import os
 from pathlib import Path
 
 from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import (
-    _TASK_PRESETS,
+    MAX_IMAGES_PER_REQUEST,
     build_prompt_tokens,
     resolve_stop_token_ids,
+    resolve_sys_type,
 )
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniPromptType
 
-# Default deploy configs are absolute so this example works from any cwd.
 _REPO_ROOT = Path(__file__).resolve().parents[3]
 _DEFAULT_DEPLOY_CONFIG = str(_REPO_ROOT / "vllm_omni" / "deploy" / "hunyuan_image3.yaml")
 _DEFAULT_AR_DEPLOY_CONFIG = str(_REPO_ROOT / "vllm_omni" / "deploy" / "hunyuan_image3_ar.yaml")
 
+_MODALITY_TASK_MAP: dict[str, tuple[str, str | None]] = {
+    "text2img": ("t2i", "think"),
+    "img2img": ("it2i", "think"),
+    "img2text": ("i2t", None),
+    "text2text": ("t2t", None),
+}
+
 _MODALITY_DEFAULT_DEPLOY_CONFIG = {
     "text2img": _DEFAULT_DEPLOY_CONFIG,
     "img2img": _DEFAULT_DEPLOY_CONFIG,
@@ -45,73 +41,37 @@ _MODALITY_MODE = {
     "text2text": "text-to-text",
 }
 
-_MODALITY_TASK_MAP = {
-    "text2img": "t2i",
-    "img2img": "it2i",
-    "img2text": "i2t",
-    "text2text": "t2t",
-}
-
 
 def parse_args():
     parser = argparse.ArgumentParser(description="HunyuanImage-3.0-Instruct end-to-end inference.")
-    parser.add_argument(
+    parser.add_argument("--model", default="tencent/HunyuanImage-3.0-Instruct", help="Model name or local path.")
-        "--model",
-        default="tencent/HunyuanImage-3.0-Instruct",
-        help="Model name or local path.",
-    )
     parser.add_argument(
         "--modality",
         default="text2img",
-        choices=["text2img", "img2img", "img2text", "text2text"],
+        choices=list(_MODALITY_TASK_MAP),
-        help="Modality mode to control stage execution.",
     )
     parser.add_argument("--prompts", nargs="+", default=None, help="Input text prompts.")
     parser.add_argument(
         "--image-path",
         type=str,
         default=None,
-        help="Path to input image (for img2img/img2text).",
+        help="Input image path(s) for img2img/img2text. Comma-separated for multi-image (up to 3).",
     )
-    parser.add_argument(
+    parser.add_argument("--output", type=str, default=".", help="Output directory to save results.")
-        "--output",
-        type=str,
-        default=".",
-        help="Output directory to save results.",
-    )
-
-    # Generation parameters
     parser.add_argument("--steps", type=int, default=50, help="Number of inference steps.")
     parser.add_argument("--guidance-scale", type=float, default=5.0, help="Classifier-free guidance scale.")
     parser.add_argument("--seed", type=int, default=42, help="Random seed.")
     parser.add_argument("--height", type=int, default=1024, help="Output image height.")
     parser.add_argument("--width", type=int, default=1024, help="Output image width.")
-    parser.add_argument(
+    parser.add_argument("--vae-use-tiling", action="store_true", help="Enable VAE tiling.")
-        "--vae-use-tiling",
-        action="store_true",
-        help="Enable VAE tiling for memory optimization.",
-    )
-
-    # Prompt configuration
     parser.add_argument(
         "--bot-task",
         type=str,
-        default="auto",
-        choices=["auto", "think", "recaption", "think_recaption", "vanilla"],
-        help=(
-            "Prompt behavior. 'auto' selects the default for the modality; "
-            "'think' adds <think>; 'recaption' adds <recaption>; "
-            "'vanilla' uses the t2i pretrain template."
-        ),
-    )
-    parser.add_argument(
-        "--sys-type",
-        type=str,
         default=None,
-        help="Override system prompt type (e.g. en_unified, en_vanilla).",
+        choices=["none", "think", "recaption", "think_recaption", "vanilla"],
+        help="Override prompt mode. Default: auto from --modality.",
     )
-
+    parser.add_argument("--sys-type", type=str, default=None, help="Override system prompt type.")
-    # Omni init args
     parser.add_argument("--deploy-config", type=str, default=None, help="Custom deploy YAML path.")
     parser.add_argument("--stage-configs-path", type=str, default=None, help="Custom legacy stage config YAML path.")
     parser.add_argument("--log-stats", action="store_true", default=False)
@@ -157,22 +117,13 @@ def main():
     os.makedirs(args.output, exist_ok=True)
     additional_config = parse_additional_config(args.additional_config)
 
-    # Determine task for prompt formatting from modality + bot behavior.
+    task, default_bot_task = _MODALITY_TASK_MAP[args.modality]
-    task = _MODALITY_TASK_MAP[args.modality]
+    if args.bot_task is None:
-    assert task is not None
+        bot_task: str | None = default_bot_task
-    bot_task = args.bot_task
+    elif args.bot_task == "none":
-    if bot_task != "auto":
+        bot_task = None
-        task = task + "_" + bot_task
+    else:
-    if task not in _TASK_PRESETS:
+        bot_task = args.bot_task
-        valid_bot_tasks = {
-            "text2img": ["think", "recaption", "vanilla"],
-            "img2img": ["think", "recaption", "think_recaption"],
-            "img2text": ["auto"],
-            "text2text": ["auto"],
-        }[args.modality]
-        raise ValueError(
-            f"--bot-task {bot_task!r} is not supported for {args.modality}. Choose from: {valid_bot_tasks}"
-        )
 
     if args.deploy_config is not None and args.stage_configs_path is not None:
         raise ValueError("--deploy-config and --stage-configs-path are mutually exclusive.")
@@ -182,13 +133,13 @@ def main():
     if deploy_config is None and stage_configs_path is None:
         deploy_config = _MODALITY_DEFAULT_DEPLOY_CONFIG[args.modality]
 
-    # Build Omni
     omni_kwargs = {
         "model": args.model,
         "vae_use_tiling": args.vae_use_tiling,
         "log_stats": args.log_stats,
         "init_timeout": args.init_timeout,
         "enforce_eager": args.enforce_eager,
+        "mode": _MODALITY_MODE[args.modality],
     }
 
     if additional_config is not None:
@@ -197,71 +148,67 @@ def main():
         omni_kwargs["deploy_config"] = deploy_config
     else:
         omni_kwargs["stage_configs_path"] = stage_configs_path
-    omni_kwargs["mode"] = _MODALITY_MODE[args.modality]
 
     omni = Omni(**omni_kwargs)
 
-    # Prepare prompts
     prompts = args.prompts or ["A cute cat"]
-    if not prompts:
+    input_images: list = []
-        print("[Info] No prompts provided, using default.")
-        prompts = ["A cute cat"]
-
-    # Load image if needed
-    input_image = None
     if args.modality in ("img2img", "img2text"):
-        if not args.image_path or not os.path.exists(args.image_path):
+        if not args.image_path:
             raise ValueError(f"--image-path required for {args.modality}, got: {args.image_path}")
         from PIL import Image
 
-        input_image = Image.open(args.image_path).convert("RGB")
+        image_paths = [p.strip() for p in args.image_path.split(",") if p.strip()]
+        if len(image_paths) > MAX_IMAGES_PER_REQUEST:
+            raise ValueError(
+                f"--image-path accepts at most {MAX_IMAGES_PER_REQUEST} images for "
+                f"HunyuanImage-3.0 IT2I, got {len(image_paths)}: {args.image_path}"
+            )
+        for image_path in image_paths:
+            if not os.path.exists(image_path):
+                raise ValueError(f"Image path does not exist: {image_path}")
+            input_images.append(Image.open(image_path).convert("RGB"))
+        if not input_images:
+            raise ValueError(f"--image-path produced no usable paths: {args.image_path!r}")
 
-    # Load tokenizer for segment-wise prompt tokenization (matches HF
-    # apply_chat_template byte-for-byte; see build_prompt_tokens docstring).
     from transformers import AutoTokenizer
 
     tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
+    mm_image_payload = (input_images[0] if len(input_images) == 1 else input_images) if input_images else None
 
-    # Format prompts
     formatted_prompts: list[OmniPromptType] = []
-    for p in prompts:
+    for prompt in prompts:
-        result = build_prompt_tokens(p, tokenizer, task=task, sys_type=args.sys_type)
+        build_kwargs: dict = {"task": task, "bot_task": bot_task, "sys_type": args.sys_type}
+        if input_images:
+            build_kwargs["num_images"] = len(input_images)
+        result = build_prompt_tokens(prompt, tokenizer, **build_kwargs)
         token_ids = result.token_ids
-        effective_sys_type = result.system_prompt_type
+        effective_sys_type = args.sys_type or resolve_sys_type(bot_task)
 
-        # `prompt_token_ids` drives the AR stage (matches HF byte-for-byte).
-        # `prompt` and `use_system_prompt` are forwarded by ar2diffusion to
-        # the DiT stage so the diffusion pipeline can rebuild the same
-        # system prefix when constructing its model inputs.
         prompt_dict: dict = {
             "prompt_token_ids": token_ids,
-            "prompt": p,
+            "prompt": prompt,
             "use_system_prompt": effective_sys_type,
         }
-
         if args.modality == "text2img":
             prompt_dict["modalities"] = ["image"]
         elif args.modality == "img2img":
             prompt_dict["modalities"] = ["image"]
-            prompt_dict["multi_modal_data"] = {"image": input_image}
+            prompt_dict["multi_modal_data"] = {"image": mm_image_payload}
-            prompt_dict["height"] = input_image.height
+            prompt_dict["height"] = input_images[0].height
-            prompt_dict["width"] = input_image.width
+            prompt_dict["width"] = input_images[0].width
         elif args.modality == "img2text":
             prompt_dict["modalities"] = ["text"]
-            prompt_dict["multi_modal_data"] = {"image": input_image}
+            prompt_dict["multi_modal_data"] = {"image": mm_image_payload}
-        elif args.modality == "text2text":
+        else:
             prompt_dict["modalities"] = ["text"]
-
         formatted_prompts.append(prompt_dict)
 
-    # Build sampling params from defaults
     params_list = list(omni.default_sampling_params_list)
 
-    # Override diffusion params if applicable
     from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
     ar_stop_token_ids = resolve_stop_token_ids(task=task, bot_task=bot_task, tokenizer=tokenizer)
-    assert ar_stop_token_ids is not None
     for sp in params_list:
         if isinstance(sp, OmniDiffusionSamplingParams):
             sp.num_inference_steps = args.steps
@@ -269,13 +216,12 @@ def main():
             sp.guidance_scale_provided = True
             if args.seed is not None:
                 sp.seed = args.seed
-            if args.modality in ("text2img",):
+            if args.modality == "text2img":
                 sp.height = args.height
                 sp.width = args.width
         elif hasattr(sp, "stop_token_ids"):
             sp.stop_token_ids = ar_stop_token_ids
 
-    # Print configuration
     print(f"\n{'=' * 60}")
     print("HunyuanImage-3.0 Generation Configuration:")
     print(f"  Model: {args.model}")
@@ -300,13 +246,9 @@ def main():
     print(f"  Prompts: {prompts}")
     print(f"{'=' * 60}\n")
 
-    # Generate
     omni_outputs = list(omni.generate(prompts=formatted_prompts, sampling_params_list=params_list))
-
-    # Process outputs
     img_idx = 0
     for req_output in omni_outputs:
-        # Text output (AR stage or text-only)
         ro = getattr(req_output, "request_output", None)
         txt = ""
         if ro and getattr(ro, "outputs", None):
@@ -320,11 +262,9 @@ def main():
         if txt:
             print(f"[Output] Text:\n{txt}")
 
-        # Image output (DiT stage)
         images = getattr(req_output, "images", None)
         if not images and ro and hasattr(ro, "images"):
             images = ro.images
-
         if images:
             for j, img in enumerate(images):
                 save_path = os.path.join(args.output, f"output_{img_idx}_{j}.png")

examples/offline_inference/image_to_video/image_to_video.py

@@ -103,22 +103,22 @@ def parse_args() -> argparse.Namespace:
         help="Sampling solver for Wan2.2 pipelines. Use 'euler' for Lightning/Distill setups.",
     )
     parser.add_argument(
-        "--kv-cache-dtype",
+        "--diffusion-kv-cache-dtype",
         type=str,
         default=None,
-        help="Config-level KV cache dtype (e.g. float8_e4m3fn).",
+        help="Diffusion attention KV cache dtype (e.g. float8_e4m3fn). Separate from vLLM --kv-cache-dtype.",
     )
     parser.add_argument(
-        "--kv-cache-skip-steps",
+        "--diffusion-kv-cache-skip-steps",
         type=str,
         default=None,
-        help="Config-level KV-cache quantization skip-step selector, e.g. '0-9,20,25-30'.",
+        help="Diffusion KV-cache quantization skip-step selector, e.g. '0-9,20,25-30'.",
     )
     parser.add_argument(
-        "--kv-cache-skip-layers",
+        "--diffusion-kv-cache-skip-layers",
         type=str,
         default=None,
-        help="Config-level KV-cache quantization skip-layer selector, e.g. '0,1,4-8'.",
+        help="Diffusion KV-cache quantization skip-layer selector, e.g. '0,1,4-8'.",
     )
     parser.add_argument("--output", type=str, default="i2v_output.mp4", help="Path to save the video (mp4).")
     parser.add_argument("--fps", type=int, default=None, help="Frames per second for the output video.")
@@ -334,9 +334,9 @@ def main():
         vae_use_tiling=args.vae_use_tiling,
         boundary_ratio=args.boundary_ratio,
         flow_shift=args.flow_shift,
-        kv_cache_dtype=args.kv_cache_dtype,
+        diffusion_kv_cache_dtype=args.diffusion_kv_cache_dtype,
-        kv_cache_skip_steps=args.kv_cache_skip_steps,
+        diffusion_kv_cache_skip_steps=args.diffusion_kv_cache_skip_steps,
-        kv_cache_skip_layers=args.kv_cache_skip_layers,
+        diffusion_kv_cache_skip_layers=args.diffusion_kv_cache_skip_layers,
         enable_cpu_offload=args.enable_cpu_offload,
         parallel_config=parallel_config,
         enforce_eager=args.enforce_eager,
@@ -361,9 +361,9 @@ def main():
     print(f"  Inference steps: {args.num_inference_steps}")
     print(f"  Frames: {args.num_frames}")
     print(f"  Solver: {args.sample_solver}")
-    print(f"  kv_cache_dtype(config): {args.kv_cache_dtype}")
+    print(f"  diffusion_kv_cache_dtype(config): {args.diffusion_kv_cache_dtype}")
-    print(f"  kv_cache_skip_steps(config): {args.kv_cache_skip_steps}")
+    print(f"  diffusion_kv_cache_skip_steps(config): {args.diffusion_kv_cache_skip_steps}")
-    print(f"  kv_cache_skip_layers(config): {args.kv_cache_skip_layers}")
+    print(f"  diffusion_kv_cache_skip_layers(config): {args.diffusion_kv_cache_skip_layers}")
     print(
         f"  Parallel configuration: cfg_parallel_size={args.cfg_parallel_size},"
         f" tensor_parallel_size={args.tensor_parallel_size}, vae_patch_parallel_size={args.vae_patch_parallel_size}"

examples/offline_inference/text_to_speech/README.md

@@ -73,13 +73,15 @@ python examples/offline_inference/text_to_speech/cosyvoice3/end2end.py \
 ```
 
 ### Voice cloning
-Pass a reference audio. Note that CosyVoice3's `--prompt-text` is a system-style prompt for the GPT stage, not a reference transcript:
+If `--ref-audio` is omitted, the script downloads the upstream
+[`zero_shot_prompt.wav`](https://github.com/FunAudioLLM/CosyVoice/blob/main/asset/zero_shot_prompt.wav) from the CosyVoice repo into the current directory.
+To use your own clip, pass `--ref-audio /path/to/reference.wav`, and modify `--prompt-text` correspondingly.
 ```bash
 python examples/offline_inference/text_to_speech/cosyvoice3/end2end.py \
     --model pretrained_models/Fun-CosyVoice3-0.5B \
     --tokenizer pretrained_models/Fun-CosyVoice3-0.5B/CosyVoice-BlankEN \
-    --ref-audio prompt.wav \
+    --ref-audio /path/to/reference.wav \
-    --prompt-text "You are a helpful assistant.<|endofprompt|>Testing my voices. Why should I not?"
+    --prompt-text "You are a helpful assistant.<|endofprompt|>Trascript in your ref audio clip"
 ```
 
 ### Streaming
@@ -317,44 +319,39 @@ pip install voxcpm soundfile
 export VLLM_OMNI_VOXCPM_CODE_PATH=/path/to/VoxCPM/src
 ```
 
-If the native VoxCPM `config.json` does not contain HF metadata such as `model_type`, prepare a persistent HF-compatible config directory and point the stage configs to it via `VLLM_OMNI_VOXCPM_HF_CONFIG_PATH`:
-
-```bash
-export VOXCPM_MODEL=/path/to/voxcpm-model
-export VLLM_OMNI_VOXCPM_HF_CONFIG_PATH=/tmp/voxcpm_hf_config
-mkdir -p "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH"
-cp "$VOXCPM_MODEL/config.json" "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH/config.json"
-cp "$VOXCPM_MODEL/generation_config.json" "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH/generation_config.json" 2>/dev/null || true
-python3 -c 'import json, os; p=os.path.join(os.environ["VLLM_OMNI_VOXCPM_HF_CONFIG_PATH"], "config.json"); cfg=json.load(open(p, "r", encoding="utf-8")); cfg["model_type"]="voxcpm"; cfg.setdefault("architectures", ["VoxCPMForConditionalGeneration"]); json.dump(cfg, open(p, "w", encoding="utf-8"), indent=2, ensure_ascii=False)'
-```
-
 ### Quick start
 ```bash
 python examples/offline_inference/text_to_speech/voxcpm/end2end.py \
-    --model "$VOXCPM_MODEL" \
+    --model openbmb/VoxCPM-0.5B \
     --text "This is a split-stage VoxCPM synthesis example running on vLLM Omni."
 ```
 
 ### Voice cloning
 ```bash
 python examples/offline_inference/text_to_speech/voxcpm/end2end.py \
-    --model "$VOXCPM_MODEL" \
+    --model openbmb/VoxCPM-0.5B \
     --text "This sentence is synthesized with a cloned voice." \
     --ref-audio /path/to/reference.wav \
     --ref-text  "The exact transcript spoken in reference.wav."
 ```
 
 ### Streaming
-Pass the async-chunk stage config:
+Pass `--streaming` together with the legacy async-chunk stage config:
 ```bash
 python examples/offline_inference/text_to_speech/voxcpm/end2end.py \
-    --model "$VOXCPM_MODEL" \
+    --model openbmb/VoxCPM-0.5B \
+    --streaming \
     --stage-configs-path vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml \
     --text "This is a split-stage VoxCPM streaming example running on vLLM Omni."
 ```
 
+### Persistent HF config (optional)
+The engine auto-patches a missing `model_type` in the checkpoint's `config.json` in-memory at startup (see `vllm_omni/engine/arg_utils.py`), so no manual setup is required.
+
+If you'd rather maintain a persistent patched config dir (e.g. to share across runs or to add fields the auto-patcher doesn't set), export `VLLM_OMNI_VOXCPM_HF_CONFIG_PATH` to its location — both the bundled deploy and async-chunk yamls read it via `${oc.env:VLLM_OMNI_VOXCPM_HF_CONFIG_PATH,}`.
+
 ### Notes
-- `voxcpm.yaml` is the default non-streaming stage config; `voxcpm_async_chunk.yaml` enables streaming.
+- Non-streaming auto-loads `vllm_omni/deploy/voxcpm.yaml` from the model directory name / HF `model_type`. Streaming uses the legacy `voxcpm_async_chunk.yaml` passed via `--stage-configs-path`.
 - Streaming is currently single-request oriented.
 - `--ref-text` must be the real transcript of `--ref-audio`; mismatched text degrades quality.
 - For online serving, see the [VoxCPM section in the online hub](../../online_serving/text_to_speech/README.md#voxcpm). For benchmark reporting, see [`benchmarks/voxcpm`](../../../benchmarks/voxcpm/README.md).

examples/offline_inference/text_to_speech/cosyvoice3/end2end.py

@@ -2,18 +2,32 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import argparse
 import os
+import urllib.request
+from pathlib import Path
 
 import numpy as np
 import soundfile as sf
 from vllm import SamplingParams
-from vllm.assets.audio import AudioAsset
 from vllm.multimodal.media.audio import load_audio
 
 from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
 from vllm_omni.entrypoints.omni import Omni
-from vllm_omni.model_executor.models.cosyvoice3.config import CosyVoice3Config
 from vllm_omni.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokenizer
 from vllm_omni.model_executor.models.cosyvoice3.utils import extract_text_token
+from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config
+
+# Upstream zero-shot reference clip
+ZERO_SHOT_PROMPT_URL = "https://raw.githubusercontent.com/FunAudioLLM/CosyVoice/main/asset/zero_shot_prompt.wav"
+
+
+def _default_ref_audio() -> str:
+    # Download the upstream zero_shot_prompt.wav into the current dir
+    dest = Path("zero_shot_prompt.wav")
+    if not dest.exists() or dest.stat().st_size == 0:
+        print(f"Downloading default reference audio to {dest}")
+        urllib.request.urlretrieve(ZERO_SHOT_PROMPT_URL, dest)
+
+    return str(dest)
 
 
 def run_e2e():
@@ -36,9 +50,15 @@ def run_e2e():
     parser.add_argument(
         "--prompt-text",
         type=str,
-        default="You are a helpful assistant.<|endofprompt|>Testing my voices. Why should I not?",
+        default="You are a helpful assistant.<|endofprompt|>希望你以后，能够做的比我还好呦!",
+    )
+    parser.add_argument(
+        "--ref-audio",
+        type=str,
+        default=None,
+        help="Path to reference audio for voice cloning. "
+        "If unset, downloads the upstream CosyVoice3 zero-shot prompt audio clip",
     )
-    parser.add_argument("--ref-audio", type=str, default="prompt.wav")
     parser.add_argument(
         "--tokenizer",
         type=str,
@@ -66,24 +86,22 @@ def run_e2e():
     sampling_cfg = {"top_p": 0.8, "top_k": 25, "eos_token_id": 6561 + 1}
 
     print("Model initialized. Preparing inputs...")
-    if args.ref_audio:
+    ref_audio_path = args.ref_audio or _default_ref_audio()
-        if not os.path.exists(args.ref_audio):
+    if not os.path.exists(ref_audio_path):
-            raise FileNotFoundError(f"Audio file not found: {args.ref_audio}")
+        raise FileNotFoundError(f"Audio file not found: {ref_audio_path}")
-        # Load at native sample rate
+    # Load at native sample rate
-        audio_signal, sr = load_audio(args.ref_audio, sr=None)
+    audio_signal, sr = load_audio(ref_audio_path, sr=None)
 
-        # Validate sample rate before processing (similar to original CosyVoice)
+    # Validate sample rate before processing (similar to original CosyVoice)
-        min_sr = 16000
+    min_sr = 16000
-        if sr < min_sr:
+    if sr < min_sr:
-            raise ValueError(
+        raise ValueError(
-                f"Audio sample rate {sr} Hz is too low. "
+            f"Audio sample rate {sr} Hz is too low. "
-                f"Minimum required: {min_sr} Hz. "
+            f"Minimum required: {min_sr} Hz. "
-                f"Please provide audio with sample rate >= {min_sr} Hz."
+            f"Please provide audio with sample rate >= {min_sr} Hz."
-            )
+        )
 
-        audio_data = (audio_signal.astype(np.float32), sr)
+    audio_data = (audio_signal.astype(np.float32), sr)
-    else:
-        audio_data = AudioAsset("mary_had_lamb").audio_and_sample_rate
 
     prompts = {
         "prompt": args.text,

examples/offline_inference/text_to_speech/voxcpm/end2end.py

@@ -11,12 +11,9 @@ import soundfile as sf
 import torch
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 
-from tests.helpers.stage_config import get_deploy_config_path
 from vllm_omni import AsyncOmni, Omni
 from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
 
-DEFAULT_SYNC_STAGE_CONFIG = get_deploy_config_path("voxcpm.yaml")
-
 
 def _build_prompt(args) -> dict[str, Any]:
     additional_information: dict[str, list[Any]] = {
@@ -59,10 +56,6 @@ def _extract_sample_rate(mm: dict[str, Any]) -> int:
     return int(sr_raw)
 
 
-def _is_streaming_stage_config(stage_config_path: str) -> bool:
-    return "async_chunk" in Path(stage_config_path).stem
-
-
 def _save_audio(audio: torch.Tensor, sample_rate: int, output_dir: Path, request_id: str) -> Path:
     output_dir.mkdir(parents=True, exist_ok=True)
     output_path = output_dir / f"output_{request_id}.wav"
@@ -87,6 +80,7 @@ async def _run_streaming(args) -> Path:
     omni = AsyncOmni(
         model=args.model,
         stage_configs_path=args.stage_configs_path,
+        deploy_config=args.deploy_config,
         log_stats=args.log_stats,
         stage_init_timeout=args.stage_init_timeout,
     )
@@ -133,6 +127,7 @@ def _run_sync(args) -> Path:
     omni = Omni(
         model=args.model,
         stage_configs_path=args.stage_configs_path,
+        deploy_config=args.deploy_config,
         log_stats=args.log_stats,
         stage_init_timeout=args.stage_init_timeout,
     )
@@ -164,11 +159,27 @@ def _run_sync(args) -> Path:
 def parse_args():
     parser = FlexibleArgumentParser(description="Minimal offline VoxCPM example for vLLM Omni.")
     parser.add_argument("--model", type=str, required=True, help="Local VoxCPM model directory.")
+    parser.add_argument(
+        "--deploy-config",
+        type=str,
+        default=None,
+        help=("Override the deploy config path."),
+    )
     parser.add_argument(
         "--stage-configs-path",
         type=str,
-        default=DEFAULT_SYNC_STAGE_CONFIG,
+        default=None,
-        help=("Stage config path. Use voxcpm.yaml for non-streaming or voxcpm_async_chunk.yaml for streaming."),
+        help=(
+            "Legacy stage_args yaml path. Required for streaming "
+            "(vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml); "
+            "leave unset for non-streaming to use the auto-resolved deploy config."
+        ),
+    )
+    parser.add_argument(
+        "--streaming",
+        action="store_true",
+        default=False,
+        help="Stream audio chunks as they arrive via AsyncOmni.",
     )
     parser.add_argument("--text", type=str, required=True, help="Input text for synthesis.")
     parser.add_argument("--ref-audio", type=str, default=None, help="Reference audio path for voice cloning.")
@@ -194,9 +205,10 @@ def parse_args():
 
 
 def main(args) -> None:
-    route = "streaming" if _is_streaming_stage_config(args.stage_configs_path) else "sync"
+    route = "streaming" if args.streaming else "sync"
     print(f"Model: {args.model}")
-    print(f"Stage config: {args.stage_configs_path}")
+    print(f"Deploy config: {args.deploy_config or '<auto from HF model_type>'}")
+    print(f"Stage configs path: {args.stage_configs_path or '<unused>'}")
     print(f"Route: {route}")
     if route == "streaming":
         asyncio.run(_run_streaming(args))

examples/online_serving/text_to_speech/qwen3_tts/streaming_speech_client.py

@@ -7,6 +7,12 @@ Usage:
     # Send full text at once
     python streaming_speech_client.py --text "Hello world. How are you? I am fine."
 
+    # Pick a built-in speaker for CustomVoice models
+    python streaming_speech_client.py \
+        --text "打开电子枪，关闭扫描，切换到低倍模式。" \
+        --speaker "Serena" \
+        --language "Chinese"
+
     # Simulate STT: send text word-by-word with delay
     python streaming_speech_client.py \
         --text "Hello world. How are you? I am fine." \
@@ -211,7 +217,9 @@ def main():
 
     args = parser.parse_args()
 
-    # Build session config (only include non-None values)
+    # Build session config (only include non-None values).
+    # Server canonical field is `voice`; `speaker` is accepted as an alias
+    # (see StreamingSpeechSessionConfig.voice in protocol/audio.py).
     config = {}
     for key in [
         "model",

tests/core/sched/test_omni_scheduler_input_coordinator_cleanup.py

@@ -0,0 +1,94 @@
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+from vllm.v1.request import RequestStatus
+
+import vllm_omni.core.sched.omni_ar_scheduler as ar_sched_mod
+import vllm_omni.core.sched.omni_generation_scheduler as gen_sched_mod
+from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
+from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class FakeInputCoordinator:
+    def __init__(self) -> None:
+        self.freed: list[str] = []
+
+    def free_finished_request(self, request_id: str) -> None:
+        self.freed.append(request_id)
+
+
+class DummyRequest:
+    request_id = "req-free"
+    client_index = 0
+    additional_information = None
+
+    def is_finished(self) -> bool:
+        return True
+
+
+@pytest.mark.parametrize(
+    ("scheduler_cls", "scheduler_mod"),
+    [
+        (OmniARScheduler, ar_sched_mod),
+        (OmniGenerationScheduler, gen_sched_mod),
+    ],
+)
+def test_finish_requests_cleans_input_coordinator_for_finished_ids(
+    monkeypatch: pytest.MonkeyPatch,
+    scheduler_cls,
+    scheduler_mod,
+) -> None:
+    coordinator = FakeInputCoordinator()
+    scheduler = scheduler_cls.__new__(scheduler_cls)
+    scheduler.chunk_transfer_adapter = None
+    scheduler.input_coordinator = coordinator
+
+    def fake_finish_requests(self, request_ids, finished_status):
+        assert request_ids == ["req-a", "req-b"]
+        assert finished_status == RequestStatus.FINISHED_STOPPED
+        return [("req-a", 0), ("req-b", 1)]
+
+    monkeypatch.setattr(scheduler_mod.VLLMScheduler, "finish_requests", fake_finish_requests)
+
+    result = scheduler_cls.finish_requests(
+        scheduler,
+        ["req-a", "req-b"],
+        RequestStatus.FINISHED_STOPPED,
+    )
+
+    assert result == [("req-a", 0), ("req-b", 1)]
+    assert coordinator.freed == ["req-a", "req-b"]
+
+
+def test_ar_free_request_cleans_input_coordinator_on_normal_free() -> None:
+    coordinator = FakeInputCoordinator()
+    scheduler = OmniARScheduler.__new__(OmniARScheduler)
+    scheduler.input_coordinator = coordinator
+    scheduler._omits_kv_transfer_cache = {"req-free": True}
+    scheduler.encoder_cache_manager = SimpleNamespace(free=lambda request: None)
+    scheduler.finished_req_ids = set()
+    scheduler.finished_req_ids_dict = None
+    scheduler._new_prompt_len_snapshot = {"req-free": 3}
+    scheduler._connector_finished = lambda request: (False, None)
+    scheduler._should_transfer_kv_for_request = lambda request_id: False
+    scheduler._free_blocks = lambda request: None
+
+    result = OmniARScheduler._free_request(scheduler, DummyRequest())
+
+    assert result is None
+    assert coordinator.freed == ["req-free"]
+    assert scheduler._omits_kv_transfer_cache == {}
+    assert scheduler._new_prompt_len_snapshot == {}
+
+
+def test_generation_free_request_cleans_input_coordinator(monkeypatch: pytest.MonkeyPatch) -> None:
+    coordinator = FakeInputCoordinator()
+    scheduler = OmniGenerationScheduler.__new__(OmniGenerationScheduler)
+    scheduler.input_coordinator = coordinator
+
+    def fake_free_request(self, request, delay_free_blocks=False):
+        assert delay_free_blocks is True

tests/core/sched/test_omni_scheduling_coordinator.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Unit tests for OmniSchedulingCoordinator (formerly ChunkSchedulingCoordinator).
+"""Unit tests for OmniSchedulingCoordinator.
 
 These tests use mock request objects and mock queues.  They do not require
 GPU, vLLM runtime, or any connector.
@@ -11,10 +11,12 @@ from __future__ import annotations
 import unittest
 from types import SimpleNamespace
 
+import torch
+
 import vllm_omni.core.sched.omni_scheduling_coordinator as coord_mod
 from vllm_omni.core.sched.omni_scheduling_coordinator import (
-    ChunkSchedulingCoordinator,
     OmniSchedulingCoordinator,
+    uses_qwen3_omni_full_payload_input_coordinator,
 )
 
 # ------------------------------------------------------------------ #
@@ -89,11 +91,55 @@ class MockQueue:
 # ------------------------------------------------------------------ #
 
 
+class TestFullPayloadCoordinatorSelection(unittest.TestCase):
+    def test_qwen3_omni_talker_and_code2wav_use_full_payload_input_coordinator(self):
+        for model_stage in ("talker", "code2wav"):
+            model_config = SimpleNamespace(
+                stage_id=1,
+                async_chunk=False,
+                model_arch="Qwen3OmniMoeForConditionalGeneration",
+                model_stage=model_stage,
+            )
+
+            self.assertTrue(uses_qwen3_omni_full_payload_input_coordinator(model_config))
+
+    def test_async_chunk_and_non_qwen3_omni_do_not_use_full_payload_input_coordinator(self):
+        cases = [
+            SimpleNamespace(
+                stage_id=1,
+                async_chunk=True,
+                model_arch="Qwen3OmniMoeForConditionalGeneration",
+                model_stage="talker",
+            ),
+            SimpleNamespace(
+                stage_id=1,
+                async_chunk=False,
+                model_arch="Qwen3TTSForConditionalGeneration",
+                model_stage="code2wav",
+            ),
+            SimpleNamespace(
+                stage_id=1,
+                async_chunk=False,
+                model_arch="Qwen2_5OmniForConditionalGeneration",
+                model_stage="talker",
+            ),
+            SimpleNamespace(
+                stage_id=0,
+                async_chunk=False,
+                model_arch="Qwen3OmniMoeForConditionalGeneration",
+                model_stage="thinker",
+            ),
+        ]
+
+        for model_config in cases:
+            self.assertFalse(uses_qwen3_omni_full_payload_input_coordinator(model_config))
+
+
 class TestChunkCoordinatorStateTransition(unittest.TestCase):
     """Test 5: process_pending_chunks transitions WAITING_FOR_CHUNK → target."""
 
     def test_ready_request_transitions_to_waiting(self):
-        coord = ChunkSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1, async_chunk=True)
+        coord = OmniSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1, async_chunk=True)
 
         req = _make_request("r1", status=RequestStatus.WAITING_FOR_CHUNK)
         waiting = MockQueue([req])
@@ -110,7 +156,7 @@ class TestChunkCoordinatorStateTransition(unittest.TestCase):
         self.assertIn("r1", coord.requests_with_ready_chunks)
 
     def test_non_ready_stays_waiting_for_chunk(self):
-        coord = ChunkSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1, async_chunk=True)
+        coord = OmniSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1, async_chunk=True)
 
         req = _make_request("r1", status=RequestStatus.WAITING_FOR_CHUNK)
         waiting = MockQueue([req])
@@ -126,7 +172,7 @@ class TestChunkCoordinatorStateTransition(unittest.TestCase):
         self.assertEqual(req.status, RequestStatus.WAITING_FOR_CHUNK)
 
     def test_stage_0_is_noop(self):
-        coord = ChunkSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=0)
+        coord = OmniSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=0)
         req = _make_request("r1")
         waiting = MockQueue([req])
         running: list = []
@@ -144,7 +190,7 @@ class TestChunkCoordinatorRestoreQueues(unittest.TestCase):
     """Test 6: restore_queues returns waiting-for-chunk requests."""
 
     def test_restore(self):
-        coord = ChunkSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1)
+        coord = OmniSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1)
 
         r1 = _make_request("r1")
         r2 = _make_request("r2")
@@ -166,7 +212,7 @@ class TestChunkCoordinatorFinishedSignal(unittest.TestCase):
     """Test 8: chunk_finished_req_ids → finished_requests."""
 
     def test_finished_signal(self):
-        coord = ChunkSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1, async_chunk=True)
+        coord = OmniSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1, async_chunk=True)
 
         req = _make_request("r1", status=RequestStatus.WAITING_FOR_CHUNK)
         waiting = MockQueue([req])
@@ -187,7 +233,7 @@ class TestChunkCoordinatorUpdateRequestMetadata(unittest.TestCase):
 
     def test_ar_mode_no_longer_sets_additional_information(self):
         """AR mode only processes scheduling metadata, not full payloads."""
-        coord = ChunkSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1)
+        coord = OmniSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1)
 
         req = _make_request("r1")
         requests = {"r1": req}
@@ -204,10 +250,14 @@ class TestChunkCoordinatorUpdateRequestMetadata(unittest.TestCase):
         self.assertIsNone(getattr(req, "additional_information", None))
 
     def test_generation_mode(self):
-        coord = ChunkSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1)
+        coord = OmniSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1)
 
         req = _make_request("r1")
         req.prompt_token_ids = [0, 0, 0]
+        req.num_prompt_tokens = 3
+        req.num_computed_tokens = 3
+        req._all_token_ids = [0, 0, 0, 99]
+        req._output_token_ids = [99]
         requests = {"r1": req}
 
         request_metadata = {
@@ -220,16 +270,61 @@ class TestChunkCoordinatorUpdateRequestMetadata(unittest.TestCase):
         coord.update_request_metadata(requests, request_metadata, model_mode="generation")
 
         self.assertEqual(req.prompt_token_ids, [10, 20, 30])
+        self.assertEqual(req.num_prompt_tokens, 3)
         self.assertEqual(req.num_computed_tokens, 0)
+        self.assertEqual(req._all_token_ids, [10, 20, 30])
+        self.assertEqual(req._output_token_ids, [])
         self.assertIsNone(req.additional_information)
         self.assertEqual(req._omni_initial_model_buffer, {"meta": {"left_context_size": 25}})
 
+    def test_generation_mode_flattens_tensor_code_predictor_codes(self):
+        coord = OmniSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1)
+
+        req = _make_request("r1")
+        req.prompt_token_ids = [9]
+        req.num_prompt_tokens = 1
+        req._all_token_ids = [9, 8]
+        req._output_token_ids = [8]
+        requests = {"r1": req}
+
+        coord.update_request_metadata(
+            requests,
+            {"r1": {"code_predictor_codes": torch.tensor([[1, 2, 3]], dtype=torch.long)}},
+            model_mode="generation",
+        )
+
+        self.assertEqual(req.prompt_token_ids, [1, 2, 3])
+        self.assertEqual(req.num_prompt_tokens, 3)
+        self.assertEqual(req._all_token_ids, [1, 2, 3])
+        self.assertEqual(req._output_token_ids, [])
+
+    def test_generation_mode_flattens_nested_code_predictor_codes(self):
+        coord = OmniSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1)
+
+        req = _make_request("r1")
+        req.prompt_token_ids = [9]
+        req.num_prompt_tokens = 1
+        req._all_token_ids = [9, 8]
+        req._output_token_ids = [8]
+        requests = {"r1": req}
+
+        coord.update_request_metadata(
+            requests,
+            {"r1": {"code_predictor_codes": [[1, 2], [3, 4]]}},
+            model_mode="generation",
+        )
+
+        self.assertEqual(req.prompt_token_ids, [1, 2, 3, 4])
+        self.assertEqual(req.num_prompt_tokens, 4)
+        self.assertEqual(req._all_token_ids, [1, 2, 3, 4])
+        self.assertEqual(req._output_token_ids, [])
+
 
 class TestChunkCoordinatorPostprocess(unittest.TestCase):
     """Test postprocess_scheduler_output clears ready chunks."""
 
     def test_clear_ready(self):
-        coord = ChunkSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1)
+        coord = OmniSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1)
         coord.requests_with_ready_chunks = {"r1", "r2"}
 
         new_req = SimpleNamespace(req_id="r1")
@@ -366,6 +461,33 @@ class TestWaitingForInputTransition(unittest.TestCase):
         self.assertEqual(len(coord.pending_input_registrations), 1)
         self.assertEqual(coord.pending_input_registrations[0].request_id, "r1")
 
+    def test_idle_cycles_retain_received_marker_before_request_appears(self):
+        coord = OmniSchedulingCoordinator(
+            scheduler_max_num_seqs=10,
+            stage_id=1,
+            async_chunk=False,
+        )
+        coord._full_payload_input_received.add("late")
+        coord.finished_requests.add("late")
+
+        waiting = MockQueue()
+        running: list = []
+
+        coord.process_pending_full_payload_inputs(waiting, running, stage_recv_req_ids=set())
+
+        self.assertIn("late", coord._full_payload_input_received)
+        self.assertIn("late", coord.finished_requests)
+
+        late_req = _make_request("late", status=RequestStatus.WAITING)
+        waiting.add_request(late_req)
+
+        coord.process_pending_full_payload_inputs(waiting, running, stage_recv_req_ids=set())
+
+        self.assertEqual(late_req.status, RequestStatus.WAITING)
+        self.assertEqual(coord.pending_input_registrations, [])
+        self.assertIn("late", coord._full_payload_input_received)
+        self.assertIn("late", coord.finished_requests)
+
 
 class TestTimeoutDetection(unittest.TestCase):
     """Regression tests for orphaned pending-recv timeout detection.
@@ -508,14 +630,18 @@ class TestTimeoutDetection(unittest.TestCase):
         self.assertEqual(len(coord._waiting_for_input), 0)
 
     def test_free_finished_request_clears_waiting_since(self):
-        """free_finished_request clears _waiting_since."""
+        """free_finished_request clears coordinator lifecycle markers."""
         coord = OmniSchedulingCoordinator(
             scheduler_max_num_seqs=10,
             stage_id=1,
         )
         coord._waiting_since["r1"] = 0.0
+        coord._full_payload_input_received.add("r1")
+        coord.finished_requests.add("r1")
         coord.free_finished_request("r1")
         self.assertNotIn("r1", coord._waiting_since)
+        self.assertNotIn("r1", coord._full_payload_input_received)
+        self.assertNotIn("r1", coord.finished_requests)
 
     def test_timeout_from_running_queue_full_lifecycle(self):
         """End-to-end: request from running → WAITING_FOR_CHUNK → restore →

tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json

@@ -10,16 +10,16 @@
                 "dataset_name": "random",
                 "backend": "openai-chat-omni",
                 "endpoint": "/v1/chat/completions",
-                "num_prompts": [4, 16, 32, 64],
+                "num_prompts": [4, 16, 32, 64, 128],
-                "max_concurrency": [1, 4, 8, 16],
+                "max_concurrency": [1, 4, 8, 16, 32],
                 "random_input_len": 2500,
                 "random_output_len": 900,
                 "ignore_eos": true,
                 "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
                 "baseline": {
-                    "mean_ttft_ms": [1000, 3000, 5000, 7000],
+                    "mean_ttft_ms": [1000, 3000, 5000, 7000, 9000],
-                    "mean_audio_ttfp_ms": [1000, 3000, 5000, 7000],
+                    "mean_audio_ttfp_ms": [1000, 3000, 5000, 7000, 9000],
-                    "mean_audio_rtf": [0.2, 0.35, 0.6, 0.85]
+                    "mean_audio_rtf": [0.2, 0.35, 0.6, 0.85, 0.9]
                 }
             },
             {

tests/dfx/perf/tests/test_tts.json

@@ -2,16 +2,19 @@
   {
     "test_name": "test_qwen3_tts_base",
     "server_params": {
-      "model": "Qwen/Qwen3-TTS-12Hz-1.7B-Base"
+      "model": "Qwen/Qwen3-TTS-12Hz-1.7B-Base",
+      "stage_overrides": {
+        "1": {"devices": "1"}
+      }
     },
     "benchmark_params": [
       {
         "task": "voice_clone",
         "eval_phase": "latency",
-        "enabled": false,
         "dataset_name": "seed-tts",
         "backend": "openai-audio-speech",
         "endpoint": "/v1/audio/speech",
+        "dataset_path": "linyueqian/seed-tts-eval-subset",
         "num_prompts": [20],
         "max_concurrency": [1],
         "seed_tts_locale": "en",
@@ -24,18 +27,18 @@
       {
         "task": "voice_clone",
         "eval_phase": "throughput",
-        "enabled": false,
         "dataset_name": "seed-tts",
         "backend": "openai-audio-speech",
         "endpoint": "/v1/audio/speech",
-        "num_prompts": [80],
+        "dataset_path": "linyueqian/seed-tts-eval-subset",
-        "max_concurrency": [8],
+        "num_prompts": [80, 128, 128],
+        "max_concurrency": [8, 16, 64],
         "seed_tts_locale": "en",
         "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
         "baseline": {
-          "median_audio_ttfp_ms": [3500],
+          "median_audio_ttfp_ms": [3500, 30000, 7000],
-          "median_audio_rtf": [0.75],
+          "median_audio_rtf": [0.75, 10.0, 1.5],
-          "audio_throughput": [10.0]
+          "audio_throughput": [10.0, 1.0, 40.0]
         }
       },
       {
@@ -45,6 +48,7 @@
         "dataset_name": "seed-tts",
         "backend": "openai-audio-speech",
         "endpoint": "/v1/audio/speech",
+        "dataset_path": "linyueqian/seed-tts-eval-subset",
         "num_prompts": [200],
         "max_concurrency": [4],
         "seed_tts_locale": "en",
@@ -59,7 +63,10 @@
   {
     "test_name": "test_qwen3_tts_customvoice",
     "server_params": {
-      "model": "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice"
+      "model": "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
+      "stage_overrides": {
+        "1": {"devices": "1"}
+      }
     },
     "benchmark_params": [
       {
@@ -86,15 +93,15 @@
         "backend": "openai-audio-speech",
         "endpoint": "/v1/audio/speech",
         "dataset_path": "benchmarks/build_dataset/seed_tts_smoke",
-        "num_prompts": [80],
+        "num_prompts": [80, 128, 128],
-        "max_concurrency": [8],
+        "max_concurrency": [8, 16, 64],
         "seed_tts_locale": "en",
         "extra_body": {"voice": "Vivian", "language": "English", "task_type": "CustomVoice"},
         "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
         "baseline": {
-          "median_audio_ttfp_ms": [1500],
+          "median_audio_ttfp_ms": [1500, 30000, 3000],
-          "median_audio_rtf": [0.30],
+          "median_audio_rtf": [0.30, 10.0, 0.60],
-          "audio_throughput": [30.0]
+          "audio_throughput": [30.0, 1.0, 100.0]
         }
       },
       {
@@ -157,15 +164,15 @@
         "backend": "openai-audio-speech",
         "endpoint": "/v1/audio/speech",
         "dataset_path": "benchmarks/build_dataset/seed_tts_design",
-        "num_prompts": [80],
+        "num_prompts": [80, 128, 128],
-        "max_concurrency": [8],
+        "max_concurrency": [8, 16, 64],
         "seed_tts_locale": "en",
         "extra_body": {"task_type": "VoiceDesign", "language": "English"},
         "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
         "baseline": {
-          "median_audio_ttfp_ms": [1500],
+          "median_audio_ttfp_ms": [1500, 30000, 3500],
-          "median_audio_rtf": [0.35],
+          "median_audio_rtf": [0.35, 10.0, 0.70],
-          "audio_throughput": [25.0]
+          "audio_throughput": [25.0, 1.0, 80.0]
         }
       },
       {

tests/diffusion/attention/test_attention_config.py

@@ -394,6 +394,9 @@ class TestAttentionInitUsesCurrentDiffusionConfig:
                 per_role={"cross": AttentionSpec(backend="TORCH_SDPA", extra={"block_size": 128})},
             ),
             parallel_config=SimpleNamespace(ring_degree=2),
+            diffusion_kv_cache_dtype=None,
+            diffusion_kv_cache_skip_step_indices=None,
+            diffusion_kv_cache_skip_layer_indices=None,
         )
 
         with set_current_diffusion_config(od_config):

tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_it2i_multi_image.py

@@ -0,0 +1,255 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Multi-image input regression for HunyuanImage3 IT2I prompt construction.
+
+The official HunyuanImage-3.0-Instruct supports up to 3 reference images
+per IT2I request ("Multi-Image Fusion"; see hunyuan3.0_ins/README.md
+section 200-216 + line 500). Each cond image becomes its own user-role
+message and `apply_general_template` concatenates successive user
+messages back-to-back inside ONE user_prefix/user_suffix wrap (see
+hunyuan3.0_ins/tokenization_hunyuan_image_3.py:1399-1400, 1499-1515).
+The lightweight `<img>` + `multi_modal_data` builder used by the example
+flow must match that contract: N consecutive `<img>` placeholders sit
+between `User: ` and the user prompt, with no separator between them.
+
+This file pins:
+  1. N consecutive `<img>` placeholders for N=1/2/3 across both the
+     string builder (`build_prompt`) and the token builder
+     (`build_prompt_tokens`).
+  2. The N=1 path stays bit-identical to the legacy single-image builder
+     (regression guard so default callers don't notice).
+  3. N=2 / N=3 token sequences differ from N=1 by exactly (N-1) extra
+     `<img>` ids inserted between `User: ` and `user_prompt`.
+  4. Validation: N<1 and N>3 raise ValueError (hard cap N<=3 mirrors
+     official upstream).
+  5. Text-only tasks ignore `num_images` (no validation, no extra ids).
+"""
+
+from __future__ import annotations
+
+import os
+
+import pytest
+
+from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import (
+    MAX_IMAGES_PER_REQUEST,
+    build_prompt,
+    build_prompt_tokens,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class FakeTokenizer:
+    """Recording fake tokenizer mirroring the one in test_prompt_utils.
+
+    Special token ids: `<|startoftext|>`=1, `<img>`=2, `<think>`=3,
+    `<recaption>`=4. encode() returns one id per character starting at
+    100, so substring-position assertions are stable.
+    """
+
+    SPECIAL = {
+        "<|startoftext|>": 1,
+        "<img>": 2,
+        "<think>": 3,
+        "<recaption>": 4,
+    }
+
+    def __init__(self) -> None:
+        self.encode_calls: list[str] = []
+
+    def convert_tokens_to_ids(self, tok: str) -> int:
+        return self.SPECIAL.get(tok, 0)
+
+    def encode(self, text: str, add_special_tokens: bool = False) -> list[int]:
+        self.encode_calls.append(text)
+        return list(range(100, 100 + len(text)))
+
+
+_IMAGE_TASK_COMBOS = (
+    ("i2t", None),
+    ("it2i", "think"),
+    ("it2i", "recaption"),
+)
+_TEXT_ONLY_TASK_COMBOS = (("t2t", None),)
+
+
+# -------------------- string builder --------------------
+
+
+@pytest.mark.parametrize("task,bot_task", _IMAGE_TASK_COMBOS)
+@pytest.mark.parametrize("num_images", [1, 2, 3])
+def test_build_prompt_emits_N_consecutive_img_placeholders(task: str, bot_task: str | None, num_images: int):
+    """N=1/2/3 -> exactly N `<img>` substrings appear consecutively
+    between `User: ` and the user prompt, with no separator between them."""
+    s = build_prompt("HELLO", task=task, bot_task=bot_task, num_images=num_images)
+    assert s.count("<img>") == num_images, (
+        f"task={task} bot_task={bot_task} num_images={num_images}: expected {num_images} <img> "
+        f"placeholders, found {s.count('<img>')} -- prompt was: {s!r}"
+    )
+
+    # All `<img>` placeholders must form one contiguous run "<img><img>..."
+    # immediately after `User: ` and before HELLO.
+    user_idx = s.index("User: ") + len("User: ")
+    hello_idx = s.index("HELLO")
+    between = s[user_idx:hello_idx]
+    assert between == "<img>" * num_images, (
+        f"region between `User: ` and prompt must be exactly N <img> placeholders; got {between!r}"
+    )
+
+
+def test_build_prompt_default_num_images_matches_legacy():
+    """num_images default = 1 must produce a string bit-identical to the
+    pre-multi-image behavior (single `<img>` placeholder)."""
+    legacy = build_prompt("HELLO", task="it2i", bot_task="think")
+    explicit = build_prompt("HELLO", task="it2i", bot_task="think", num_images=1)
+    assert legacy == explicit, "default num_images=1 must match legacy single-image output"
+
+
+# -------------------- token builder --------------------
+
+
+@pytest.mark.parametrize("task,bot_task", _IMAGE_TASK_COMBOS)
+def test_build_prompt_tokens_inserts_N_img_ids(task: str, bot_task: str | None):
+    """N=1/2/3 -> the resulting id sequence contains exactly N copies of
+    img_id (=2) sitting consecutively after the `User: ` segment."""
+    tok = FakeTokenizer()
+    ids_n1 = build_prompt_tokens("hi", tok, task=task, bot_task=bot_task, num_images=1).token_ids
+    tok = FakeTokenizer()
+    ids_n2 = build_prompt_tokens("hi", tok, task=task, bot_task=bot_task, num_images=2).token_ids
+    tok = FakeTokenizer()
+    ids_n3 = build_prompt_tokens("hi", tok, task=task, bot_task=bot_task, num_images=3).token_ids
+
+    assert ids_n1.count(2) == 1
+    assert ids_n2.count(2) == 2
+    assert ids_n3.count(2) == 3
+
+    # Each additional image must extend the sequence by exactly one img_id,
+    # not shift other tokens around.
+    assert len(ids_n2) == len(ids_n1) + 1
+    assert len(ids_n3) == len(ids_n1) + 2
+
+    # The img_ids must be CONSECUTIVE (no other token between successive
+    # `<img>` placeholders -- mirrors the official `process_successive_message`
+    # wrapping where successive user messages share one user_prefix/suffix).
+    for ids, n in [(ids_n2, 2), (ids_n3, 3)]:
+        first = ids.index(2)
+        for k in range(n):
+            assert ids[first + k] == 2, (
+                f"img_ids must be consecutive starting at position {first} for n={n}; got {ids[first : first + n]!r}"
+            )
+
+
+def test_build_prompt_tokens_default_num_images_matches_legacy():
+    """num_images default = 1 must produce the same id sequence as
+    omitting the parameter (regression guard for existing single-image
+    callers)."""
+    tok_a = FakeTokenizer()
+    legacy = build_prompt_tokens("hi", tok_a, task="it2i", bot_task="think").token_ids
+    tok_b = FakeTokenizer()
+    explicit = build_prompt_tokens("hi", tok_b, task="it2i", bot_task="think", num_images=1).token_ids
+    assert legacy == explicit
+    # Also: encode() must have been called on the same set of segments,
+    # so segment boundaries are preserved.
+    assert tok_a.encode_calls == tok_b.encode_calls
+
+
+# -------------------- validation --------------------
+
+
+@pytest.mark.parametrize("task,bot_task", _IMAGE_TASK_COMBOS)
+@pytest.mark.parametrize("bad", [0, -1, MAX_IMAGES_PER_REQUEST + 1, 99])
+def test_build_prompt_rejects_out_of_range_num_images(task: str, bot_task: str | None, bad: int):
+    with pytest.raises(ValueError, match="num_images must be in"):
+        build_prompt("hi", task=task, bot_task=bot_task, num_images=bad)
+    with pytest.raises(ValueError, match="num_images must be in"):
+        build_prompt_tokens("hi", FakeTokenizer(), task=task, bot_task=bot_task, num_images=bad)
+
+
+@pytest.mark.parametrize("task,bot_task", _TEXT_ONLY_TASK_COMBOS)
+@pytest.mark.parametrize("num_images", [0, 1, 2, 99])
+def test_text_only_tasks_ignore_num_images(task: str, bot_task: str | None, num_images: int):
+    """Validation only kicks in for image-input tasks; t2t et al. accept
+    any num_images and emit zero `<img>` placeholders."""
+    s = build_prompt("hi", task=task, bot_task=bot_task, num_images=num_images)
+    assert "<img>" not in s
+    ids = build_prompt_tokens("hi", FakeTokenizer(), task=task, bot_task=bot_task, num_images=num_images).token_ids
+    assert 2 not in ids
+
+
+# -------------------- real HF tokenizer regression --------------------
+
+_HUNYUAN_MODEL_ID = "tencent/HunyuanImage-3.0-Instruct"
+
+
+def _hf_cached(model_id: str) -> bool:
+    hf_home = os.environ.get("HF_HOME") or os.path.expanduser("~/.cache/huggingface")
+    snap_dir = os.path.join(hf_home, "hub", f"models--{model_id.replace('/', '--')}", "snapshots")
+    return os.path.isdir(snap_dir) and any(os.scandir(snap_dir))
+
+
+@pytest.mark.skipif(not _hf_cached(_HUNYUAN_MODEL_ID), reason=f"{_HUNYUAN_MODEL_ID} tokenizer not in HF cache")
+@pytest.mark.parametrize("num_images", [1, 2, 3])
+def test_real_tokenizer_emits_n_consecutive_img_ids(num_images: int):
+    """Real `AutoTokenizer.from_pretrained(...)` (the production path) must
+    encode N=1/2/3 prompts to a sequence with exactly N consecutive `<img>`
+    token-ids in the right place — proves the placeholder layout from
+    `build_prompt_tokens` survives a real BPE tokenizer, not just FakeTokenizer.
+    """
+    from transformers import AutoTokenizer
+
+    tok = AutoTokenizer.from_pretrained(_HUNYUAN_MODEL_ID, trust_remote_code=True)
+    img_id = tok.convert_tokens_to_ids("<img>")
+    assert img_id is not None and img_id >= 0, f"<img> not in tokenizer vocab; got id={img_id}"
+
+    ids = build_prompt_tokens("hi", tok, task="it2i", bot_task="think", num_images=num_images).token_ids
+
+    # Exactly N copies of <img> id, all consecutive.
+    img_positions = [i for i, x in enumerate(ids) if x == img_id]
+    assert len(img_positions) == num_images, (
+        f"expected {num_images} <img> ids, got {len(img_positions)} at positions {img_positions}"
+    )
+    assert img_positions == list(range(img_positions[0], img_positions[0] + num_images)), (
+        f"<img> ids must be contiguous; got positions {img_positions}"
+    )
+
+
+@pytest.mark.skipif(not _hf_cached(_HUNYUAN_MODEL_ID), reason=f"{_HUNYUAN_MODEL_ID} tokenizer not in HF cache")
+def test_real_tokenizer_n_plus_one_extends_by_exactly_one_img_id():
+    """Going from N to N+1 images must extend the encoded id sequence by
+    exactly one extra `<img>` token-id and shift nothing else. Catches
+    accidental separator tokens between successive `<img>` placeholders
+    that a FakeTokenizer (deterministic encode) can't surface."""
+    from transformers import AutoTokenizer
+
+    tok = AutoTokenizer.from_pretrained(_HUNYUAN_MODEL_ID, trust_remote_code=True)
+    img_id = tok.convert_tokens_to_ids("<img>")
+
+    ids_n1 = build_prompt_tokens("hi", tok, task="it2i", bot_task="think", num_images=1).token_ids
+    ids_n2 = build_prompt_tokens("hi", tok, task="it2i", bot_task="think", num_images=2).token_ids
+    ids_n3 = build_prompt_tokens("hi", tok, task="it2i", bot_task="think", num_images=3).token_ids
+
+    assert len(ids_n2) == len(ids_n1) + 1, f"N=2 should be N=1 + 1 token; got {len(ids_n2)} vs {len(ids_n1)}"
+    assert len(ids_n3) == len(ids_n1) + 2, f"N=3 should be N=1 + 2 tokens; got {len(ids_n3)} vs {len(ids_n1)}"
+
+    # Insert one img_id at the existing position; everything else unchanged.
+    p1 = ids_n1.index(img_id)
+    assert ids_n2[: p1 + 1] == ids_n1[: p1 + 1] + [], "prefix before extra <img> must match N=1"
+    assert ids_n2[p1] == img_id and ids_n2[p1 + 1] == img_id, "two consecutive <img> ids at the insertion point"
+    assert ids_n2[p1 + 2 :] == ids_n1[p1 + 1 :], "tail after the extra <img> must match N=1's tail"
+    # N=3 same pattern, three in a row.
+    assert ids_n3[p1 : p1 + 3] == [img_id, img_id, img_id]
+    assert ids_n3[p1 + 3 :] == ids_n1[p1 + 1 :]
+
+
+@pytest.mark.skipif(not _hf_cached(_HUNYUAN_MODEL_ID), reason=f"{_HUNYUAN_MODEL_ID} tokenizer not in HF cache")
+def test_real_tokenizer_default_n1_byte_identical_to_legacy():
+    """Default `num_images=1` must produce the exact same id sequence as
+    omitting the parameter — pins the legacy single-image regression
+    against the real tokenizer (not just FakeTokenizer)."""
+    from transformers import AutoTokenizer
+
+    tok = AutoTokenizer.from_pretrained(_HUNYUAN_MODEL_ID, trust_remote_code=True)
+    legacy = build_prompt_tokens("hi", tok, task="it2i", bot_task="think").token_ids
+    explicit = build_prompt_tokens("hi", tok, task="it2i", bot_task="think", num_images=1).token_ids
+    assert legacy == explicit, "real tokenizer: default num_images=1 must be byte-identical to legacy"

tests/diffusion/models/hunyuan_image3/test_prompt_utils.py

@@ -1,20 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Regression tests for HunyuanImage3 prompt construction (PR #3243).
-
-Two layers:
-  1. Pure-logic tests with a recording fake tokenizer -- protect the
-     prompt template structure (BOS, User:/Assistant: framing, trigger
-     placement, image placeholder position) and protect the segment-
-     by-segment tokenization contract (each segment must hit
-     `tokenizer.encode` in isolation).
-  2. Real-tokenizer regression -- run when the HunyuanImage3-Instruct
-     tokenizer is in the local HF cache. Asserts the segment-tokenized
-     output diverges from the naive full-string encode, which is the
-     bug-tripping fixture for the cross-segment BPE merge fix
-     (commit 7bd429ed).
-"""
-
 from __future__ import annotations
 
 import ast
@@ -24,7 +9,9 @@ import pathlib
 import pytest
 
 from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import (
+    _TASK_PRESETS,
     HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS,
+    available_bot_tasks,
     available_tasks,
     build_prompt,
     build_prompt_tokens,
@@ -34,18 +21,7 @@ from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import (
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
-# -------------------- Pure-logic structural tests --------------------
-
-
 class FakeTokenizer:
-    """Minimal tokenizer stub that records every encode() call.
-
-    Returns deterministic ids from convert_tokens_to_ids while
-    encode() returns one id per character starting at 100. This lets
-    tests both verify segmentation (by inspecting `encode_calls`) and
-    locate substrings inside the returned id list.
-    """
-
     SPECIAL = {
         "<|startoftext|>": 1,
         "<img>": 2,
@@ -72,85 +48,131 @@ class FakeTokenizer:
 
 
 def test_available_tasks_covers_all_modalities():
-    tasks = set(available_tasks())
+    assert set(available_tasks()) == {"t2t", "i2t", "it2i", "t2i"}
-    assert tasks >= {
+
-        "t2t",
+
-        "i2t",
+def test_available_bot_tasks_covers_all_modes():
+    assert set(available_bot_tasks()) == {None, "think", "recaption", "think_recaption", "vanilla"}
+
+
+def test_legacy_task_presets_still_available():
+    assert {
         "it2i_think",
         "it2i_recaption",
         "it2i_think_recaption",
         "t2i_think",
         "t2i_recaption",
         "t2i_vanilla",
+    } <= set(_TASK_PRESETS)
+
+
+def test_legacy_base_task_omitted_bot_task_keeps_plain_mode():
+    prompt = build_prompt("HELLO", task="i2t")
+    assert prompt.endswith("Assistant: ")
+    assert not prompt.endswith("<think>")
+
+    result = build_prompt_tokens("hi", FakeTokenizer(), task="i2t")
+    assert result.system_prompt_type == "en_unified"
+    assert result.token_ids[-1] not in {
+        FakeTokenizer.SPECIAL["<think>"],
+        FakeTokenizer.SPECIAL["<recaption>"],
     }
 
 
-def test_resolve_stop_token_ids_uses_answer_for_generation_tasks():
+def test_legacy_composite_task_with_none_bot_task_keeps_encoded_mode():
+    prompt = build_prompt("HELLO", task="it2i_think", bot_task=None)
+    assert prompt.endswith("Assistant: <think>")
+
+    result = build_prompt_tokens("hi", FakeTokenizer(), task="it2i_recaption", bot_task=None)
+    assert result.token_ids[-1] == FakeTokenizer.SPECIAL["<recaption>"]
+
+
+def test_default_prompt_still_uses_it2i_think_mode():
+    prompt = build_prompt("HELLO")
+    assert prompt.endswith("Assistant: <think>")
+
+    result = build_prompt_tokens("hi", FakeTokenizer())
+    assert result.system_prompt_type == "en_unified"
+    assert result.token_ids[-1] == FakeTokenizer.SPECIAL["<think>"]
+
+
+def test_resolve_stop_token_ids_image_tasks_stop_on_ratio_range():
+    """Image-output tasks stop on any ``<img_ratio_*>`` token.
+
+    Mirrors upstream ``modeling_hunyuan_image_3.py::generate_image``
+    (line 3289-3303): when ``need_ratio`` is true,
+    ``final_stop_tokens = list(range(start_ratio, end_ratio + 1)) +
+    ratio_token_other_slices``. AR stops AT the ratio token sampled
+    after ``<img_size_*>``; the bridge then strips the trailing ratio
+    token before passing the cot to DiT.
+    """
     tok = FakeTokenizer()
 
+    start = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<img_ratio_0>"]
+    end = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<img_ratio_32>"]
+    other_start = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<img_ratio_33>"]
+    other_end = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<img_ratio_36>"]
+    expected = list(range(start, end + 1)) + list(range(other_start, other_end + 1))
+
+    # Image-output: t2i / it2i stop on the full ratio token range.
+    for bot in ("think", "recaption", "think_recaption", "vanilla"):
+        assert resolve_stop_token_ids(task="t2i", bot_task=bot, tokenizer=tok) == expected
+        assert resolve_stop_token_ids(task="it2i", bot_task=bot, tokenizer=tok) == expected
+
+    # Text-output: i2t / t2t comprehension stops on <answer> (response sits inside).
     answer_id = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<answer>"]
-    assert resolve_stop_token_ids(task="t2i_think", tokenizer=tok) == [answer_id]
+    assert resolve_stop_token_ids(task="i2t", bot_task=None, tokenizer=tok) == [answer_id]
-    assert resolve_stop_token_ids(task="t2i_recaption", tokenizer=tok) == [answer_id]
+    assert resolve_stop_token_ids(task="t2t", bot_task=None, tokenizer=tok) == [answer_id]
 
 
 @pytest.mark.parametrize(
-    "task",
+    "task,bot_task",
     [
-        "t2t",
+        ("t2t", None),
-        "i2t",
+        ("i2t", None),
-        "it2i_think",
+        ("it2i", "think"),
-        "it2i_recaption",
+        ("it2i", "recaption"),
-        "it2i_think_recaption",
+        ("it2i", "think_recaption"),
-        "t2i_think",
+        ("t2i", "think"),
-        "t2i_recaption",
+        ("t2i", "recaption"),
+        ("t2i", "think_recaption"),
     ],
 )
-def test_build_prompt_string_structure_chat_template(task: str):
+def test_build_prompt_string_structure_chat_template(task: str, bot_task: str | None):
-    """Chat-template tasks must produce <|startoftext|>...User: ...Assistant: ...
+    s = build_prompt("HELLO", task=task, bot_task=bot_task)
-    with image placeholder (when applicable) and trigger tag AFTER `Assistant: `."""
-    s = build_prompt("HELLO", task=task)
-
     assert s.startswith("<|startoftext|>")
     assert "User: " in s
     assert "Assistant: " in s
     assert s.index("User: ") < s.index("HELLO") < s.index("Assistant: ")
 
-    if task.startswith(("i2t", "it2i")):
+    if task in ("i2t", "it2i"):
-        assert s.index("User: ") < s.index("<img>") < s.index("HELLO"), (
+        assert s.index("User: ") < s.index("<img>") < s.index("HELLO")
-            "<img> placeholder must sit between `User: ` and the user prompt"
-        )
     else:
         assert "<img>" not in s
 
-    # Trigger tag must be the FINAL token of the prompt (after `Assistant: `).
+    if bot_task in ("think", "think_recaption"):
-    # Note: the system prompt itself mentions <think>/<recaption> as mode
+        assert s.endswith("Assistant: <think>")
-    # documentation, so substring index() catches the wrong occurrence -- use
+    elif bot_task == "recaption":
-    # endswith() which directly captures "trigger is at the tail" (the Part A
+        assert s.endswith("Assistant: <recaption>")
-    # fix: trigger goes AFTER `Assistant: `, not before user_prompt).
+    elif bot_task is None:
-    if task in ("it2i_think", "t2i_think", "it2i_think_recaption"):
+        assert s.endswith("Assistant: ")
-        assert s.endswith("Assistant: <think>"), (
-            f"Trigger <think> must be appended right after `Assistant: ` (Part A fix). Got tail: ...{s[-40:]!r}"
-        )
-    if task in ("it2i_recaption", "t2i_recaption"):
-        assert s.endswith("Assistant: <recaption>"), (
-            f"Trigger <recaption> must be appended right after `Assistant: ` (Part A fix). Got tail: ...{s[-40:]!r}"
-        )
-    if task in ("t2t", "i2t"):
-        assert s.endswith("Assistant: "), "Plain (no-trigger) task must end at `Assistant: ` with no trailing tag."
 
 
 def test_build_prompt_vanilla_uses_pretrain_template():
-    """t2i_vanilla is the only task that bypasses chat structure -- direct
+    s = build_prompt("HELLO", task="t2i", bot_task="vanilla")
-    text->image generation driven by the vanilla system prompt."""
-    s = build_prompt("HELLO", task="t2i_vanilla")
     assert s.startswith("<|startoftext|>")
     assert "User: " not in s
     assert "Assistant: " not in s
-    assert "<think>" not in s
-    assert "<recaption>" not in s
     assert s.endswith("HELLO")
 
 
+def test_build_prompt_vanilla_rejects_non_t2i_task():
+    with pytest.raises(ValueError, match="bot_task='vanilla'"):
+        build_prompt("x", task="it2i", bot_task="vanilla")
+    with pytest.raises(ValueError, match="bot_task='vanilla'"):
+        build_prompt_tokens("x", FakeTokenizer(), task="i2t", bot_task="vanilla")
+
+
 def test_build_prompt_unknown_task_raises():
     with pytest.raises(ValueError, match="Unknown task"):
         build_prompt("x", task="bogus")
@@ -158,127 +180,83 @@ def test_build_prompt_unknown_task_raises():
         build_prompt_tokens("x", FakeTokenizer(), task="bogus")
 
 
+def test_build_prompt_unknown_bot_task_raises():
+    with pytest.raises(ValueError, match="Unknown bot_task"):
+        build_prompt("x", task="t2i", bot_task="bogus")
+    with pytest.raises(ValueError, match="Unknown bot_task"):
+        build_prompt_tokens("x", FakeTokenizer(), task="t2i", bot_task="bogus")
+
+
 def test_build_prompt_tokens_segments_each_boundary():
-    """Regression for cross-segment BPE merge bug (commit 7bd429ed):
-    each template segment must hit tokenizer.encode() independently;
-    user_prompt MUST NOT be concatenated with the following separator
-    in the same encode() call."""
     tok = FakeTokenizer()
-    build_prompt_tokens("写诗。", tok, task="i2t")
+    build_prompt_tokens("写诗。", tok, task="i2t", bot_task=None)
-
-    # Each canonical segment is encoded in its own call.
     assert "User: " in tok.encode_calls
-    assert "写诗。" in tok.encode_calls, (
+    assert "写诗。" in tok.encode_calls
-        "user_prompt must be encoded alone -- if it is concatenated with the "
-        "trailing separator, BPE will merge across the boundary (the PR-#3243 bug)."
-    )
     assert "\n\nAssistant: " in tok.encode_calls
-
-    # No call must contain user_prompt glued to neighboring text.
     for call in tok.encode_calls:
         if call != "写诗。":
-            assert "写诗。" not in call, f"user_prompt leaked into a multi-segment encode call: {call!r}"
+            assert "写诗。" not in call
 
 
 def test_build_prompt_tokens_image_placeholder_present_for_image_tasks():
     tok = FakeTokenizer()
-    result = build_prompt_tokens("hi", tok, task="i2t")
+    result = build_prompt_tokens("hi", tok, task="i2t", bot_task=None)
     ids = result.token_ids
-    assert ids[0] == FakeTokenizer.SPECIAL["<|startoftext|>"], "BOS (<|startoftext|>) must be the first token"
+    assert ids[0] == FakeTokenizer.SPECIAL["<|startoftext|>"]
-    assert FakeTokenizer.SPECIAL["<img>"] in ids, "<img> placeholder must be present for i2t/it2i tasks"
+    assert FakeTokenizer.SPECIAL["<img>"] in ids
 
 
 def test_build_prompt_tokens_no_image_for_text_only_tasks():
     tok = FakeTokenizer()
-    result = build_prompt_tokens("hi", tok, task="t2t")
+    result = build_prompt_tokens("hi", tok, task="t2t", bot_task=None)
     ids = result.token_ids
-    assert FakeTokenizer.SPECIAL["<img>"] not in ids, "<img> must NOT appear for text-only tasks"
+    assert FakeTokenizer.SPECIAL["<img>"] not in ids
 
 
 @pytest.mark.parametrize(
-    "task,trigger_id",
+    "task,bot_task,trigger_id",
     [
-        ("it2i_think", FakeTokenizer.SPECIAL["<think>"]),
+        ("it2i", "think", FakeTokenizer.SPECIAL["<think>"]),
-        ("t2i_think", FakeTokenizer.SPECIAL["<think>"]),
+        ("t2i", "think", FakeTokenizer.SPECIAL["<think>"]),
-        ("it2i_recaption", FakeTokenizer.SPECIAL["<recaption>"]),
+        ("t2i", "think_recaption", FakeTokenizer.SPECIAL["<think>"]),
-        ("t2i_recaption", FakeTokenizer.SPECIAL["<recaption>"]),
+        ("it2i", "recaption", FakeTokenizer.SPECIAL["<recaption>"]),
+        ("t2i", "recaption", FakeTokenizer.SPECIAL["<recaption>"]),
+        ("it2i_think", None, FakeTokenizer.SPECIAL["<think>"]),
+        ("it2i_recaption", None, FakeTokenizer.SPECIAL["<recaption>"]),
     ],
 )
-def test_build_prompt_tokens_trigger_is_last_token(task: str, trigger_id: int):
+def test_build_prompt_tokens_trigger_is_last_token(task: str, bot_task: str | None, trigger_id: int):
-    """Trigger tag id must be the LAST token (after `Assistant: ` segment)."""
     tok = FakeTokenizer()
-    result = build_prompt_tokens("hi", tok, task=task)
+    result = build_prompt_tokens("hi", tok, task=task, bot_task=bot_task)
-    ids = result.token_ids
+    assert result.token_ids[-1] == trigger_id
-    assert ids[-1] == trigger_id
 
 
 def test_build_prompt_tokens_no_trigger_for_plain_tasks():
-    """Tasks without trigger_tag (t2t / i2t) must NOT append a trigger id."""
     tok = FakeTokenizer()
-    result = build_prompt_tokens("hi", tok, task="t2t")
+    result = build_prompt_tokens("hi", tok, task="t2t", bot_task=None)
-    ids = result.token_ids
+    assert result.token_ids[-1] not in {
-    assert ids[-1] not in {
         FakeTokenizer.SPECIAL["<think>"],
         FakeTokenizer.SPECIAL["<recaption>"],
     }
 
 
-# -------------------- end2end.py wiring guard --------------------
-
-
 def _repo_root() -> pathlib.Path:
-    # tests/diffusion/models/hunyuan_image3/test_prompt_utils.py -> repo root
     return pathlib.Path(__file__).resolve().parents[4]
 
 
 def test_end2end_routes_through_shared_prompt_utils():
-    """Regression for the *delivery vector* of PR #3243.
-
-    Background: the wrong-template bug that PR #3243 fixes was introduced
-    when end2end.py grew its own hand-rolled prompt builder that diverged
-    from the canonical instruct chat template. To prevent that exact
-    failure mode from recurring, end2end.py MUST:
-      1. Import the prompt builders from the shared prompt_utils module.
-      2. NOT redefine `build_prompt` or `build_prompt_tokens` locally.
-
-    A local redefinition is precisely how a future merge can silently
-    re-introduce a pretrain-style template (trigger BEFORE user_prompt,
-    no User:/Assistant: framing, etc.) without touching prompt_utils,
-    bypassing every other test in this file.
-    """
     end2end_path = _repo_root() / "examples" / "offline_inference" / "hunyuan_image3" / "end2end.py"
-    assert end2end_path.is_file(), f"end2end.py not found at {end2end_path}"
-
     tree = ast.parse(end2end_path.read_text(encoding="utf-8"))
 
     local_func_names = {n.name for n in ast.walk(tree) if isinstance(n, ast.FunctionDef)}
-    forbidden = {"build_prompt", "build_prompt_tokens"}
+    assert not (local_func_names & {"build_prompt", "build_prompt_tokens"})
-    redefined = local_func_names & forbidden
-    assert not redefined, (
-        f"end2end.py defines {sorted(redefined)} locally. This is exactly how "
-        "the wrong prompt template re-entered the example before PR #3243. "
-        "Use the shared `vllm_omni.diffusion.models.hunyuan_image3.prompt_utils` "
-        "helpers instead."
-    )
 
     imported_from_prompt_utils: set[str] = set()
     for node in ast.walk(tree):
         if isinstance(node, ast.ImportFrom) and node.module and node.module.endswith("hunyuan_image3.prompt_utils"):
             imported_from_prompt_utils.update(alias.name for alias in node.names)
-    expected_imports = {
+    expected_imports = {"build_prompt_tokens", "resolve_stop_token_ids", "resolve_sys_type"}
-        "_TASK_PRESETS",
+    assert expected_imports <= imported_from_prompt_utils
-        "build_prompt_tokens",
-        "resolve_stop_token_ids",
-    }
-    assert expected_imports <= imported_from_prompt_utils, (
-        "end2end.py must import the HunyuanImage3 prompt and stop-token helpers from "
-        "vllm_omni.diffusion.models.hunyuan_image3.prompt_utils -- the shared "
-        "module is the single source of truth for the AR-prefill template and "
-        "bot_task-derived AR stop token ids."
-    )
-
-
-# -------------------- Real-tokenizer regression --------------------
 
 
 _HUNYUAN_MODEL_ID = "tencent/HunyuanImage-3.0-Instruct"
@@ -290,41 +268,14 @@ def _hf_cached(model_id: str) -> bool:
     return os.path.isdir(snap_dir) and any(os.scandir(snap_dir))
 
 
-@pytest.mark.skipif(
+@pytest.mark.skipif(not _hf_cached(_HUNYUAN_MODEL_ID), reason=f"{_HUNYUAN_MODEL_ID} tokenizer not in HF cache")
-    not _hf_cached(_HUNYUAN_MODEL_ID),
-    reason=f"{_HUNYUAN_MODEL_ID} tokenizer not in HF cache",
-)
 def test_segment_tokenize_diverges_from_full_string_encode():
-    """Regression for PR #3243 segment-tokenization fix.
-
-    The naive `tokenizer.encode(build_prompt(...))` lets BPE merge tokens
-    across segment boundaries (notably `。\\n\\n` -> a single id), which
-    drifts the AR prefill away from HF's apply_chat_template output. The
-    segment-by-segment build_prompt_tokens must produce a STRICTLY
-    DIFFERENT id sequence on a prompt that triggers the merge.
-
-    If someone "simplifies" build_prompt_tokens to call encode() on the
-    full string, this assertion fires.
-    """
     from transformers import AutoTokenizer
 
     tok = AutoTokenizer.from_pretrained(_HUNYUAN_MODEL_ID, trust_remote_code=True)
-
     user_prompt = "写一首关于夜的诗。"
-    result = build_prompt_tokens(user_prompt, tok, task="i2t")
+    result = build_prompt_tokens(user_prompt, tok, task="i2t", bot_task=None)
     seg_ids = result.token_ids
-    full_ids = tok.encode(build_prompt(user_prompt, task="i2t"), add_special_tokens=False)
+    full_ids = tok.encode(build_prompt(user_prompt, task="i2t", bot_task=None), add_special_tokens=False)
-
+    assert seg_ids != full_ids
-    assert seg_ids != full_ids, (
+    assert len(seg_ids) >= len(full_ids)
-        "build_prompt_tokens output equals naive full-string encode -- "
-        "the BPE-merge-bypass behavior is no longer exercised. This means "
-        "the segment-by-segment fix from PR #3243 has been silently undone."
-    )
-
-    # Segmenting prevents merges, so the segment id list should have AT LEAST
-    # as many tokens as the merged version (a merge consumes 2+ ids -> 1).
-    assert len(seg_ids) >= len(full_ids), (
-        f"segment-encoded length ({len(seg_ids)}) shorter than full-string "
-        f"merged length ({len(full_ids)}) -- impossible if segmentation is "
-        f"genuinely bypassing merges."
-    )

tests/diffusion/offloader/test_diffusion_cpu_offload.py

@@ -5,6 +5,7 @@ import pytest
 import torch
 from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
 
+from tests.helpers import skip_if_gated_repo_inaccessible
 from tests.helpers.env import DeviceMemoryMonitor
 from tests.helpers.mark import hardware_test
 from tests.helpers.runtime import OmniRunner
@@ -21,6 +22,11 @@ IMAGE_VIDEO_MODELS = {
 
 MODELS = {**AUDIO_MODEL, **IMAGE_VIDEO_MODELS}
 
+_GATED_MODELS = {"stabilityai/stable-audio-open-1.0"}
+
+# Aliased for backward compatibility (imported by test_diffusion_layerwise_offload.py).
+_skip_if_gated_repo_inaccessible = skip_if_gated_repo_inaccessible
+
 AUDIO_MODEL_PARAMS = {
     "runner_params": {},
     "sampler_params": {},
@@ -94,10 +100,23 @@ def check_audio_determinism(audio1, audio2, atol=1e-2):
 @hardware_test(res={"cuda": "L4", "rocm": "MI325"})
 @pytest.mark.parametrize("model_name", list(MODELS.keys()))
 def test_cpu_offload_diffusion_model(model_name: str):
+    if model_name in _GATED_MODELS:
+        _skip_if_gated_repo_inaccessible(model_name)
     try:
         offload_peak_memory, output_offload = inference(model_name, offload=True)
         cleanup_dist_env_and_memory()
         no_offload_peak_memory, output_no_offload = inference(model_name, offload=False)
+    except ValueError as exc:
+        # omni_snapshot_download wraps GatedRepoError in a ValueError.
+        # If the pre-flight guard above did not catch it (e.g. partial
+        # HF_TOKEN where config.json is accessible but weight shards are
+        # blocked), skip instead of failing.
+        if "Access to model" in str(exc) and "is restricted" in str(exc):
+            pytest.skip(
+                f"Skipping: gated HF repo {model_name!r} inaccessible "
+                f"({exc}). See docs/contributing/ci/hf_credentials.md."
+            )
+        pytest.fail(f"Inference failed: {exc}")
     except Exception:
         pytest.fail("Inference failed")
     print(f"Offload peak memory: {offload_peak_memory} MB")

tests/diffusion/offloader/test_diffusion_layerwise_offload.py

@@ -3,6 +3,10 @@ import pytest
 import torch
 from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
 
+from tests.diffusion.offloader.test_diffusion_cpu_offload import (
+    _GATED_MODELS,
+    _skip_if_gated_repo_inaccessible,
+)
 from tests.helpers.env import DeviceMemoryMonitor
 from tests.helpers.runtime import OmniRunner
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
@@ -96,6 +100,8 @@ def test_layerwise_offload_diffusion_model(model_name: str):
     offloader keeps only a single transformer block on GPU at a time, with
     prefetching for compute-memory overlap.
     """
+    if model_name in _GATED_MODELS:
+        _skip_if_gated_repo_inaccessible(model_name)
     try:
         # Run without layerwise offloading (baseline)
         no_offload_peak_memory, output_no_offload = run_inference(model_name, layerwise_offload=False)
@@ -104,6 +110,17 @@ def test_layerwise_offload_diffusion_model(model_name: str):
         # Run with layerwise offloading (1 layer on device)
         layerwise_offload_peak_memory, output_offload = run_inference(model_name, layerwise_offload=True)
         cleanup_dist_env_and_memory()
+    except ValueError as exc:
+        # omni_snapshot_download wraps GatedRepoError in a ValueError.
+        # If the pre-flight guard above did not catch it (e.g. partial
+        # HF_TOKEN where config.json is accessible but weight shards are
+        # blocked), skip instead of failing.
+        if "Access to model" in str(exc) and "is restricted" in str(exc):
+            pytest.skip(
+                f"Skipping: gated HF repo {model_name!r} inaccessible "
+                f"({exc}). See docs/contributing/ci/hf_credentials.md."
+            )
+        pytest.fail(f"Inference failed: {exc}")
     except Exception:
         pytest.fail("Inference failed")
 

tests/diffusion/test_inline_stage_diffusion_client.py

@@ -5,6 +5,7 @@ from types import SimpleNamespace
 from unittest.mock import MagicMock, patch
 
 import pytest
+from vllm.v1.engine.exceptions import EngineDeadError
 
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.inline_stage_diffusion_client import InlineStageDiffusionClient
@@ -19,6 +20,10 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 def mock_engine():
     with patch("vllm_omni.diffusion.inline_stage_diffusion_client.DiffusionEngine") as mock:
         engine_instance = MagicMock()
+        engine_instance.executor = SimpleNamespace(
+            register_failure_callback=MagicMock(),
+            check_health=MagicMock(),
+        )
         mock.make_engine.return_value = engine_instance
         yield engine_instance
 
@@ -97,6 +102,41 @@ def test_inline_shutdown(client, mock_engine):
     mock_engine.close.assert_called_once()
 
 
+def test_inline_registers_executor_failure_callback(client, mock_engine):
+    mock_engine.executor.register_failure_callback.assert_called_once()
+
+
+def test_inline_executor_failure_marks_engine_dead(client, mock_engine):
+    callback = mock_engine.executor.register_failure_callback.call_args.args[0]
+
+    callback()
+
+    assert client._engine_dead is True
+    with pytest.raises(EngineDeadError, match="inline diffusion engine is dead"):
+        client.get_diffusion_output_nowait()
+
+
+def test_inline_returns_queued_output_before_engine_dead(client, mock_engine):
+    callback = mock_engine.executor.register_failure_callback.call_args.args[0]
+    output = OmniRequestOutput.from_diffusion(request_id="req-queued", images=[])
+    client._output_queue.put_nowait(output)
+
+    callback()
+
+    assert client.get_diffusion_output_nowait() is output
+    with pytest.raises(EngineDeadError, match="inline diffusion engine is dead"):
+        client.get_diffusion_output_nowait()
+
+
+def test_inline_check_health_marks_engine_dead(client, mock_engine):
+    mock_engine.executor.check_health.side_effect = EngineDeadError()
+
+    with pytest.raises(EngineDeadError):
+        client.check_health()
+
+    assert client._engine_dead is True
+
+
 def test_inline_client_requires_replica_id(mock_engine):
     metadata = SimpleNamespace(
         stage_id=0,

tests/distributed/omni_connectors/test_chunk_transfer_adapter.py

@@ -199,37 +199,53 @@ def test_send_single_request_cleans_up_after_finished_payload(build_adapter, mon
     assert args[1] == "ext-finished"
 
 
-def test_update_request_payload(build_adapter):
+def test_load_poll_non_ar_merges_into_existing_additional_information(build_adapter):
-    adapter, _ = build_adapter()
+    adapter, connector = build_adapter(stage_id=2, model_mode="diffusion")
-
+    request = _req("req-non-ar", RequestStatus.WAITING, external_req_id="ext-non-ar")
-    first: OmniPayload = {
+    request.additional_information = {
         "hidden_states": {"output": torch.tensor([[1.0]])},
-        "codes": {"audio": [1]},
+        "ids": {"prompt": [11, 12]},
-        "meta": {"finished": torch.tensor(False, dtype=torch.bool)},
+        "meta": {"finished": torch.tensor(False, dtype=torch.bool), "step": 1},
     }
-    adapter._update_request_payload("ext", first)
+    request.num_computed_tokens = 9
-    second: OmniPayload = {
-        "hidden_states": {"output": torch.tensor([[2.0]])},
-        "codes": {"audio": [2]},
-        "meta": {"finished": torch.tensor(True, dtype=torch.bool)},
-    }
-    merged = adapter._update_request_payload("ext", second)
 
-    assert torch.equal(merged["hidden_states"]["output"], torch.tensor([[1.0], [2.0]]))
+    payload: OmniPayload = {
-    assert merged["codes"]["audio"] == [1, 2]
+        "hidden_states": {"output": torch.tensor([[2.0]])},
-    assert merged["meta"]["finished"].item() is True
+        "ids": {"all": [21, 22]},
+        "codes": {"audio": torch.tensor([7, 8], dtype=torch.long)},
+        "meta": {"finished": torch.tensor(True, dtype=torch.bool), "phase": "decode"},
+        "kv_metadata": {"foo": "bar"},
+    }
+    connector.get.return_value = (payload, 8)
+
+    assert adapter._poll_single_request(request) is True
+
+    assert request.prompt_token_ids == [7, 8]
+    assert request.num_computed_tokens == 0
+    assert torch.equal(
+        request.additional_information["hidden_states"]["output"],
+        torch.tensor([[2.0]]),
+    )
+    assert request.additional_information["ids"]["prompt"] == [11, 12]
+    assert request.additional_information["ids"]["all"] == [21, 22]
+    # non-ar merge path intentionally doesn't overwrite meta.finished.
+    assert request.additional_information["meta"]["finished"].item() is False
+    assert request.additional_information["meta"]["phase"] == "decode"
+    assert request.additional_information["kv_metadata"] == {"foo": "bar"}
+    assert "req-non-ar" in adapter._finished_load_reqs
+    assert "req-non-ar" in adapter.finished_requests
 
 
 def test_load_poll_ar_request_additional_information_concats_tensors(build_adapter):
     adapter, connector = build_adapter(stage_id=2, model_mode="ar")
     request = _req("req-merged", RequestStatus.WAITING, external_req_id="ext-merged")
-
+    request.additional_information = {
-    adapter.request_ids_mapping["req-merged"] = "ext-merged"
-    adapter.request_payload["ext-merged"] = {
         "hidden_states": {"output": torch.tensor([[1.0]])},
         "ids": {"prompt": [11, 12]},
         "meta": {"finished": torch.tensor(False, dtype=torch.bool)},
     }
+
+    adapter.request_ids_mapping["req-merged"] = "ext-merged"
     payload: OmniPayload = {
         "hidden_states": {"output": torch.tensor([[2.0]])},
         "meta": {"finished": torch.tensor(True, dtype=torch.bool)},
@@ -238,12 +254,8 @@ def test_load_poll_ar_request_additional_information_concats_tensors(build_adapt
 
     adapter._poll_single_request(request)
 
-    assert torch.equal(
+    # AR mode now forwards the latest payload directly.
-        request.additional_information["hidden_states"]["output"],
+    assert request.additional_information == payload
-        torch.tensor([[1.0], [2.0]]),
-    )
-    # Keys absent from the new chunk are dropped (matches main's behavior).
-    assert "ids" not in request.additional_information
     assert request.additional_information["meta"]["finished"].item() is True
 
 
@@ -528,6 +540,7 @@ def test_generation_scheduler_calls_cleanup_on_finished(monkeypatch, mocker: Moc
         kv_connector_output=None,
         cudagraph_stats=None,
         req_id_to_index={"req-s1": 0},
+        routed_experts_dict=None,
     )
 
     OmniGenerationScheduler.update_from_output(scheduler, scheduler_output, model_runner_output)
@@ -613,6 +626,7 @@ def test_ar_scheduler_defers_cleanup_and_queues_save_on_finished(mocker: MockerF
         cudagraph_stats=None,
         req_id_to_index={"req-ar": 0},
         kv_extracted_req_ids=None,
+        routed_experts_dict=None,
     )
 
     OmniARScheduler.update_from_output(scheduler, scheduler_output, model_runner_output)
@@ -686,3 +700,232 @@ def test_wire_round_trip_struct_to_dict_contract():
     assert set(decoded.keys()) == set(expected.keys())
     assert set(decoded["meta"].keys()) == set(expected["meta"].keys())
     assert set(decoded["codes"].keys()) == set(expected["codes"].keys())
+
+
+# ---------------------------------------------------------------
+# Deferred finish for upstream-completed requests
+# ---------------------------------------------------------------
+
+
+def _build_deferred_finish_scheduler(mocker, *, running, pending_finish_reqs):
+    """Build a mock scheduler with requests queued for deferred finish."""
+    adapter_mock = mocker.MagicMock()
+    adapter_mock.finished_requests = {r.request_id for r in pending_finish_reqs}
+    cleanup_calls = []
+    adapter_mock.cleanup = lambda *a, **kw: cleanup_calls.append((a, kw))
+
+    scheduler = mocker.MagicMock()
+    scheduler.chunk_transfer_adapter = adapter_mock
+    scheduler.connector = None
+    scheduler.ec_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.take_events.return_value = None
+    scheduler.kv_event_publisher = mocker.MagicMock()
+    scheduler._pending_finish_reqs = list(pending_finish_reqs)
+
+    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 = list(running)
+    scheduler.waiting = mocker.MagicMock()
+    scheduler.waiting.remove_requests = mocker.MagicMock()
+    scheduler.make_stats = mocker.MagicMock(return_value=None)
+    scheduler.requests = {r.request_id: r for r in running}
+
+    scheduler_output = SimpleNamespace(
+        num_scheduled_tokens={},
+        scheduled_spec_decode_tokens={},
+        num_invalid_spec_tokens=0,
+    )
+    model_runner_output = SimpleNamespace(
+        sampled_token_ids=None,
+        logprobs=None,
+        prompt_logprobs_dict={},
+        pooler_output=None,
+        num_nans_in_logits=None,
+        kv_connector_output=None,
+        cudagraph_stats=None,
+        req_id_to_index={},
+    )
+    return scheduler, scheduler_output, model_runner_output, cleanup_calls
+
+
+def test_deferred_finish_emits_finished_output(mocker: MockerFixture):
+    """A request whose upstream completed with no remaining tokens should
+    emit a FINISHED EngineCoreOutput, free resources, and clean up adapter state."""
+    from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
+
+    request = _HashableRequest(
+        request_id="req-df1",
+        external_req_id="ext-df1",
+        status=RequestStatus.RUNNING,
+        is_finished=lambda: False,
+        num_computed_tokens=16,
+        num_prompt_tokens=16,
+        prompt_token_ids=list(range(1, 17)),
+        num_output_placeholders=0,
+        sampling_params=None,
+        pooling_params=None,
+        stop_reason=None,
+        client_index=0,
+        take_events=lambda: [],
+        trace_headers=None,
+        has_encoder_inputs=False,
+        take_prefill_stats=lambda: None,
+        num_nans_in_logits=0,
+        get_finished_reason=lambda: "stop",
+    )
+    scheduler, sched_out, model_out, cleanup_calls = _build_deferred_finish_scheduler(
+        mocker,
+        running=[request],
+        pending_finish_reqs=[request],
+    )
+    scheduler._free_request.return_value = {"mock": "kv_params"}
+
+    result = OmniGenerationScheduler.update_from_output(scheduler, sched_out, model_out)
+
+    assert request.status == RequestStatus.FINISHED_STOPPED
+    scheduler._handle_stopped_request.assert_called_once_with(request)
+    scheduler._free_request.assert_called_once_with(request)
+    assert len(cleanup_calls) == 1
+    assert cleanup_calls[0][0] == ("req-df1", "ext-df1")
+
+    eco = result[0]
+    assert len(eco.outputs) == 1
+    assert eco.outputs[0].request_id == "req-df1"
+    assert eco.outputs[0].finish_reason == "stop"
+    assert eco.outputs[0].kv_transfer_params == {"mock": "kv_params"}
+    assert scheduler._pending_finish_reqs == []
+
+
+def test_deferred_finish_empty_prompt(mocker: MockerFixture):
+    """A request that never received any tokens (finished immediately upstream)
+    should still emit a FINISHED output and clean up."""
+    from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
+
+    request = _HashableRequest(
+        request_id="req-df2",
+        external_req_id="ext-df2",
+        status=RequestStatus.WAITING,
+        is_finished=lambda: False,
+        num_computed_tokens=0,
+        num_prompt_tokens=0,
+        prompt_token_ids=[],
+        num_output_placeholders=0,
+        sampling_params=None,
+        pooling_params=None,
+        stop_reason=None,
+        client_index=0,
+        take_events=lambda: [],
+        trace_headers=None,
+        has_encoder_inputs=False,
+        take_prefill_stats=lambda: None,
+        num_nans_in_logits=0,
+        get_finished_reason=lambda: "stop",
+    )
+    scheduler, sched_out, model_out, cleanup_calls = _build_deferred_finish_scheduler(
+        mocker,
+        running=[],
+        pending_finish_reqs=[request],
+    )
+
+    result = OmniGenerationScheduler.update_from_output(scheduler, sched_out, model_out)
+
+    assert request.status == RequestStatus.FINISHED_STOPPED
+    scheduler._free_request.assert_called_once_with(request)
+    assert len(cleanup_calls) == 1
+    eco = result[0]
+    assert len(eco.outputs) == 1
+    assert eco.outputs[0].finish_reason == "stop"
+    assert scheduler._pending_finish_reqs == []
+
+
+def test_deferred_finish_skips_already_finished(mocker: MockerFixture):
+    """A request that was aborted between schedule() and update_from_output()
+    should be skipped without emitting output or freeing resources twice."""
+    from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
+
+    request = _HashableRequest(
+        request_id="req-df3",
+        external_req_id="ext-df3",
+        status=RequestStatus.FINISHED_ABORTED,
+        is_finished=lambda: True,
+        num_computed_tokens=0,
+        num_prompt_tokens=0,
+        prompt_token_ids=[],
+        num_output_placeholders=0,
+        sampling_params=None,
+        pooling_params=None,
+        stop_reason=None,
+        client_index=0,
+        take_events=lambda: [],
+        trace_headers=None,
+        has_encoder_inputs=False,
+        take_prefill_stats=lambda: None,
+        num_nans_in_logits=0,
+        get_finished_reason=lambda: "stop",
+    )
+    scheduler, sched_out, model_out, cleanup_calls = _build_deferred_finish_scheduler(
+        mocker,
+        running=[],
+        pending_finish_reqs=[request],
+    )
+
+    result = OmniGenerationScheduler.update_from_output(scheduler, sched_out, model_out)
+
+    scheduler._handle_stopped_request.assert_not_called()
+    scheduler._free_request.assert_not_called()
+    assert len(cleanup_calls) == 0
+    assert 0 not in result
+    assert scheduler._pending_finish_reqs == []
+
+
+def test_deferred_finish_not_finished_still_emits_output(mocker: MockerFixture):
+    """When _handle_stopped_request returns False (resumable request), the
+    output must still be emitted so the client stream doesn't hang."""
+    from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
+
+    request = _HashableRequest(
+        request_id="req-df4",
+        external_req_id="ext-df4",
+        status=RequestStatus.RUNNING,
+        is_finished=lambda: False,
+        num_computed_tokens=16,
+        num_prompt_tokens=16,
+        prompt_token_ids=list(range(1, 17)),
+        num_output_placeholders=0,
+        sampling_params=None,
+        pooling_params=None,
+        stop_reason=None,
+        client_index=0,
+        take_events=lambda: [],
+        trace_headers=None,
+        has_encoder_inputs=False,
+        take_prefill_stats=lambda: None,
+        num_nans_in_logits=0,
+        get_finished_reason=lambda: "stop",
+    )
+    scheduler, sched_out, model_out, cleanup_calls = _build_deferred_finish_scheduler(
+        mocker,
+        running=[request],
+        pending_finish_reqs=[request],
+    )
+    scheduler._handle_stopped_request.return_value = False
+
+    result = OmniGenerationScheduler.update_from_output(scheduler, sched_out, model_out)
+
+    scheduler._handle_stopped_request.assert_called_once_with(request)
+    scheduler._free_request.assert_not_called()
+    assert len(cleanup_calls) == 0
+
+    eco = result[0]
+    assert len(eco.outputs) == 1
+    assert eco.outputs[0].request_id == "req-df4"
+    assert eco.outputs[0].finish_reason == "stop"
+    assert eco.outputs[0].kv_transfer_params is None
+    assert scheduler._pending_finish_reqs == []

tests/distributed/omni_connectors/test_mooncake_transfer_engine_rdma.py

@@ -25,7 +25,11 @@ from vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_c
 )
 
 # All tests in this file require Mooncake TransferEngine and an RDMA environment.
-pytestmark = [pytest.mark.parallel, pytest.mark.core_model]
+pytestmark = [
+    pytest.mark.parallel,
+    pytest.mark.core_model,
+    pytest.mark.skipif(TransferEngine is None, reason="Mooncake TransferEngine not installed"),
+]
 
 # ---------------------------------------------------------------------------
 # Shared helpers

tests/distributed/omni_connectors/test_tp_rank_aware.py

@@ -115,39 +115,43 @@ class TestConnectorKeyFormat:
 class TestRankMapping:
     """Verify get_kv_target_ranks and get_kv_source_ranks for various TP configs."""
 
-    def test_homogeneous_tp2_rank0(self):
+    @pytest.mark.parametrize(
-        topo = KVTPTopology(source_tp_size=2, target_tp_size=2, local_rank=0)
+        "tp_size,local_rank",
-        assert get_kv_target_ranks(topo) == [0]
+        [
-        assert get_kv_source_ranks(topo) == [0]
+            (2, 0),
+            (2, 1),
+            (4, 3),
+        ],
+    )
+    def test_homogeneous_tp(self, tp_size, local_rank):
+        """Homogeneous TP: each rank sends/receives to/from itself."""
+        topo = KVTPTopology(source_tp_size=tp_size, target_tp_size=tp_size, local_rank=local_rank)
+        assert get_kv_target_ranks(topo) == [local_rank]
+        assert get_kv_source_ranks(topo) == [local_rank]
 
-    def test_homogeneous_tp2_rank1(self):
+    @pytest.mark.parametrize(
-        topo = KVTPTopology(source_tp_size=2, target_tp_size=2, local_rank=1)
+        "local_rank,expected_source_ranks",
-        assert get_kv_target_ranks(topo) == [1]
+        [
-        assert get_kv_source_ranks(topo) == [1]
+            (0, [0, 1]),  # rank 0 receives from sender ranks 0 and 1
+            (1, [2, 3]),  # rank 1 receives from sender ranks 2 and 3
+        ],
+    )
+    def test_sender_gt_receiver_tp4_to_tp2(self, local_rank, expected_source_ranks):
+        """Receiver with source_tp=4, target_tp=2."""
+        topo = KVTPTopology(source_tp_size=4, target_tp_size=2, local_rank=local_rank)
+        assert get_kv_source_ranks(topo) == expected_source_ranks
 
-    def test_homogeneous_tp4_rank3(self):
+    @pytest.mark.parametrize(
-        topo = KVTPTopology(source_tp_size=4, target_tp_size=4, local_rank=3)
+        "local_rank,expected_target_ranks",
-        assert get_kv_target_ranks(topo) == [3]
+        [
-        assert get_kv_source_ranks(topo) == [3]
+            (0, [0, 1]),  # sender rank 0 sends to receiver ranks 0 and 1
-
+            (1, [2, 3]),  # sender rank 1 sends to receiver ranks 2 and 3
-    def test_sender_gt_receiver_tp4_to_tp2_rank0(self):
+        ],
-        """Receiver rank 0 should receive from sender rank 0 and 1."""
+    )
-        topo = KVTPTopology(source_tp_size=4, target_tp_size=2, local_rank=0)
+    def test_sender_lt_receiver_tp2_to_tp4(self, local_rank, expected_target_ranks):
-        assert get_kv_source_ranks(topo) == [0, 1]
+        """Sender with source_tp=2, target_tp=4."""
-
+        topo = KVTPTopology(source_tp_size=2, target_tp_size=4, local_rank=local_rank)
-    def test_sender_gt_receiver_tp4_to_tp2_rank1(self):
+        assert get_kv_target_ranks(topo) == expected_target_ranks
-        """Receiver rank 1 should receive from sender rank 2 and 3."""
-        topo = KVTPTopology(source_tp_size=4, target_tp_size=2, local_rank=1)
-        assert get_kv_source_ranks(topo) == [2, 3]
-
-    def test_sender_lt_receiver_tp2_to_tp4_rank0(self):
-        """Sender rank 0 should send to receiver ranks 0 and 1."""
-        topo = KVTPTopology(source_tp_size=2, target_tp_size=4, local_rank=0)
-        assert get_kv_target_ranks(topo) == [0, 1]
-
-    def test_sender_lt_receiver_tp2_to_tp4_rank1(self):
-        topo = KVTPTopology(source_tp_size=2, target_tp_size=4, local_rank=1)
-        assert get_kv_target_ranks(topo) == [2, 3]
 
     def test_receiver_lt_sender_source_ranks(self):
         """Receiver rank 0 with tp2_to_tp4 should source from rank 0 only."""
@@ -174,21 +178,21 @@ class TestBuildRankAwareRecvKeys:
         assert key == "omni_stage0_to_stage1_kv_cache_req-1"
         assert rank is None
 
-    def test_homogeneous_tp2_rank0(self):
+    @pytest.mark.parametrize(
-        topo = KVTPTopology(source_tp_size=2, target_tp_size=2, local_rank=0)
+        "local_rank,expected_key,expected_rank",
+        [
+            (0, "req-1_stage0_0_0_0", 0),
+            (1, "req-1_stage0_0_1_1", 1),
+        ],
+    )
+    def test_homogeneous_tp2(self, local_rank, expected_key, expected_rank):
+        """Homogeneous TP=2: each rank gets one key."""
+        topo = KVTPTopology(source_tp_size=2, target_tp_size=2, local_rank=local_rank)
         pairs = build_rank_aware_recv_keys("req-1", "stage0", "stage1", topo)
         assert len(pairs) == 1
         key, rank = pairs[0]
-        assert key == "req-1_stage0_0_0_0"
+        assert key == expected_key
-        assert rank == 0
+        assert rank == expected_rank
-
-    def test_homogeneous_tp2_rank1(self):
-        topo = KVTPTopology(source_tp_size=2, target_tp_size=2, local_rank=1)
-        pairs = build_rank_aware_recv_keys("req-1", "stage0", "stage1", topo)
-        assert len(pairs) == 1
-        key, rank = pairs[0]
-        assert key == "req-1_stage0_0_1_1"
-        assert rank == 1
 
     def test_heterogeneous_tp4_to_tp2_rank0_gets_two_keys(self):
         """Receiver rank 0 with source_tp=4, target_tp=2 should get 2 keys."""
@@ -311,9 +315,16 @@ class TestReceiveConstructsMetadata:
         assert len(calls) > 0
         assert calls[0]["metadata"] is None
 
-    def test_homogeneous_tp2_rank0_passes_metadata(self):
+    @pytest.mark.parametrize(
-        """TP=2 rank 0: metadata should point to sender rank 0's port."""
+        "local_rank,expected_port_offset",
-        mgr = _make_manager(from_tp=2, to_tp=2, local_rank=0, recv_timeout=0.05)
+        [
+            (0, 0),
+            (1, 1),
+        ],
+    )
+    def test_homogeneous_tp2_passes_metadata(self, local_rank, expected_port_offset):
+        """TP=2: metadata should point to sender rank's port with correct offset."""
+        mgr = _make_manager(from_tp=2, to_tp=2, local_rank=local_rank, recv_timeout=0.05)
         mgr.update_sender_info({"host": "10.0.0.1", "zmq_port": 50151})
 
         calls = []
@@ -330,24 +341,7 @@ class TestReceiveConstructsMetadata:
         meta = calls[0]["metadata"]
         assert meta is not None
         assert meta["source_host"] == "10.0.0.1"
-        assert meta["source_port"] == 50151 + 0 * KV_RANK_PORT_STRIDE
+        assert meta["source_port"] == 50151 + expected_port_offset * KV_RANK_PORT_STRIDE
-
-    def test_homogeneous_tp2_rank1_passes_metadata_with_offset(self):
-        mgr = _make_manager(from_tp=2, to_tp=2, local_rank=1, recv_timeout=0.05)
-        mgr.update_sender_info({"host": "10.0.0.1", "zmq_port": 50151})
-
-        calls = []
-
-        class _Connector:
-            def get(self, from_stage, to_stage, get_key, metadata=None):
-                calls.append({"key": get_key, "metadata": metadata})
-                return None
-
-        mgr._connector = _Connector()
-        mgr.receive_kv_cache_for_request("req-1")
-
-        meta = calls[0]["metadata"]
-        assert meta["source_port"] == 50151 + 1 * KV_RANK_PORT_STRIDE
 
     def test_heterogeneous_tp4_to_tp2_rank0_multiple_metadata(self):
         """Receiver rank 0 with source_tp=4, target_tp=2 should call get() with
@@ -677,6 +671,230 @@ class TestDistributedReceive:
 
         mgr.receive_multi_kv_cache.assert_called_once_with(req, None, torch.device("cpu"))
 
+    # ── SP-only scenarios ────────────────────────────────────────────
+
+    @pytest.mark.parametrize(
+        "sp_rank,world_rank_in_group,is_owner",
+        [
+            (0, 0, True),  # owner receives and broadcasts
+            (1, 1, False),  # follower receives via broadcast
+        ],
+    )
+    def test_sp_only_scenarios(self, sp_rank, world_rank_in_group, is_owner):
+        """SP-only (sp_size=2, cfg_size=1): Test owner and follower behavior."""
+        mgr = _make_manager(from_tp=2, to_tp=4, local_rank=0)
+        req = SimpleNamespace(request_id="req-1", sampling_params=SimpleNamespace())
+        world_group = _MockBroadcastGroup(world_size=4, rank_in_group=world_rank_in_group)
+        sp_payload = {
+            "past_key_values": SimpleNamespace(key_cache=[torch.tensor([1.0])]),
+            "kv_metadata": {"source": "owner"},
+            "sp.past_key_values": SimpleNamespace(key_cache=[torch.tensor([1.0])]),
+            "sp.kv_metadata": {"source": "owner"},
+        }
+        sp_group = _MockBroadcastGroup(world_size=2, rank_in_group=sp_rank, broadcast_value=sp_payload)
+
+        def _receive(req_obj, cfg_func, target_device):
+            req_obj.past_key_values = SimpleNamespace(key_cache=[torch.tensor([1.0])])
+            req_obj.kv_metadata = {"source": "owner"}
+            req_obj.sampling_params.past_key_values = req_obj.past_key_values
+            req_obj.sampling_params.kv_metadata = req_obj.kv_metadata
+            return True
+
+        mgr.receive_multi_kv_cache = MagicMock(side_effect=_receive)
+        with (
+            patch("vllm_omni.diffusion.distributed.parallel_state.get_world_group", return_value=world_group),
+            patch(
+                "vllm_omni.diffusion.distributed.parallel_state.get_classifier_free_guidance_world_size",
+                return_value=1,
+            ),
+            patch(
+                "vllm_omni.diffusion.distributed.parallel_state.get_classifier_free_guidance_rank",
+                return_value=0,
+            ),
+            patch("vllm_omni.diffusion.distributed.parallel_state.get_cfg_group", return_value=None),
+            patch(
+                "vllm_omni.diffusion.distributed.parallel_state.get_sequence_parallel_world_size",
+                return_value=2,
+            ),
+            patch(
+                "vllm_omni.diffusion.distributed.parallel_state.get_sequence_parallel_rank",
+                return_value=sp_rank,
+            ),
+            patch("vllm_omni.diffusion.distributed.parallel_state.get_sp_group", return_value=sp_group),
+        ):
+            assert mgr.receive_multi_kv_cache_distributed(req) is True
+
+        # Verify behavior based on role
+        if is_owner:
+            mgr.receive_multi_kv_cache.assert_called_once()
+            assert mgr.receive_multi_kv_cache.call_args.args[2] == torch.device("cpu")
+            assert sp_group.broadcast_calls[0][1] == 0
+        else:
+            mgr.receive_multi_kv_cache.assert_not_called()
+
+        assert len(sp_group.broadcast_calls) == 1
+        assert torch.equal(req.past_key_values.key_cache[0], torch.tensor([1.0]))
+        assert req.sampling_params.kv_metadata == {"source": "owner"}
+
+    # ── CFG + SP mixed scenarios ─────────────────────────────────────
+
+    @pytest.mark.parametrize(
+        "cfg_rank,sp_rank,world_rank_in_group,role",
+        [
+            (0, 0, 0, "owner"),  # owner: receives, cfg scatter, sp broadcast
+            (1, 0, 2, "cfg_follower_sp_leader"),  # cfg follower + sp leader: cfg recv, sp broadcast
+            (0, 1, 1, "sp_follower"),  # sp follower: only sp broadcast
+        ],
+    )
+    def test_cfg_sp_mixed_scenarios(self, cfg_rank, sp_rank, world_rank_in_group, role):
+        """CFG+SP (cfg_size=2, sp_size=2): Test various rank combinations."""
+        mgr = _make_manager(from_tp=2, to_tp=4, local_rank=0)
+        req = SimpleNamespace(request_id="req-1", sampling_params=SimpleNamespace())
+        world_group = _MockBroadcastGroup(world_size=8, rank_in_group=world_rank_in_group)
+
+        # Setup payloads based on role
+        if role == "owner":
+            owner_sp_payload = {
+                "past_key_values": SimpleNamespace(key_cache=[torch.tensor([1.0])]),
+                "kv_metadata": {"source": "owner"},
+                "sp.past_key_values": SimpleNamespace(key_cache=[torch.tensor([1.0])]),
+                "sp.kv_metadata": {"source": "owner"},
+                "sp.cfg_branch_roles": ["cfg_text"],
+                "sp.cfg_active_branch": None,
+            }
+            cfg_group = _MockBroadcastGroup(world_size=2, rank_in_group=cfg_rank)
+            sp_group = _MockBroadcastGroup(world_size=2, rank_in_group=sp_rank, broadcast_value=owner_sp_payload)
+
+            def _receive(req_obj, cfg_func, target_device):
+                req_obj.past_key_values = SimpleNamespace(key_cache=[torch.tensor([1.0])])
+                req_obj.kv_metadata = {"source": "owner"}
+                req_obj.sampling_params.past_key_values = req_obj.past_key_values
+                req_obj.sampling_params.kv_metadata = req_obj.kv_metadata
+                req_obj.sampling_params.cfg_text_past_key_values = SimpleNamespace(key_cache=[torch.tensor([2.0])])
+                req_obj.sampling_params.cfg_text_kv_metadata = {"source": "cfg_text"}
+                return True
+
+            mgr.receive_multi_kv_cache = MagicMock(side_effect=_receive)
+        else:
+            cfg_payload = {
+                "past_key_values": SimpleNamespace(key_cache=[torch.tensor([1.0])]),
+                "kv_metadata": {"source": "main"},
+                "sp.past_key_values": SimpleNamespace(key_cache=[torch.tensor([1.0])]),
+                "sp.kv_metadata": {"source": "main"},
+                "sp.cfg_active_branch": "cfg_text" if role == "cfg_follower_sp_leader" else None,
+                "sp.cfg_branch_roles": ["cfg_text"],
+                "sp.cfg_branch_past_key_values": {
+                    "cfg_text": SimpleNamespace(key_cache=[torch.tensor([2.0])]),
+                }
+                if role == "cfg_follower_sp_leader"
+                else {},
+            }
+            if role == "cfg_follower_sp_leader":
+                cfg_payload["sp.cfg_text_past_key_values"] = SimpleNamespace(key_cache=[torch.tensor([2.0])])
+
+            cfg_group = _MockBroadcastGroup(world_size=2, rank_in_group=cfg_rank, recv_value=cfg_payload)
+            sp_group = _MockBroadcastGroup(world_size=2, rank_in_group=sp_rank, broadcast_value=cfg_payload)
+            mgr.receive_multi_kv_cache = MagicMock(return_value=True)
+
+        with (
+            patch("vllm_omni.diffusion.distributed.parallel_state.get_world_group", return_value=world_group),
+            patch(
+                "vllm_omni.diffusion.distributed.parallel_state.get_classifier_free_guidance_world_size",
+                return_value=2,
+            ),
+            patch(
+                "vllm_omni.diffusion.distributed.parallel_state.get_classifier_free_guidance_rank",
+                return_value=cfg_rank,
+            ),
+            patch("vllm_omni.diffusion.distributed.parallel_state.get_cfg_group", return_value=cfg_group),
+            patch(
+                "vllm_omni.diffusion.distributed.parallel_state.get_sequence_parallel_world_size",
+                return_value=2,
+            ),
+            patch(
+                "vllm_omni.diffusion.distributed.parallel_state.get_sequence_parallel_rank",
+                return_value=sp_rank,
+            ),
+            patch("vllm_omni.diffusion.distributed.parallel_state.get_sp_group", return_value=sp_group),
+        ):
+            assert mgr.receive_multi_kv_cache_distributed(req) is True
+
+        # Verify behavior based on role
+        if role == "owner":
+            mgr.receive_multi_kv_cache.assert_called_once()
+            assert mgr.receive_multi_kv_cache.call_args.args[2] == torch.device("cpu")
+            assert [dst for dst, _ in cfg_group.send_calls] == [1]
+            assert cfg_group.send_calls[0][1]["sp.cfg_active_branch"] == "cfg_text"
+            assert len(sp_group.broadcast_calls) == 1
+            assert sp_group.broadcast_calls[0][1] == 0
+        elif role == "cfg_follower_sp_leader":
+            mgr.receive_multi_kv_cache.assert_not_called()
+            assert cfg_group.recv_calls == [0]
+            assert len(sp_group.broadcast_calls) == 1
+            assert sp_group.broadcast_calls[0][1] == 0
+            assert req.sampling_params.cfg_active_branch == "cfg_text"
+            assert torch.equal(
+                req.sampling_params.cfg_branch_past_key_values["cfg_text"].key_cache[0],
+                torch.tensor([2.0]),
+            )
+        else:  # sp_follower
+            mgr.receive_multi_kv_cache.assert_not_called()
+            assert cfg_group.recv_calls == []
+            assert cfg_group.send_calls == []
+            assert len(sp_group.broadcast_calls) == 1
+            assert torch.equal(req.past_key_values.key_cache[0], torch.tensor([1.0]))
+            assert req.kv_metadata == {"source": "main"}
+
+    @pytest.mark.parametrize(
+        "has_cfg,cfg_size",
+        [
+            (False, 1),  # SP-only: owner receive failure
+            (True, 2),  # CFG+SP: owner receive failure with deadlock prevention
+        ],
+    )
+    def test_owner_receive_failure_scenarios(self, has_cfg, cfg_size):
+        """Test receive failure handling in SP-only and CFG+SP scenarios."""
+        mgr = _make_manager(from_tp=2, to_tp=4, local_rank=0)
+        req = SimpleNamespace(request_id="req-1", sampling_params=SimpleNamespace())
+        world_group = _MockBroadcastGroup(world_size=4 if not has_cfg else 8, rank_in_group=0)
+        cfg_group = _MockBroadcastGroup(world_size=2, rank_in_group=0) if has_cfg else None
+        sp_group = _MockBroadcastGroup(world_size=2, rank_in_group=0, broadcast_value=None)
+
+        mgr.receive_multi_kv_cache = MagicMock(return_value=False)
+        with (
+            patch("vllm_omni.diffusion.distributed.parallel_state.get_world_group", return_value=world_group),
+            patch(
+                "vllm_omni.diffusion.distributed.parallel_state.get_classifier_free_guidance_world_size",
+                return_value=cfg_size,
+            ),
+            patch(
+                "vllm_omni.diffusion.distributed.parallel_state.get_classifier_free_guidance_rank",
+                return_value=0,
+            ),
+            patch("vllm_omni.diffusion.distributed.parallel_state.get_cfg_group", return_value=cfg_group),
+            patch(
+                "vllm_omni.diffusion.distributed.parallel_state.get_sequence_parallel_world_size",
+                return_value=2,
+            ),
+            patch(
+                "vllm_omni.diffusion.distributed.parallel_state.get_sequence_parallel_rank",
+                return_value=0,
+            ),
+            patch("vllm_omni.diffusion.distributed.parallel_state.get_sp_group", return_value=sp_group),
+        ):
+            assert mgr.receive_multi_kv_cache_distributed(req) is False
+
+        # Owner attempted to receive but failed
+        mgr.receive_multi_kv_cache.assert_called_once()
+        # SP broadcast still called but with None
+        assert len(sp_group.broadcast_calls) == 1
+        assert sp_group.broadcast_calls[0][0] is None
+
+        # CFG+SP: owner must send None to cfg followers to prevent deadlock
+        if has_cfg:
+            assert len(cfg_group.send_calls) == 1
+            assert cfg_group.send_calls[0] == (1, None)
+
 
 # ── TP auto-detect ───────────────────────────────────────────────────
 

tests/e2e/accuracy/conftest.py

@@ -185,6 +185,12 @@ def accuracy_artifact_root() -> Path:
     return root
 
 
+@pytest.fixture(scope="session")
+def accuracy_assets_root() -> Path:
+    root = Path(__file__).resolve().parent / "assets"
+    return root
+
+
 @pytest.fixture(scope="session")
 def qwen_bear_image(accuracy_artifact_root: Path):
     """Download the Qwen bear image from the URL and save it to the accuracy artifact root."""

tests/e2e/accuracy/helpers.py

@@ -1,9 +1,11 @@
 import os
+from io import BytesIO
 from pathlib import Path
 from typing import Any
 
 import numpy as np
 import pytest
+import requests
 import torch
 from PIL import Image
 from torchmetrics.image import PeakSignalNoiseRatio, StructuralSimilarityIndexMeasure
@@ -125,6 +127,74 @@ def assert_image_sequence_similarity(
         )
 
 
+def assert_images_pixel_close(
+    *,
+    model_name: str,
+    vllm_image: Image.Image,
+    diffusers_image: Image.Image,
+    mean_threshold: float,
+    p99_threshold: float,
+    compare_mode: str = "RGB",
+) -> None:
+    """Assert full-image pixel closeness between online serving and diffusers output.
+
+    Match the threshold style used by diffusion parallelism tests: convert both
+    images to float32 RGB tensors in the [0, 1] range, then compare mean and p99
+    absolute channel differences.
+
+    Accuracy improves as both values move lower:
+    - ``mean_abs_diff`` tracks global average drift across the whole image.
+    - ``p99_abs_diff`` tracks tail drift while ignoring the noisiest 1% of
+      channel samples.
+
+    The printed mismatch ratios are diagnostics only. ``pixel_ratio`` counts
+    pixels where any channel exceeds a tolerance, while ``channel_ratio`` counts
+    individual channel samples above that tolerance.
+    """
+    assert vllm_image.size == diffusers_image.size, (
+        f"Online and diffusers output sizes mismatch: online={vllm_image.size}, diffusers={diffusers_image.size}"
+    )
+
+    vllm_array = np.asarray(vllm_image.convert(compare_mode), dtype=np.float32) / 255.0
+    diffusers_array = np.asarray(diffusers_image.convert(compare_mode), dtype=np.float32) / 255.0
+    channel_abs_diff = np.abs(vllm_array - diffusers_array)
+    mean_abs_diff = float(channel_abs_diff.mean())
+    p99_abs_diff = float(np.quantile(channel_abs_diff, 0.99))
+    percentiles = {
+        percentile: float(np.quantile(channel_abs_diff, percentile / 100.0)) for percentile in (50, 90, 95, 99, 99.9)
+    }
+
+    print(f"{model_name} pixel metrics:")
+    print(
+        f"  mean_abs_diff: value={mean_abs_diff:.6e}, threshold<={mean_threshold:.6e}, "
+        "range=[0, 1], lower_is_better=True"
+    )
+    print(
+        f"  p99_abs_diff: value={p99_abs_diff:.6e}, threshold<={p99_threshold:.6e}, range=[0, 1], lower_is_better=True"
+    )
+    print("  abs_diff_percentiles:")
+    for percentile, value in percentiles.items():
+        print(f"    p{percentile}: value={value:.6e}, range=[0, 1], lower_is_better=True")
+    print("  mismatch_ratios_by_channel_threshold:")
+    for tolerance in (0, 1, 2, 4, 8, 16, 32, 64, 128):
+        normalized_tolerance = tolerance / 255.0
+        pixel_mismatch = np.any(channel_abs_diff > normalized_tolerance, axis=-1)
+        pixel_mismatch_ratio = float(np.count_nonzero(pixel_mismatch) / pixel_mismatch.size)
+        channel_mismatch_ratio = float(
+            np.count_nonzero(channel_abs_diff > normalized_tolerance) / channel_abs_diff.size
+        )
+        print(
+            f"    threshold>{normalized_tolerance:.6e} ({tolerance}/255): pixel_ratio={pixel_mismatch_ratio:.8f}, "
+            f"channel_ratio={channel_mismatch_ratio:.8f}"
+        )
+
+    assert mean_abs_diff <= mean_threshold and p99_abs_diff <= p99_threshold, (
+        f"Image diff exceeded threshold for {model_name}: mean_abs_diff={mean_abs_diff:.6e}, "
+        f"p99_abs_diff={p99_abs_diff:.6e} (thresholds: mean<={mean_threshold:.6e}, "
+        f"p99<={p99_threshold:.6e})"
+    )
+
+
 def compute_image_ssim_psnr(
     *,
     prediction: Image.Image,
@@ -161,6 +231,16 @@ class CLIPScorer:
         similarity = (img_emb * txt_emb).sum(dim=-1)
         return float(similarity.item() * 100)
 
+    def image_image_score(self, image1: Image.Image, image2: Image.Image) -> float:
+        inputs = self._processor(text=[""], images=[image1, image2], return_tensors="pt", padding=True)
+        with torch.no_grad():
+            outputs = self._model(**inputs)
+        img_emb1, img_emb2 = outputs.image_embeds.chunk(2)
+        img_emb1 = img_emb1 / img_emb1.norm(p=2, dim=-1, keepdim=True)
+        img_emb2 = img_emb2 / img_emb2.norm(p=2, dim=-1, keepdim=True)
+        similarity = (img_emb1 * img_emb2).sum(dim=-1)
+        return float(similarity.item() * 100)
+
     def assert_score(self, *, model_name: str, image: Image.Image, text: str, threshold: float) -> None:
         value = self.score(image, text)
         print(f"{model_name} CLIP score:")
@@ -180,3 +260,73 @@ def _pil_to_batched_tensor(image: Image.Image, *, compare_mode: str) -> torch.Te
     array = np.asarray(image.convert(compare_mode), dtype=np.float32) / 255.0
     tensor = torch.from_numpy(array).permute(2, 0, 1).unsqueeze(0)
     return tensor
+
+
+class SemanticSimilarityScorer:
+    """Semantic similarity scorer for text-text and image-text comparisons"""
+
+    def __init__(
+        self,
+        text_model_name: str = "BAAI/bge-m3",
+    ):
+        """
+        Initialize semantic similarity scorer
+
+        Args:
+            text_model_name: Model for text-text similarity (BGE-M3)
+        """
+        self._text_model = None
+        self._text_model_name = text_model_name
+
+    def _load_text_model(self):
+        """Lazy load text similarity model (BGE-M3)"""
+        if self._text_model is None:
+            try:
+                from FlagEmbedding import BGEM3FlagModel
+
+                print(f"Loading text similarity model: {self._text_model_name}")
+                self._text_model = BGEM3FlagModel(self._text_model_name, use_fp16=True)
+            except ImportError:
+                raise ImportError("FlagEmbedding not installed. Install with: pip install FlagEmbedding")
+        return self._text_model
+
+    def text_similarity(self, text1: str, text2: str) -> float:
+        """
+        Compute semantic similarity between two texts using BGE-M3
+        """
+        results = {}
+        # 1. Text prefix matching (maximum matching character length)
+        match_count = 0
+        for char_vllm, char_ref in zip(text1, text2):
+            if char_vllm == char_ref:
+                match_count += 1
+            else:
+                break
+        results["text_prefix_match_count"] = match_count
+
+        # 2. Full CoT semantic similarity comparison
+        model = self._load_text_model()
+        emb1 = model.encode([text1], return_dense=True, return_sparse=False, return_colbert_vecs=False)["dense_vecs"][0]
+        emb2 = model.encode([text2], return_dense=True, return_sparse=False, return_colbert_vecs=False)["dense_vecs"][0]
+        similarity = np.dot(emb1, emb2) / (np.linalg.norm(emb1) * np.linalg.norm(emb2))
+        results["cot_semantic_sim"] = float(similarity)
+        return results
+
+    def release_models(self):
+        """Release loaded models to free memory"""
+        if self._text_model is not None:
+            del self._text_model
+            self._text_model = None
+
+
+def download_images(urls: list[str]) -> list[Image.Image]:
+    """Download and cache multiple images"""
+    images = []
+    for i, url in enumerate(urls):
+        print(f"Downloading image {i} from {url}")
+        response = requests.get(url, timeout=30)
+        response.raise_for_status()
+        image = Image.open(BytesIO(response.content)).convert("RGB")
+        images.append(image)
+
+    return images

tests/e2e/accuracy/test_diffusers_backend_similarity.py

@@ -284,7 +284,6 @@ def test_diffusers_backend_t2i_matches_diffusers(model_id: str, accuracy_artifac
     [
         pytest.param(
             "Wan-AI/Wan2.2-I2V-A14B-Diffusers",
-            marks=pytest.mark.skip(reason="#3488"),
         ),
     ],
 )

tests/e2e/accuracy/test_hunyuan_image3.py

@@ -5,6 +5,8 @@ from __future__ import annotations
 
 import copy
 import gc
+import importlib
+import os
 import tempfile
 import time
 from pathlib import Path
@@ -14,28 +16,59 @@ import torch
 import yaml
 from PIL import Image
 
-from tests.e2e.accuracy.helpers import CLIPScorer, assert_similarity, model_output_dir
+from tests.e2e.accuracy.helpers import (
+    CLIPScorer,
+    SemanticSimilarityScorer,
+    compute_image_ssim_psnr,
+    download_images,
+    model_output_dir,
+)
 from tests.helpers.runtime import OmniRunner
-from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import build_prompt_tokens
+from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import build_prompt_tokens, resolve_stop_token_ids
+
+os.environ["DIFFUSION_ATTENTION_BACKEND"] = "TORCH_SDPA"
 
 pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
 
+# ============================================================================
+# Configurable Parameters
+# ============================================================================
+AR_DEVICES = "0,1"
+DIT_DEVICES = "2,3"
 MODEL_NAME = "tencent/HunyuanImage-3.0-Instruct"
-SEED = 42
 NUM_INFERENCE_STEPS = 50
-GUIDANCE_SCALE = 0
+GUIDANCE_SCALE = 2.5
-HEIGHT, WIDTH = 1024, 1024
-PSNR_THRESHOLD = 0.0
-SSIM_THRESHOLD = 0.0
-CLIP_SCORE_THRESHOLD = 20.0
-PROMPT = "A brown and white dog is running on the grass."
 
+# ============================================================================
+# Constants
+# ============================================================================
+MODEL_PATH = os.environ.get("HUNYUAN_MODEL_PATH", MODEL_NAME)
+# Test input
+PROMPT = "基于图一的logo，参考图二中冰箱贴的材质，制作一个新的冰箱贴"
+TEST_IMAGE_URLS = [
+    "https://raw.githubusercontent.com/Tencent-Hunyuan/HunyuanImage-3.0/main/assets/demo_instruct_imgs/input_1_0.png",
+    "https://raw.githubusercontent.com/Tencent-Hunyuan/HunyuanImage-3.0/main/assets/demo_instruct_imgs/input_1_1.png",
+]
+SEED = 42
+AR_TP_SIZE = len(AR_DEVICES.split(","))
+DIT_TP_SIZE = len(DIT_DEVICES.split(","))
+
+# Precision thresholds
+THRESHOLDS = {
+    # AR text comparison
+    "text_prefix_match": 10,  # First 10 characters must match exactly
+    "cot_semantic_sim": 0.9,  # Full CoT semantic similarity
+    # Image comparison
+    "clip_score": 85,  # CLIP image semantic similarity
+    "ssim": 0.20,  # Structural similarity
+    "psnr": 11.0,  # Peak signal-to-noise ratio (dB)
+}
 # fmt: off
 _BASE_CONFIG = {
     "stage_args": [
         {
             "stage_id": 0, "stage_type": "llm",
-            "runtime": {"process": True, "devices": "0,1", "max_batch_size": 1, "requires_multimodal_data": True},
+            "runtime": {"process": True, "devices": AR_DEVICES, "max_batch_size": 1, "requires_multimodal_data": True},
             "engine_args": {
                 "model_stage": "AR", "model_arch": "HunyuanImage3ForCausalMM",
                 "worker_cls": "vllm_omni.worker.gpu_ar_worker.GPUARWorker",
@@ -50,11 +83,11 @@ _BASE_CONFIG = {
                 "temperature": 0.0, "top_p": 1, "top_k": -1, "max_tokens": 8192,
                 "stop_token_ids": [128025], "detokenize": True, "skip_special_tokens": False,
             },
-            "output_connectors": {"to_stage_1": "rdma_connector"},
+            "output_connectors": {"to_stage_1": "shared_memory_connector"},
         },
         {
             "stage_id": 1, "stage_type": "diffusion",
-            "runtime": {"process": True, "devices": "2,3", "max_batch_size": 1, "requires_multimodal_data": True},
+            "runtime": {"process": True, "devices": DIT_DEVICES, "max_batch_size": 1, "requires_multimodal_data": True},
             "engine_args": {
                 "model_stage": "dit", "model_arch": "HunyuanImage3ForCausalMM",
                 "enforce_eager": True, "trust_remote_code": True, "distributed_executor_backend": "mp",
@@ -64,20 +97,32 @@ _BASE_CONFIG = {
             "custom_process_input_func": "vllm_omni.model_executor.stage_input_processors.hunyuan_image3.ar2diffusion",
             "final_output": True, "final_output_type": "image",
             "default_sampling_params": {"num_inference_steps": NUM_INFERENCE_STEPS, "guidance_scale": GUIDANCE_SCALE},
-            "input_connectors": {"from_stage_0": "rdma_connector"},
+            "input_connectors": {"from_stage_0": "shared_memory_connector"},
         },
     ],
     "runtime": {
         "enabled": True,
-        "connectors": {"rdma_connector": {
+        "connectors": {"shared_memory_connector": {
-            "name": "MooncakeTransferEngineConnector",
+            "name": "SharedMemoryConnector",
-            "extra": {"host": "auto", "zmq_port": 50051, "protocol": "rdma", "device_name": "", "memory_pool_size": 4294967296, "memory_pool_device": "cpu"},
+            "extra": {"shm_threshold_bytes": 65536}
         }},
         "edges": [{"from": 0, "to": 1}],
     },
 }
 # fmt: on
 
+# fmt: off
+COT_REF = ("首先，我分析所有输入图像：图像1是一个圆形的logo，设计现代且抽象。它由不同色调的蓝色（深蓝、中蓝、浅蓝）和白色构成，这些色块以流畅的曲线相互交织，形成一个动态的、类似旋涡或波浪的图案。整个logo是扁平化的矢量图形，背景为纯黑色。图像2展示了四个并排摆放的卡通动物造型冰箱贴，"
+           "它们被放置在灰色的织物背景上。这些冰箱贴的关键特征是其材质：它们具有光滑、高光的珐琅或烤漆质感，边缘有明显的金属包边，整体呈现出一种立体的、有厚度的实体感。用户的指令是“基于图一的logo，参考图二中冰箱贴的材质，制作一个新的冰箱贴”。这个指令要求将一个二维的平面设计（logo）"
+           "转化为一个具有特定物理属性（材质和立体感）的三维物体。核心任务是保留logo的视觉识别性，同时赋予其冰箱贴的实体质感。为了构建答案图像，我会将图一的圆形logo作为基础形状。然后，我会将图二中冰箱贴的材质特性应用到这个logo上。具体来说，logo中的每一个色块（深蓝、中蓝、浅蓝、白色）"
+           "都会被渲染成具有高光泽度的珐琅质感，表面会反射出柔和的环境光，形成自然的高光。logo中不同颜色区域之间的分界线，将被处理成纤细的、带有金属光泽的凸起边缘，这既能清晰地勾勒出图案，也符合珐琅工艺品的典型特征。整个冰箱贴会呈现出轻微的厚度和圆润的边缘，使其看起来像一个真实的、可触摸的物体。"
+           "最后，将这个制作完成的冰箱贴放置在图二所示的灰色织物背景上，并为其添加一个微妙的、柔和的阴影，以增强其立体感和与背景的融合度，最终呈现出一个精致、逼真的产品展示图。</think><recaption>这幅图像以产品摄影的精致风格，呈现了一枚根据`image_1`标志定制的圆形珐琅冰箱贴。最终图像使用`image_2`的分辨率。"
+           "冰箱贴居中放置在`image_2`的灰色织物背景上，其设计完美复刻了`image_1`中由深蓝、中蓝、浅蓝和白色构成的动态旋涡图案。整个冰箱贴被赋予了`image_2`中冰箱贴特有的高级质感：表面覆盖着一层光滑如镜的珐琅釉面，反射出柔和而清晰的高光；图案的每一个色块边缘都由纤细的抛光金属边框精确勾勒，增强了立体感。"
+           "柔和的顶光在冰箱贴的弧形边缘上形成平滑的过渡，并在其下方投下淡淡的、轮廓模糊的阴影，使其与织物背景无缝融合，营造出一种真实、静谧的视觉效果。<relation_1>最终图像完整保留了`image_1`中标志的全部设计元素。这包括其完美的圆形轮廓，以及内部由深蓝、中蓝、浅蓝和白色组成的精确旋涡状图案布局、形状和色彩关系。"
+           "</relation_1><relation_2>最终图像的分辨率、背景和材质均来自`image_2`。背景中灰色织物的纹理和质感被完整保留。冰箱贴的材质被完美重构，精确复刻了`image_2`中冰箱贴所展示的光滑珐琅质感、抛光金属边框的视觉效果，以及整体柔和、均匀的布光环境和由此产生的自然阴影。</relation_2></recaption><answer><boi>"
+           "<img_size_1024><img_ratio_36><timestep>[<img>]{3600}<eoi></answer>")
+# fmt: on
+
 
 def _make_config(enable_kv_reuse: bool, path: Path) -> None:
     config = copy.deepcopy(_BASE_CONFIG)
@@ -86,18 +131,25 @@ def _make_config(enable_kv_reuse: bool, path: Path) -> None:
     path.write_text(yaml.dump(config, default_flow_style=False, sort_keys=False))
 
 
-def _run(stage_config_path: str, output_path: Path) -> tuple[Image.Image, str, float]:
+def _run_offline(stage_configs_path: str, output_path: Path) -> tuple[Image.Image, str, float]:
     from transformers import AutoTokenizer
 
     from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType
     from vllm_omni.platforms import current_omni_platform
 
+    build_kwargs: dict = {"task": "it2i", "bot_task": "think_recaption", "sys_type": "en_unified", "num_images": 2}
+
     tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
-    result = build_prompt_tokens(PROMPT, tokenizer, task="it2i_recaption", sys_type="en_unified")
+    result = build_prompt_tokens(
+        PROMPT,
+        tokenizer,
+        **build_kwargs,
+    )
     token_ids = result.token_ids
     system_prompt_type = result.system_prompt_type
 
-    with OmniRunner(MODEL_NAME, stage_configs_path=stage_config_path) as runner:
+    ar_stop_token_ids = resolve_stop_token_ids(task="it2i", bot_task="think_recaption", tokenizer=tokenizer)
+    with OmniRunner(MODEL_NAME, stage_configs_path=stage_configs_path) as runner:
         params_list = list(runner.omni.default_sampling_params_list)
         for sp in params_list:
             if isinstance(sp, OmniDiffusionSamplingParams):
@@ -105,14 +157,17 @@ def _run(stage_config_path: str, output_path: Path) -> tuple[Image.Image, str, f
                 sp.guidance_scale = GUIDANCE_SCALE
                 sp.seed = SEED
                 sp.generator = torch.Generator(device=current_omni_platform.device_type or "cuda").manual_seed(SEED)
+            elif hasattr(sp, "stop_token_ids"):
+                sp.stop_token_ids = ar_stop_token_ids
 
+        images = download_images(TEST_IMAGE_URLS)
         prompts: list[OmniPromptType] = [
             {
                 "prompt_token_ids": token_ids,
                 "prompt": PROMPT,
                 "use_system_prompt": system_prompt_type,
-                "height": HEIGHT,
+                "modalities": ["image"],
-                "width": WIDTH,
+                "multi_modal_data": {"image": images},
             }
         ]
         t0 = time.perf_counter()
@@ -125,9 +180,13 @@ def _run(stage_config_path: str, output_path: Path) -> tuple[Image.Image, str, f
     for out in outputs:
         ro = getattr(out, "request_output", None)
         if ro and getattr(ro, "outputs", None):
-            txt = "".join(getattr(o, "text", "") or "" for o in ro.outputs)
+            cot_text = "".join(getattr(o, "text", "") or "" for o in ro.outputs)
-            if txt:
+        if not cot_text:
-                cot_text = txt
+            ar_text = getattr(out, "custom_output", {}).get("ar_generated_text")
+            if isinstance(ar_text, list):
+                cot_text = "\n".join(text for text in ar_text if text)
+            else:
+                cot_text = ar_text or ""
 
         imgs = getattr(out, "images", None)
         if not imgs and ro and hasattr(ro, "images"):
@@ -136,9 +195,11 @@ def _run(stage_config_path: str, output_path: Path) -> tuple[Image.Image, str, f
             images = imgs
 
     assert images, "Pipeline output had no images"
+    cot_text = cot_text.lstrip("\n")
 
     image = images[0].convert("RGB")
-    image.save(output_path)
+    image.save(output_path / "image_offline.png")
+    (output_path / "cot_offline.txt").write_text(cot_text, encoding="utf-8")
     gc.collect()
     if torch.accelerator.is_available():
         torch.accelerator.empty_cache()
@@ -146,71 +207,44 @@ def _run(stage_config_path: str, output_path: Path) -> tuple[Image.Image, str, f
 
 
 @pytest.mark.skipif(torch.accelerator.device_count() < 4, reason="Needs 4+ GPUs (2 AR + 2 DiT)")
-def test_text_to_image_alignment(accuracy_artifact_root: Path) -> None:
+def test_image_to_image_alignment(accuracy_artifact_root: Path, accuracy_assets_root: Path) -> None:
-    """KV reuse ON vs OFF: same pipeline, same seed → PSNR >= 40 dB."""
+    if importlib.util.find_spec("FlagEmbedding") is None:
-    output_dir = model_output_dir(accuracy_artifact_root, MODEL_NAME + "-kv-reuse")
+        raise ImportError("Missing dependency: FlagEmbedding\nInstall with: pip install FlagEmbedding")
+    from tabulate import tabulate  # lazy import
+
+    """KV reuse ON vs OFF: same pipeline, same seed → PSNR >= 10 dB."""
+    output_dir = model_output_dir(accuracy_artifact_root, MODEL_NAME + "-offline-kv-reuse")
 
     with tempfile.TemporaryDirectory() as tmpdir:
         tmp = Path(tmpdir)
-        _make_config(False, tmp / "off.yaml")
-        image_no_reuse_1, cot_no_reuse_1, time_no_reuse_1 = _run(
-            str(tmp / "off.yaml"), output_dir / "without_kv_reuse_1.png"
-        )
-        image_no_reuse_2, cot_no_reuse_2, time_no_reuse_2 = _run(
-            str(tmp / "off.yaml"), output_dir / "without_kv_reuse_2.png"
-        )
         _make_config(True, tmp / "on.yaml")
-        image_reuse, cot_reuse, time_reuse = _run(str(tmp / "on.yaml"), output_dir / "with_kv_reuse.png")
+        omni_image, omni_cot, time_reuse = _run_offline(str(tmp / "on.yaml"), output_dir)
-
-    print("\n--- End-to-end time ---")
-    print(f"  WITHOUT KV reuse 1: {time_no_reuse_1:.2f}s")
-    print(f"  WITHOUT KV reuse 2: {time_no_reuse_2:.2f}s")
-    print(f"  WITH KV reuse:    {time_reuse:.2f}s")
-
-    (output_dir / "cot_with_kv_reuse.txt").write_text(cot_reuse, encoding="utf-8")
-    (output_dir / "cot_without_kv_reuse_1.txt").write_text(cot_no_reuse_1, encoding="utf-8")
-    (output_dir / "cot_without_kv_reuse_2.txt").write_text(cot_no_reuse_2, encoding="utf-8")
-
-    print(f"\n--- CoT WITH KV reuse (len={len(cot_reuse)}) ---\n{cot_reuse}")
-    print(f"\n--- CoT WITHOUT KV reuse 1 (len={len(cot_no_reuse_1)}) ---\n{cot_no_reuse_1}")
-    print(f"\n--- CoT WITHOUT KV reuse 2 (len={len(cot_no_reuse_2)}) ---\n{cot_no_reuse_2}")
 
+    scorer = SemanticSimilarityScorer()
     clip_scorer = CLIPScorer()
-    clip_scorer.assert_score(
+    cot_results = scorer.text_similarity(omni_cot, COT_REF)
-        model_name=f"{MODEL_NAME} with-reuse",
+    image_ref = Image.open(str(accuracy_assets_root / "hunyuan_image_ref.png")).convert("RGB")
-        image=image_reuse,
+    image_clip_score = clip_scorer.image_image_score(omni_image, image_ref)
-        text=PROMPT,
+    ssim_value, psnr_value = compute_image_ssim_psnr(prediction=omni_image, reference=image_ref, compare_mode="RGB")
-        threshold=CLIP_SCORE_THRESHOLD,
-    )
-    clip_scorer.assert_score(
-        model_name=f"{MODEL_NAME} without-reuse 1",
-        image=image_no_reuse_1,
-        text=PROMPT,
-        threshold=CLIP_SCORE_THRESHOLD,
-    )
-    clip_scorer.assert_score(
-        model_name=f"{MODEL_NAME} without-reuse 2",
-        image=image_no_reuse_2,
-        text=PROMPT,
-        threshold=CLIP_SCORE_THRESHOLD,
-    )
 
-    assert_similarity(
+    table = [
-        model_name=f"{MODEL_NAME} non-reuse run1 vs run2",
+        ["COT similarity to reference", f"{cot_results['cot_semantic_sim']:.4f}", 0.9644],
-        vllm_image=image_no_reuse_1,
+        ["COT prefix match", f"{cot_results['text_prefix_match_count']:.4f}", 29],
-        diffusers_image=image_no_reuse_2,
+        ["Image-Image similarity", f"{image_clip_score:.4f}", 94.5538],
-        ssim_threshold=SSIM_THRESHOLD,
+        ["SSIM", f"{ssim_value:.4f}", 0.242],
-        psnr_threshold=PSNR_THRESHOLD,
+        ["PSNR (dB)", f"{psnr_value:.2f}", 14.1],
-        width=WIDTH,
+    ]
-        height=HEIGHT,
-    )
 
-    assert_similarity(
+    print(tabulate(table, headers=["Metric", "Value", "L20x Reference"], tablefmt="grid"))
-        model_name=f"{MODEL_NAME} KV-reuse vs no-reuse",
+
-        vllm_image=image_reuse,
+    assert cot_results["cot_semantic_sim"] >= THRESHOLDS["cot_semantic_sim"], (
-        diffusers_image=image_no_reuse_1,
+        f"COT semantic similarity {cot_results['cot_semantic_sim']:.4f} is below threshold {THRESHOLDS['cot_semantic_sim']}"
-        ssim_threshold=SSIM_THRESHOLD,
-        psnr_threshold=PSNR_THRESHOLD,
-        width=WIDTH,
-        height=HEIGHT,
     )
+    assert cot_results["text_prefix_match_count"] >= THRESHOLDS["text_prefix_match"], (
+        f"COT prefix match count {cot_results['text_prefix_match_count']} is below threshold {THRESHOLDS['text_prefix_match']}"
+    )
+    assert image_clip_score >= THRESHOLDS["clip_score"], (
+        f"Image-Image similarity{image_clip_score:.4f} is below threshold {THRESHOLDS['clip_score']}"
+    )
+    assert ssim_value >= THRESHOLDS["ssim"], f"SSIM {ssim_value:.4f} is below threshold {THRESHOLDS['ssim']}"
+    assert psnr_value >= THRESHOLDS["psnr"], f"PSNR {psnr_value:.2f} dB is below threshold {THRESHOLDS['psnr']} dB"

tests/e2e/accuracy/test_qwen_image.py

@@ -12,7 +12,7 @@ import torch
 from diffusers.pipelines.pipeline_utils import DiffusionPipeline
 from PIL import Image
 
-from tests.e2e.accuracy.helpers import assert_similarity, model_output_dir
+from tests.e2e.accuracy.helpers import assert_images_pixel_close, assert_similarity, model_output_dir
 from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup
 from tests.helpers.mark import hardware_test
 from tests.helpers.runtime import OmniServer
@@ -32,11 +32,38 @@ SEED = 42
 SSIM_THRESHOLD = 0.97
 PSNR_THRESHOLD = 30.0
 
+MODEL_2512_ID = "Qwen/Qwen-Image-2512"
+MODEL_2512_ENV_VAR = "QWEN_IMAGE_2512_MODEL"
+PROMPT_2512 = (
+    "A 20-year-old East Asian girl with delicate, charming features and large, bright brown eyes—expressive and "
+    "lively, with a cheerful or subtly smiling expression. Her naturally wavy long hair is either loose or tied in "
+    "twin ponytails. She has fair skin and light makeup accentuating her youthful freshness. She wears a modern, "
+    "cute dress or relaxed outfit in bright, soft colors—lightweight fabric, minimalist cut. She stands indoors at "
+    "an anime convention, surrounded by banners, posters, or stalls. Lighting is typical indoor illumination—no "
+    "staged lighting—and the image resembles a casual iPhone snapshot: unpretentious composition, yet brimming "
+    "with vivid, fresh, youthful charm."
+)
+NEGATIVE_PROMPT_2512 = (
+    "低分辨率，低画质，肢体畸形，手指畸形，画面过饱和，蜡像感，人脸无细节，过度光滑，画面具有AI感。"
+    "构图混乱。文字模糊，扭曲。"
+)
+WIDTH_2512 = 1664
+HEIGHT_2512 = 928
+NUM_INFERENCE_STEPS_2512 = 50
+TRUE_CFG_SCALE_2512 = 4.0
+SEED_2512 = 42
+MEAN_ABS_DIFF_THRESHOLD_2512 = 3e-2
+P99_ABS_DIFF_THRESHOLD_2512 = 4e-1
+
 
 def _model_name() -> str:
     return os.environ.get(MODEL_ENV_VAR, MODEL_ID)
 
 
+def _model_2512_name() -> str:
+    return os.environ.get(MODEL_2512_ENV_VAR, MODEL_2512_ID)
+
+
 def _local_files_only(model: str) -> bool:
     return Path(model).exists()
 
@@ -102,6 +129,67 @@ def _run_diffusers_qwen_image(*, model: str, output_path: Path) -> Image.Image:
         run_post_test_cleanup()
 
 
+def _run_vllm_omni_qwen_image_2512(*, model: str, output_path: Path) -> Image.Image:
+    server_args = ["--num-gpus", "1", "--stage-init-timeout", "300", "--init-timeout", "900"]
+    with OmniServer(model, server_args, use_omni=True) as omni_server:
+        response = requests.post(
+            f"http://{omni_server.host}:{omni_server.port}/v1/images/generations",
+            json={
+                "model": omni_server.model,
+                "prompt": PROMPT_2512,
+                "size": f"{WIDTH_2512}x{HEIGHT_2512}",
+                "n": 1,
+                "response_format": "b64_json",
+                "negative_prompt": NEGATIVE_PROMPT_2512,
+                "num_inference_steps": NUM_INFERENCE_STEPS_2512,
+                "true_cfg_scale": TRUE_CFG_SCALE_2512,
+                "seed": SEED_2512,
+            },
+            timeout=1200,
+        )
+        response.raise_for_status()
+        payload = response.json()
+        assert len(payload["data"]) == 1
+        image_bytes = base64.b64decode(payload["data"][0]["b64_json"])
+        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+        image.load()
+        image.save(output_path)
+        return image
+
+
+def _run_diffusers_qwen_image_2512(*, model: str, output_path: Path) -> Image.Image:
+    run_pre_test_cleanup()
+    pipe: DiffusionPipeline | None = None
+    try:
+        pipe = DiffusionPipeline.from_pretrained(
+            model,
+            torch_dtype=torch.bfloat16,
+            trust_remote_code=True,
+            local_files_only=_local_files_only(model),
+        ).to("cuda")
+        generator = torch.Generator(device="cuda").manual_seed(SEED_2512)
+        result = pipe(  # pyright: ignore[reportCallIssue]
+            prompt=PROMPT_2512,
+            negative_prompt=NEGATIVE_PROMPT_2512,
+            width=WIDTH_2512,
+            height=HEIGHT_2512,
+            num_inference_steps=NUM_INFERENCE_STEPS_2512,
+            true_cfg_scale=TRUE_CFG_SCALE_2512,
+            generator=generator,
+        )
+        output_image = result.images[0].convert("RGB")
+        output_image.save(output_path)
+        return output_image
+    finally:
+        if pipe is not None and hasattr(pipe, "maybe_free_model_hooks"):
+            pipe.maybe_free_model_hooks()
+        del pipe
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.accelerator.empty_cache()
+        run_post_test_cleanup()
+
+
 @pytest.mark.benchmark
 @hardware_test(res={"cuda": "H100"}, num_cards=1)
 def test_qwen_image_matches_diffusers(accuracy_artifact_root: Path) -> None:
@@ -120,3 +208,27 @@ def test_qwen_image_matches_diffusers(accuracy_artifact_root: Path) -> None:
         ssim_threshold=SSIM_THRESHOLD,
         psnr_threshold=PSNR_THRESHOLD,
     )
+
+
+@pytest.mark.benchmark
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+def test_qwen_image_2512_matches_diffusers_pixelwise(accuracy_artifact_root: Path) -> None:
+    model = _model_2512_name()
+    output_dir = model_output_dir(accuracy_artifact_root, MODEL_2512_ID)
+    vllm_output_path = output_dir / "vllm_omni_2512.png"
+    diffusers_output_path = output_dir / "diffusers_2512.png"
+
+    vllm_output = _run_vllm_omni_qwen_image_2512(model=model, output_path=vllm_output_path)
+    diffusers_output = _run_diffusers_qwen_image_2512(model=model, output_path=diffusers_output_path)
+
+    print(f"{MODEL_2512_ID} generated images:")
+    print(f"  vllm_omni: {vllm_output_path}")
+    print(f"  diffusers: {diffusers_output_path}")
+
+    assert_images_pixel_close(
+        model_name=MODEL_2512_ID,
+        vllm_image=vllm_output,
+        diffusers_image=diffusers_output,
+        mean_threshold=MEAN_ABS_DIFF_THRESHOLD_2512,
+        p99_threshold=P99_ABS_DIFF_THRESHOLD_2512,
+    )

tests/e2e/accuracy/test_qwen_image_edit.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import gc
+import os
 from pathlib import Path
 
 import pytest
@@ -10,7 +11,7 @@ from diffusers import QwenImageEditPipeline, QwenImageEditPlusPipeline
 from PIL import Image
 
 from benchmarks.accuracy.common import decode_base64_image, pil_to_png_bytes
-from tests.e2e.accuracy.helpers import assert_similarity, model_output_dir
+from tests.e2e.accuracy.helpers import assert_images_pixel_close, assert_similarity, model_output_dir
 from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup
 from tests.helpers.mark import hardware_test
 from tests.helpers.runtime import OmniServer
@@ -28,12 +29,43 @@ SEED = 42
 SSIM_THRESHOLD = 0.94
 PSNR_THRESHOLD = 28.0
 
+EDIT_2511_MODEL = "Qwen/Qwen-Image-Edit-2511"
+EDIT_2511_MODEL_ENV_VAR = "QWEN_IMAGE_EDIT_2511_MODEL"
+EDIT_2511_WIDTH = 1190
+EDIT_2511_HEIGHT = 2032
+EDIT_2511_NUM_INFERENCE_STEPS = 40
+EDIT_2511_TRUE_CFG_SCALE = 4.0
+EDIT_2511_GUIDANCE_SCALE = 1.0
+EDIT_2511_SEED = 42
+EDIT_2511_MEAN_ABS_DIFF_THRESHOLD = 3e-2
+EDIT_2511_P99_ABS_DIFF_THRESHOLD = 4e-1
+EDIT_2511_NEGATIVE_PROMPT = " "
+EDIT_2511_PROMPT = (
+    "将第二张图中人脸/猫狗脸转换为3D卡通形象，质量要求：毛发纹理自然电影级色彩校准，注意保留原图的外貌、"
+    "肤色、发型特征，仅替换第一张财神爷/猫狗财神的用绿框框住的头部区域，耳朵样式和原图的3D卡通形象一致， "
+    "（帽子紧贴头顶，未贴合部分可调整帽子摆放角度达到完全贴合）， 头部与身体过渡必须自然，无拼接痕迹，"
+    "无额外元素出现在头部四周，纯白背景，不要出现绿色圆圈"
+)
+EDIT_2511_ASSET_DIR = Path(__file__).resolve().parent / "assets"
+EDIT_2511_IMAGE_PATHS = [
+    EDIT_2511_ASSET_DIR / "qwen_image_edit_2511_test1.png",
+    EDIT_2511_ASSET_DIR / "qwen_image_edit_2511_test2.png",
+]
+
 PROMPT_SINGLE_IMAGE = "The input is a 2D cartoon bear mascot. Restyle it into a painterly oil artwork with warm colors while preserving the main structure."
 PROMPT_MULTIPLE_IMAGE = "Put the cartoon bear mascot and the furry rabbit into one coherent scene with a painterly oil artwork style and consistent lighting."
 NEGATIVE_PROMPT = "low quality, blurry, artifacts, distortion"
 SERVER_ARGS = ["--num-gpus", "1", "--stage-init-timeout", "300", "--init-timeout", "900"]
 
 
+def _edit_2511_model_name() -> str:
+    return os.environ.get(EDIT_2511_MODEL_ENV_VAR, EDIT_2511_MODEL)
+
+
+def _local_files_only(model: str) -> bool:
+    return Path(model).exists()
+
+
 def _run_vllm_omni_image_edit(
     *,
     omni_server: OmniServer,
@@ -79,17 +111,15 @@ def _run_diffusers_image_edit(
 ) -> Image.Image:
     run_pre_test_cleanup()
     pipe: QwenImageEditPipeline | QwenImageEditPlusPipeline | None = None
-    device = torch.device("cuda:0")
-    torch.cuda.set_device(device)
     try:
         images = input_images[0] if len(input_images) == 1 else input_images
         pipe = pipeline_class.from_pretrained(
             model,
             torch_dtype=torch.bfloat16,
             trust_remote_code=True,
-        ).to(device)
+        ).to("cuda")
         pipe.set_progress_bar_config(disable=False)
-        generator = torch.Generator(device=device).manual_seed(SEED)
+        generator = torch.Generator(device="cuda").manual_seed(SEED)
         result = pipe(  # pyright: ignore[reportCallIssue]
             prompt=prompt,
             image=images,
@@ -113,6 +143,88 @@ def _run_diffusers_image_edit(
         run_post_test_cleanup()
 
 
+def _load_edit_2511_input_images() -> list[Image.Image]:
+    return [Image.open(path).convert("RGB") for path in EDIT_2511_IMAGE_PATHS]
+
+
+def _run_vllm_omni_image_edit_2511(
+    *,
+    model: str,
+    input_images: list[Image.Image],
+    output_path: Path,
+) -> Image.Image:
+    with OmniServer(model=model, serve_args=SERVER_ARGS) as server:
+        response = requests.post(
+            f"http://{server.host}:{server.port}/v1/images/edits",
+            data={
+                "model": server.model,
+                "prompt": EDIT_2511_PROMPT,
+                "size": f"{EDIT_2511_WIDTH}x{EDIT_2511_HEIGHT}",
+                "n": 1,
+                "response_format": "b64_json",
+                "negative_prompt": EDIT_2511_NEGATIVE_PROMPT,
+                "num_inference_steps": EDIT_2511_NUM_INFERENCE_STEPS,
+                "true_cfg_scale": EDIT_2511_TRUE_CFG_SCALE,
+                "guidance_scale": EDIT_2511_GUIDANCE_SCALE,
+                "seed": EDIT_2511_SEED,
+            },
+            files=[
+                ("image", (f"image_{index}.png", pil_to_png_bytes(image), "image/png"))
+                for index, image in enumerate(input_images)
+            ],
+            timeout=1200,
+        )
+        response.raise_for_status()
+        payload = response.json()
+        assert len(payload["data"]) == 1
+        image = decode_base64_image(payload["data"][0]["b64_json"])
+        image.load()
+        image.save(output_path)
+        return image
+
+
+def _run_diffusers_image_edit_2511(
+    *,
+    model: str,
+    input_images: list[Image.Image],
+    output_path: Path,
+) -> Image.Image:
+    run_pre_test_cleanup()
+    pipe: QwenImageEditPlusPipeline | None = None
+    try:
+        pipe = QwenImageEditPlusPipeline.from_pretrained(
+            model,
+            torch_dtype=torch.bfloat16,
+            trust_remote_code=True,
+            local_files_only=_local_files_only(model),
+        ).to("cuda")
+        pipe.set_progress_bar_config(disable=False)
+        generator = torch.Generator(device="cuda").manual_seed(EDIT_2511_SEED)
+        result = pipe(  # pyright: ignore[reportCallIssue]
+            image=input_images,
+            prompt=EDIT_2511_PROMPT,
+            generator=generator,
+            true_cfg_scale=EDIT_2511_TRUE_CFG_SCALE,
+            negative_prompt=EDIT_2511_NEGATIVE_PROMPT,
+            num_inference_steps=EDIT_2511_NUM_INFERENCE_STEPS,
+            guidance_scale=EDIT_2511_GUIDANCE_SCALE,
+            num_images_per_prompt=1,
+            width=EDIT_2511_WIDTH,
+            height=EDIT_2511_HEIGHT,
+        )
+        output_image = result.images[0].convert("RGB")  # pyright: ignore[reportAttributeAccessIssue]
+        output_image.save(output_path)
+        return output_image
+    finally:
+        if pipe is not None and hasattr(pipe, "maybe_free_model_hooks"):
+            pipe.maybe_free_model_hooks()
+        del pipe
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.accelerator.empty_cache()
+        run_post_test_cleanup()
+
+
 def _vllm_omni_output_single_image(
     accuracy_artifact_root: Path,
     qwen_bear_image: Image.Image,
@@ -226,3 +338,43 @@ def test_qwen_image_edit_multiple_matches_diffusers(
         ssim_threshold=SSIM_THRESHOLD,
         psnr_threshold=PSNR_THRESHOLD,
     )
+
+
+@pytest.mark.benchmark
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+def test_qwen_image_edit_2511_matches_diffusers_pixelwise(accuracy_artifact_root: Path) -> None:
+    model = _edit_2511_model_name()
+    output_dir = model_output_dir(accuracy_artifact_root, EDIT_2511_MODEL)
+    input_images = _load_edit_2511_input_images()
+    input_paths: list[Path] = []
+    for index, image in enumerate(input_images, start=1):
+        input_path = output_dir / f"edit_2511_input_{index}.png"
+        image.save(input_path)
+        input_paths.append(input_path)
+    vllm_output_path = output_dir / "vllm_omni_edit_2511.png"
+    diffusers_output_path = output_dir / "diffusers_edit_2511.png"
+
+    vllm_image = _run_vllm_omni_image_edit_2511(
+        model=model,
+        input_images=input_images,
+        output_path=vllm_output_path,
+    )
+    diffusers_image = _run_diffusers_image_edit_2511(
+        model=model,
+        input_images=input_images,
+        output_path=diffusers_output_path,
+    )
+
+    print(f"{EDIT_2511_MODEL} generated images:")
+    for index, input_path in enumerate(input_paths, start=1):
+        print(f"  input_{index}: {input_path}")
+    print(f"  vllm_omni: {vllm_output_path}")
+    print(f"  diffusers: {diffusers_output_path}")
+
+    assert_images_pixel_close(
+        model_name=EDIT_2511_MODEL,
+        vllm_image=vllm_image,
+        diffusers_image=diffusers_image,
+        mean_threshold=EDIT_2511_MEAN_ABS_DIFF_THRESHOLD,
+        p99_threshold=EDIT_2511_P99_ABS_DIFF_THRESHOLD,
+    )

tests/e2e/offline_inference/test_cosyvoice3_expansion.py

@@ -27,9 +27,9 @@ from tests.helpers.mark import hardware_test
 from tests.helpers.media import get_asset_path
 from tests.helpers.runtime import OmniRunner
 from tests.helpers.stage_config import get_deploy_config_path
-from vllm_omni.model_executor.models.cosyvoice3.config import CosyVoice3Config
 from vllm_omni.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokenizer
 from vllm_omni.outputs import OmniRequestOutput
+from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config
 
 MODEL = "FunAudioLLM/Fun-CosyVoice3-0.5B-2512"
 MODEL_DIR_ENV = "VLLM_OMNI_COSYVOICE3_MODEL_DIR"

tests/e2e/offline_inference/test_moss_tts_nano_expansion.py

@@ -106,18 +106,21 @@ def _build_request(
 
 def _collect_audio(omni: Omni, request: dict) -> tuple[torch.Tensor, int]:
     """Run a single request and return (waveform, sample_rate)."""
-    # Omni.generate returns list[OmniRequestOutput] by default (py_generator=False).
+    for stage_outputs in omni.generate(request, DEFAULT_SAMPLING):
-    for omni_out in omni.generate(request, DEFAULT_SAMPLING):
+        req_output = stage_outputs.request_output
-        mm = omni_out.multimodal_output
+        if req_output is not None:
-        assert mm is not None, "Expected multimodal_output to be non-None"
+            mm = req_output.outputs[0].multimodal_output
-        audio = mm.get("audio")
+            assert mm is not None, "Expected multimodal_output to be non-None"
-        sr = mm.get("sr")
+            audio = mm.get("audio")
-        assert audio is not None, "Expected 'audio' key in multimodal_output"
+            sr = mm.get("sr")
-        assert isinstance(audio, torch.Tensor), f"audio should be Tensor, got {type(audio)}"
+            assert audio is not None, "Expected 'audio' key in multimodal_output"
-        return audio.cpu(), int(sr.item()) if sr is not None else SAMPLE_RATE
+            assert isinstance(audio, torch.Tensor), f"audio should be Tensor, got {type(audio)}"
+            return audio.cpu(), int(sr.item()) if sr is not None else SAMPLE_RATE
     raise AssertionError("No stage outputs received")
 
 
+@pytest.mark.advanced_model
+@pytest.mark.omni
 @hardware_test(res={"cuda": "L4"})
 def test_moss_tts_nano_english(omni_runner: OmniRunner, ref_audio_path) -> None:
     """English TTS produces non-empty 48 kHz stereo audio."""
@@ -129,6 +132,8 @@ def test_moss_tts_nano_english(omni_runner: OmniRunner, ref_audio_path) -> None:
     assert not torch.all(audio == 0), "Audio should not be all-zeros (silence)"
 
 
+@pytest.mark.advanced_model
+@pytest.mark.omni
 @hardware_test(res={"cuda": "L4"})
 def test_moss_tts_nano_chinese(omni_runner: OmniRunner, ref_audio_path) -> None:
     """Chinese TTS produces non-empty audio."""
@@ -140,6 +145,8 @@ def test_moss_tts_nano_chinese(omni_runner: OmniRunner, ref_audio_path) -> None:
     assert not torch.all(audio == 0)
 
 
+@pytest.mark.advanced_model
+@pytest.mark.omni
 @hardware_test(res={"cuda": "L4"})
 def test_moss_tts_nano_deterministic(omni_runner: OmniRunner, ref_audio_path) -> None:
     """Same seed produces identical waveforms."""
@@ -151,6 +158,8 @@ def test_moss_tts_nano_deterministic(omni_runner: OmniRunner, ref_audio_path) ->
     assert torch.allclose(audio1, audio2, atol=1e-4), "Waveforms should match with same seed"
 
 
+@pytest.mark.advanced_model
+@pytest.mark.omni
 @hardware_test(res={"cuda": "L4"})
 def test_moss_tts_nano_batch(omni_runner: OmniRunner, ref_audio_path) -> None:
     """Batch of two requests returns audio for each."""
@@ -159,11 +168,13 @@ def test_moss_tts_nano_batch(omni_runner: OmniRunner, ref_audio_path) -> None:
         _build_request("Second request.", ref_audio_path),
     ]
     results = []
-    # Single-stage pipeline: one SamplingParams object is broadcast to all prompts.
+    # Single-stage model (num_stages=1): one sampling param for all requests.
-    for omni_out in omni_runner.omni.generate(requests, DEFAULT_SAMPLING):
+    for stage_outputs in omni_runner.omni.generate(requests, [DEFAULT_SAMPLING]):
-        mm = omni_out.multimodal_output
+        req_output = stage_outputs.request_output
-        assert mm is not None
+        if req_output is not None:
-        results.append(mm["audio"].cpu())
+            mm = req_output.outputs[0].multimodal_output
+            assert mm is not None
+            results.append(mm["audio"].cpu())
 
     assert len(results) == 2, f"Expected 2 outputs, got {len(results)}"
     for i, audio in enumerate(results):

tests/e2e/offline_inference/test_stable_audio_expansion.py

@@ -15,6 +15,7 @@ import numpy as np
 import pytest
 import torch
 
+from tests.helpers import skip_if_gated_repo_inaccessible
 from tests.helpers.assertions import assert_audio_valid
 from tests.helpers.mark import hardware_test
 from vllm_omni import Omni
@@ -24,6 +25,8 @@ from vllm_omni.platforms import current_omni_platform
 
 pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
 
+_MODEL_REPO = "stabilityai/stable-audio-open-1.0"
+
 _SAMPLE_RATE = 44100
 _CLIP_DURATION_S = 2.0
 
@@ -77,6 +80,7 @@ def test_stable_audio_quantization_and_teacache() -> None:
 
     CI should provide ``HF_TOKEN`` if the checkpoint is gated.
     """
+    skip_if_gated_repo_inaccessible(_MODEL_REPO)
     # ``model_class_name`` must be passed explicitly: the default-stage-cfg
     # factory in ``async_omni_engine.py`` reads it out of ``kwargs`` when
     # deciding ``final_output_type`` (#2077), and at construction time the

tests/e2e/online_serving/test_qwen3_omni_expansion.py

@@ -44,7 +44,7 @@ def get_batch_token_config(default_path):
     return modify_stage_config(
         default_path,
         updates={
-            "stages": {1: {"max_num_batched_tokens": 64}},
+            "stages": {0: {"max_num_batched_tokens": 64}, 1: {"max_num_batched_tokens": 64}},
         },
     )
 
@@ -95,7 +95,12 @@ test_params = [
 
 test_token_params = [
     pytest.param(
-        OmniServerParams(model=model, stage_config_path=get_batch_token_config(default_path), use_stage_cli=True),
+        OmniServerParams(
+            model=model,
+            stage_config_path=get_batch_token_config(default_path),
+            use_stage_cli=True,
+            server_args=["--async-chunk"],
+        ),
         id="batch_token_64",
     )
 ]

tests/engine/test_async_omni_engine_input.py

@@ -5,13 +5,14 @@ from vllm.v1.engine import EngineCoreRequest
 
 from vllm_omni.engine import OmniEngineCoreRequest
 from vllm_omni.engine.async_omni_engine import AsyncOmniEngine, StageRuntimeInfo
+from vllm_omni.engine.stage_pool import StagePool
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
-def _make_engine_core_request() -> EngineCoreRequest:
+def _make_engine_core_request(request_id: str = "req-1") -> EngineCoreRequest:
     return EngineCoreRequest(
-        request_id="req-1",
+        request_id=request_id,
         prompt_token_ids=[1, 1, 1],
         mm_features=None,
         sampling_params=SamplingParams(max_tokens=8),
@@ -88,3 +89,86 @@ def test_build_add_request_message_with_resumable_streaming(mocker: MockerFixtur
     assert msg.type == "streaming_update"
     input_processor.process_inputs.assert_called_once()
     assert input_processor.process_inputs.call_args.kwargs["resumable"] is True
+
+
+class _FakeStageClient:
+    stage_type = "llm"
+    final_output = False
+
+
+def test_build_add_request_message_scopes_mm_uuids_to_selected_stage0_replica(mocker: MockerFixture):
+    engine = object.__new__(AsyncOmniEngine)
+    params = SamplingParams(max_tokens=8)
+    engine.model = "test-model"
+    engine.default_sampling_params_list = [params]
+    engine.stage_metadata = [StageRuntimeInfo(final_output=False, final_output_type=None, stage_type="llm")]
+    engine.supported_tasks = ("generate",)
+    engine.stage_pools = [StagePool(0, [_FakeStageClient(), _FakeStageClient()])]
+
+    seen_uuids: list[str] = []
+
+    def process_inputs(**kwargs):
+        prompt = kwargs["prompt"]
+        seen_uuids.append(prompt["multi_modal_uuids"]["image"][0])
+        return _make_engine_core_request(kwargs["request_id"])
+
+    input_processor = mocker.Mock()
+    input_processor.process_inputs.side_effect = process_inputs
+    engine.input_processor = input_processor
+
+    for request_id in ("req-1", "req-2"):
+        engine._build_add_request_message(
+            request_id=request_id,
+            prompt={
+                "prompt": "describe",
+                "multi_modal_data": {"image": "same-image"},
+            },
+            sampling_params_list=[params],
+            final_stage_id=0,
+        )
+
+    assert seen_uuids[0].startswith("stage0:rep0:")
+    assert seen_uuids[1].startswith("stage0:rep1:")
+    assert seen_uuids[0].removeprefix("stage0:rep0:") == seen_uuids[1].removeprefix("stage0:rep1:")
+
+
+def test_stage_pool_replica_count_falls_back_to_clients():
+    class PoolWithoutLiveNumReplicas:
+        clients = [object(), None, object()]
+
+    assert AsyncOmniEngine._stage_pool_replica_count(PoolWithoutLiveNumReplicas()) == 2
+
+
+def test_stage_pool_is_distributed_falls_back_to_hub():
+    class PoolWithoutIsDistributed:
+        _hub = object()
+
+    assert AsyncOmniEngine._stage_pool_is_distributed(PoolWithoutIsDistributed()) is True
+
+
+def test_build_add_request_message_releases_preselected_replica_on_preprocess_error(mocker: MockerFixture):
+    engine = object.__new__(AsyncOmniEngine)
+    params = SamplingParams(max_tokens=8)
+    engine.model = "test-model"
+    engine.default_sampling_params_list = [params]
+    engine.stage_metadata = [StageRuntimeInfo(final_output=False, final_output_type=None, stage_type="llm")]
+    engine.supported_tasks = ("generate",)
+    stage_pool = StagePool(0, [_FakeStageClient(), _FakeStageClient()])
+    engine.stage_pools = [stage_pool]
+
+    input_processor = mocker.Mock()
+    input_processor.process_inputs.side_effect = RuntimeError("boom")
+    engine.input_processor = input_processor
+
+    with pytest.raises(RuntimeError, match="boom"):
+        engine._build_add_request_message(
+            request_id="req-error",
+            prompt={
+                "prompt": "describe",
+                "multi_modal_data": {"image": "same-image"},
+            },
+            sampling_params_list=[params],
+            final_stage_id=0,
+        )
+
+    assert stage_pool.get_bound_replica_id("req-error") is None

tests/engine/test_orchestrator.py

@@ -79,7 +79,7 @@ class FakeStageClient:
         self._diffusion_outputs = queue.Queue()
 
     # Orchestrator-facing interface.
-    async def add_request_async(self, *args, **_kwargs) -> None:
+    async def add_request_async(self, *args, **kwargs) -> None:
         self.add_request_calls.append(args)
 
     async def get_output_async(self):
@@ -448,7 +448,7 @@ def test_stage_engine_core_client_shutdown_cleans_children_if_base_shutdown_fail
             assert recursive
             return [fake_child]
 
-    def fail_base_shutdown(self):
+    def fail_base_shutdown(self, **kwargs):
         raise RuntimeError("base shutdown failed")
 
     monkeypatch.setattr(psutil, "Process", FakePsutilProcess)
@@ -480,7 +480,7 @@ def test_stage_engine_core_client_shutdown_kills_stubborn_children(monkeypatch):
 
     monkeypatch.setattr(psutil, "Process", FakePsutilProcess)
     monkeypatch.setattr(psutil, "wait_procs", lambda procs, timeout: ([], list(procs)))
-    monkeypatch.setattr(AsyncMPClient, "shutdown", lambda self: None)
+    monkeypatch.setattr(AsyncMPClient, "shutdown", lambda self, **kwargs: None)
 
     client = object.__new__(StageEngineCoreClient)
     client._proc = fake_proc
@@ -994,7 +994,7 @@ async def test_handle_streaming_update_passes_prompt_text_to_stage_pool() -> Non
 @pytest.mark.asyncio
 async def test_stage_pool_submit_initial_rolls_back_output_processor_when_client_submit_fails() -> None:
     class FailingStageClient(FakeStageClient):
-        async def add_request_async(self, *args, **_kwargs) -> None:
+        async def add_request_async(self, *args, **kwargs) -> None:
             raise RuntimeError("submit failed")
 
     class TrackingOutputProcessor(FakeOutputProcessor):

tests/entrypoints/openai_api/test_image_server.py

@@ -1349,8 +1349,16 @@ def test_image_edit_parameter_default(async_omni_test_client):
     engine = async_omni_test_client.app.state.engine_client
     captured_sampling_params = engine.captured_sampling_params_list[-1]
 
-    assert captured_sampling_params.width == 24
+    # size="auto" on multi-stage pipelines deliberately leaves the diffusion
-    assert captured_sampling_params.height == 16
+    # stages sampling_params width/height unset so AR-driven pipelines (e.g.
+    # HunyuanImage-3.0) can let ar2diffusion override the final bucket from
+    # the AR-predicted ratio token; see
+    # test_image_edits_size_auto_preserves_bridge_size for the contract.
+    # Single-stage diffusion (test_image_edit_parameter_default_single_stage)
+    # still pins width/height to the input image size via api_servers
+    # gen_params, which is unchanged.
+    assert captured_sampling_params.width is None
+    assert captured_sampling_params.height is None
     assert captured_sampling_params.num_outputs_per_prompt == 1
     assert captured_sampling_params.num_inference_steps == 4
     assert captured_sampling_params.guidance_scale == 7.5
@@ -1649,3 +1657,59 @@ def test_extract_images_from_result():
     assert len(images) == 1
     assert isinstance(images[0], Image.Image)
     assert images[0].size == (32, 32)
+
+
+def test_image_edits_size_auto_preserves_bridge_size(async_omni_stage_configs_only_client):
+    """size=auto must NOT pin the diffusion stage sampling_params.height/width.
+
+    Regression: prior to the fix, edit_images resolved size=auto to the
+    first input image dimensions and forwarded them through gen_params +
+    extra_body to the diffusion stages sampling_params. AR-driven
+    pipelines (e.g. HunyuanImage-3.0) rely on ar2diffusions
+    bridge to override the final bucket via the AR-predicted ratio token,
+    and the DiT pre_process_func only fills sampling_params from the
+    bridge value when sampling_params.width is None (see
+    pipeline_hunyuan_image3.py:290). Non-None width from the input image
+    silently suppressed the AR decision, producing the wrong bucket
+    (e.g. 1024x1024 square instead of the AR-decided 1280x720 landscape
+    for multi-image fusion).
+
+    Cross-pins the multi-image fix at the API level: 2 reference images
+    with bot_task=think must produce 2 <img> placeholders in the captured
+    AR prompt (build_prompt called with num_images=2).
+    """
+    img_a = make_test_image_bytes((32, 32))
+    img_b = make_test_image_bytes((128, 64))
+    response = async_omni_stage_configs_only_client.post(
+        "/v1/images/edits",
+        files=[("image", img_a), ("image", img_b)],
+        data={
+            "prompt": "fuse",
+            "size": "auto",
+            "bot_task": "think",
+        },
+    )
+    assert response.status_code == 200, response.text
+
+    engine = async_omni_stage_configs_only_client.app.state.engine_client
+    captured = engine.captured_sampling_params_list
+    assert captured is not None
+    assert len(captured) == 2
+
+    diffusion_params = captured[1]
+    assert diffusion_params.height is None, (
+        f"size=auto leaked into diffusion sampling_params.height={diffusion_params.height}; "
+        "must stay None so AR-driven pipelines can apply the bridges decision."
+    )
+    assert diffusion_params.width is None, (
+        f"size=auto leaked into diffusion sampling_params.width={diffusion_params.width}; "
+        "must stay None so AR-driven pipelines can apply the bridges decision."
+    )
+
+    KEY = "prompt"
+    IMG = "<img>"
+    captured_prompt = engine.captured_prompt
+    if isinstance(captured_prompt, dict) and isinstance(captured_prompt.get("prompt"), str):
+        assert captured_prompt["prompt"].count("<img>") == 2, (
+            f"N=2 reference images must emit 2 <img> placeholders in AR prompt; got {captured_prompt[KEY].count(IMG)} -- prompt: {captured_prompt[KEY]!r}"
+        )

tests/entrypoints/openai_api/test_serving_audio_generate.py

@@ -389,7 +389,7 @@ class TestErrorHandling:
 
         async def empty_gen(*args, **kwargs):
             return
-            yield  # noqa: unreachable – makes this an async generator
+            yield  # unreachable – makes this an async generator
 
         engine.generate = MagicMock(side_effect=empty_gen)
         server = _make_server(engine_client=engine)
@@ -448,7 +448,7 @@ class TestErrorHandling:
 
         async def gen_value_error(*args, **kwargs):
             raise ValueError("bad value")
-            yield  # noqa: unreachable
+            yield  # unreachable
 
         engine.generate = MagicMock(side_effect=gen_value_error)
         server = _make_server(engine_client=engine)
@@ -463,7 +463,7 @@ class TestErrorHandling:
 
         async def gen_runtime_error(*args, **kwargs):
             raise RuntimeError("something went wrong")
-            yield  # noqa: unreachable
+            yield  # unreachable
 
         engine.generate = MagicMock(side_effect=gen_runtime_error)
         server = _make_server(engine_client=engine)

tests/entrypoints/openai_api/test_serving_chat_multistage_generation.py

@@ -91,3 +91,275 @@ def test_build_multistage_generation_inputs_applies_stage_specific_overrides(ser
     assert engine.default_sampling_params_list[1].lora_request is None
     assert engine.default_sampling_params_list[2].resolution == 640
     assert engine.default_sampling_params_list[2].lora_request is None
+
+
+def test_build_multistage_generation_inputs_multi_image_emits_n_img_placeholders(serving_chat):
+    """N reference images with bot_task set must emit N <img> placeholders.
+
+    Regression: prior to the multi-image online fix, build_prompt was
+    called without num_images, defaulting to 1. A 2-image edit request
+    would only get a single <img> placeholder in the AR prompt; vLLMs
+    _process_multimodal then raised
+    AssertionError(Failed to apply prompt replacement for mm_items[image][1])
+    when trying to replace the second image (no placeholder left for it).
+
+    Pins the contract that build_prompt() is invoked with the actual image
+    count so multi-image IT2I is wired correctly through the online
+    /v1/images/edits path.
+    """
+    from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
+
+    engine = SimpleNamespace(
+        stage_configs=[
+            SimpleNamespace(stage_type="llm", is_comprehension=True),
+            SimpleNamespace(stage_type="diffusion", is_comprehension=False),
+        ],
+        default_sampling_params_list=[
+            SamplingParams(temperature=0.0),
+            OmniDiffusionSamplingParams(),
+        ],
+    )
+    IMG = "<img>"
+    images = [Image.new("RGB", (32, 32), color="red") for _ in range(3)]
+
+    for n in (1, 2, 3):
+        engine_prompt, _ = OmniOpenAIServingChat._build_multistage_generation_inputs(
+            serving_chat,
+            engine=engine,
+            prompt="edit me",
+            extra_body={"bot_task": "think"},
+            reference_images=images[:n],
+            gen_params=OmniDiffusionSamplingParams(),
+        )
+        prompt_str = engine_prompt["prompt"]
+        assert prompt_str.count("<img>") == n, (
+            f"N={n}: expected {n} <img> placeholders, got {prompt_str.count(IMG)} -- prompt: {prompt_str!r}"
+        )
+
+
+def test_build_multistage_generation_inputs_tokenizer_path_emits_prompt_token_ids(serving_chat):
+    """When a tokenizer is provided, the helper must emit HF byte-for-byte
+    prompt_token_ids and forward use_system_prompt to the engine prompt.
+
+    Regression: prior to the HF-byte-equivalent fix, online IT2I always
+    passed the prompt as a single string. The engine then BPE-merged across
+    chat-template segment boundaries (e.g. user_prompt-ending punctuation
+    plus the trailing \n\n before \"Assistant: \") producing a token
+    sequence that differs from HF apply_chat_template / offline
+    end2end.py. AR generated different cot_text (706 tokens / 1190 chars
+    vs offline 661 / 1118 for the same inputs) and DiT produced a visually
+    different image (yin-yang on brushed-metal vs three-blue swirl on
+    canvas) under the same seed.
+
+    Pins:
+      1. engine_prompt[\"prompt_token_ids\"] is set when tokenizer is passed.
+      2. engine_prompt[\"prompt\"] stays as the raw user prompt -- the DiT
+         side rebuilds its own system prefix via use_system_prompt.
+      3. engine_prompt[\"use_system_prompt\"] == \"en_unified\" so
+         ar2diffusion forwards the matching system prompt to DiT.
+      4. N reference images emit N <img> token ids in the AR sequence.
+    """
+    from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
+
+    # Minimal FakeTokenizer mirroring tests/diffusion/.../test_hunyuan_image3_it2i_multi_image.py
+    class FakeTokenizer:
+        SPECIAL = {
+            "<|startoftext|>": 1,
+            "<img>": 2,
+            "<think>": 3,
+            "<recaption>": 4,
+        }
+
+        def convert_tokens_to_ids(self, tok: str) -> int:
+            return self.SPECIAL.get(tok, 0)
+
+        def encode(self, text: str, add_special_tokens: bool = False) -> list[int]:
+            return list(range(100, 100 + len(text)))
+
+    engine = SimpleNamespace(
+        stage_configs=[
+            SimpleNamespace(stage_type="llm", is_comprehension=True),
+            SimpleNamespace(stage_type="diffusion", is_comprehension=False),
+        ],
+        default_sampling_params_list=[
+            SamplingParams(temperature=0.0),
+            OmniDiffusionSamplingParams(),
+        ],
+    )
+    PROMPT_KEY = "prompt"
+    USP_KEY = "use_system_prompt"
+    images = [Image.new("RGB", (32, 32), color="red") for _ in range(3)]
+
+    for n in (1, 2, 3):
+        tok = FakeTokenizer()
+        engine_prompt, _ = OmniOpenAIServingChat._build_multistage_generation_inputs(
+            serving_chat,
+            engine=engine,
+            prompt="edit me",
+            extra_body={"bot_task": "think"},
+            reference_images=images[:n],
+            gen_params=OmniDiffusionSamplingParams(),
+            tokenizer=tok,
+        )
+        # (1) prompt_token_ids must be set and non-empty
+        assert "prompt_token_ids" in engine_prompt, f"N={n}: prompt_token_ids missing"
+        token_ids = engine_prompt["prompt_token_ids"]
+        assert isinstance(token_ids, list) and len(token_ids) > 0, f"N={n}: prompt_token_ids empty"
+        # (2) raw prompt preserved (DiT bridge needs raw user text)
+        assert engine_prompt["prompt"] == "edit me", (
+            f"N={n}: prompt must stay raw user text, got {engine_prompt[PROMPT_KEY]!r}"
+        )
+        # (3) use_system_prompt forwarded for ar2diffusion bridge
+        assert engine_prompt.get("use_system_prompt") == "en_unified", (
+            f"N={n}: use_system_prompt must be en_unified, got {engine_prompt.get(USP_KEY)!r}"
+        )
+        # (4) N <img> token ids (id=2 in FakeTokenizer)
+        img_count = token_ids.count(2)
+        assert img_count == n, f"N={n}: expected {n} <img> token ids in prompt_token_ids, got {img_count}"
+
+
+def test_build_multistage_generation_inputs_bot_task_semantic_changes_trigger_and_sys(serving_chat):
+    """Passing bot_task=think_recaption (vs default "think") must flip the
+    resolved sys_type to en_think_recaption (and trigger tag is still
+    <think>). Pins that the API actually plumbs the bot_task semantic
+    through to build_prompt rather than ignoring it.
+    """
+    from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
+
+    engine = SimpleNamespace(
+        stage_configs=[
+            SimpleNamespace(stage_type="llm", is_comprehension=True),
+            SimpleNamespace(stage_type="diffusion", is_comprehension=False),
+        ],
+        default_sampling_params_list=[
+            SamplingParams(temperature=0.0),
+            OmniDiffusionSamplingParams(),
+        ],
+    )
+    images = [Image.new("RGB", (32, 32), color="red")]
+
+    # Default bot_task (think) -> en_unified system prompt baked into the
+    # legacy string path. Use legacy build_prompt (tokenizer=None) so the
+    # rendered prompt is a string we can grep.
+    think_prompt, _ = OmniOpenAIServingChat._build_multistage_generation_inputs(
+        serving_chat,
+        engine=engine,
+        prompt="edit me",
+        extra_body={"task": "it2i", "bot_task": "think"},
+        reference_images=images,
+        gen_params=OmniDiffusionSamplingParams(),
+    )
+    # think_recaption -> en_think_recaption system prompt (different content).
+    recap_prompt, _ = OmniOpenAIServingChat._build_multistage_generation_inputs(
+        serving_chat,
+        engine=engine,
+        prompt="edit me",
+        extra_body={"task": "it2i", "bot_task": "think_recaption"},
+        reference_images=images,
+        gen_params=OmniDiffusionSamplingParams(),
+    )
+    assert think_prompt["prompt"] != recap_prompt["prompt"], (
+        "bot_task semantic must change the rendered system prompt: "
+        f"think/think_recaption produced identical strings (len={len(think_prompt['prompt'])})"
+    )
+
+
+def test_build_multistage_generation_inputs_sys_type_override(serving_chat):
+    """Caller-supplied sys_type must override the bot_task-derived default.
+    Mirrors offline `--bot-task think_recaption --sys-type en_unified`
+    where the user wants think_recaptions trigger but the unified system
+    prompt body.
+    """
+    from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
+
+    engine = SimpleNamespace(
+        stage_configs=[
+            SimpleNamespace(stage_type="llm", is_comprehension=True),
+            SimpleNamespace(stage_type="diffusion", is_comprehension=False),
+        ],
+        default_sampling_params_list=[
+            SamplingParams(temperature=0.0),
+            OmniDiffusionSamplingParams(),
+        ],
+    )
+    images = [Image.new("RGB", (32, 32), color="red")]
+
+    # think_recaption defaults sys_type -> en_think_recaption.
+    default_sys, _ = OmniOpenAIServingChat._build_multistage_generation_inputs(
+        serving_chat,
+        engine=engine,
+        prompt="edit me",
+        extra_body={"task": "it2i", "bot_task": "think_recaption"},
+        reference_images=images,
+        gen_params=OmniDiffusionSamplingParams(),
+    )
+    # sys_type=en_unified overrides -> same system body as bot_task=think.
+    overridden, _ = OmniOpenAIServingChat._build_multistage_generation_inputs(
+        serving_chat,
+        engine=engine,
+        prompt="edit me",
+        extra_body={"task": "it2i", "bot_task": "think_recaption", "sys_type": "en_unified"},
+        reference_images=images,
+        gen_params=OmniDiffusionSamplingParams(),
+    )
+    plain_think, _ = OmniOpenAIServingChat._build_multistage_generation_inputs(
+        serving_chat,
+        engine=engine,
+        prompt="edit me",
+        extra_body={"task": "it2i", "bot_task": "think"},
+        reference_images=images,
+        gen_params=OmniDiffusionSamplingParams(),
+    )
+
+    # Override must (a) differ from the no-override default, and (b) equal
+    # the prompt that bot_task=think produces (both end up with
+    # en_unified system body + <think> trigger).
+    assert overridden["prompt"] != default_sys["prompt"], (
+        "sys_type override must change the rendered prompt body vs the bot_task default"
+    )
+    assert overridden["prompt"] == plain_think["prompt"], (
+        "sys_type=en_unified + bot_task=think_recaption must produce the same prompt as "
+        "bot_task=think (both = en_unified system body + <think> trigger)"
+    )
+
+
+def test_build_multistage_generation_inputs_custom_system_prompt(serving_chat):
+    """`extra_body["system_prompt"]` must reach build_prompt as
+    `custom_system_prompt`, enabling sys_type="custom" callers to inject
+    a verbatim system body. Without this plumbing the sys_type="custom"
+    branch in get_system_prompt() returns None and silently drops the
+    user-supplied content.
+    """
+    from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
+
+    engine = SimpleNamespace(
+        stage_configs=[
+            SimpleNamespace(stage_type="llm", is_comprehension=True),
+            SimpleNamespace(stage_type="diffusion", is_comprehension=False),
+        ],
+        default_sampling_params_list=[
+            SamplingParams(temperature=0.0),
+            OmniDiffusionSamplingParams(),
+        ],
+    )
+    images = [Image.new("RGB", (32, 32), color="red")]
+
+    marker = "ZZZ_CUSTOM_SYSTEM_PROMPT_MARKER_ZZZ"
+
+    out, _ = OmniOpenAIServingChat._build_multistage_generation_inputs(
+        serving_chat,
+        engine=engine,
+        prompt="edit me",
+        extra_body={
+            "task": "it2i",
+            "bot_task": "think",
+            "sys_type": "custom",
+            "system_prompt": marker,
+        },
+        reference_images=images,
+        gen_params=OmniDiffusionSamplingParams(),
+    )
+    assert marker in out["prompt"], (
+        f"custom system_prompt content must reach the rendered prompt; "
+        f"marker {marker!r} not found in prompt of length {len(out['prompt'])}"
+    )

tests/entrypoints/openai_api/test_serving_chat_speaker.py

@@ -22,6 +22,7 @@ def serving_chat():
 
     instance = object.__new__(OmniOpenAIServingChat)
     instance._supported_speakers = None
+    instance.has_kv_connector = False
     return instance
 
 

tests/entrypoints/test_omni_entrypoints.py

@@ -943,6 +943,20 @@ def test_omni_base_errored_false_when_alive():
     assert base.errored is False
 
 
+def test_omni_base_is_running_false_when_stage_engine_dead():
+    base = _make_base()
+    base.engine.is_alive.return_value = True
+    base.engine.stage_clients = [SimpleNamespace(_engine_dead=True)]
+    assert base.is_running is False
+
+
+def test_omni_base_is_running_false_when_stage_resources_engine_dead():
+    base = _make_base()
+    base.engine.is_alive.return_value = True
+    base.engine.stage_clients = [SimpleNamespace(resources=SimpleNamespace(engine_dead=True))]
+    assert base.is_running is False
+
+
 def test_omni_base_errored_true_when_orchestrator_dead():
     base = _make_base()
     base.engine.is_alive.return_value = False

tests/helpers/__init__.py

@@ -6,3 +6,32 @@ refactor only inside the old monolithic ``conftest``; a greedy ``__init__``
 changes import order and can affect in-process Omni (``OmniRunner`` / offline
 e2e) vs subprocess-based ``OmniServer`` tests.
 """
+
+import pytest
+
+
+def skip_if_gated_repo_inaccessible(repo_id: str) -> None:
+    """Skip the test if a gated HuggingFace repo is not accessible.
+
+    Tries to download the model's config.json via ``hf_hub_download``,
+    which performs an actual file-access check (unlike ``HfApi().model_info()``
+    that only checks metadata).  If the token has metadata access but not
+    file-download access, ``hf_hub_download`` will raise ``GatedRepoError``
+    and we skip cleanly.
+    """
+    try:
+        from huggingface_hub import hf_hub_download
+        from huggingface_hub.errors import GatedRepoError, RepositoryNotFoundError
+    except Exception:
+        return
+    try:
+        hf_hub_download(repo_id=repo_id, filename="config.json")
+    except GatedRepoError as exc:
+        pytest.skip(
+            f"Skipping: gated HF repo {repo_id!r} inaccessible to the current "
+            f"HF_TOKEN ({exc}). See docs/contributing/ci/hf_credentials.md."
+        )
+    except RepositoryNotFoundError as exc:
+        pytest.skip(f"Skipping: HF repo {repo_id!r} not found ({exc}).")
+    except Exception:
+        return

tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py

@@ -48,6 +48,7 @@ def _build_mock_modules(mocker: MockerFixture) -> dict[str, object]:
     """Build the dict of modules to inject into sys.modules."""
     platforms_mock = mocker.MagicMock()
     platforms_mock.current_omni_platform.supports_torch_inductor.return_value = False
+    platforms_mock.current_omni_platform.is_npu.return_value = False
 
     logger_mock = mocker.MagicMock()
     logger_mock.init_logger = lambda name: mocker.MagicMock()
@@ -211,6 +212,10 @@ class TestCodePredictorDtypeAlignment:
         predictor._model_dtype = torch.float16
         predictor._compiled_model_fwd = predictor.model.forward
 
+        # Ensure NPU path is not taken on non-NPU hardware
+        common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"]
+        mocker.patch.object(common_mod.current_omni_platform, "is_npu", return_value=False)
+
         # _warmup_buckets should fix the dtype mismatch
         predictor._warmup_buckets()
 
@@ -230,6 +235,9 @@ class TestCodePredictorDtypeAlignment:
         predictor = predictor.to(torch.float16)
 
         assert predictor._model_dtype is None
+        # Ensure NPU path is not taken on non-NPU hardware
+        common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"]
+        mocker.patch.object(common_mod.current_omni_platform, "is_npu", return_value=False)
         predictor._setup_compile()
         assert predictor._model_dtype == torch.float16
 

tests/model_executor/stage_input_processors/test_hunyuan_image3_bridge.py

@@ -0,0 +1,92 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for HunyuanImage3 stage input processor."""
+
+import builtins
+from types import SimpleNamespace
+
+import pytest
+
+from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import (
+    HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS,
+)
+from vllm_omni.model_executor.stage_input_processors.hunyuan_image3 import (
+    _extract_ratio_index,
+    _truncate_at_cot_end,
+    ar2diffusion,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _source_output(token_ids: list[int], text: str = ""):
+    return SimpleNamespace(
+        outputs=[
+            SimpleNamespace(
+                token_ids=token_ids,
+                cumulative_token_ids=token_ids,
+                text=text,
+            )
+        ],
+        multimodal_output=None,
+    )
+
+
+def test_extract_ratio_index_uses_fixed_special_token_ids():
+    ratio_33 = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<img_ratio_33>"]
+    ratio_36 = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<img_ratio_36>"]
+
+    assert _extract_ratio_index([1, ratio_33, 2]) == 33
+    assert _extract_ratio_index([1, ratio_33, 2, ratio_36]) == 36
+
+
+def test_truncate_at_cot_end_strips_tail_after_recaption_marker():
+    text = _truncate_at_cot_end("body text</recaption><answer><boi><img_size_1024><img_ratio_0>")
+    assert text == "body text</recaption>"
+
+
+def test_ar2diffusion_applies_ratio_and_truncates_tail_without_tokenizer(monkeypatch: pytest.MonkeyPatch):
+    real_import = builtins.__import__
+
+    def _block_transformers_import(name, *args, **kwargs):
+        if name == "transformers" or name.startswith("transformers."):
+            raise AssertionError("ar2diffusion must not import transformers on the bridge path")
+        return real_import(name, *args, **kwargs)
+
+    monkeypatch.setattr(builtins, "__import__", _block_transformers_import)
+
+    end_recaption = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["</recaption>"]
+    answer = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<answer>"]
+    boi = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<boi>"]
+    size = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<img_size_1024>"]
+    ratio_0 = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<img_ratio_0>"]
+    token_ids = [100, 101, end_recaption, answer, boi, size, ratio_0]
+
+    result = ar2diffusion(
+        [_source_output(token_ids, text="decoded without special tokens")],
+        prompt=[{"prompt": "edit", "height": 64, "width": 64}],
+    )
+
+    assert len(result) == 1
+    assert (result[0]["height"], result[0]["width"]) == (512, 2048)
+    assert result[0]["extra"]["ar_generated_text"] == "decoded without special tokens"
+    assert "ar_token_ids" not in result[0]["extra"]
+
+
+def test_ar2diffusion_forwards_custom_system_prompt_body():
+    end_think = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["</think>"]
+    marker = "CUSTOM_SYSTEM_BODY"
+
+    result = ar2diffusion(
+        [_source_output([100, end_think], text="thought</think>")],
+        prompt=[
+            {
+                "prompt": "edit",
+                "use_system_prompt": "custom",
+                "system_prompt": marker,
+            }
+        ],
+    )
+
+    assert result[0]["use_system_prompt"] == "custom"
+    assert result[0]["system_prompt"] == marker

tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py

@@ -7,6 +7,7 @@ from __future__ import annotations
 from types import SimpleNamespace
 
 import pytest
+import torch
 
 import vllm_omni.model_executor.stage_input_processors.qwen3_omni as q3
 
@@ -79,3 +80,110 @@ def test_get_streaming_codec_delta_len_increments_and_finishes(_streaming_contex
     assert d3 == 1
     state = q3._get_qwen3_streaming_state("c1", _streaming_context)
     assert state.talker2code2wav_last_seq_len == 0
+
+
+def test_talker2code2wav_full_payload_filters_by_output_token_ids() -> None:
+    request = SimpleNamespace(
+        request_id="codec",
+        output_token_ids=[4197, 1, 2, 4198, -1, 2048],
+    )
+    rows = torch.tensor(
+        [
+            [100, 101, 102],
+            [10, 11, 12],
+            [20, 21, 22],
+            [30, 31, 32],
+            [40, 41, 42],
+            [50, 51, 52],
+        ],
+        dtype=torch.long,
+    )
+
+    payload = q3.talker2code2wav_full_payload(None, {"codes.audio": rows}, request)
+
+    assert payload is not None
+    assert payload["codes"]["audio"] == [10, 20, 11, 21, 12, 22]
+    assert payload["code_predictor_codes"] == payload["codes"]["audio"]
+
+
+def test_talker2code2wav_full_payload_drops_count_matched_terminal_row() -> None:
+    request = SimpleNamespace(
+        request_id="codec_terminal_row",
+        output_token_ids=[0, 4198],
+    )
+    rows = torch.tensor(
+        [
+            [10, 11, 12],
+        ],
+        dtype=torch.long,
+    )
+
+    payload = q3.talker2code2wav_full_payload(None, {"codes.audio": rows}, request)
+
+    assert payload is None
+
+
+def test_talker2code2wav_full_payload_drops_rows_aligned_to_non_codec_ids() -> None:
+    request = SimpleNamespace(
+        request_id="codec_invalid_ids",
+        output_token_ids=[4197, 0, 4198, 4196, -1, 2048],
+    )
+    rows = torch.tensor(
+        [
+            [91, 92, 93],
+            [0, 0, 0],
+            [81, 82, 83],
+            [71, 72, 73],
+            [61, 62, 63],
+            [51, 52, 53],
+        ],
+        dtype=torch.long,
+    )
+
+    payload = q3.talker2code2wav_full_payload(None, {"codes.audio": rows}, request)
+
+    assert payload is not None
+    assert payload["codes"]["audio"] == [0, 0, 0]
+    assert payload["code_predictor_codes"] == payload["codes"]["audio"]
+
+
+def test_talker2code2wav_full_payload_keeps_all_zero_codec_rows() -> None:
+    request = SimpleNamespace(
+        request_id="codec_zero",
+        output_token_ids=[0, 1],
+    )
+    rows = torch.tensor(
+        [
+            [0, 0, 0],
+            [7, 8, 9],
+        ],
+        dtype=torch.long,
+    )
+
+    payload = q3.talker2code2wav_full_payload(None, {"codes.audio": rows}, request)
+
+    assert payload is not None
+    assert payload["codes"]["audio"] == [0, 7, 0, 8, 0, 9]
+    assert payload["code_predictor_codes"] == payload["codes"]["audio"]
+
+
+def test_thinker2talker_full_payload_packs_complete_tensors() -> None:
+    request = SimpleNamespace(
+        request_id="thinker",
+        prompt_token_ids=[151644, 872],
+        output_token_ids=[3],
+        all_token_ids=[151644, 872, 3],
+    )
+    pooling_output = {
+        "hidden_states.layer_0": torch.ones(3, 2),
+        "hidden_states.layer_24": torch.full((3, 2), 2.0),
+        "embed.tts_bos": torch.zeros(1, 2),
+    }
+
+    payload = q3.thinker2talker_full_payload(None, pooling_output, request)
+
+    assert payload is not None
+    assert payload["ids"]["all"] == [151644, 872, 3]
+    assert payload["embed"]["prefill"].device.type == "cpu"
+    assert payload["hidden_states"]["output"].device.type == "cpu"
+    assert payload["next_stage_prompt_len"] > 0

tests/test_config_factory.py

@@ -4,27 +4,40 @@
 Unit tests for StageConfigFactory and related classes.
 """
 
+import importlib
+import inspect
 import warnings
-from dataclasses import dataclass
+from dataclasses import dataclass, fields
 from pathlib import Path
+from unittest.mock import patch
 
 import pytest
 
+from tests.helpers.stage_config import get_deploy_config_path
 from vllm_omni.config.stage_config import (
     _PIPELINE_REGISTRY,
+    DeployConfig,
     ModelPipeline,
     PipelineConfig,
     StageConfig,
     StageConfigFactory,
+    StageDeployConfig,
     StageExecutionType,
     StagePipelineConfig,
     StageType,
+    _apply_platform_overrides,
+    _deep_merge_stage,
     _resolve_scheduler,
     build_stage_runtime_overrides,
+    deploy_override_field_names,
+    load_deploy_config,
+    merge_pipeline_deploy,
     register_pipeline,
     strip_parent_engine_args,
 )
-from vllm_omni.engine.arg_utils import SHARED_FIELDS, internal_blacklist_keys
+from vllm_omni.engine.arg_utils import SHARED_FIELDS, EngineArgs, internal_blacklist_keys
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
 class TestStageType:
@@ -99,8 +112,6 @@ class TestStageConfig:
 
     def test_to_omegaconf_with_runtime_overrides(self):
         """Test that runtime overrides are applied to OmegaConf output."""
-        import warnings
-
         config = StageConfig(
             stage_id=0,
             model_stage="thinker",
@@ -123,8 +134,6 @@ class TestStageConfig:
 
     def test_to_omegaconf_max_batch_size_deprecation(self):
         """Test that runtime.max_batch_size emits a FutureWarning."""
-        import warnings
-
         config = StageConfig(
             stage_id=0,
             model_stage="thinker",
@@ -149,7 +158,6 @@ class TestStageConfig:
 
     def test_to_omegaconf_omits_none_deploy_overrides_for_engine_args(self):
         """None deploy overrides must fall through to EngineArgs defaults."""
-        from vllm_omni.config.stage_config import deploy_override_field_names
 
         config = StageConfig(
             stage_id=0,
@@ -370,11 +378,7 @@ class TestStageResolutionHelpers:
         assert "parallel_config" not in overrides
 
     def test_strip_parent_engine_args_reports_only_surprising_parent_overrides(self):
-        from dataclasses import fields as dc_fields
+        parent_fields = {f.name: f for f in fields(EngineArgs)}
-
-        from vllm.engine.arg_utils import EngineArgs
-
-        parent_fields = {f.name: f for f in dc_fields(EngineArgs)}
         filtered, overridden = strip_parent_engine_args(
             {
                 "model": "some/model",
@@ -396,11 +400,7 @@ class TestStageResolutionHelpers:
         assert overridden == ["tensor_parallel_size"]
 
     def test_strip_parent_engine_args_keeps_allowed_media_access_controls(self):
-        from dataclasses import fields as dc_fields
+        parent_fields = {f.name: f for f in fields(EngineArgs)}
-
-        from vllm.engine.arg_utils import EngineArgs
-
-        parent_fields = {f.name: f for f in dc_fields(EngineArgs)}
         filtered, overridden = strip_parent_engine_args(
             {
                 "model": "some/model",
@@ -826,8 +826,6 @@ class TestPipelineRegistry:
 
 class TestDeployConfigLoading:
     def test_deploy_override_fields_include_deploy_schema_fields(self):
-        from vllm_omni.config.stage_config import deploy_override_field_names
-
         expected_fields = {
             "async_chunk",
             "async_scheduling",
@@ -864,10 +862,6 @@ class TestDeployConfigLoading:
         )
 
     def test_load_qwen3_omni_moe_deploy_config(self):
-        from pathlib import Path
-
-        from vllm_omni.config.stage_config import load_deploy_config
-
         deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
         deploy = load_deploy_config(deploy_path)
         assert len(deploy.stages) == 3
@@ -876,10 +870,6 @@ class TestDeployConfigLoading:
         assert deploy.platforms is not None
 
     def test_load_voxtral_tts_deploy_config_schema_fields(self):
-        from pathlib import Path
-
-        from vllm_omni.config.stage_config import load_deploy_config
-
         deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "voxtral_tts.yaml"
         deploy = load_deploy_config(deploy_path)
         assert deploy.stages[0].config_format == "mistral"
@@ -890,21 +880,12 @@ class TestDeployConfigLoading:
         )
 
     def test_load_ming_flash_omni_deploy_config_schema_fields(self):
-        from pathlib import Path
-
-        from vllm_omni.config.stage_config import load_deploy_config
-
         deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "ming_flash_omni.yaml"
         deploy = load_deploy_config(deploy_path)
         assert deploy.stages[0].compilation_config == {"pass_config": {"fuse_allreduce_rms": False}}
         assert "compilation_config" not in deploy.stages[0].engine_extras
 
     def test_merge_pipeline_deploy(self):
-        from pathlib import Path
-
-        import vllm_omni.model_executor.models.qwen3_omni.pipeline  # noqa: F401
-        from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy
-
         pipeline = _PIPELINE_REGISTRY["qwen3_omni_moe"]
         deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
         if not deploy_path.exists():
@@ -921,11 +902,6 @@ class TestDeployConfigLoading:
         assert s0.yaml_extras["default_sampling_params"]["detokenize"] is True
 
     def test_merge_pipeline_deploy_preserves_num_replicas(self, tmp_path):
-        from pathlib import Path
-
-        import vllm_omni.model_executor.models.qwen3_omni.pipeline  # noqa: F401
-        from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy
-
         pipeline = _PIPELINE_REGISTRY["qwen3_omni_moe"]
         base = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
         if not base.exists():
@@ -942,15 +918,6 @@ class TestDeployConfigLoading:
         assert stages[1].yaml_runtime["num_replicas"] == 2
 
     def test_merge_pipeline_deploy_preserves_requires_multimodal_data(self):
-        from vllm_omni.config.stage_config import (
-            DeployConfig,
-            PipelineConfig,
-            StageDeployConfig,
-            StageExecutionType,
-            StagePipelineConfig,
-            merge_pipeline_deploy,
-        )
-
         pipeline = PipelineConfig(
             model_type="test_mm",
             model_arch="TestModel",
@@ -969,11 +936,186 @@ class TestDeployConfigLoading:
 
         assert stages[0].yaml_runtime["requires_multimodal_data"] is True
 
+    def test_mixed_schema_preserves_flat_fields(self):
+        """Ensure flat fields are not dropped when engine_args are present."""
+        fake_config = {
+            "stages": [
+                {
+                    "stage_id": 0,
+                    "gpu_memory_utilization": 0.7,
+                    "max_num_seqs": 8,
+                    "engine_args": {
+                        "tensor_parallel_size": 4,
+                        "enforce_eager": True,
+                    },
+                },
+            ]
+        }
+
+        with patch("vllm_omni.config.stage_config.resolve_deploy_yaml", return_value=fake_config):
+            deploy = load_deploy_config("dummy.yaml")
+
+        assert len(deploy.stages) == 1
+        stage = deploy.stages[0]
+        # Check that the engine args are set
+        assert stage.tensor_parallel_size == 4
+        assert stage.enforce_eager is True
+        # Check that the other settings are also preserved
+        assert stage.gpu_memory_utilization == 0.7
+        assert stage.max_num_seqs == 8
+
+    def test_engine_parse_engine_fields(self):
+        """Test that we correctly parse & recursively merge stage deploy fields."""
+        fake_config = {
+            "stages": [
+                {
+                    "stage_id": 0,
+                    "compilation_config": {
+                        "encoder_cudagraph_token_budgets": [1024, 2048],
+                        "pass_config": {"fuse_norm_quant": True},
+                    },
+                    "engine_args": {
+                        "compilation_config": {
+                            "cudagraph_mm_encoder": True,
+                            "pass_config": {"fuse_allreduce_rms": False},
+                        },
+                    },
+                },
+            ]
+        }
+
+        with patch("vllm_omni.config.stage_config.resolve_deploy_yaml", return_value=fake_config):
+            deploy = load_deploy_config("dummy.yaml")
+
+        assert len(deploy.stages) == 1
+        stage = deploy.stages[0]
+        assert stage.compilation_config == {
+            "encoder_cudagraph_token_budgets": [1024, 2048],
+            "pass_config": {
+                "fuse_norm_quant": True,
+                "fuse_allreduce_rms": False,
+            },
+            "cudagraph_mm_encoder": True,
+        }
+
+    def test_engine_extras_deep_merges_dicts_simple(self):
+        """Ensure dictionary valued keys merge properly for top level dicts."""
+        fake_config = {
+            "stages": [
+                {
+                    "stage_id": 0,
+                    "foo": {"a": 111, "b": 1},
+                    "engine_args": {
+                        "foo": {"b": 2, "c": 3},
+                    },
+                },
+            ]
+        }
+
+        with patch("vllm_omni.config.stage_config.resolve_deploy_yaml", return_value=fake_config):
+            deploy = load_deploy_config("dummy.yaml")
+
+        assert len(deploy.stages) == 1
+        stage = deploy.stages[0]
+        assert "foo" in stage.engine_extras
+        assert stage.engine_extras["foo"] == {"a": 111, "b": 2, "c": 3}
+
+    def test_engine_extras_dict_type_mismatch(self):
+        """Ensure that we handle type mismatches with nested dicts correctly."""
+        fake_config = {
+            "stages": [
+                {
+                    "stage_id": 0,
+                    "foo": {"b": {1: 1}},
+                    "engine_args": {
+                        "foo": {"b": 2},
+                    },
+                },
+            ]
+        }
+
+        with patch("vllm_omni.config.stage_config.resolve_deploy_yaml", return_value=fake_config):
+            deploy = load_deploy_config("dummy.yaml")
+
+        assert len(deploy.stages) == 1
+        stage = deploy.stages[0]
+        assert "foo" in stage.engine_extras
+        assert stage.engine_extras["foo"] == {"b": 2}
+
+    def test_mixed_engine_extras_deep_merges_dicts(self):
+        """Ensure dictionary valued keys merge properly for nested dicts."""
+        fake_config = {
+            "stages": [
+                {
+                    "stage_id": 0,
+                    "foo": {"a": 111, "b": {"e": 199}},
+                    "engine_args": {
+                        "foo": {"b": {"d": 9}, "c": 3},
+                    },
+                },
+            ]
+        }
+
+        with patch("vllm_omni.config.stage_config.resolve_deploy_yaml", return_value=fake_config):
+            deploy = load_deploy_config("dummy.yaml")
+
+        assert len(deploy.stages) == 1
+        stage = deploy.stages[0]
+        assert "foo" in stage.engine_extras
+        assert stage.engine_extras["foo"] == {"a": 111, "b": {"d": 9, "e": 199}, "c": 3}
+
+    def test_deep_merge_does_not_mutate_inputs(self):
+        """Merging engine_args must not mutate the base stage dict."""
+        fake_config = {
+            "stages": [
+                {
+                    "stage_id": 0,
+                    "foo": {"a": 1, "b": {"x": 10, "y": 20}},
+                    "engine_args": {
+                        "foo": {"b": {"y": 99, "z": 30}, "c": 3},
+                    },
+                },
+            ]
+        }
+        original_foo = fake_config["stages"][0]["foo"]
+        original_b = original_foo["b"]
+
+        with patch("vllm_omni.config.stage_config.resolve_deploy_yaml", return_value=fake_config):
+            deploy = load_deploy_config("dummy.yaml")
+
+        # Values should merge into the new dict correctly
+        assert deploy.stages[0].engine_extras["foo"] == {
+            "a": 1,
+            "b": {"x": 10, "y": 99, "z": 30},
+            "c": 3,
+        }
+        # But the original nested dicts foo / b should be unchanged, since
+        # we recursively shallow copy to avoid mutating in place.
+        assert original_foo == {"a": 1, "b": {"x": 10, "y": 20}}
+        assert original_b == {"x": 10, "y": 20}
+
+    def test_mixed_schema_engine_args_wins_scalars(self):
+        """engine_args takes precedence over flat fields for scalar conflicts."""
+        fake_config = {
+            "stages": [
+                {
+                    "stage_id": 0,
+                    "gpu_memory_utilization": 0.7,
+                    "engine_args": {
+                        "gpu_memory_utilization": 0.5,
+                    },
+                },
+            ]
+        }
+
+        with patch("vllm_omni.config.stage_config.resolve_deploy_yaml", return_value=fake_config):
+            deploy = load_deploy_config("dummy.yaml")
+
+        assert deploy.stages[0].gpu_memory_utilization == 0.5
+
 
 class TestQwen3OmniPipeline:
     def test_registered(self):
-        import vllm_omni.model_executor.models.qwen3_omni.pipeline  # noqa: F401
-
         p = _PIPELINE_REGISTRY.get("qwen3_omni_moe")
         assert p is not None
         assert p.model_arch == "Qwen3OmniMoeForConditionalGeneration"
@@ -981,8 +1123,6 @@ class TestQwen3OmniPipeline:
         assert p.validate() == []
 
     def test_thinker(self):
-        import vllm_omni.model_executor.models.qwen3_omni.pipeline  # noqa: F401
-
         s = _PIPELINE_REGISTRY["qwen3_omni_moe"].get_stage(0)
         assert s.model_stage == "thinker"
         assert s.execution_type == StageExecutionType.LLM_AR
@@ -991,8 +1131,6 @@ class TestQwen3OmniPipeline:
         assert s.sampling_constraints["detokenize"] is True
 
     def test_talker(self):
-        import vllm_omni.model_executor.models.qwen3_omni.pipeline  # noqa: F401
-
         s = _PIPELINE_REGISTRY["qwen3_omni_moe"].get_stage(1)
         assert s.input_sources == (0,)
         assert s.sampling_constraints["stop_token_ids"] == [2150]
@@ -1000,8 +1138,6 @@ class TestQwen3OmniPipeline:
         assert s.custom_process_next_stage_input_func is not None
 
     def test_code2wav(self):
-        import vllm_omni.model_executor.models.qwen3_omni.pipeline  # noqa: F401
-
         s = _PIPELINE_REGISTRY["qwen3_omni_moe"].get_stage(2)
         assert s.execution_type == StageExecutionType.LLM_GENERATION
         assert s.final_output_type == "audio"
@@ -1010,8 +1146,6 @@ class TestQwen3OmniPipeline:
 
 class TestQwen2_5OmniPipeline:
     def test_registered(self):
-        import vllm_omni.model_executor.models.qwen2_5_omni.pipeline  # noqa: F401
-
         p = _PIPELINE_REGISTRY.get("qwen2_5_omni")
         assert p is not None
         assert p.model_arch == "Qwen2_5OmniForConditionalGeneration"
@@ -1019,8 +1153,6 @@ class TestQwen2_5OmniPipeline:
         assert p.validate() == []
 
     def test_thinker(self):
-        import vllm_omni.model_executor.models.qwen2_5_omni.pipeline  # noqa: F401
-
         s = _PIPELINE_REGISTRY["qwen2_5_omni"].get_stage(0)
         assert s.model_stage == "thinker"
         assert s.execution_type == StageExecutionType.LLM_AR
@@ -1029,16 +1161,12 @@ class TestQwen2_5OmniPipeline:
         assert s.requires_multimodal_data is True
 
     def test_talker(self):
-        import vllm_omni.model_executor.models.qwen2_5_omni.pipeline  # noqa: F401
-
         s = _PIPELINE_REGISTRY["qwen2_5_omni"].get_stage(1)
         assert s.input_sources == (0,)
         assert s.sampling_constraints["stop_token_ids"] == [8294]
         assert s.custom_process_input_func is not None
 
     def test_code2wav(self):
-        import vllm_omni.model_executor.models.qwen2_5_omni.pipeline  # noqa: F401
-
         s = _PIPELINE_REGISTRY["qwen2_5_omni"].get_stage(2)
         assert s.execution_type == StageExecutionType.LLM_GENERATION
         assert s.final_output_type == "audio"
@@ -1047,8 +1175,6 @@ class TestQwen2_5OmniPipeline:
 
 class TestQwen3TTSPipeline:
     def test_registered(self):
-        import vllm_omni.model_executor.models.qwen3_tts.pipeline  # noqa: F401
-
         p = _PIPELINE_REGISTRY.get("qwen3_tts")
         assert p is not None
         assert p.model_arch == "Qwen3TTSTalkerForConditionalGeneration"
@@ -1056,8 +1182,6 @@ class TestQwen3TTSPipeline:
         assert p.validate() == []
 
     def test_talker_stage(self):
-        import vllm_omni.model_executor.models.qwen3_tts.pipeline  # noqa: F401
-
         s = _PIPELINE_REGISTRY["qwen3_tts"].get_stage(0)
         assert s.model_stage == "qwen3_tts"
         assert s.execution_type == StageExecutionType.LLM_AR
@@ -1068,8 +1192,6 @@ class TestQwen3TTSPipeline:
         assert s.model_arch is None
 
     def test_code2wav_stage_has_per_stage_model_arch(self):
-        import vllm_omni.model_executor.models.qwen3_tts.pipeline  # noqa: F401
-
         s = _PIPELINE_REGISTRY["qwen3_tts"].get_stage(1)
         assert s.execution_type == StageExecutionType.LLM_GENERATION
         assert s.final_output_type == "audio"
@@ -1081,9 +1203,6 @@ class TestQwen3TTSPipeline:
 
     def test_per_stage_model_arch_flows_through_merge(self, tmp_path):
         """Verify the new ps.model_arch override survives merge_pipeline_deploy."""
-        import vllm_omni.model_executor.models.qwen3_tts.pipeline  # noqa: F401
-        from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy
-
         deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_tts.yaml"
         if not deploy_path.exists():
             pytest.skip("qwen3_tts deploy yaml not found")
@@ -1099,8 +1218,6 @@ class TestQwen3TTSPipeline:
 
     def test_subtalker_sampling_params_deep_merge_preserves_base_keys(self):
         """Verify subtalker sampling params participate in stage deep-merge."""
-        from vllm_omni.config.stage_config import _deep_merge_stage
-
         base = {
             "stage_id": 0,
             "subtalker_sampling_params": {
@@ -1130,8 +1247,6 @@ class TestQwen3TTSPipeline:
 
 class TestMingFlashOmniPipeline:
     def test_registered(self):
-        import vllm_omni.model_executor.models.ming_flash_omni.pipeline  # noqa: F401
-
         p = _PIPELINE_REGISTRY.get("ming_flash_omni")
         assert p is not None
         assert p.model_arch == "MingFlashOmniForConditionalGeneration"
@@ -1139,8 +1254,6 @@ class TestMingFlashOmniPipeline:
         assert p.validate() == []
 
     def test_thinker_stage(self):
-        import vllm_omni.model_executor.models.ming_flash_omni.pipeline  # noqa: F401
-
         s = _PIPELINE_REGISTRY["ming_flash_omni"].get_stage(0)
         assert s.model_stage == "thinker"
         assert s.execution_type == StageExecutionType.LLM_AR
@@ -1151,8 +1264,6 @@ class TestMingFlashOmniPipeline:
         assert s.sampling_constraints["detokenize"] is True
 
     def test_talker_stage(self):
-        import vllm_omni.model_executor.models.ming_flash_omni.pipeline  # noqa: F401
-
         s = _PIPELINE_REGISTRY["ming_flash_omni"].get_stage(1)
         assert s.model_stage == "ming_tts"
         assert s.execution_type == StageExecutionType.LLM_GENERATION
@@ -1171,18 +1282,12 @@ class TestMingFlashOmniPipeline:
         Lazy string references only fail at first inference otherwise — this
         catches typos in the pipeline declaration at import / registration time.
         """
-        import importlib
-
-        import vllm_omni.model_executor.models.ming_flash_omni.pipeline  # noqa: F401
-
         s = _PIPELINE_REGISTRY["ming_flash_omni"].get_stage(1)
         module_path, _, attr = s.custom_process_input_func.rpartition(".")
         module = importlib.import_module(module_path)
         assert callable(getattr(module, attr))
 
     def test_tts_pipeline_registered(self):
-        import vllm_omni.model_executor.models.ming_flash_omni.pipeline  # noqa: F401
-
         p = _PIPELINE_REGISTRY.get("ming_flash_omni_tts")
         assert p is not None
         assert p.model_arch == "MingFlashOmniTalkerForConditionalGeneration"
@@ -1190,8 +1295,6 @@ class TestMingFlashOmniPipeline:
         assert p.validate() == []
 
     def test_tts_stage(self):
-        import vllm_omni.model_executor.models.ming_flash_omni.pipeline  # noqa: F401
-
         s = _PIPELINE_REGISTRY["ming_flash_omni_tts"].get_stage(0)
         assert s.model_stage == "ming_tts"
         assert s.execution_type == StageExecutionType.LLM_GENERATION
@@ -1204,9 +1307,6 @@ class TestMingFlashOmniPipeline:
 
     def test_full_yaml_loads_and_merges(self):
         """deploy/ming_flash_omni.yaml parses and merges with the registered pipeline."""
-        import vllm_omni.model_executor.models.ming_flash_omni.pipeline  # noqa: F401
-        from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy
-
         deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "ming_flash_omni.yaml"
         if not deploy_path.exists():
             pytest.skip("ming_flash_omni deploy yaml not found")
@@ -1226,9 +1326,6 @@ class TestMingFlashOmniPipeline:
 
     def test_tts_yaml_loads_and_merges(self):
         """deploy/ming_flash_omni_tts.yaml parses and routes to the TTS-only pipeline."""
-        import vllm_omni.model_executor.models.ming_flash_omni.pipeline  # noqa: F401
-        from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy
-
         deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "ming_flash_omni_tts.yaml"
         if not deploy_path.exists():
             pytest.skip("ming_flash_omni_tts deploy yaml not found")
@@ -1243,8 +1340,6 @@ class TestMingFlashOmniPipeline:
         assert stages[0].yaml_engine_args["model_arch"] == "MingFlashOmniTalkerForConditionalGeneration"
 
     def test_thinker_only_pipeline_registered(self):
-        import vllm_omni.model_executor.models.ming_flash_omni.pipeline  # noqa: F401
-
         p = _PIPELINE_REGISTRY.get("ming_flash_omni_thinker_only")
         assert p is not None
         assert p.model_arch == "MingFlashOmniForConditionalGeneration"
@@ -1252,8 +1347,6 @@ class TestMingFlashOmniPipeline:
         assert p.validate() == []
 
     def test_thinker_only_stage(self):
-        import vllm_omni.model_executor.models.ming_flash_omni.pipeline  # noqa: F401
-
         s = _PIPELINE_REGISTRY["ming_flash_omni_thinker_only"].get_stage(0)
         assert s.model_stage == "thinker"
         assert s.execution_type == StageExecutionType.LLM_AR
@@ -1267,9 +1360,6 @@ class TestMingFlashOmniPipeline:
 
     def test_thinker_only_yaml_loads_and_merges(self):
         """deploy/ming_flash_omni_thinker_only.yaml parses and routes to the thinker-only pipeline."""
-        import vllm_omni.model_executor.models.ming_flash_omni.pipeline  # noqa: F401
-        from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy
-
         deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "ming_flash_omni_thinker_only.yaml"
         if not deploy_path.exists():
             pytest.skip("ming_flash_omni_thinker_only deploy yaml not found")
@@ -1288,9 +1378,6 @@ class TestBaseConfigInheritance:
     """Test deploy YAML base_config inheritance."""
 
     def test_ci_inherits_from_main(self):
-        from tests.helpers.stage_config import get_deploy_config_path
-        from vllm_omni.config.stage_config import load_deploy_config
-
         ci_path = Path(get_deploy_config_path("ci/qwen3_omni_moe.yaml"))
         if not ci_path.exists():
             pytest.skip("CI deploy config not found")
@@ -1298,7 +1385,7 @@ class TestBaseConfigInheritance:
         deploy = load_deploy_config(ci_path)
         assert len(deploy.stages) == 3
         # CI overrides
-        assert deploy.stages[0].load_format == "dummy"
+        assert deploy.stages[0].load_format is None
         assert "load_format" not in deploy.stages[0].engine_extras
         assert deploy.stages[0].max_num_seqs == 5
         # Inherited from base
@@ -1311,9 +1398,6 @@ class TestBaseConfigInheritance:
         assert deploy.async_chunk is False
 
     def test_ci_sampling_merge(self):
-        from tests.helpers.stage_config import get_deploy_config_path
-        from vllm_omni.config.stage_config import load_deploy_config
-
         ci_path = Path(get_deploy_config_path("ci/qwen3_omni_moe.yaml"))
         if not ci_path.exists():
             pytest.skip("CI deploy config not found")
@@ -1328,8 +1412,6 @@ class TestBaseConfigInheritance:
 
     def test_pure_inheritance_overlay(self, tmp_path):
         """An overlay with only ``base_config`` inherits everything."""
-        from vllm_omni.config.stage_config import load_deploy_config
-
         base = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
         if not base.exists():
             pytest.skip("Base deploy config not found")
@@ -1343,8 +1425,6 @@ class TestBaseConfigInheritance:
 
     def test_single_field_overlay(self, tmp_path):
         """An overlay overriding one stage field merges with the base."""
-        from vllm_omni.config.stage_config import load_deploy_config
-
         base = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
         if not base.exists():
             pytest.skip("Base deploy config not found")
@@ -1362,10 +1442,6 @@ class TestPlatformOverrides:
     """Test platform-specific deploy config overrides."""
 
     def test_npu_overrides(self):
-        from pathlib import Path
-
-        from vllm_omni.config.stage_config import _apply_platform_overrides, load_deploy_config
-
         deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
         if not deploy_path.exists():
             pytest.skip("Deploy config not found")
@@ -1380,10 +1456,6 @@ class TestPlatformOverrides:
         assert deploy.stages[2].enforce_eager is True
 
     def test_xpu_overrides(self):
-        from pathlib import Path
-
-        from vllm_omni.config.stage_config import _apply_platform_overrides, load_deploy_config
-
         deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
         if not deploy_path.exists():
             pytest.skip("Deploy config not found")
@@ -1396,10 +1468,6 @@ class TestPlatformOverrides:
         assert deploy.stages[0].engine_extras.get("max_cudagraph_capture_size") == 0
 
     def test_unknown_platform_noop(self):
-        from pathlib import Path
-
-        from vllm_omni.config.stage_config import _apply_platform_overrides, load_deploy_config
-
         deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
         if not deploy_path.exists():
             pytest.skip("Deploy config not found")
@@ -1411,8 +1479,6 @@ class TestPlatformOverrides:
 
     def test_platforms_deep_merge_inheritance(self, tmp_path):
         """Overlay's platforms: block layers onto base's, per-stage."""
-        from vllm_omni.config.stage_config import _apply_platform_overrides, load_deploy_config
-
         base = tmp_path / "base.yaml"
         base.write_text(
             "stages:\n"
@@ -1447,11 +1513,6 @@ class TestCLIOverrideFlow:
     """Test --stage-overrides JSON merge into StageConfig."""
 
     def test_stage_overrides_merge(self):
-        from pathlib import Path
-
-        import vllm_omni.model_executor.models.qwen3_omni.pipeline  # noqa: F401
-        from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy
-
         pipeline = _PIPELINE_REGISTRY["qwen3_omni_moe"]
         deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
         if not deploy_path.exists():
@@ -1466,11 +1527,6 @@ class TestCLIOverrideFlow:
         assert stages[0].runtime_overrides["gpu_memory_utilization"] == 0.5
 
     def test_global_override_applies_to_all(self):
-        from pathlib import Path
-
-        import vllm_omni.model_executor.models.qwen3_omni.pipeline  # noqa: F401
-        from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy
-
         pipeline = _PIPELINE_REGISTRY["qwen3_omni_moe"]
         deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
         if not deploy_path.exists():
@@ -1490,8 +1546,6 @@ class TestSentinelDefaultPrecedence:
     to YAML / dataclass defaults (#3035)."""
 
     def _stages(self, cli_overrides):
-        import vllm_omni.model_executor.models.qwen3_omni.pipeline  # noqa: F401
-
         return StageConfigFactory._create_from_registry(
             "qwen3_omni_moe",
             cli_overrides=cli_overrides,
@@ -1504,7 +1558,7 @@ class TestSentinelDefaultPrecedence:
     def test_none_value_skipped_yaml_wins(self):
         stages = self._stages({"max_num_seqs": None})
         assert stages[2].runtime_overrides.get("max_num_seqs") is None
-        assert "max_num_seqs" not in stages[2].yaml_engine_args
+        assert stages[2].yaml_engine_args.get("max_num_seqs") == 64
 
     def test_empty_kwargs_yaml_only(self):
         stages = self._stages({})
@@ -1550,15 +1604,11 @@ class TestSentinelDefaultPrecedence:
             assert stage.runtime_overrides == {}
 
     def test_create_from_registry_no_cli_explicit_keys_param(self):
-        import inspect
-
         sig = inspect.signature(StageConfigFactory._create_from_registry)
         named = [p for p in sig.parameters.values() if p.kind != p.VAR_KEYWORD]
         assert "cli_explicit_keys" not in {p.name for p in named}
 
     def test_cli_explicit_keys_kwarg_emits_deprecation(self):
-        import vllm_omni.model_executor.models.qwen3_omni.pipeline  # noqa: F401
-
         with warnings.catch_warnings(record=True) as w:
             warnings.simplefilter("always")
             StageConfigFactory._create_from_registry(
@@ -1572,13 +1622,6 @@ class TestSentinelDefaultPrecedence:
         """A single ``qwen3_tts`` pipeline picks per-chunk vs end-to-end
         processors based on ``deploy.async_chunk``, without needing a
         separate variant pipeline registration."""
-        import vllm_omni.model_executor.models.qwen3_tts.pipeline  # noqa: F401
-        from vllm_omni.config.stage_config import (
-            _PIPELINE_REGISTRY,
-            DeployConfig,
-            merge_pipeline_deploy,
-        )
-
         pipeline = _PIPELINE_REGISTRY["qwen3_tts"]
 
         # async_chunk=True → stage 0's per-chunk processor wires up, stage 1
@@ -1598,16 +1641,46 @@ class TestSentinelDefaultPrecedence:
         assert sync_stages[1].custom_process_input_func is not None
         assert sync_stages[1].custom_process_input_func.endswith("talker2code2wav")
 
+    def test_async_chunk_dispatches_qwen3_omni_processors(self):
+        import runpy
+        from pathlib import Path
+
+        from vllm_omni.config.stage_config import DeployConfig, merge_pipeline_deploy
+
+        pipeline_path = (
+            Path(__file__).parent.parent / "vllm_omni" / "model_executor" / "models" / "qwen3_omni" / "pipeline.py"
+        )
+        pipeline = runpy.run_path(str(pipeline_path))["QWEN3_OMNI_PIPELINE"]
+
+        async_stages = merge_pipeline_deploy(pipeline, DeployConfig(async_chunk=True))
+        assert (
+            async_stages[0]
+            .yaml_engine_args["custom_process_next_stage_input_func"]
+            .endswith("thinker2talker_async_chunk")
+        )
+        assert (
+            async_stages[1]
+            .yaml_engine_args["custom_process_next_stage_input_func"]
+            .endswith("talker2code2wav_async_chunk")
+        )
+
+        sync_stages = merge_pipeline_deploy(pipeline, DeployConfig(async_chunk=False))
+        assert (
+            sync_stages[0]
+            .yaml_engine_args["custom_process_next_stage_input_func"]
+            .endswith("thinker2talker_full_payload")
+        )
+        assert (
+            sync_stages[1]
+            .yaml_engine_args["custom_process_next_stage_input_func"]
+            .endswith("talker2code2wav_full_payload")
+        )
+
 
 class TestSamplingConstraintsPrecedence:
     """Test that pipeline sampling_constraints override deploy defaults."""
 
     def test_constraints_win(self):
-        from pathlib import Path
-
-        import vllm_omni.model_executor.models.qwen3_omni.pipeline  # noqa: F401
-        from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy
-
         pipeline = _PIPELINE_REGISTRY["qwen3_omni_moe"]
         deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
         if not deploy_path.exists():

tests/worker/test_gpu_generation_model_runner.py

@@ -35,6 +35,7 @@ def _make_runner(multimodal_outputs):
     runner.device = torch.device("cpu")
     runner.supports_mm_inputs = False
     runner.speculative_config = None
+    runner.routed_experts_initialized = False
     return runner
 
 

tests/worker/test_omni_gpu_model_runner.py

@@ -5,6 +5,7 @@ import pytest
 import torch
 
 from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
+from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -330,6 +331,117 @@ def test_update_intermediate_buffer_skips_unknown_req_id():
     assert "unknown_req" not in runner.model_intermediate_buffer
 
 
+def _make_full_payload_accumulation_runner(
+    model_arch="Qwen3OmniMoeForConditionalGeneration",
+    model_stage="talker",
+    async_chunk=False,
+):
+    runner = object.__new__(OmniConnectorModelRunnerMixin)
+    runner.model_config = SimpleNamespace(
+        model_arch=model_arch,
+        model_stage=model_stage,
+        async_chunk=async_chunk,
+    )
+    runner._custom_process_func = object()
+    runner._pending_full_payload_send = {}
+    runner._stage_id = 1
+    return runner
+
+
+def test_accumulate_full_payload_output_preserves_aligned_all_zero_qwen3_omni_codec_rows():
+    runner = _make_full_payload_accumulation_runner()
+    request = SimpleNamespace(output_token_ids=[0, 1])
+    codes = torch.zeros((2, 3), dtype=torch.long)
+
+    OmniConnectorModelRunnerMixin.accumulate_full_payload_output(runner, "r1", {"codes.audio": codes}, request)
+
+    stored, _ = OmniConnectorModelRunnerMixin._materialize_full_payload_entry(runner._pending_full_payload_send["r1"])
+    assert torch.equal(stored["codes.audio"], codes)
+
+
+def test_accumulate_full_payload_output_keeps_misaligned_all_zero_qwen3_omni_codec_rows():
+    # After removing the sender-side zero filter, the accumulator keeps every
+    # codec row including misaligned all-zero rows. The downstream consumer
+    # (_extract_qwen3_full_payload_codec_rows) is the authoritative crop and
+    # filters by output_token_ids.
+    runner = _make_full_payload_accumulation_runner()
+    request = SimpleNamespace(output_token_ids=[0, 1])
+    codes = torch.zeros((1, 3), dtype=torch.long)
+
+    OmniConnectorModelRunnerMixin.accumulate_full_payload_output(runner, "r1", {"codes.audio": codes}, request)
+
+    stored, _ = OmniConnectorModelRunnerMixin._materialize_full_payload_entry(runner._pending_full_payload_send["r1"])
+    assert "codes.audio" in stored
+    assert torch.equal(stored["codes.audio"], codes)
+
+
+def test_accumulate_full_payload_output_preserves_incremental_aligned_all_zero_qwen3_omni_codec_rows():
+    runner = _make_full_payload_accumulation_runner()
+    request = SimpleNamespace(output_token_ids=[0, 1])
+    runner._pending_full_payload_send["r1"] = (
+        {"codes.audio": torch.ones((1, 3), dtype=torch.long)},
+        request,
+    )
+    codes = torch.zeros((1, 3), dtype=torch.long)
+
+    OmniConnectorModelRunnerMixin.accumulate_full_payload_output(runner, "r1", {"codes.audio": codes}, request)
+
+    stored, _ = OmniConnectorModelRunnerMixin._materialize_full_payload_entry(runner._pending_full_payload_send["r1"])
+    assert stored["codes.audio"].shape == (2, 3)
+    assert torch.equal(stored["codes.audio"][1], torch.zeros(3, dtype=torch.long))
+
+
+def test_accumulate_full_payload_output_keeps_all_zero_qwen3_omni_prefill_placeholder():
+    # Prefill placeholder rows (output_token_ids empty) are no longer dropped
+    # at the sender. The consumer-side crop trims them off using
+    # output_token_ids, so the end-to-end semantics are unchanged.
+    runner = _make_full_payload_accumulation_runner()
+    request = SimpleNamespace(output_token_ids=[])
+    codes = torch.zeros((2, 3), dtype=torch.long)
+
+    OmniConnectorModelRunnerMixin.accumulate_full_payload_output(runner, "r1", {"codes.audio": codes}, request)
+
+    stored, _ = OmniConnectorModelRunnerMixin._materialize_full_payload_entry(runner._pending_full_payload_send["r1"])
+    assert "codes.audio" in stored
+    assert torch.equal(stored["codes.audio"], codes)
+
+
+def test_full_payload_output_accumulation_hook_matrix():
+    assert _make_full_payload_accumulation_runner(model_stage="thinker")._should_accumulate_full_payload_output()
+    assert _make_full_payload_accumulation_runner(model_stage="talker")._should_accumulate_full_payload_output()
+    assert not _make_full_payload_accumulation_runner(model_stage="code2wav")._should_accumulate_full_payload_output()
+    assert not _make_full_payload_accumulation_runner(
+        model_stage="talker", async_chunk=True
+    )._should_accumulate_full_payload_output()
+    assert not _make_full_payload_accumulation_runner(
+        model_arch="Qwen3TTSForConditionalGeneration"
+    )._should_accumulate_full_payload_output()
+    assert not _make_full_payload_accumulation_runner(
+        model_arch="Qwen2_5OmniForConditionalGeneration"
+    )._should_accumulate_full_payload_output()
+
+
+def test_sync_local_stage_payloads_retains_payload_until_request_is_active():
+    runner = object.__new__(OmniGPUModelRunner)
+    payload = {"codes": {"audio": [1, 2, 3]}}
+    runner._local_stage_payload_cache = {"late": payload}
+    runner._full_payload_pending_broadcast_req_ids = set()
+    runner.requests = {}
+    runner.model_intermediate_buffer = {}
+
+    OmniGPUModelRunner._sync_local_stage_payloads(runner)
+
+    assert runner._local_stage_payload_cache == {"late": payload}
+    assert runner.model_intermediate_buffer == {}
+
+    runner.requests = {"late": DummyReqState()}
+    OmniGPUModelRunner._sync_local_stage_payloads(runner)
+
+    assert runner._local_stage_payload_cache == {}
+    assert runner.model_intermediate_buffer["late"] == payload
+    assert runner.requests["late"].additional_information_cpu == payload
+
+
 def test_maybe_attach_mimo_audio_req_infos_enriches_dict():
     runner = _make_runner_for_mimo()
     req_id = "r_mimo"

vllm_omni/config/stage_config.py

@@ -473,7 +473,7 @@ class DeployConfig:
     pipeline_parallel_size: int | None = None
 
 
-_STAGE_NON_ENGINE_KEYS = frozenset(
+_STAGE_RESERVED_KEYS = frozenset(
     {
         "stage_id",
         "devices",
@@ -482,24 +482,31 @@ _STAGE_NON_ENGINE_KEYS = frozenset(
         "input_connectors",
         "default_sampling_params",
         "engine_extras",
+        "engine_args",
+        "runtime",
     }
 )
 
 # Fields on StageDeployConfig that are populated from engine_args dict
-_STAGE_DEPLOY_FIELDS = {f.name: f for f in fields(StageDeployConfig) if f.name not in _STAGE_NON_ENGINE_KEYS}
+_STAGE_DEPLOY_FIELDS = {f.name: f for f in fields(StageDeployConfig) if f.name not in _STAGE_RESERVED_KEYS}
 
 
 def _parse_stage_deploy(stage_data: dict[str, Any]) -> StageDeployConfig:
     """Parse a single stage entry from deploy YAML into StageDeployConfig."""
+    # Get the non-reserved keys for this stage
+    flat_args = {k: v for k, v in stage_data.items() if k not in _STAGE_RESERVED_KEYS}
+    runtime_cfg = dict(stage_data.get("runtime", {}))
+    devices = runtime_cfg.get("devices", stage_data.get("devices"))
+    num_replicas = runtime_cfg.get("num_replicas", stage_data.get("num_replicas", 1))
+
     if "engine_args" in stage_data:
-        runtime_cfg = dict(stage_data.get("runtime", {}))
+        for k, v in stage_data["engine_args"].items():
-        engine_args = dict(stage_data["engine_args"])
+            existing = flat_args.get(k)
-        devices = stage_data.get("runtime", {}).get("devices", stage_data.get("devices"))
+            # If we have multiple dictionaries, merge recursively.
-        num_replicas = runtime_cfg.get("num_replicas", stage_data.get("num_replicas", 1))
+            if isinstance(v, dict) and isinstance(existing, dict):
-    else:
+                flat_args[k] = _get_recursively_merged_dict(existing, v)
-        engine_args = {k: v for k, v in stage_data.items() if k not in _STAGE_NON_ENGINE_KEYS and k != "stage_id"}
+            else:
-        devices = stage_data.get("devices")
+                flat_args[k] = v
-        num_replicas = stage_data.get("num_replicas", stage_data.get("runtime", {}).get("num_replicas", 1))
 
     kwargs: dict[str, Any] = {
         "stage_id": stage_data["stage_id"],
@@ -507,13 +514,13 @@ def _parse_stage_deploy(stage_data: dict[str, Any]) -> StageDeployConfig:
         "num_replicas": int(num_replicas),
     }
     for name, f in _STAGE_DEPLOY_FIELDS.items():
-        if name in engine_args:
+        if name in flat_args:
-            kwargs[name] = engine_args.pop(name)
+            kwargs[name] = flat_args.pop(name)
 
     kwargs["output_connectors"] = stage_data.get("output_connectors")
     kwargs["input_connectors"] = stage_data.get("input_connectors")
     kwargs["default_sampling_params"] = stage_data.get("default_sampling_params")
-    kwargs["engine_extras"] = engine_args
+    kwargs["engine_extras"] = flat_args
     return StageDeployConfig(**kwargs)
 
 
@@ -522,24 +529,34 @@ _DEEP_MERGE_KEYS = frozenset({"default_sampling_params", "subtalker_sampling_par
 
 def _deep_merge_stage(base: dict, overlay: dict) -> dict:
     """Deep-merge ``_DEEP_MERGE_KEYS`` so thin overlays don't drop base keys."""
-    merged = dict(base)
+    # Deep merge _DEEP_MERGE_KEYS recursively
-    for k, v in overlay.items():
+    base_merge_dict = {k: v for k, v in base.items() if k in _DEEP_MERGE_KEYS}
-        if k in _DEEP_MERGE_KEYS:
+    overlay_merge_dict = {k: v for k, v in overlay.items() if k in _DEEP_MERGE_KEYS}
-            base_val = merged.get(k)
+
-            if isinstance(v, dict) and isinstance(base_val, dict):
+    # Get the merge dict; priority is base < overlay < merged sub
-                merged[k] = {**base_val, **v}
+    merged_subdict = _get_recursively_merged_dict(original=base_merge_dict, update=overlay_merge_dict)
-                continue
+    merged_dict = {**base, **overlay, **merged_subdict}
-            # Deep-merge key but at least one side isn't a dict: surface the
+    return merged_dict
-            # silent clobber so mismatched YAML types don't get past review.
+
-            if base_val is not None:
+
+def _get_recursively_merged_dict(original: dict, update: dict) -> dict:
+    """Recursively merge two dicts, returning a new dict."""
+    merged = original.copy()
+    for k, update_v in update.items():
+        orig_v = merged.get(k)
+        if isinstance(orig_v, dict) and isinstance(update_v, dict):
+            merged[k] = _get_recursively_merged_dict(orig_v, update_v)
+        else:
+            if orig_v is not None and (isinstance(orig_v, dict) != isinstance(update_v, dict)):
                 logger.warning(
                     "Deep-merge key %r has non-dict value (base=%s, overlay=%s); "
                     "overlay will fully replace base instead of merging.",
                     k,
-                    type(base_val).__name__,
+                    type(orig_v).__name__,
-                    type(v).__name__,
+                    type(update_v).__name__,
                 )
-        merged[k] = v
+
+            merged[k] = update_v
     return merged
 
 
@@ -779,7 +796,7 @@ def _build_engine_args(
     # Per-stage StageDeployConfig values override pipeline-wide settings.
     if ds is not None:
         for k, v in asdict(ds).items():
-            if k in _STAGE_NON_ENGINE_KEYS or v is None:
+            if k in _STAGE_RESERVED_KEYS or v is None:
                 continue
             engine_args[k] = v
         engine_args.update(ds.engine_extras)

vllm_omni/core/sched/omni_ar_scheduler.py

@@ -20,11 +20,16 @@ from vllm.v1.request import Request, RequestStatus, StreamingUpdate
 from vllm.v1.spec_decode.metrics import SpecDecodingStats
 
 from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
+from vllm_omni.core.sched.omni_scheduling_coordinator import (
+    OmniSchedulingCoordinator,
+    uses_qwen3_omni_full_payload_input_coordinator,
+)
 from vllm_omni.core.sched.output import OmniSchedulerOutput
 from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import (
     OmniChunkTransferAdapter,
 )
 from vllm_omni.engine.serialization import deserialize_additional_information
+from vllm_omni.outputs import OmniConnectorOutput
 
 logger = init_logger(__name__)
 
@@ -75,6 +80,14 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         self.chunk_transfer_adapter = None
         if getattr(model_config, "async_chunk", False):
             self.chunk_transfer_adapter = OmniChunkTransferAdapter(self.vllm_config)
+        self.input_coordinator: OmniSchedulingCoordinator | None = None
+        if uses_qwen3_omni_full_payload_input_coordinator(model_config):
+            self.input_coordinator = OmniSchedulingCoordinator(
+                scheduler_max_num_seqs=self.vllm_config.scheduler_config.max_num_seqs,
+                stage_id=getattr(model_config, "stage_id", 0),
+                async_chunk=False,
+            )
+        self._latest_omni_connector_output: OmniConnectorOutput | None = None
         # Snapshot prompt length for each streaming input update
         self._new_prompt_len_snapshot: dict[str, int] = {}
 
@@ -190,6 +203,27 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         return False
 
     def schedule(self) -> SchedulerOutput:  # type: ignore[override]
+        # Remove FINISHED_ABORTED requests before the upstream scheduler sees
+        # them. Upstream vllm raises RuntimeError on this status; omni allows
+        # async abort (e.g. client disconnect during TTS streaming) to leave
+        # requests in the waiting/running queues temporarily.
+        for queue in (self.waiting, self.running):
+            for req in list(queue):
+                if getattr(req, "status", None) == RequestStatus.FINISHED_ABORTED:
+                    queue.remove(req)
+        connector_output = self._latest_omni_connector_output
+        self._latest_omni_connector_output = None
+        if self.input_coordinator:
+            if connector_output and connector_output.request_metadata:
+                self.input_coordinator.update_request_metadata(
+                    self.requests, connector_output.request_metadata, model_mode="ar"
+                )
+            self.input_coordinator.process_pending_full_payload_inputs(
+                self.waiting,
+                self.running,
+                connector_output.stage_recv_req_ids if connector_output else set(),
+            )
+
         if self.chunk_transfer_adapter:
             self.chunk_transfer_adapter.process_pending_chunks(self.waiting, self.running)
 
@@ -199,6 +233,8 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
             if self.chunk_transfer_adapter:
                 # Add request waiting for chunk to the waiting and running queue
                 self.chunk_transfer_adapter.restore_queues(self.waiting, self.running)
+            if self.input_coordinator:
+                self.input_coordinator.restore_queues(self.waiting, self.running)
         try:
             # Late import to avoid circulars in some launch modes
             from .output import OmniNewRequestData
@@ -222,6 +258,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                     lora_request=nr.lora_request,
                     # Enrich with omni payloads from the live request object
                     prompt_embeds=(getattr(request, "prompt_embeds", None) if request else None),
+                    prompt_is_token_ids=nr.prompt_is_token_ids,
                     additional_information=(getattr(request, "additional_information", None) if request else None),
                 )
                 new_list.append(omni_nr)
@@ -239,9 +276,11 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         # Wrap in omni scheduler output to carry transfer metadata.
         base_fields = SchedulerOutput.__dataclass_fields__.keys()
         base_data = {name: getattr(scheduler_output, name) for name in base_fields}
+        input_regs = self.input_coordinator.pending_input_registrations if self.input_coordinator else []
         return OmniSchedulerOutput(
             **base_data,
             finished_requests_needing_kv_transfer=finished_reqs,
+            pending_input_registrations=input_regs,
         )
 
     def update_from_output(
@@ -389,7 +428,11 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                     stopped = True
 
             if stopped:
-                routed_experts = self._get_routed_experts(request)
+                if (
+                    model_runner_output.routed_experts_dict is not None
+                    and req_id in model_runner_output.routed_experts_dict
+                ):
+                    routed_experts = model_runner_output.routed_experts_dict[req_id]
 
                 # Capture finish_reason BEFORE _handle_stopped_request, which may
                 # reset the status to WAITING for streaming requests that continue.
@@ -534,6 +577,16 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                 engine_core_outputs[0] = eco = EngineCoreOutputs()
             eco.scheduler_stats = stats
 
+        omni_output = getattr(model_runner_output, "omni_connector_output", None)
+        if omni_output is not None:
+            self._latest_omni_connector_output = omni_output
+            if self.input_coordinator and omni_output.request_metadata:
+                self.input_coordinator.update_request_metadata(
+                    self.requests,
+                    omni_output.request_metadata,
+                    model_mode="ar",
+                )
+
         # Free blocks that were held for transfer (kv_ready and
         # active_kv_transfers updates already done before the per-request loop).
         if kv_extracted_ids:
@@ -572,7 +625,12 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         if self.chunk_transfer_adapter:
             self.chunk_transfer_adapter.finish_requests(request_ids, finished_status, self.requests)
 
-        return super().finish_requests(request_ids, finished_status)
+        finished = super().finish_requests(request_ids, finished_status)
+        input_coordinator = getattr(self, "input_coordinator", None)
+        if input_coordinator is not None:
+            for request_id, _ in finished:
+                self._free_input_coordinator_request(request_id)
+        return finished
 
     def _update_request_as_session(self, session: Request, update: StreamingUpdate) -> None:
         """
@@ -616,10 +674,14 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                     # It triggered but hasn't finished yet. We MUST wait.
                     logger.debug(f"[Omni] Request {request_id} finished but transfer is still ACTIVE. Waiting.")
                     self.waiting_for_transfer_free.add(request_id)
+                    if self.input_coordinator is not None:
+                        self._free_input_coordinator_request(request_id)
                     kv_xfer_params = None
                     return kv_xfer_params
                 elif request_id in self.waiting_for_transfer_free:
                     # Blocks held until KV extraction completes in a future step.
+                    if self.input_coordinator is not None:
+                        self._free_input_coordinator_request(request_id)
                     return None
                 else:
                     logger.debug(
@@ -654,10 +716,14 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                     if isinstance(add_info, dict):
                         add_info.update(kv_xfer_params)
 
+                if self.input_coordinator is not None:
+                    self._free_input_coordinator_request(request_id)
                 return kv_xfer_params
 
         # 3. Standard Freeing
         delay_free_blocks |= connector_delay_free_blocks
+        if self.input_coordinator is not None:
+            self._free_input_coordinator_request(request_id)
         if not delay_free_blocks:
             self._free_blocks(request)
 

vllm_omni/core/sched/omni_generation_scheduler.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import time
 from collections import defaultdict
+from typing import Any
 
 from vllm.compilation.cuda_graph import CUDAGraphStat
 from vllm.distributed.kv_events import KVEventBatch
@@ -23,11 +24,15 @@ from vllm.v1.request import Request, RequestStatus, StreamingUpdate
 from vllm.v1.spec_decode.metrics import SpecDecodingStats
 
 from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
-from vllm_omni.core.sched.output import OmniCachedRequestData, OmniNewRequestData
+from vllm_omni.core.sched.omni_scheduling_coordinator import (
+    OmniSchedulingCoordinator,
+    uses_qwen3_omni_full_payload_input_coordinator,
+)
+from vllm_omni.core.sched.output import OmniCachedRequestData, OmniNewRequestData, OmniSchedulerOutput
 from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import (
     OmniChunkTransferAdapter,
 )
-from vllm_omni.outputs import OmniModelRunnerOutput
+from vllm_omni.outputs import OmniConnectorOutput, OmniModelRunnerOutput
 
 logger = init_logger(__name__)
 
@@ -39,6 +44,15 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
         self.chunk_transfer_adapter = None
         if getattr(model_config, "async_chunk", False):
             self.chunk_transfer_adapter = OmniChunkTransferAdapter(self.vllm_config)
+        self._pending_finish_reqs: list[Request] = []
+        self.input_coordinator: OmniSchedulingCoordinator | None = None
+        if uses_qwen3_omni_full_payload_input_coordinator(model_config):
+            self.input_coordinator = OmniSchedulingCoordinator(
+                scheduler_max_num_seqs=self.vllm_config.scheduler_config.max_num_seqs,
+                stage_id=getattr(model_config, "stage_id", 0),
+                async_chunk=False,
+            )
+        self._latest_omni_connector_output: OmniConnectorOutput | None = None
 
     def schedule(self) -> SchedulerOutput:
         """Diffusion fast path:
@@ -68,6 +82,18 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
         # Temporary queue: preserve waiting order, do not disturb non-diffusion requests
         skipped_waiting_requests = create_request_queue(self.policy)
         req_index = 0
+        connector_output = self._latest_omni_connector_output
+        self._latest_omni_connector_output = None
+        if self.input_coordinator:
+            if connector_output and connector_output.request_metadata:
+                self.input_coordinator.update_request_metadata(
+                    self.requests, connector_output.request_metadata, model_mode="generation"
+                )
+            self.input_coordinator.process_pending_full_payload_inputs(
+                self.waiting,
+                self.running,
+                connector_output.stage_recv_req_ids if connector_output else set(),
+            )
         if self.chunk_transfer_adapter:
             self.chunk_transfer_adapter.process_pending_chunks(self.waiting, self.running)
 
@@ -86,23 +112,14 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
 
             num_computed_tokens = request.num_computed_tokens
             required_tokens = len(request.prompt_token_ids) - num_computed_tokens
-            # async_chunk: don't schedule placeholder tokens when no new chunk is available.
             if required_tokens <= 0:
                 if (
                     self.chunk_transfer_adapter is not None
                     and request.request_id in self.chunk_transfer_adapter.finished_requests
                 ):
-                    # Upstream may finish with no terminal tokens; append one pad token so we can emit FINISHED.
+                    self._pending_finish_reqs.append(request)
-                    if len(request.prompt_token_ids) <= num_computed_tokens:
+                req_index += 1
-                        request.prompt_token_ids.append(0)
+                continue
-                        try:
-                            request._all_token_ids.append(0)  # type: ignore[attr-defined]
-                        except Exception:
-                            pass
-                    required_tokens = len(request.prompt_token_ids) - num_computed_tokens
-                else:
-                    req_index += 1
-                    continue
             num_new_tokens = min(required_tokens, token_budget)
             new_blocks = self.kv_cache_manager.allocate_slots(
                 request,
@@ -149,11 +166,9 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
             # async_chunk: wait for the first upstream chunk (don't start with placeholders).
             if self.chunk_transfer_adapter is not None and len(request.prompt_token_ids) == 0:
                 if request.request_id in self.chunk_transfer_adapter.finished_requests:
-                    request.prompt_token_ids.append(0)
+                    self.waiting.pop_request()
-                    try:
+                    self._pending_finish_reqs.append(request)
-                        request._all_token_ids.append(0)  # type: ignore[attr-defined]
+                    continue
-                    except Exception:
-                        pass
                 else:
                     self.waiting.pop_request()
                     skipped_waiting_requests.prepend_request(request)
@@ -201,7 +216,16 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
                 self.chunk_transfer_adapter.restore_queues(self.waiting, self.running)
             else:
                 res = super().schedule()
-                return res
+                if self.input_coordinator:
+                    self.input_coordinator.restore_queues(self.waiting, self.running)
+                base_fields = SchedulerOutput.__dataclass_fields__.keys()
+                base_data = {name: getattr(res, name) for name in base_fields}
+                return OmniSchedulerOutput(
+                    **base_data,
+                    pending_input_registrations=(
+                        self.input_coordinator.pending_input_registrations if self.input_coordinator else []
+                    ),
+                )
 
         # Compute common prefix blocks (aligned with v1)
         num_common_prefix_blocks = [0] * len(self.kv_cache_config.kv_cache_groups)
@@ -303,6 +327,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
                     lora_request=nr.lora_request,
                     # Enrich with omni payloads from the live request object
                     prompt_embeds=(getattr(request, "prompt_embeds", None) if request else None),
+                    prompt_is_token_ids=nr.prompt_is_token_ids,
                     additional_information=(getattr(request, "additional_information", None) if request else None),
                 )
                 new_list.append(omni_nr)
@@ -321,8 +346,17 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
             # internal deques and never scheduled again.
             if self.chunk_transfer_adapter:
                 self.chunk_transfer_adapter.restore_queues(self.waiting, self.running)
+            if self.input_coordinator:
+                self.input_coordinator.restore_queues(self.waiting, self.running)
 
-        return scheduler_output
+        base_fields = SchedulerOutput.__dataclass_fields__.keys()
+        base_data = {name: getattr(scheduler_output, name) for name in base_fields}
+        return OmniSchedulerOutput(
+            **base_data,
+            pending_input_registrations=(
+                self.input_coordinator.pending_input_registrations if self.input_coordinator else []
+            ),
+        )
 
     def finish_requests(self, request_ids, finished_status: RequestStatus) -> list[tuple[str, int]]:
         """Handles the finish signal from outside the scheduler.
@@ -340,7 +374,20 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
         if self.chunk_transfer_adapter:
             self.chunk_transfer_adapter.finish_requests(request_ids, finished_status, self.requests)
 
-        return super().finish_requests(request_ids, finished_status)
+        finished = super().finish_requests(request_ids, finished_status)
+        if self.input_coordinator is not None:
+            for request_id, _ in finished:
+                self._free_input_coordinator_request(request_id)
+        return finished
+
+    def _free_request(self, request: Request, delay_free_blocks: bool = False) -> dict[str, Any] | None:
+        if self.input_coordinator is None:
+            return super()._free_request(request, delay_free_blocks)
+
+        try:
+            return super()._free_request(request, delay_free_blocks)
+        finally:
+            self._free_input_coordinator_request(request.request_id)
 
     """
     Scheduler for the diffusion model.
@@ -464,7 +511,11 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
                 stopped = True
 
             if stopped:
-                routed_experts = self._get_routed_experts(request)
+                if (
+                    model_runner_output.routed_experts_dict is not None
+                    and req_id in model_runner_output.routed_experts_dict
+                ):
+                    routed_experts = model_runner_output.routed_experts_dict[req_id]
                 finish_reason = request.get_finished_reason()
                 finished = self._handle_stopped_request(request)
                 if finished:
@@ -520,6 +571,36 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
                 # Invariant: EngineCore returns no partial prefill outputs.
                 assert not prompt_logprobs_tensors
 
+        # Finish async_chunk requests that schedule() collected because their
+        # upstream completed with no remaining codec tokens.
+        for request in self._pending_finish_reqs:
+            if request.is_finished():
+                continue
+            request.status = RequestStatus.FINISHED_STOPPED
+            finish_reason = request.get_finished_reason()
+            finished = self._handle_stopped_request(request)
+            kv_transfer_params = None
+            if finished:
+                kv_transfer_params = self._free_request(request)
+                if self.chunk_transfer_adapter is not None:
+                    self.chunk_transfer_adapter.cleanup(
+                        request.request_id,
+                        getattr(request, "external_req_id", None),
+                    )
+            outputs[request.client_index].append(
+                EngineCoreOutput(
+                    request_id=request.request_id,
+                    new_token_ids=[],
+                    finish_reason=finish_reason,
+                    stop_reason=request.stop_reason,
+                    events=request.take_events(),
+                    kv_transfer_params=kv_transfer_params,
+                    trace_headers=request.trace_headers,
+                )
+            )
+            stopped_running_reqs.add(request)
+        self._pending_finish_reqs.clear()
+
         # Remove the stopped requests from the running and waiting queues.
         if stopped_running_reqs:
             self.running = remove_all(self.running, stopped_running_reqs)
@@ -589,6 +670,16 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
                 engine_core_outputs[0] = eco = EngineCoreOutputs()
             eco.scheduler_stats = stats
 
+        omni_output = getattr(model_runner_output, "omni_connector_output", None)
+        if omni_output is not None:
+            self._latest_omni_connector_output = omni_output
+            if self.input_coordinator and omni_output.request_metadata:
+                self.input_coordinator.update_request_metadata(
+                    self.requests,
+                    omni_output.request_metadata,
+                    model_mode="generation",
+                )
+
         return engine_core_outputs
 
     def _update_request_as_session(self, session: Request, update: StreamingUpdate) -> None:

vllm_omni/core/sched/omni_scheduler_mixin.py

@@ -7,6 +7,12 @@ from vllm.v1.request import Request, RequestStatus, StreamingUpdate
 class OmniSchedulerMixin:
     """Shared scheduler helpers for omni-specific request handling."""
 
+    def _free_input_coordinator_request(self, request_id: str) -> None:
+        """Prune full-payload coordinator state for a completed request."""
+        input_coordinator = getattr(self, "input_coordinator", None)
+        if input_coordinator is not None:
+            input_coordinator.free_finished_request(request_id)
+
     def _replace_session_with_streaming_update(
         self,
         session: Request,

vllm_omni/core/sched/omni_scheduling_coordinator.py

@@ -22,6 +22,15 @@ from vllm.v1.request import Request, RequestStatus
 logger = init_logger(__name__)
 
 
+def uses_qwen3_omni_full_payload_input_coordinator(model_config: Any) -> bool:
+    return (
+        getattr(model_config, "stage_id", 0) > 0
+        and not getattr(model_config, "async_chunk", False)
+        and getattr(model_config, "model_arch", None) == "Qwen3OmniMoeForConditionalGeneration"
+        and getattr(model_config, "model_stage", None) in {"talker", "code2wav"}
+    )
+
+
 class OmniSchedulingCoordinator:
     """Pure-scheduling coordinator for chunk and full_payload input waiting.
 
@@ -166,8 +175,10 @@ class OmniSchedulingCoordinator:
                         to_remove.append(request)
                         self._waiting_for_input.append(request)
                         self.pending_input_registrations.append(request)
-            for request in to_remove:
+            if to_remove:
-                waiting_queue.remove(request)
+                # Use the bulk-remove helper: one O(N) sweep instead of N
+                # repeated O(N) removes from a list-backed queue.
+                waiting_queue.remove_requests(to_remove)
 
     def process_pending_full_payload_inputs_legacy(
         self,
@@ -255,6 +266,30 @@ class OmniSchedulingCoordinator:
             waiting_queue.add_request(request)
         self._waiting_for_input = deque()
 
+    @staticmethod
+    def _flatten_prompt_token_ids(value: Any) -> list[int]:
+        """Normalize connector metadata into flat prompt token ids."""
+        if value is None:
+            return []
+        if hasattr(value, "detach") and hasattr(value, "cpu") and hasattr(value, "tolist"):
+            value = value.detach().cpu().tolist()
+        elif hasattr(value, "tolist") and not isinstance(value, (list, tuple)):
+            value = value.tolist()
+
+        if isinstance(value, (list, tuple)):
+            flattened: list[int] = []
+            for item in value:
+                if hasattr(item, "detach") and hasattr(item, "cpu") and hasattr(item, "tolist"):
+                    item = item.detach().cpu().tolist()
+                elif hasattr(item, "tolist") and not isinstance(item, (list, tuple)):
+                    item = item.tolist()
+                if isinstance(item, (list, tuple)):
+                    flattened.extend(int(token_id) for token_id in item)
+                else:
+                    flattened.append(int(item))
+            return flattened
+        return [int(value)]
+
     def update_request_metadata(
         self,
         requests: dict[str, Request],
@@ -310,7 +345,7 @@ class OmniSchedulingCoordinator:
                             )
 
             if model_mode != "ar":
-                new_ids = metadata.get("code_predictor_codes", [])
+                new_ids = self._flatten_prompt_token_ids(metadata.get("code_predictor_codes"))
                 runtime_seed = None
                 if "left_context_size" in metadata:
                     runtime_seed = {
@@ -319,6 +354,10 @@ class OmniSchedulingCoordinator:
                 request._omni_initial_model_buffer = runtime_seed
                 if new_ids:
                     request.prompt_token_ids = new_ids
+                    request.num_prompt_tokens = len(new_ids)
+                    request._all_token_ids.clear()
+                    request._all_token_ids.extend(new_ids)
+                    request._output_token_ids.clear()
                     request.num_computed_tokens = 0
 
     def postprocess_scheduler_output(
@@ -374,7 +413,3 @@ class OmniSchedulingCoordinator:
         if scheduler_output.scheduled_cached_reqs:
             for req_id in scheduler_output.scheduled_cached_reqs.req_ids:
                 self.requests_with_ready_chunks.discard(req_id)
-
-
-# Backward-compatible alias
-ChunkSchedulingCoordinator = OmniSchedulingCoordinator

vllm_omni/core/sched/output.py

@@ -1,4 +1,5 @@
 from dataclasses import dataclass, field
+from typing import Any
 
 from vllm.v1.core.sched.output import CachedRequestData, NewRequestData, SchedulerOutput
 from vllm.v1.request import Request
@@ -53,6 +54,7 @@ class OmniNewRequestData(NewRequestData):
             num_computed_tokens=request.num_computed_tokens,
             lora_request=request.lora_request,
             prompt_embeds=getattr(request, "prompt_embeds", None),
+            prompt_is_token_ids=getattr(request, "prompt_is_token_ids", None),
             prefill_token_ids=prefill_token_ids,
             additional_information=getattr(request, "additional_information", None),
         )
@@ -75,3 +77,4 @@ class OmniSchedulerOutput(SchedulerOutput):
     """Scheduler output with omni-specific transfer metadata."""
 
     finished_requests_needing_kv_transfer: dict[str, dict] = field(default_factory=dict)
+    pending_input_registrations: list[Any] = field(default_factory=list)

vllm_omni/deploy/hunyuan_image3.yaml

@@ -22,6 +22,13 @@ connectors:
 
 stages:
   - stage_id: 0
+    # ``is_comprehension`` in vllm-omni names the tokenizer-owning AR stage
+    # (see config/stage_config.py + serving_chat AR-stage lookup), independent
+    # of whether the AR's task is comprehension (i2t/t2t) or generation
+    # (it2i/t2i). HunyuanImage-3.0's stage-0 owns the tokenizer and emits the
+    # cot+ratio token sequence consumed by stage-1, so it must be marked True
+    # for the serving path to set AR seed/stop_token_ids on this stage.
+    is_comprehension: true
     final_output: true
     final_output_type: text
     max_num_seqs: 1

vllm_omni/deploy/mimo_audio.yaml

@@ -23,7 +23,7 @@ connectors:
 stages:
   - stage_id: 0
     max_num_seqs: 1
-    gpu_memory_utilization: 0.3
+    gpu_memory_utilization: 0.5
     enforce_eager: true
     trust_remote_code: true
     enable_prefix_caching: false
@@ -42,7 +42,7 @@ stages:
 
   - stage_id: 1
     max_num_seqs: 1
-    gpu_memory_utilization: 0.2
+    gpu_memory_utilization: 0.3
     enforce_eager: true
     trust_remote_code: true
     enable_prefix_caching: false

vllm_omni/diffusion/attention/layer.py

@@ -171,7 +171,7 @@ class Attention(nn.Module):
     def _init_kv_cache_quantization(self, config) -> None:
         if config is None:
             return
-        dtype = config.kv_cache_dtype
+        dtype = config.diffusion_kv_cache_dtype
         if dtype:
             if config.parallel_config.ring_degree > 1:
                 raise ValueError(
@@ -190,8 +190,8 @@ class Attention(nn.Module):
                 )
                 dtype = None
         self._kv_cache_dtype = dtype
-        self._kv_cache_skip_steps = config.kv_cache_skip_step_indices
+        self._kv_cache_skip_steps = config.diffusion_kv_cache_skip_step_indices
-        self._kv_cache_skip_layers = config.kv_cache_skip_layer_indices
+        self._kv_cache_skip_layers = config.diffusion_kv_cache_skip_layer_indices
 
     def _should_apply_kv_cache_quant(self) -> bool:
         skip_steps = self._kv_cache_skip_steps

vllm_omni/diffusion/cache/cache_dit_backend.py

@@ -1383,6 +1383,68 @@ def enable_cache_for_ernie_image(pipeline: Any, cache_config: Any) -> Callable[[
     return refresh_cache_context
 
 
+def enable_cache_for_helios(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
+    """Enable cache-dit for Helios pipeline.
+
+    Helios extends Wan2.2 with multi-term memory patches and guidance cross-attention.
+    Its transformer blocks have the same single-output signature as Wan2.2 blocks
+    (returns hidden_states only), so we use ForwardPattern.Pattern_2.
+
+    Args:
+        pipeline: The HeliosPipeline instance.
+        cache_config: DiffusionCacheConfig instance with cache configuration.
+
+    Returns:
+        A refresh function that can be called to update cache context with new num_inference_steps.
+    """
+    db_cache_config = _build_db_cache_config(cache_config)
+
+    calibrator_config = None
+    if cache_config.enable_taylorseer:
+        taylorseer_order = cache_config.taylorseer_order
+        calibrator_config = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order)
+        logger.info(f"TaylorSeer enabled with order={taylorseer_order}")
+
+    cache_dit.enable_cache(
+        BlockAdapter(
+            transformer=pipeline.transformer,
+            blocks=[pipeline.transformer.blocks],
+            forward_pattern=[ForwardPattern.Pattern_2],
+            params_modifiers=[
+                ParamsModifier(cache_config=db_cache_config, calibrator_config=calibrator_config),
+            ],
+            has_separate_cfg=False,
+        ),
+        cache_config=db_cache_config,
+        calibrator_config=calibrator_config,
+    )
+
+    logger.info(
+        f"Enabling cache-dit on Helios transformer: "
+        f"Fn={db_cache_config.Fn_compute_blocks}, "
+        f"Bn={db_cache_config.Bn_compute_blocks}, "
+        f"W={db_cache_config.max_warmup_steps}, "
+    )
+
+    def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
+        if cache_config.scm_steps_mask_policy is None:
+            cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose)
+        else:
+            cache_dit.refresh_context(
+                pipeline.transformer,
+                cache_config=DBCacheConfig().reset(
+                    num_inference_steps=num_inference_steps,
+                    steps_computation_mask=cache_dit.steps_mask(
+                        mask_policy=cache_config.scm_steps_mask_policy, total_steps=num_inference_steps
+                    ),
+                    steps_computation_policy=cache_config.scm_steps_policy,
+                ),
+                verbose=verbose,
+            )
+
+    return refresh_cache_context
+
+
 def enable_cache_for_hunyuan_video_15(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
     """Enable cache-dit for HunyuanVideo 1.5 pipeline.
 
@@ -1463,6 +1525,7 @@ CUSTOM_DIT_ENABLERS.update(
         "ErnieImagePipeline": enable_cache_for_ernie_image,
         "HunyuanVideo15Pipeline": enable_cache_for_hunyuan_video_15,
         "HunyuanVideo15I2VPipeline": enable_cache_for_hunyuan_video_15,
+        "HeliosPipeline": enable_cache_for_helios,
     }
 )
 

vllm_omni/diffusion/data.py

@@ -568,18 +568,18 @@ class OmniDiffusionConfig:
     # has already resolved to vLLM's ModelOpt FP8 linear method.
     force_cutlass_fp8: bool = False
 
-    # KV cache dtype for attention. Aligned with upstream vLLM's --kv-cache-dtype.
+    # Diffusion attention KV cache dtype (not vLLM's --kv-cache-dtype for AR models).
     # None = native dtype (no quantization).
     # "fp8" = dynamic FP8 (float8_e4m3fn) quantization per forward pass.
     # On Hopper+FA3: native FP8 attention (memory + compute savings).
     # On other backends: no benefit, backends skip quantization.
-    kv_cache_dtype: str | None = None
+    diffusion_kv_cache_dtype: str | None = None
     # Optional skip selectors for KV-cache quantization. Format: "0-9,20,25-30".
     # Listed steps/layers skip quantization; others keep quantized execution.
-    kv_cache_skip_steps: str | None = None
+    diffusion_kv_cache_skip_steps: str | None = None
-    kv_cache_skip_layers: str | None = None
+    diffusion_kv_cache_skip_layers: str | None = None
-    kv_cache_skip_step_indices: set[int] | None = None
+    diffusion_kv_cache_skip_step_indices: set[int] | None = None
-    kv_cache_skip_layer_indices: set[int] | None = None
+    diffusion_kv_cache_skip_layer_indices: set[int] | None = None
 
     # Diffusion pipeline Profiling config
     enable_diffusion_pipeline_profiler: bool = False
@@ -726,8 +726,8 @@ class OmniDiffusionConfig:
         # Match vLLM's config flow: parse entrypoint shorthands before the
         # config object is built, and keep a single runtime truth source.
         self.diffusion_attention_config = build_attention_config(self.diffusion_attention_config)
-        self.kv_cache_skip_step_indices = parse_kv_cache_skip_selector(self.kv_cache_skip_steps)
+        self.diffusion_kv_cache_skip_step_indices = parse_kv_cache_skip_selector(self.diffusion_kv_cache_skip_steps)
-        self.kv_cache_skip_layer_indices = parse_kv_cache_skip_selector(self.kv_cache_skip_layers)
+        self.diffusion_kv_cache_skip_layer_indices = parse_kv_cache_skip_selector(self.diffusion_kv_cache_skip_layers)
 
         if self.max_cpu_loras is None:
             self.max_cpu_loras = 1
@@ -814,6 +814,15 @@ class OmniDiffusionConfig:
                 # (non-DiT models don't have a separate transformer folder/config)
                 if self.diffusion_load_format == "diffusers":
                     self.set_tf_model_config(TransformerConfig())
+                    try:
+                        diffusers_pipeline_cls_name = config_dict["_class_name"]
+                        self.diffusers_pipeline_cls = getattr(diffusers, diffusers_pipeline_cls_name)
+                    except (KeyError, AttributeError) as exc:
+                        logger.warning(
+                            "Could not find valid _class_name for diffusers pipeline in model_index.json: %s. "
+                            "Without the underlying pipeline class the dummy run may omit required inputs.",
+                            exc,
+                        )
                 else:
                     tf_config_dict = get_hf_file_to_dict("transformer/config.json", self.model)
                     self.set_tf_model_config(TransformerConfig.from_dict(tf_config_dict))
@@ -824,7 +833,13 @@ class OmniDiffusionConfig:
             # (non-DiT models don't have a separate transformer folder/config)
             if self.diffusion_load_format == "diffusers":
                 self.set_tf_model_config(TransformerConfig())
-                self.update_multimodal_support()
+                logger.warning(
+                    "Could not find valid model_index.json per diffusers format. "
+                    "This model is likely unsupported by the diffusers backend. "
+                    "Also, without knowing the underlying diffusers pipeline class from model_index.json, "
+                    "the dummy run will input only text prompt, which may cause errors for pipelines "
+                    "that require additional inputs."
+                )
             else:
                 cfg = get_hf_file_to_dict("config.json", self.model)
                 if cfg is None:
@@ -874,6 +889,20 @@ class OmniDiffusionConfig:
         else:
             kwargs.pop("quantization", None)
 
+        # Renamed from kv_cache_* to avoid clashing with vLLM's --kv-cache-dtype.
+        if kwargs.get("diffusion_kv_cache_dtype") is None and "kv_cache_dtype" in kwargs:
+            kwargs["diffusion_kv_cache_dtype"] = kwargs.pop("kv_cache_dtype")
+        else:
+            kwargs.pop("kv_cache_dtype", None)
+        if kwargs.get("diffusion_kv_cache_skip_steps") is None and "kv_cache_skip_steps" in kwargs:
+            kwargs["diffusion_kv_cache_skip_steps"] = kwargs.pop("kv_cache_skip_steps")
+        else:
+            kwargs.pop("kv_cache_skip_steps", None)
+        if kwargs.get("diffusion_kv_cache_skip_layers") is None and "kv_cache_skip_layers" in kwargs:
+            kwargs["diffusion_kv_cache_skip_layers"] = kwargs.pop("kv_cache_skip_layers")
+        else:
+            kwargs.pop("kv_cache_skip_layers", None)
+
         # Handle "diffusion_attention_backend" shorthand: merge into
         # diffusion_attention_config before field filtering.
         diffusion_attn_backend = kwargs.pop("diffusion_attention_backend", None)

vllm_omni/diffusion/executor/abstract.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
+from collections.abc import Callable
 from typing import TYPE_CHECKING, Any
 
 from vllm.utils.import_utils import resolve_obj_by_qualname
@@ -101,6 +102,14 @@ class DiffusionExecutor(ABC):
         """Check if the executor and workers are healthy."""
         pass
 
+    def register_failure_callback(self, callback: Callable[[], None]) -> None:
+        """Register a callback invoked when the executor fatally fails.
+
+        Executors without a background failure monitor can keep the default
+        no-op implementation.
+        """
+        return None
+
     @abstractmethod
     def shutdown(self) -> None:
         """Shutdown the executor and release resources."""

vllm_omni/diffusion/inline_stage_diffusion_client.py

@@ -64,8 +64,11 @@ class InlineStageDiffusionClient(StageClientBase):
 
         self._output_queue: asyncio.Queue[OmniRequestOutput] = asyncio.Queue()
         self._tasks: dict[str, asyncio.Task] = {}
+        self._engine_dead = False
         self._shutting_down = False
 
+        self._engine.executor.register_failure_callback(self._mark_engine_dead)
+
         logger.info(
             "[InlineStageDiffusionClient] stage-%s [rep-%s] initialized inline (batch_size=%d)",
             self.stage_id,
@@ -77,6 +80,16 @@ class InlineStageDiffusionClient(StageClientBase):
         """Load model metadata from HuggingFace and populate od_config fields."""
         self.od_config.enrich_config()
 
+    def _mark_engine_dead(self) -> None:
+        if self._engine_dead:
+            return
+        self._engine_dead = True
+        logger.error(
+            "[InlineStageDiffusionClient] stage-%s [rep-%s] diffusion executor died unexpectedly.",
+            self.stage_id,
+            self.replica_id,
+        )
+
     # ------------------------------------------------------------------
     # Request processing
     # ------------------------------------------------------------------
@@ -249,6 +262,8 @@ class InlineStageDiffusionClient(StageClientBase):
         try:
             return self._output_queue.get_nowait()
         except asyncio.QueueEmpty:
+            if self._engine_dead:
+                raise EngineDeadError(f"Stage-{self.stage_id} inline diffusion engine is dead")
             return None
 
     async def abort_requests_async(self, request_ids: list[str]) -> None:
@@ -351,7 +366,11 @@ class InlineStageDiffusionClient(StageClientBase):
         """Check if the inline diffusion engine and its workers are healthy."""
         if self._shutting_down:
             raise EngineDeadError("InlineStageDiffusionClient is shutting down")
-        self._engine.executor.check_health()
+        try:
+            self._engine.executor.check_health()
+        except EngineDeadError:
+            self._mark_engine_dead()
+            raise
 
     def shutdown(self) -> None:
         self._shutting_down = True

vllm_omni/diffusion/model_metadata.py

@@ -13,6 +13,8 @@ class DiffusionModelMetadata:
 
 
 QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES = 4
+# Upstream HunyuanImage-3.0 "Multi-Image Fusion" caps reference images at 3.
+HUNYUAN_IMAGE3_MAX_INPUT_IMAGES = 3
 
 
 _DIFFUSION_MODEL_METADATA: dict[str, DiffusionModelMetadata] = {
@@ -20,6 +22,10 @@ _DIFFUSION_MODEL_METADATA: dict[str, DiffusionModelMetadata] = {
         supports_multimodal_inputs=True,
         max_multimodal_image_inputs=QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES,
     ),
+    "HunyuanImage3Pipeline": DiffusionModelMetadata(
+        supports_multimodal_inputs=True,
+        max_multimodal_image_inputs=HUNYUAN_IMAGE3_MAX_INPUT_IMAGES,
+    ),
 }
 
 

vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py

@@ -471,8 +471,21 @@ class Resolution:
         return f"{self.h}x{self.w}"
 
 
+# Baked-in extras matching the official model's
+# `HunyuanImage3ImageProcessor.vae_reso_group` (image_processor.py:147-152).
+# These four aspect buckets sit at ratio_token indices 33-36 in the trained
+# model and the AR was trained to address them, so any deviation breaks the
+# ratio-token vocab → output-shape lookup.
+HUNYUAN_IMAGE3_EXTRA_RESOLUTIONS: tuple[str, ...] = (
+    "1024x768",
+    "1280x720",
+    "768x1024",
+    "720x1280",
+)
+
+
 class ResolutionGroup:
-    def __init__(self, base_size=None, step=None, align=1):
+    def __init__(self, base_size=None, step=None, align=1, extra_resolutions=None):
         self.align = align
         self.base_size = base_size
         assert base_size % align == 0, f"base_size {base_size} is not divisible by align {align}"
@@ -486,6 +499,11 @@ class ResolutionGroup:
         self.step = step
         self.data = self._calc_by_step()
 
+        if extra_resolutions is not None:
+            for er in extra_resolutions:
+                if not any(r.ratio == er.ratio for r in self.data):
+                    self.data.append(er)
+
         self.ratio = np.array([x.ratio for x in self.data])
         self.attr = ["" for _ in range(len(self.data))]
         self.prefix_space = 0
@@ -1351,7 +1369,10 @@ class HunyuanImage3ImageProcessor:
     def __init__(self, config):
         self.config = config
 
-        self.reso_group = ResolutionGroup(base_size=config.image_base_size)
+        self.reso_group = ResolutionGroup(
+            base_size=config.image_base_size,
+            extra_resolutions=[Resolution(s) for s in HUNYUAN_IMAGE3_EXTRA_RESOLUTIONS],
+        )
         self.vae_processor = transforms.Compose(
             [
                 transforms.ToTensor(),

vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py

@@ -283,11 +283,13 @@ def get_hunyuan_image_3_pre_process_func(od_config: OmniDiffusionConfig):
                 cond_image_infos = [_build_cond_joint_image(image) for image in image_list]
                 prompt["additional_information"]["batch_cond_image_info"] = cond_image_infos
 
+                bridge_h = prompt.get("height") if isinstance(prompt, dict) else None
+                bridge_w = prompt.get("width") if isinstance(prompt, dict) else None
                 first_image_w, first_image_h = _to_pil_image(image_list[0]).size
                 if request.sampling_params.width is None:
-                    request.sampling_params.width = int(first_image_w)
+                    request.sampling_params.width = int(bridge_w or first_image_w)
                 if request.sampling_params.height is None:
-                    request.sampling_params.height = int(first_image_h)
+                    request.sampling_params.height = int(bridge_h or first_image_h)
 
             request.prompts[i] = prompt
 
@@ -539,7 +541,12 @@ class HunyuanImage3Pipeline(
         timestep_scatter_index: BatchRaggedTensor,
     ):
         batch_size, seq_len, n_embd = x.shape
-        # batch_size x n x n_embd
+        # `_encode_cond_image` returns `t` as list[Tensor] for the
+        # multi-image branch (outer length = batch_size, currently fixed
+        # at 1 by the stage runtime `max_batch_size`); flatten to a Tensor
+        # before reshape.
+        if isinstance(t, list):
+            t = torch.cat([ti.reshape(-1) for ti in t], dim=0)
         timestep_scatter_src = self.timestep_emb(t.reshape(-1)).reshape(batch_size, -1, n_embd)
         x.scatter_(
             dim=1,
@@ -627,7 +634,11 @@ class HunyuanImage3Pipeline(
             if isinstance(vae_encode_result, torch.Tensor):
                 latents = vae_encode_result
             else:
-                latents = vae_encode_result.latent_dist.sample()
+                # Match HunyuanImage-3's cond encode path: sample the
+                # posterior, but use a fixed generator so repeated online
+                # requests are deterministic.
+                _cond_vae_gen = torch.Generator(device=image.device).manual_seed(0)
+                latents = vae_encode_result.latent_dist.sample(_cond_vae_gen)
             if hasattr(config, "shift_factor") and config.shift_factor:
                 latents.sub_(config.shift_factor)
             if hasattr(config, "scaling_factor") and config.scaling_factor:
@@ -1353,25 +1364,21 @@ class HunyuanImage3Pipeline(
         use_system_prompt = extra_args.get("use_system_prompt")
         system_prompt = extra_args.get("system_prompt")
         # Fall back to per-prompt use_system_prompt forwarded by ar2diffusion
-        if use_system_prompt is None and req.prompts:
+        if req.prompts:
             first_prompt = req.prompts[0]
             if isinstance(first_prompt, dict):
-                use_system_prompt = first_prompt.get("use_system_prompt")
+                if use_system_prompt is None:
+                    use_system_prompt = first_prompt.get("use_system_prompt")
+                if system_prompt is None:
+                    system_prompt = first_prompt.get("system_prompt")
         if use_system_prompt is not None:
             system_prompt = get_system_prompt(use_system_prompt, "image", system_prompt)
             system_prompt = system_prompt.strip() if system_prompt is not None else ""
         prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt
 
-        # Extract AR-generated CoT/recaption text from each prompt's extra dict.
+        cot_text_list = [
-        # The AR-side stage input processor (``ar2diffusion``) already prepends
+            (p.get("extra", {}).get("ar_generated_text") if isinstance(p, dict) else None) or None for p in req.prompts
-        # the trigger tag (e.g. ``<think>``) when the AR used the KV-reuse
+        ]
-        # pretrain format, so ``ar_generated_text`` is a self-contained string
-        # and ``get_cot_sections()`` can parse the think/recaption structure
-        # directly.
-        cot_text_list = []
-        for p in req.prompts:
-            extra = p.get("extra", {}) if isinstance(p, dict) else {}
-            cot_text_list.append(extra.get("ar_generated_text") or None)
         cot_text = (
             [self._normalize_cot_text(t) for t in cot_text_list] if any(t is not None for t in cot_text_list) else None
         )

vllm_omni/diffusion/models/hunyuan_image3/prompt_utils.py

@@ -11,8 +11,23 @@ The DiT pipeline (`pipeline_hunyuan_image3.py`) builds prompts through
 `JointImageInfo` objects produced by image preprocessing. The example
 flow uses an `<img>` placeholder + `multi_modal_data` instead, so it
 needs a lighter-weight builder that only requires a HF tokenizer. This
-module provides that builder; the task -> template mapping below is the
+module provides that builder; the (task, bot_task) -> template mapping
-canonical mapping for both flows.
+below is the canonical mapping for both flows.
+
+Two orthogonal axes:
+
+  * `task` selects the I/O modality combination, which only controls
+    whether `<img>` placeholders are emitted between `User: ` and the
+    user prompt: ``i2t`` / ``it2i`` produce them, ``t2t`` / ``t2i`` do
+    not.
+
+  * `bot_task` selects the prompting mode and drives both the system
+    prompt and the trigger tag appended after ``Assistant: ``. ``None``
+    (default) gives a plain Assistant turn under the unified prompt;
+    ``think`` / ``recaption`` switch the trigger tag to ``<think>`` /
+    ``<recaption>``; ``think_recaption`` swaps the system prompt for
+    the dedicated combined-mode template; ``vanilla`` drops the chat
+    structure entirely (pretrain template, ``t2i`` only).
 """
 
 from __future__ import annotations
@@ -45,143 +60,234 @@ HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS: dict[str, int] = {
     "<img_ratio_36>": 130106,
 }
 
-# task -> (sys_type, bot_task, trigger_tag)
+# bot_task -> (sys_type, trigger_tag).
+# ``vanilla`` is special-cased downstream: it bypasses the chat template
+# (no ``User:`` / ``Assistant:`` framing) and is only valid with
+# ``task='t2i'``.
+_BOT_TASK_PRESETS: dict[str | None, tuple[str, str | None]] = {
+    None: ("en_unified", None),
+    "think": ("en_unified", "<think>"),
+    "recaption": ("en_unified", "<recaption>"),
+    "think_recaption": ("en_think_recaption", "<think>"),
+    "vanilla": ("en_vanilla", None),
+}
+
+_TASKS: frozenset[str] = frozenset({"t2t", "i2t", "it2i", "t2i"})
+
+
+class _DefaultBotTask:
+    pass
+
+
+_DEFAULT_BOT_TASK = _DefaultBotTask()
+
+# Legacy composite task alias -> (task, bot_task). Keep this during rebase so
+# older callers and intermediate commits still resolve cleanly.
 _TASK_PRESETS: dict[str, tuple[str, str | None, str | None]] = {
     "t2t": ("en_unified", None, None),
     "i2t": ("en_unified", None, None),
     "it2i_think": ("en_unified", "think", "<think>"),
     "it2i_recaption": ("en_unified", "recaption", "<recaption>"),
     "it2i_think_recaption": ("en_unified", "think_recaption", "<think>"),
-    "t2i": ("en_unified", "image", None),
+    "t2i": ("en_unified", None, None),
-    "t2i_vanilla": ("en_vanilla", "image", None),
+    "t2i_vanilla": ("en_vanilla", "vanilla", None),
     "t2i_think": ("en_unified", "think", "<think>"),
     "t2i_recaption": ("en_unified", "recaption", "<recaption>"),
 }
 
+_LEGACY_COMPOSITE_TASKS: frozenset[str] = frozenset(_TASK_PRESETS) - {"t2t", "i2t", "t2i"}
+
+
+def _normalize_task_and_bot_task(
+    task: str,
+    bot_task: str | None | _DefaultBotTask,
+) -> tuple[str, str | None]:
+    bot_task_was_omitted = bot_task is _DEFAULT_BOT_TASK
+    if task in _TASK_PRESETS:
+        _, legacy_bot_task, _ = _TASK_PRESETS[task]
+        base_task = task.split("_", 1)[0]
+        if base_task == "t2i" and task == "t2i":
+            base_task = "t2i"
+        if task in ("t2t", "i2t", "t2i"):
+            base_task = task
+        if bot_task_was_omitted:
+            bot_task = legacy_bot_task
+        elif task in _LEGACY_COMPOSITE_TASKS and bot_task is None:
+            # Composite task names already encode the legacy bot_task. Keep
+            # calls like build_prompt_tokens(task="it2i_think", bot_task=None)
+            # on their historical meaning; explicit None is the plain-mode
+            # escape hatch only for the new two-axis base tasks.
+            bot_task = legacy_bot_task
+        task = base_task
+    elif bot_task_was_omitted:
+        bot_task = "think"
+    return task, bot_task
+
 
 def available_tasks() -> list[str]:
-    """Sorted list of task keys accepted by `build_prompt` / `build_prompt_tokens`."""
+    """Sorted list of `task` values accepted by the prompt builders."""
-    return sorted(_TASK_PRESETS)
+    return sorted(_TASKS)
+
+
+def available_bot_tasks() -> list[str | None]:
+    """Sorted list of `bot_task` values (with ``None`` first)."""
+    rest = sorted(k for k in _BOT_TASK_PRESETS if k is not None)
+    return [None, *rest]
+
+
+def resolve_sys_type(bot_task: str | None) -> str:
+    """Default system-prompt type for a given ``bot_task``."""
+    if bot_task not in _BOT_TASK_PRESETS:
+        raise ValueError(f"Unknown bot_task {bot_task!r}. Choose from: {available_bot_tasks()}")
+    return _BOT_TASK_PRESETS[bot_task][0]
 
 
 def resolve_stop_token_ids(
-    task: str = "it2i_think",
+    task: str = "it2i",
-    bot_task: str = "think",
+    bot_task: str | None | _DefaultBotTask = _DEFAULT_BOT_TASK,
     tokenizer: Any | None = None,
-):
+) -> list[int]:
+    """AR stop-token ids for a given (task, bot_task) generation request.
+
+    Image-output tasks (``it2i`` / ``t2i``) stop on any ``<img_ratio_*>``
+    token. Upstream ``modeling_hunyuan_image_3.py::generate_image``
+    (line 3289-3303) sets ``final_stop_tokens`` to the full ratio token
+    range when ``need_ratio`` is true, then strips the trailing ratio
+    token before passing the cot to the image stage. AR's natural
+    trajectory under ``_stage_transitions`` is
+    ``</recaption><answer><boi><img_size_base><img_ratio_X>``; stopping
+    AT the ratio token means KV ends exactly at the prefix DiT reuses,
+    and ``ar2diffusion`` can read the ratio off the last sampled token
+    without AR wasting decode steps on ``<|endoftext|>``.
+
+    Text-output tasks (``i2t`` / ``t2t``) stop on ``<answer>`` -- the AR
+    is the final stage, and the comprehension response sits inside the
+    ``<answer>`` body so the answer-open is the natural cot/recaption
+    terminator.
+    """
+    task, bot_task = _normalize_task_and_bot_task(task, bot_task)
+    if task not in _TASKS:
+        raise ValueError(f"Unknown task {task!r}. Choose from: {available_tasks()}")
+    if bot_task not in _BOT_TASK_PRESETS:
+        raise ValueError(f"Unknown bot_task {bot_task!r}. Choose from: {available_bot_tasks()}")
+    if task in ("it2i", "t2i"):
+        # Main ratio range: <img_ratio_0> .. <img_ratio_32>.
+        start = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<img_ratio_0>"]
+        end = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<img_ratio_32>"]
+        stops = list(range(start, end + 1))
+        # Other slices (upstream tokenizer ``ratio_token_other_slices``):
+        # <img_ratio_33> .. <img_ratio_36>.
+        other_start = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<img_ratio_33>"]
+        other_end = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<img_ratio_36>"]
+        stops.extend(range(other_start, other_end + 1))
+        return stops
     return [HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<answer>"]]
 
 
+# Upstream "Multi-Image Fusion" caps reference images at 3 per request.
+MAX_IMAGES_PER_REQUEST = 3
+
+
+def _validate_num_images(num_images: int) -> None:
+    if not (1 <= num_images <= MAX_IMAGES_PER_REQUEST):
+        raise ValueError(f"num_images must be in [1, {MAX_IMAGES_PER_REQUEST}], got {num_images}")
+
+
+def _resolve_preset(task: str, bot_task: str | None) -> tuple[str, str | None]:
+    """Validate (task, bot_task) and return ``(sys_type, trigger_tag)``."""
+    task, bot_task = _normalize_task_and_bot_task(task, bot_task)
+    if task not in _TASKS:
+        raise ValueError(f"Unknown task {task!r}. Choose from: {available_tasks()}")
+    if bot_task not in _BOT_TASK_PRESETS:
+        raise ValueError(f"Unknown bot_task {bot_task!r}. Choose from: {available_bot_tasks()}")
+    if bot_task == "vanilla" and task != "t2i":
+        raise ValueError(f"bot_task='vanilla' is only valid with task='t2i' (pretrain template); got task={task!r}")
+    return _BOT_TASK_PRESETS[bot_task]
+
+
 def build_prompt(
     user_prompt: str,
-    task: str = "it2i_think",
+    task: str = "it2i",
+    bot_task: str | None | _DefaultBotTask = _DEFAULT_BOT_TASK,
     sys_type: str | None = None,
     custom_system_prompt: str | None = None,
+    num_images: int = 1,
 ) -> str:
-    """Build a HunyuanImage-3.0 prompt as a string (legacy/compat path).
+    """Build a HunyuanImage-3.0 prompt as a string (legacy/compat path)."""
-
+    task, bot_task = _normalize_task_and_bot_task(task, bot_task)
-    NOTE: when this string is passed to the engine, the engine's tokenizer
+    preset_sys_type, trigger_tag = _resolve_preset(task, bot_task)
-    will run a single BPE pass over the whole string, which can merge
-    tokens across segment boundaries (e.g. `。\\n\\n` -> id 3490). For
-    inputs that need to match HF baseline byte-for-byte, use
-    `build_prompt_tokens` instead and feed the result via prompt_token_ids.
-    """
-    if task not in _TASK_PRESETS:
-        raise ValueError(f"Unknown task {task!r}. Choose from: {available_tasks()}")
-
-    preset_sys_type, preset_bot_task, trigger_tag = _TASK_PRESETS[task]
     effective_sys_type = sys_type or preset_sys_type
 
-    system_prompt = get_system_prompt(effective_sys_type, preset_bot_task, custom_system_prompt)
+    system_prompt = get_system_prompt(effective_sys_type, bot_task, custom_system_prompt)
-    sys_text = system_prompt.strip() if system_prompt else ""
+    sys_text = system_prompt or ""
 
-    has_image_input = task.startswith("i2t") or task.startswith("it2i")
+    has_image_input = task in ("i2t", "it2i")
+    if has_image_input:
+        _validate_num_images(num_images)
 
-    # t2i_vanilla: pretrain mode for direct text->image generation. The
+    if bot_task == "vanilla":
-    # vanilla system prompt drives the model with no chat structure.
-    if task == "t2i_vanilla":
         parts = ["<|startoftext|>"]
         if sys_text:
             parts.append(sys_text)
         parts.append(user_prompt)
         return "".join(parts)
 
-    # All other tasks (t2t / i2t / t2i_think / t2i_recaption /
-    # it2i_think / it2i_recaption) use HunyuanImage3 Instruct chat template:
-    #   <|startoftext|>{system?}\n\nUser: {<img>?}{user_prompt}\n\nAssistant: {trigger?}
-    # generation_config.json declares sequence_template="instruct", so the
-    # AR prefill MUST use this template -- verified to match HF's
-    # apply_chat_template output token-for-token (modulo BPE boundary merges).
-    # The trigger_tag (e.g. <think>) MUST come AFTER the `Assistant: ` prefix:
-    # if it goes BEFORE user_prompt (the old pretrain layout) the model puts
-    # the user's instructions inside the "thinking section" and collapses
-    # into repetition garbage under greedy decoding.
     parts = ["<|startoftext|>"]
     if sys_text:
         parts.append(f"{sys_text}\n\n")
     parts.append("User: ")
     if has_image_input:
-        parts.append("<img>")
+        parts.extend(["<img>"] * num_images)
     parts.append(user_prompt)
     parts.append("\n\nAssistant: ")
     if trigger_tag:
         parts.append(trigger_tag)
-
     return "".join(parts)
 
 
 @dataclass
 class PromptTokensResult:
-    token_ids: list[int]  # The tokenized prompt
+    token_ids: list[int]
-    system_prompt_type: str  # The effective system prompt type used
+    system_prompt_type: str
 
 
 def build_prompt_tokens(
     user_prompt: str,
     tokenizer,
-    task: str = "it2i_think",
+    task: str = "it2i",
+    bot_task: str | None | _DefaultBotTask = _DEFAULT_BOT_TASK,
     sys_type: str | None = None,
     custom_system_prompt: str | None = None,
+    num_images: int = 1,
 ) -> PromptTokensResult:
-    """Segment-by-segment tokenization that matches HF apply_chat_template.
+    """Segment-by-segment tokenization that matches HF apply_chat_template."""
-
+    task, bot_task = _normalize_task_and_bot_task(task, bot_task)
-    Calling tokenizer.encode(build_prompt(...)) on the full string lets BPE
+    preset_sys_type, trigger_tag = _resolve_preset(task, bot_task)
-    merge tokens across segment boundaries (e.g. user_prompt ends with `。`
-    and the next segment is `\\n\\n` -> they merge into a single token id
-    3490 instead of HF's [1811, 271]). HF's apply_chat_template tokenizes
-    each segment independently and concatenates token_ids, so no cross-
-    boundary merge happens. We replicate that here and feed the result to
-    Omni via OmniTokensPrompt (prompt_token_ids).
-
-    Returns:
-        PromptTokensResult
-    """
-    if task not in _TASK_PRESETS:
-        raise ValueError(f"Unknown task {task!r}. Choose from: {available_tasks()}")
-
-    preset_sys_type, preset_bot_task, trigger_tag = _TASK_PRESETS[task]
     effective_sys_type = sys_type or preset_sys_type
 
     bos_id = tokenizer.convert_tokens_to_ids("<|startoftext|>")
     img_id = tokenizer.convert_tokens_to_ids("<img>")
     trig_id = tokenizer.convert_tokens_to_ids(trigger_tag) if trigger_tag else None
 
-    has_image_input = task.startswith("i2t") or task.startswith("it2i")
+    has_image_input = task in ("i2t", "it2i")
+    if has_image_input:
+        _validate_num_images(num_images)
 
-    # t2i_vanilla uses pretrain template with no chat structure; the vanilla
+    if bot_task == "vanilla":
-    # system prompt drives the model directly. No segment boundaries to
+        s = build_prompt(
-    # protect, fall back to whole-string encode.
+            user_prompt,
-    if task == "t2i_vanilla":
+            task=task,
-        s = build_prompt(user_prompt, task, sys_type, custom_system_prompt)
+            bot_task=bot_task,
+            sys_type=sys_type,
+            custom_system_prompt=custom_system_prompt,
+        )
         token_ids = tokenizer.encode(s, add_special_tokens=False)
         return PromptTokensResult(
             token_ids=token_ids,
             system_prompt_type=effective_sys_type,
         )
 
-    system_prompt = get_system_prompt(effective_sys_type, preset_bot_task, custom_system_prompt)
+    system_prompt = get_system_prompt(effective_sys_type, bot_task, custom_system_prompt)
-    # Do NOT strip -- HF apply_chat_template keeps the system prompt's
-    # natural trailing newline; stripping it would shift one token id.
     sys_text = system_prompt or ""
 
     ids: list[int] = [bos_id]
@@ -190,7 +296,7 @@ def build_prompt_tokens(
         ids += tokenizer.encode("\n\n", add_special_tokens=False)
     ids += tokenizer.encode("User: ", add_special_tokens=False)
     if has_image_input:
-        ids += [img_id]
+        ids += [img_id] * num_images
     ids += tokenizer.encode(user_prompt, add_special_tokens=False)
     ids += tokenizer.encode("\n\nAssistant: ", add_special_tokens=False)
     if trig_id is not None:
@@ -202,4 +308,14 @@ def build_prompt_tokens(
     )
 
 
-__all__ = ["build_prompt", "build_prompt_tokens", "resolve_stop_token_ids", _TASK_PRESETS]
+__all__ = [
+    "HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS",
+    "MAX_IMAGES_PER_REQUEST",
+    "_TASK_PRESETS",
+    "available_bot_tasks",
+    "available_tasks",
+    "build_prompt",
+    "build_prompt_tokens",
+    "resolve_stop_token_ids",
+    "resolve_sys_type",
+]

vllm_omni/diffusion/models/omnivoice/pipeline_omnivoice.py

@@ -26,10 +26,10 @@ from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.models.interface import SupportAudioOutput
 from vllm_omni.diffusion.request import OmniDiffusionRequest
-from vllm_omni.model_executor.models.omnivoice.config import OmniVoiceConfig
 from vllm_omni.model_executor.models.omnivoice.duration import RuleDurationEstimator
 from vllm_omni.model_executor.models.omnivoice.omnivoice_decoder import OmniVoiceDecoder
 from vllm_omni.model_executor.models.omnivoice.omnivoice_generator import OmniVoiceGenerator
+from vllm_omni.transformers_utils.configs.omnivoice import OmniVoiceConfig
 from vllm_omni.utils.speaker_cache import get_speaker_cache
 
 try:

vllm_omni/diffusion/stage_diffusion_client.py

@@ -391,6 +391,8 @@ class StageDiffusionClient(StageClientBase):
             return self._output_queue.get_nowait()
         except asyncio.QueueEmpty:
             if self._engine_dead:
+                if self._shutting_down:
+                    return None
                 raise EngineDeadError()
             if not self._shutting_down and self._owns_process and self._proc is not None and not self._proc.is_alive():
                 self._engine_dead = True

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -97,7 +97,7 @@ def _make_diffusion_vllm_model_config(od_config: OmniDiffusionConfig) -> _Diffus
 
 @contextmanager
 def _force_cutlass_fp8_linear_kernel(quant_config: object | None) -> Iterator[None]:
-    from vllm.model_executor.layers.quantization import modelopt as vllm_modelopt
+    import vllm.model_executor.layers.quantization.modelopt as vllm_modelopt
 
     linear_method_cls = getattr(quant_config, "LinearMethodCls", None)
     if linear_method_cls in {
@@ -796,6 +796,14 @@ class WorkerProc:
         """Worker initialization and execution loops."""
         from vllm_omni.plugins import load_omni_general_plugins
 
+        # Set process title for visibility in nvidia-smi / htop (optional, non-fatal)
+        try:
+            import setproctitle
+
+            setproctitle.setproctitle(f"vLLM-Omni::StageDiffusionProc-{rank}")
+        except ImportError:
+            pass  # setproctitle not installed, skip process title setting
+
         load_omni_general_plugins()
         worker_proc = WorkerProc(
             od_config,

vllm_omni/distributed/omni_connectors/kv_transfer_manager.py

@@ -1247,6 +1247,9 @@ class OmniKVTransferManager:
             get_cfg_group,
             get_classifier_free_guidance_rank,
             get_classifier_free_guidance_world_size,
+            get_sequence_parallel_rank,
+            get_sequence_parallel_world_size,
+            get_sp_group,
             get_world_group,
         )
 
@@ -1260,6 +1263,9 @@ class OmniKVTransferManager:
         cfg_size = 1
         cfg_rank = 0
         cfg_group = None
+        sp_size = 1
+        sp_rank = 0
+        sp_group = None
         try:
             cfg_size = get_classifier_free_guidance_world_size()
             cfg_rank = get_classifier_free_guidance_rank()
@@ -1269,7 +1275,16 @@ class OmniKVTransferManager:
             cfg_rank = 0
             cfg_group = None
 
-        if tp_active and cfg_size <= 1:
+        try:
+            sp_size = get_sequence_parallel_world_size()
+            sp_rank = get_sequence_parallel_rank()
+            sp_group = get_sp_group()
+        except Exception:
+            sp_size = 1
+            sp_rank = 0
+            sp_group = None
+
+        if tp_active and (cfg_size <= 1 and sp_size <= 1):
             logger.info(
                 "Rank-aware KV receive: rank %s independently receiving (from_tp=%s, to_tp=%s)",
                 topo.local_rank,
@@ -1278,16 +1293,45 @@ class OmniKVTransferManager:
             )
             return self.receive_multi_kv_cache(req, cfg_kv_collect_func, target_device)
 
-        if tp_active and cfg_size > 1 and cfg_group is not None:
+        if tp_active and (cfg_size > 1 or sp_size > 1):
             kv_payload: dict[str, object] | None = None
-            if cfg_rank == 0:
+            is_owner = cfg_rank == 0 and sp_rank == 0
-                received = self.receive_multi_kv_cache(req, cfg_kv_collect_func, torch.device("cpu"))
+
-                rank_payloads = self._build_cfg_rank_local_payloads(req, cfg_size) if received else [None] * cfg_size
+            # step1: only owner need to receive
-                kv_payload = rank_payloads[0]
+            if is_owner:
-                for dst_rank in range(1, cfg_size):
+                received = self.receive_multi_kv_cache(
-                    cfg_group.send_object(rank_payloads[dst_rank], dst_rank)
+                    req,
-            else:
+                    cfg_kv_collect_func,
+                    torch.device("cpu"),
+                )
+                if received:
+                    if cfg_size > 1:  # build cfg-local payloads
+                        cfg_rank_payloads = self._build_cfg_rank_local_payloads(
+                            req,
+                            cfg_size,
+                        )
+                        kv_payload = cfg_rank_payloads[0]
+                        # send to other cfg
+                        for dst_cfg_rank in range(1, cfg_size):
+                            cfg_group.send_object(
+                                cfg_rank_payloads[dst_cfg_rank],
+                                dst_cfg_rank,
+                            )
+                    elif sp_size > 1:
+                        kv_payload = self._collect_request_kv_payload(req)
+                else:
+                    # Owner receive failed: send None sentinel to CFG followers to avoid deadlock
+                    if cfg_size > 1:
+                        for dst_cfg_rank in range(1, cfg_size):
+                            cfg_group.send_object(None, dst_cfg_rank)
+            elif sp_rank == 0 and cfg_size > 1:
                 kv_payload = cfg_group.recv_object(0)
+            # sp broadcast
+            if sp_size > 1 and sp_group is not None:
+                kv_payload = sp_group.broadcast_object(
+                    kv_payload,
+                    src=0,
+                )
 
             if not kv_payload:
                 return False

vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py

@@ -62,6 +62,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
         self.waiting_for_chunk_running_requests: deque[Any] = deque()
         self.requests_with_ready_chunks = set()
         self.requests_origin_status = {}
+        self.requests_num_chunks_sent: dict[str, int] = defaultdict(int)
 
     @classmethod
     def create_connector(cls, model_config: Any):
@@ -117,6 +118,17 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
             pooling_output: Partial pooling output dictionary
             request: Request object
         """
+
+        # If the request is preempted, skip the already saved chunks.
+        if request.num_computed_tokens < self.requests_num_chunks_sent.get(request.external_req_id, 0):
+            logger.warning(
+                f"Enqueue save_async for request {request.external_req_id}, "
+                f"request.num_computed_tokens={request.num_computed_tokens}, "
+                f"previous_chunks_sent={self.requests_num_chunks_sent.get(request.external_req_id, 0)}"
+            )
+            return
+
+        self.requests_num_chunks_sent[request.external_req_id] = request.num_computed_tokens
         task = {
             "pooling_output": pooling_output,
             "request": request,
@@ -155,8 +167,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
 
             meta = payload_data.get("meta", {})
             if self.model_mode == "ar":
-                merged_payload = self._update_request_payload(external_req_id, payload_data)
+                request.additional_information = payload_data
-                request.additional_information = merged_payload
                 if meta.get("finished"):
                     self.finished_requests.add(req_id)
             else:
@@ -198,42 +209,6 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
 
         return False
 
-    def _update_request_payload(self, req_id: str, payload_data: dict[str, Any]) -> dict[str, Any]:
-        """Update the stored payload for *req_id* with the latest chunk."""
-        if req_id not in self.request_payload:
-            self.request_payload[req_id] = payload_data
-            return payload_data
-        origin = self.request_payload[req_id]
-        raw_ok = payload_data.get("meta", {}).pop("override_keys", [])
-        override_keys = {tuple(k) if isinstance(k, list) else k for k in raw_ok}
-
-        for key, value in payload_data.items():
-            if isinstance(value, dict):
-                origin_sub = origin.get(key)
-                if not isinstance(origin_sub, dict):
-                    continue
-                for qual, qval in value.items():
-                    if key == "meta" and qual == "finished":
-                        continue
-                    if (key, qual) in override_keys:
-                        continue
-                    osv = origin_sub.get(qual)
-                    if isinstance(qval, torch.Tensor) and isinstance(osv, torch.Tensor):
-                        value[qual] = torch.cat([osv, qval], dim=0)
-                    elif isinstance(qval, list) and isinstance(osv, list):
-                        value[qual] = osv + qval
-            else:
-                if key in override_keys:
-                    continue
-                ov = origin.get(key)
-                if isinstance(value, torch.Tensor) and isinstance(ov, torch.Tensor):
-                    payload_data[key] = torch.cat([ov, value], dim=0)
-                elif isinstance(value, list) and isinstance(ov, list):
-                    payload_data[key] = ov + value
-
-        self.request_payload[req_id] = payload_data
-        return payload_data
-
     def _send_single_request(self, task: dict):
         raw_po = task["pooling_output"]
         pooling_output = unflatten_payload(raw_po) if isinstance(raw_po, dict) else raw_po
@@ -290,6 +265,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
 
         if is_finished:
             self.code_prompt_token_ids.pop(external_req_id, None)
+            self.requests_num_chunks_sent.pop(external_req_id, None)
             cached_ic = getattr(self, "_cached_ic", None)
             if cached_ic is not None:
                 cached_ic.pop(external_req_id, None)
@@ -327,6 +303,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
         self.put_req_chunk.pop(external_req_id, None)
         self.request_payload.pop(external_req_id, None)
         self.code_prompt_token_ids.pop(external_req_id, None)
+        self.requests_num_chunks_sent.pop(external_req_id, None)
 
         cached_ic = getattr(self, "_cached_ic", None)
         if cached_ic is not None:
@@ -399,6 +376,11 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
         Add additional info for cached requests and
         clean up ready chunks from scheduler output.
         """
+        stage_id = self.connector.stage_id
+
+        if stage_id == 0:
+            return
+
         if requests is not None:
             self.attach_cached_additional_information(scheduler_output, requests)
         self._clear_chunk_ready(scheduler_output)
@@ -414,6 +396,8 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
             request = requests.get(req_id) if req_id else None
             additional_info = getattr(request, "additional_information", None) if request else None
             cached_reqs.additional_information[req_id] = additional_info
+            if request and additional_info:
+                request.additional_information = None
 
     def _process_chunk_queue(
         self,

vllm_omni/distributed/omni_connectors/utils/initialization.py

@@ -5,6 +5,7 @@
 
 import json
 import sys
+import warnings
 from pathlib import Path
 from typing import TYPE_CHECKING, Any
 
@@ -403,7 +404,16 @@ def build_stage_connectors(
     connectors_config: dict[str, Any],
     purpose: str = "request_forwarding",
 ) -> dict[tuple[str, str], Any] | None:
-    """Instantiate OmniConnectors for a stage based on config."""
+    """Instantiate OmniConnectors for a stage based on config.
+
+    Deprecated: prefer ``get_stage_connector_config`` plus the unified
+    connector factory. Kept as a thin shim so legacy callers keep working.
+    """
+    warnings.warn(
+        "build_stage_connectors is deprecated; use get_stage_connector_config and OmniConnectorFactory instead.",
+        DeprecationWarning,
+        stacklevel=2,
+    )
     if not connectors_config:
         return {}
 

vllm_omni/engine/__init__.py

@@ -94,6 +94,7 @@ class OmniEngineCoreRequest(EngineCoreRequest):
         return cls(
             request_id=request.request_id,
             prompt_token_ids=request.prompt_token_ids,
+            prompt_is_token_ids=request.prompt_is_token_ids,
             mm_features=request.mm_features,
             sampling_params=request.sampling_params,
             pooling_params=request.pooling_params,
@@ -109,6 +110,8 @@ class OmniEngineCoreRequest(EngineCoreRequest):
             resumable=request.resumable,
             external_req_id=request.external_req_id,
             reasoning_ended=request.reasoning_ended,
+            reasoning_parser_kwargs=request.reasoning_parser_kwargs,
+            abort_immediately=request.abort_immediately,
             additional_information=additional_information,
         )
 

vllm_omni/engine/arg_utils.py

@@ -35,13 +35,9 @@ def _register_omni_hf_configs() -> None:
     try:
         from transformers import AutoConfig
 
-        from vllm_omni.model_executor.models.cosyvoice3.config import CosyVoice3Config
-        from vllm_omni.model_executor.models.omnivoice.config import OmniVoiceConfig
         from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import (
             Qwen3TTSConfig,
         )
-        from vllm_omni.transformers_utils.configs.voxcpm import VoxCPMConfig
-        from vllm_omni.transformers_utils.configs.voxcpm2 import VoxCPM2Config
     except Exception as exc:  # pragma: no cover - best-effort optional registration
         logger.warning("Skipping omni HF config registration due to import error: %s", exc)
         return
@@ -56,10 +52,6 @@ def _register_omni_hf_configs() -> None:
 
     for model_type, config_cls in [
         ("qwen3_tts", Qwen3TTSConfig),
-        ("cosyvoice3", CosyVoice3Config),
-        ("omnivoice", OmniVoiceConfig),
-        ("voxcpm", VoxCPMConfig),
-        ("voxcpm2", VoxCPM2Config),
     ]:
         try:
             AutoConfig.register(model_type, config_cls)

vllm_omni/engine/async_omni_engine.py

@@ -1199,6 +1199,158 @@ class AsyncOmniEngine:
 
     # ---- request helpers ----
 
+    @staticmethod
+    def _iter_multimodal_items(value: Any) -> list[Any]:
+        if value is None:
+            return []
+        if isinstance(value, list):
+            return value
+        return [value]
+
+    def _ensure_stage_replica_mm_uuids(
+        self,
+        prompt: Any,
+        *,
+        stage_id: int,
+        replica_id: int,
+    ) -> None:
+        """Make multimodal processor-cache keys local to a stage replica.
+
+        vLLM's frontend multimodal sender cache is process-global, while each
+        vllm-omni stage replica owns a separate EngineCore receiver cache. If
+        two requests with the same image are routed to different stage-0
+        replicas, a plain content hash can make the sender omit the tensor for
+        a replica that has never received it. Prefixing user/content UUIDs with
+        the selected replica keeps cache reuse within the receiver that owns it.
+        """
+
+        if not isinstance(prompt, dict):
+            return
+
+        mm_data = prompt.get("multi_modal_data")
+        if not isinstance(mm_data, dict) or not mm_data:
+            return
+
+        from vllm.multimodal.hasher import MultiModalHasher
+
+        existing_uuids = prompt.get("multi_modal_uuids")
+        if not isinstance(existing_uuids, dict):
+            existing_uuids = {}
+
+        model_id = str(getattr(self, "model", ""))
+        scoped_uuids: dict[str, list[str | None]] = dict(existing_uuids)
+        for modality, raw_items in mm_data.items():
+            items = self._iter_multimodal_items(raw_items)
+            if not items:
+                continue
+
+            modality_existing = existing_uuids.get(modality)
+            if not isinstance(modality_existing, list):
+                modality_existing = [modality_existing] if modality_existing is not None else []
+
+            modality_uuids: list[str | None] = []
+            for idx, item in enumerate(items):
+                user_uuid = modality_existing[idx] if idx < len(modality_existing) else None
+                if user_uuid is not None:
+                    base_uuid = str(user_uuid)
+                elif item is None:
+                    base_uuid = None
+                else:
+                    base_uuid = MultiModalHasher.hash_kwargs(
+                        model_id=model_id,
+                        **{modality: item},
+                    )
+
+                if base_uuid is None:
+                    modality_uuids.append(None)
+                else:
+                    modality_uuids.append(f"stage{stage_id}:rep{replica_id}:{base_uuid}")
+
+            scoped_uuids[modality] = modality_uuids
+
+        if scoped_uuids:
+            prompt["multi_modal_uuids"] = scoped_uuids
+
+    @staticmethod
+    def _stage_pool_replica_count(stage_pool: Any) -> int:
+        try:
+            live_num_replicas = getattr(stage_pool, "live_num_replicas", None)
+            if live_num_replicas is not None:
+                return int(live_num_replicas)
+        except Exception:
+            pass
+
+        try:
+            live_replica_ids = getattr(stage_pool, "live_replica_ids", None)
+            if callable(live_replica_ids):
+                return len(live_replica_ids())
+        except Exception:
+            pass
+
+        try:
+            clients = getattr(stage_pool, "clients", None)
+            if clients is not None:
+                return sum(1 for client in clients if client is not None)
+        except Exception:
+            pass
+
+        return int(getattr(stage_pool, "num_replicas", 1) or 1)
+
+    @staticmethod
+    def _stage_pool_is_distributed(stage_pool: Any) -> bool:
+        try:
+            is_distributed = getattr(stage_pool, "is_distributed", None)
+            if is_distributed is not None:
+                return bool(is_distributed() if callable(is_distributed) else is_distributed)
+        except Exception:
+            pass
+
+        return getattr(stage_pool, "_hub", None) is not None
+
+    def _scope_stage0_multimodal_cache_to_replica(
+        self,
+        request_id: str,
+        prompt: Any,
+    ) -> int | None:
+        stage_pools = getattr(self, "stage_pools", None)
+        if isinstance(prompt, EngineCoreRequest) or not stage_pools:
+            return None
+
+        stage0_pool = stage_pools[0]
+        # TODO: Currently only supports the ar -> dit process.
+        # Future scenarios (e.g., dit -> ar) need to be added, which will require modifications here.
+        if stage0_pool.stage_type == "diffusion" or self._stage_pool_replica_count(stage0_pool) <= 1:
+            return None
+
+        # This synchronous request path can safely pre-bind local in-process
+        # replicas. Distributed head mode still uses the async picker inside
+        # StagePool.submit_initial().
+        if self._stage_pool_is_distributed(stage0_pool):
+            logger.debug(
+                "[AsyncOmniEngine] Skipping stage-0 multimodal cache scoping for distributed routing req=%s",
+                request_id,
+            )
+            return None
+
+        prompts = prompt if isinstance(prompt, list) else [prompt]
+        if not any(isinstance(p, dict) and p.get("multi_modal_data") for p in prompts):
+            return None
+
+        replica_id = stage0_pool.select_replica_id(request_id)
+        for p in prompts:
+            self._ensure_stage_replica_mm_uuids(
+                p,
+                stage_id=0,
+                replica_id=replica_id,
+            )
+
+        logger.debug(
+            "[AsyncOmniEngine] Scoped multimodal cache keys to stage-0 replica-%s for req=%s",
+            replica_id,
+            request_id,
+        )
+        return replica_id
+
     def _build_add_request_message(
         self,
         request_id: str,
@@ -1232,6 +1384,7 @@ class AsyncOmniEngine:
         # Keep the original prompt for downstream stages (they need the raw
         # dict, e.g. for multi_modal_data).
         original_prompt = prompt
+        preselected_stage0_replica: int | None = None
 
         stage_type = self.stage_metadata[0].stage_type
         output_prompt_text: Any = None
@@ -1244,21 +1397,31 @@ class AsyncOmniEngine:
                 for item in prompt:
                     _inject_global_id(item, request_id)
 
+            preselected_stage0_replica = self._scope_stage0_multimodal_cache_to_replica(
+                request_id,
+                prompt,
+            )
+
             # Full input processing (tokenization, multimodal, etc.)
             _t_preprocess = time.perf_counter()
-            request = self.input_processor.process_inputs(
+            try:
-                request_id=request_id,
+                request = self.input_processor.process_inputs(
-                prompt=prompt,
+                    request_id=request_id,
-                params=params,
+                    prompt=prompt,
-                supported_tasks=self.supported_tasks,
+                    params=params,
-                arrival_time=arrival_time,
+                    supported_tasks=self.supported_tasks,
-                lora_request=lora_request,
+                    arrival_time=arrival_time,
-                tokenization_kwargs=tokenization_kwargs,
+                    lora_request=lora_request,
-                trace_headers=trace_headers,
+                    tokenization_kwargs=tokenization_kwargs,
-                priority=priority,
+                    trace_headers=trace_headers,
-                data_parallel_rank=data_parallel_rank,
+                    priority=priority,
-                resumable=resumable,
+                    data_parallel_rank=data_parallel_rank,
-            )
+                    resumable=resumable,
+                )
+            except Exception:
+                if preselected_stage0_replica is not None and self.stage_pools:
+                    self.stage_pools[0].release_binding(request_id)
+                raise
             _preprocess_ms = (time.perf_counter() - _t_preprocess) * 1000.0
             # TODO (Peiqi): add this for Qwen3-TTS only. Other models don't have
             # additional_information field in the prompt.
@@ -1535,9 +1698,9 @@ class AsyncOmniEngine:
             "enable_multithread_weight_load": kwargs.get("enable_multithread_weight_load", True),
             "num_weight_load_threads": kwargs.get("num_weight_load_threads", 4),
             "quantization": kwargs.get("quantization", None),
-            "kv_cache_dtype": kwargs.get("kv_cache_dtype", None),
+            "diffusion_kv_cache_dtype": kwargs.get("diffusion_kv_cache_dtype", None),
-            "kv_cache_skip_steps": kwargs.get("kv_cache_skip_steps", None),
+            "diffusion_kv_cache_skip_steps": kwargs.get("diffusion_kv_cache_skip_steps", None),
-            "kv_cache_skip_layers": kwargs.get("kv_cache_skip_layers", None),
+            "diffusion_kv_cache_skip_layers": kwargs.get("diffusion_kv_cache_skip_layers", None),
             **({"diffusion_attention_config": attention_config} if attention_config is not None else {}),
             "force_cutlass_fp8": bool(kwargs.get("force_cutlass_fp8", False)),
             "enable_diffusion_pipeline_profiler": kwargs.get("enable_diffusion_pipeline_profiler", False),
@@ -1717,24 +1880,27 @@ class AsyncOmniEngine:
                 if quantization is not None:
                     if not hasattr(cfg.engine_args, "quantization") or cfg.engine_args.quantization is None:
                         cfg.engine_args.quantization = quantization
-                kv_cache_dtype = kwargs.get("kv_cache_dtype")
+                diffusion_kv_cache_dtype = kwargs.get("diffusion_kv_cache_dtype")
-                if kv_cache_dtype is not None:
+                if diffusion_kv_cache_dtype is not None:
-                    if not hasattr(cfg.engine_args, "kv_cache_dtype") or cfg.engine_args.kv_cache_dtype is None:
-                        cfg.engine_args.kv_cache_dtype = kv_cache_dtype
-                kv_cache_skip_steps = kwargs.get("kv_cache_skip_steps")
-                if kv_cache_skip_steps is not None:
                     if (
-                        not hasattr(cfg.engine_args, "kv_cache_skip_steps")
+                        not hasattr(cfg.engine_args, "diffusion_kv_cache_dtype")
-                        or cfg.engine_args.kv_cache_skip_steps is None
+                        or cfg.engine_args.diffusion_kv_cache_dtype is None
                     ):
-                        cfg.engine_args.kv_cache_skip_steps = kv_cache_skip_steps
+                        cfg.engine_args.diffusion_kv_cache_dtype = diffusion_kv_cache_dtype
-                kv_cache_skip_layers = kwargs.get("kv_cache_skip_layers")
+                diffusion_kv_cache_skip_steps = kwargs.get("diffusion_kv_cache_skip_steps")
-                if kv_cache_skip_layers is not None:
+                if diffusion_kv_cache_skip_steps is not None:
                     if (
-                        not hasattr(cfg.engine_args, "kv_cache_skip_layers")
+                        not hasattr(cfg.engine_args, "diffusion_kv_cache_skip_steps")
-                        or cfg.engine_args.kv_cache_skip_layers is None
+                        or cfg.engine_args.diffusion_kv_cache_skip_steps is None
                     ):
-                        cfg.engine_args.kv_cache_skip_layers = kv_cache_skip_layers
+                        cfg.engine_args.diffusion_kv_cache_skip_steps = diffusion_kv_cache_skip_steps
+                diffusion_kv_cache_skip_layers = kwargs.get("diffusion_kv_cache_skip_layers")
+                if diffusion_kv_cache_skip_layers is not None:
+                    if (
+                        not hasattr(cfg.engine_args, "diffusion_kv_cache_skip_layers")
+                        or cfg.engine_args.diffusion_kv_cache_skip_layers is None
+                    ):
+                        cfg.engine_args.diffusion_kv_cache_skip_layers = diffusion_kv_cache_skip_layers
             except Exception as e:
                 logger.warning("Failed to inject LoRA config for stage: %s", e)
 

vllm_omni/engine/orchestrator.py

@@ -32,7 +32,6 @@ from vllm_omni.engine.messages import (
     EngineQueueMessage,
     ErrorMessage,
     OutputMessage,
-    ShutdownRequestMessage,
     StageMetricsMessage,
     StageSubmissionMessage,
 )
@@ -128,8 +127,8 @@ class Orchestrator:
     def __init__(
         self,
         request_async_queue: janus.AsyncQueue[EngineQueueMessage],
-        output_async_queue: janus.AsyncQueue[EngineQueueMessage],
+        output_async_queue: janus.AsyncQueue[dict[str, Any]],
-        rpc_async_queue: janus.AsyncQueue[EngineQueueMessage],
+        rpc_async_queue: janus.AsyncQueue[dict[str, Any]],
         stage_pools: list[StagePool],
         *,
         async_chunk: bool = False,
@@ -174,6 +173,13 @@ class Orchestrator:
             await asyncio.gather(request_task, output_task)
         except asyncio.CancelledError:
             raise
+        except EngineDeadError as e:
+            # EngineDeadError from _orchestration_loop means the diffusion
+            # engine died.  All pending requests were already notified and
+            # _shutdown_event was already set by the loop's handler.
+            # During teardown this is expected; the finally block handles
+            # proper cleanup.  Do not re-raise.
+            logger.info("[Orchestrator] Engine dead during shutdown: %s", e)
         except Exception:
             logger.exception("[Orchestrator] Fatal error in orchestrator tasks")
             raise
@@ -208,27 +214,32 @@ class Orchestrator:
         """Read messages from the main thread via request_async_queue."""
         while True:
             msg = await self.request_async_queue.get()
+            msg_type = msg.type
 
-            if isinstance(msg, StageSubmissionMessage):
+            if msg_type == "add_request":
-                if msg.type == "add_request":
+                await self._handle_add_request(msg)
-                    await self._handle_add_request(msg)
+            elif msg_type == "streaming_update":
-                elif msg.type == "streaming_update":
+                await self._handle_streaming_update(msg)
-                    await self._handle_streaming_update(msg)
+            elif msg_type == "add_companion_request":
-                else:
-                    logger.warning("[Orchestrator] Unknown stage submission type: %s", msg.type)
-            elif isinstance(msg, AddCompanionRequestMessage):
                 await self._handle_add_companion(msg)
-            elif isinstance(msg, AbortRequestMessage):
+            elif msg_type == "abort":
                 await self._handle_abort(msg)
-            elif isinstance(msg, CollectiveRPCRequestMessage):
+            elif msg_type == "collective_rpc":
                 await self._handle_collective_rpc(msg)
-            elif isinstance(msg, ShutdownRequestMessage):
+            elif msg_type == "shutdown":
                 logger.info("[Orchestrator] Received shutdown signal")
                 self._shutdown_event.set()
+                # Pre-mark stage clients as shutting down to prevent
+                # proc_monitor daemon threads from flagging normal
+                # process exit as EngineDeadError during teardown.
+                for pool in self.stage_pools:
+                    for client in pool.clients:
+                        if hasattr(client, "_shutting_down"):
+                            client._shutting_down = True
                 self._shutdown_stages()
                 break
             else:
-                logger.warning("[Orchestrator] Unknown message type: %s", type(msg).__name__)
+                logger.warning("[Orchestrator] Unknown message type: %s", msg_type)
 
     async def _handle_add_request(self, msg: StageSubmissionMessage) -> None:
         """Handle an add_request message from the main thread."""
@@ -387,8 +398,8 @@ class Orchestrator:
         rpc_id = msg.rpc_id
         method = msg.method
         timeout = msg.timeout
-        args = msg.args
+        args = tuple(msg.args)
-        kwargs = dict(msg.kwargs)
+        kwargs = dict(msg.kwargs or {})
         requested_stage_ids = msg.stage_ids
 
         target_pools: list[StagePool] = []
@@ -793,120 +804,78 @@ class Orchestrator:
         """Forward output from the current logical stage to the next one."""
         next_logical = src_stage_id + 1
         next_pool = self.stage_pools[next_logical]
-
+        next_client = next_pool.stage_client
-        if next_pool.stage_type == "diffusion":
-            await self._forward_to_diffusion_stage(
-                req_id,
-                src_stage_id,
-                output,
-                req_state,
-            )
-        else:
-            await self._forward_to_llm_stage(
-                req_id,
-                src_stage_id,
-                output,
-                req_state,
-                is_streaming_session=is_streaming_session,
-                is_final_update=is_final_update,
-            )
-
-    async def _forward_to_diffusion_stage(
-        self,
-        req_id: str,
-        src_stage_id: int,
-        output: Any,
-        req_state: OrchestratorRequestState,
-    ) -> None:
-        """Forward output from an LLM stage to the next diffusion stage."""
-        next_logical = src_stage_id + 1
-        next_pool = self.stage_pools[next_logical]
-        params = req_state.sampling_params_list[next_logical]
-        source_outputs = [output]
-        already_submitted = self._next_stage_already_submitted(src_stage_id, req_state)
-
-        diffusion_client = next_pool.stage_client
-        if diffusion_client.custom_process_input_func is not None:
-            _t_ar2d = _time.perf_counter()
-            diffusion_prompt = diffusion_client.custom_process_input_func(
-                source_outputs,
-                req_state.prompt,
-                getattr(diffusion_client, "requires_multimodal_data", False),
-            )
-            _dt_ar2d = (_time.perf_counter() - _t_ar2d) * 1000
-            req_state.pipeline_timings["ar2diffusion_ms"] = _dt_ar2d
-            logger.info(
-                "[Orchestrator] ar2diffusion req=%s wall_time=%.3fms stage=%d->%d",
-                req_id,
-                _dt_ar2d,
-                src_stage_id,
-                next_logical,
-            )
-            if isinstance(diffusion_prompt, list):
-                if not diffusion_prompt:
-                    error_output = OmniRequestOutput.from_error(
-                        req_id,
-                        f"Stage-{src_stage_id} produced no valid inputs for diffusion stage-{next_logical}",
-                    )
-                    logger.warning(
-                        "[Orchestrator] req=%s stage=%d produced empty diffusion inputs for stage=%d; "
-                        "routing terminal error output",
-                        req_id,
-                        src_stage_id,
-                        next_logical,
-                    )
-                    await self.output_async_queue.put(
-                        OutputMessage(
-                            request_id=req_id,
-                            stage_id=next_logical,
-                            engine_outputs=error_output,
-                            metrics=None,
-                            finished=True,
-                        )
-                    )
-                    await self._cleanup_request_ids(
-                        [req_id, *self._cfg_tracker.cleanup_parent(req_id)],
-                    )
-                    return
-                if already_submitted and len(diffusion_prompt) == 1:
-                    diffusion_prompt = diffusion_prompt[0]
-        else:
-            diffusion_prompt = req_state.prompt
-
-        if already_submitted:
-            await next_pool.submit_update(req_id, req_state, diffusion_prompt)
-        else:
-            await next_pool.submit_initial(
-                req_id,
-                req_state,
-                diffusion_prompt,
-                submit_kwargs={
-                    "kv_sender_info": self._build_kv_sender_info(
-                        list(getattr(diffusion_client, "engine_input_source", None) or [src_stage_id]),
-                        request_id=req_id,
-                    )
-                },
-                params_override=self._maybe_clone_diffusion_params_for_cfg(req_id, params),
-            )
-        req_state.stage_submit_ts[next_logical] = _time.time()
-
-    async def _forward_to_llm_stage(
-        self,
-        req_id: str,
-        src_stage_id: int,
-        output: Any,
-        req_state: OrchestratorRequestState,
-        *,
-        is_streaming_session: bool = False,
-        is_final_update: bool = False,
-    ) -> None:
-        """Forward output from the current stage to the next LLM stage."""
-        next_logical = src_stage_id + 1
-        next_pool = self.stage_pools[next_logical]
         params = req_state.sampling_params_list[next_logical]
         source_outputs = [output]
         next_stage_resumable = is_streaming_session and not is_final_update
         already_submitted = self._next_stage_already_submitted(src_stage_id, req_state)
+        requires_multimodal_data = getattr(next_client, "requires_multimodal_data", False)
+
+        if next_pool.stage_type == "diffusion":
+            if next_client.custom_process_input_func is not None:
+                _t_ar2d = _time.perf_counter()
+                diffusion_prompt = next_client.custom_process_input_func(
+                    source_outputs,
+                    req_state.prompt,
+                    requires_multimodal_data,
+                )
+                _dt_ar2d = (_time.perf_counter() - _t_ar2d) * 1000
+                req_state.pipeline_timings["ar2diffusion_ms"] = _dt_ar2d
+                logger.info(
+                    "[Orchestrator] ar2diffusion req=%s wall_time=%.3fms stage=%d->%d",
+                    req_id,
+                    _dt_ar2d,
+                    src_stage_id,
+                    next_logical,
+                )
+                if isinstance(diffusion_prompt, list):
+                    if not diffusion_prompt:
+                        error_output = OmniRequestOutput.from_error(
+                            req_id,
+                            f"Stage-{src_stage_id} produced no valid inputs for diffusion stage-{next_logical}",
+                        )
+                        logger.warning(
+                            "[Orchestrator] req=%s stage=%d produced empty diffusion inputs for stage=%d; "
+                            "routing terminal error output",
+                            req_id,
+                            src_stage_id,
+                            next_logical,
+                        )
+                        await self.output_async_queue.put(
+                            OutputMessage(
+                                request_id=req_id,
+                                stage_id=next_logical,
+                                engine_outputs=error_output,
+                                metrics=None,
+                                finished=True,
+                            )
+                        )
+                        await self._cleanup_request_ids(
+                            [req_id, *self._cfg_tracker.cleanup_parent(req_id)],
+                        )
+                        return
+                    if already_submitted and len(diffusion_prompt) == 1:
+                        diffusion_prompt = diffusion_prompt[0]
+            else:
+                diffusion_prompt = req_state.prompt
+
+            if already_submitted:
+                await next_pool.submit_update(req_id, req_state, diffusion_prompt)
+            else:
+                await next_pool.submit_initial(
+                    req_id,
+                    req_state,
+                    diffusion_prompt,
+                    submit_kwargs={
+                        "kv_sender_info": self._build_kv_sender_info(
+                            list(getattr(next_client, "engine_input_source", None) or [src_stage_id]),
+                            request_id=req_id,
+                        )
+                    },
+                    params_override=self._maybe_clone_diffusion_params_for_cfg(req_id, params),
+                )
+            req_state.stage_submit_ts[next_logical] = _time.time()
+            return
 
         # PD disaggregation: prefill → decode routing uses original prompt + KV transfer params
         if self._pd_pair is not None and (src_stage_id, next_logical) == self._pd_pair:
@@ -953,7 +922,6 @@ class Orchestrator:
                 {},
             )[req_id] = req_state.pd_prefill_multimodal_output
 
-        next_client = next_pool.llm_stage_client
         try:
             next_inputs = next_client.process_engine_inputs(
                 source_outputs,
@@ -972,7 +940,7 @@ class Orchestrator:
         for next_input in next_inputs:
             # Only AR thinker stages consume encoder mm_features; downstream
             # (talker/code2wav/…) must not see them (avoids encoder-cache misses).
-            model_stage = next_client.model_stage
+            model_stage = getattr(next_client, "model_stage", None)
             mm_features = req_state.mm_features if model_stage == "thinker" else None
             request = build_engine_core_request_from_tokens(
                 request_id=req_id,
@@ -1022,7 +990,7 @@ class Orchestrator:
                     req_state.prompt,
                     submit_kwargs={
                         "kv_sender_info": self._build_kv_sender_info(
-                            list(next_pool.stage_client.engine_input_source or [next_stage_id - 1]),
+                            list(getattr(next_pool.stage_client, "engine_input_source", None) or [next_stage_id - 1]),
                             request_id=request_id,
                         )
                     },
@@ -1069,14 +1037,14 @@ class Orchestrator:
                 continue
 
             sender_pool = self.stage_pools[sender_stage_id]
-            if sender_pool.stage_type == "diffusion":
+            sender_stage = sender_pool.get_bound_client(request_id) if request_id is not None else None
+            if sender_stage is None:
+                sender_stage = sender_pool.stage_client
+            get_sender_info = getattr(sender_stage, "get_kv_sender_info", None)
+            if not callable(get_sender_info):
                 continue
 
-            sender_stage = sender_pool.get_bound_llm_client(request_id) if request_id is not None else None
+            sender_info = get_sender_info()
-            if sender_stage is None:
-                sender_stage = sender_pool.llm_stage_client
-
-            sender_info = sender_stage.get_kv_sender_info()
             if not sender_info:
                 logger.warning(
                     "[Orchestrator] Stage-%s has no KV sender info available",
@@ -1110,7 +1078,7 @@ class Orchestrator:
                 msg = self.request_async_queue.get_nowait()
             except Exception:
                 break
-            if isinstance(msg, StageSubmissionMessage) and msg.type == "add_request":
+            if msg.type == "add_request":
                 req_id = msg.request_id
                 await self.output_async_queue.put(
                     ErrorMessage(