Enable twine upload to PyPI in release pipeline

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
[CI/Build] Unify release pipeline with NIGHTLY=1 option, add x86_64/aarch64 image builds (#3428 )
2026-05-17 03:59:46 -07:00 · 2026-05-17 03:42:31 -07:00 · 2026-05-16 18:18:43 +08:00 · 2026-05-15 23:35:44 +08:00 · 2026-05-15 22:50:03 +08:00 · 2026-05-15 09:55:09 +00:00
105 changed files with 4554 additions and 1258 deletions
@@ -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"
@@ -78,15 +78,19 @@ steps:
      echo "--- Image pushed successfully"
      echo "$${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG}"

-  # - label: "Upload Ready Pipeline"
-  #   depends_on: image-build
-  #   key: upload-ready-pipeline
-  #   if: build.branch != "main" && build.pull_request.labels includes "npu-test"
-  #   agents:
-  #     queue: "ascend-a2b3"
-  #     resource_class: "npu-1"
-  #   command: |
-  #     buildkite-agent pipeline upload .buildkite/test-npu-ready.yaml
+  - label: "Upload Ready Pipeline"
+    depends_on: image-build
+    key: upload-ready-pipeline
+    if: >-
+      build.branch != "main" && (
+        build.pull_request.labels includes "npu-test" ||
+        build.pull_request.labels includes "ready"
+      )
+    agents:
+      queue: "ascend-a2b3"
+      resource_class: "npu-1"
+    command: |
+      buildkite-agent pipeline upload .buildkite/test-npu-ready.yaml

  # L4 Test — main+NIGHTLY=1 (scheduled), or PR with specific label (e.g. add label then Rebuild)
  - label: "Upload Nightly Pipeline"
@@ -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"
@@ -1,38 +1,110 @@
 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"
+  - input: "Provide Release version here"
+    id: input-release-version
+    fields:
+      - text: "What is the release version?"
+        key: release-version

-  # L2 Test
-  - label: "Upload Ready Pipeline"
-    depends_on: image-build
-    key: upload-ready-pipeline
-    commands:
-      - buildkite-agent pipeline upload .buildkite/test-ready.yml
+  - label: "Build and upload wheel"
+    depends_on: ~
+    key: "build-wheel"
    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: "Upload Merge Pipeline"
-    depends_on: image-build
-    key: upload-merge-pipeline
-    commands:
-      - buildkite-agent pipeline upload .buildkite/test-merge.yml
+  - label: "Generate and upload wheel indices"
+    depends_on: "build-wheel"
+    allow_dependency_failure: true
    agents:
-      queue: "cpu_queue_premerge"
+      queue: small_cpu_queue_release
+    commands:
+      - "bash .buildkite/scripts/generate-and-upload-nightly-index.sh"

-  # 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
+  - block: "Upload release wheels to PyPI"
+    depends_on: "build-wheel"
+    key: block-upload-release-wheels
+
+  - label: "Upload release wheels to PyPI"
+    depends_on:
+      - block-upload-release-wheels
+      - input-release-version
+    key: upload-release-wheels
    agents:
-      queue: "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"
@@ -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}"
@@ -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"
@@ -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
@@ -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
@@ -1039,7 +1040,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:
@@ -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'
@@ -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
@@ -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

@@ -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.

 ---
@@ -56,9 +56,8 @@ vLLM-Omni is flexible and easy to use with:

 vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, including:

- **Omni-modality models** (e.g. Qwen3-Omni, Cosmos, HunyuanImage, BAGEL)
- **TTS models** (e.g. Qwen3-TTS, VoxCPM2, Ming-Omni-TTS, CosyVoice3)
- **Diffusion models** — image, video, and audio generation (e.g. Qwen-Image, Wan2.2, FLUX)
+- Omni-modality models (e.g. Qwen-Omni)
+- Multi-modality generation models (e.g. Qwen-Image)

 ## Getting Started

@@ -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.
@@ -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.
@@ -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}
@@ -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 && \
@@ -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 && \
@@ -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"

@@ -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)
-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
+#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.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
 ```

@@ -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+rocm721 --extra-index-url https://wheels.vllm.ai/rocm/0.21.0/rocm721
 ```

 #### 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+rocm721 --extra-index-url https://wheels.vllm.ai/rocm/0.21.0/rocm721
 ```

-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.1. 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]
@@ -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

@@ -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
-used by vLLM's language-model KV-cache quantization. In vLLM-Omni diffusion
-pipelines, however, it is a runtime FA path: Q/K/V tensors are 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.
+This feature is configured through `diffusion_kv_cache_dtype` on
+`OmniDiffusionConfig` (CLI: `--diffusion-kv-cache-dtype`). It is intentionally
+**not** the same as vLLM's `--kv-cache-dtype`, which controls autoregressive
+language-model KV cache storage and defaults to `"auto"`. Diffusion FA
+quantization uses the dedicated diffusion flags so omni serve does not inherit
+that default.

-If `kv_cache_dtype` is not set, behavior is unchanged and attention runs in the
-native dtype.
+In vLLM-Omni diffusion pipelines, this is a runtime FA path: Q/K/V tensors are
+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
-back to native dtype execution.
+backends do not support `diffusion_kv_cache_dtype="fp8"` for diffusion attention
+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
-stage and model.
+`diffusion_kv_cache_dtype` may apply in theory; validate quality and latency for
+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 \
-    --kv-cache-skip-steps "0,1" \
-    --kv-cache-skip-layers "0-2"
+    --diffusion-kv-cache-dtype fp8 \
+    --diffusion-kv-cache-skip-steps "0,1" \
+    --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"
-      kv_cache_skip_steps: "0,1"
-      kv_cache_skip_layers: "0-2"
+      diffusion_kv_cache_dtype: "fp8"
+      diffusion_kv_cache_skip_steps: "0,1"
+      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 |
-| `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_dtype` | str \| None | `None` | Set to `"fp8"` to enable dynamic FP8 FA on supported attention backends |
+| `diffusion_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_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
@@ -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_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_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`.
 - 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.
- `t2i_vanilla` is the only task that uses the bare pretrain template without chat structure.
- 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.
+- System prompt: Auto-selected from `task` and `bot_task`.
+- `bot_task='vanilla'` with `task='t2i'` uses the bare pretrain template.

 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 |
@@ -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(
-        "--model",
-        default="tencent/HunyuanImage-3.0-Instruct",
-        help="Model name or local path.",
-    )
+    parser.add_argument("--model", default="tencent/HunyuanImage-3.0-Instruct", help="Model name or local path.")
    parser.add_argument(
        "--modality",
        default="text2img",
-        choices=["text2img", "img2img", "img2text", "text2text"],
-        help="Modality mode to control stage execution.",
+        choices=list(_MODALITY_TASK_MAP),
    )
    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(
-        "--output",
-        type=str,
-        default=".",
-        help="Output directory to save results.",
-    )
-
-    # Generation parameters
+    parser.add_argument("--output", type=str, default=".", help="Output directory to save results.")
    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(
-        "--vae-use-tiling",
-        action="store_true",
-        help="Enable VAE tiling for memory optimization.",
-    )
-
-    # Prompt configuration
+    parser.add_argument("--vae-use-tiling", action="store_true", help="Enable VAE tiling.")
    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.",
    )
-
-    # Omni init args
+    parser.add_argument("--sys-type", type=str, default=None, help="Override system prompt type.")
    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 = _MODALITY_TASK_MAP[args.modality]
-    assert task is not None
-    bot_task = args.bot_task
-    if bot_task != "auto":
-        task = task + "_" + bot_task
-    if task not in _TASK_PRESETS:
-        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}"
-        )
+    task, default_bot_task = _MODALITY_TASK_MAP[args.modality]
+    if args.bot_task is None:
+        bot_task: str | None = default_bot_task
+    elif args.bot_task == "none":
+        bot_task = None
+    else:
+        bot_task = args.bot_task

    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:
-        print("[Info] No prompts provided, using default.")
-        prompts = ["A cute cat"]
-
-    # Load image if needed
-    input_image = None
+    input_images: list = []
    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:
-        result = build_prompt_tokens(p, tokenizer, task=task, sys_type=args.sys_type)
+    for prompt in prompts:
+        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["height"] = input_image.height
-            prompt_dict["width"] = input_image.width
+            prompt_dict["multi_modal_data"] = {"image": mm_image_payload}
+            prompt_dict["height"] = input_images[0].height
+            prompt_dict["width"] = input_images[0].width
        elif args.modality == "img2text":
            prompt_dict["modalities"] = ["text"]
-            prompt_dict["multi_modal_data"] = {"image": input_image}
-        elif args.modality == "text2text":
+            prompt_dict["multi_modal_data"] = {"image": mm_image_payload}
+        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")
@@ -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,
-        kv_cache_skip_steps=args.kv_cache_skip_steps,
-        kv_cache_skip_layers=args.kv_cache_skip_layers,
+        diffusion_kv_cache_dtype=args.diffusion_kv_cache_dtype,
+        diffusion_kv_cache_skip_steps=args.diffusion_kv_cache_skip_steps,
+        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"  kv_cache_skip_steps(config): {args.kv_cache_skip_steps}")
-    print(f"  kv_cache_skip_layers(config): {args.kv_cache_skip_layers}")
+    print(f"  diffusion_kv_cache_dtype(config): {args.diffusion_kv_cache_dtype}")
+    print(f"  diffusion_kv_cache_skip_steps(config): {args.diffusion_kv_cache_skip_steps}")
+    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}"
@@ -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",
@@ -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
@@ -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 →
@@ -10,16 +10,16 @@
                "dataset_name": "random",
                "backend": "openai-chat-omni",
                "endpoint": "/v1/chat/completions",
-                "num_prompts": [4, 16, 32, 64],
-                "max_concurrency": [1, 4, 8, 16],
+                "num_prompts": [4, 16, 32, 64, 128],
+                "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_audio_ttfp_ms": [1000, 3000, 5000, 7000],
-                    "mean_audio_rtf": [0.2, 0.35, 0.6, 0.85]
+                    "mean_ttft_ms": [1000, 3000, 5000, 7000, 9000],
+                    "mean_audio_ttfp_ms": [1000, 3000, 5000, 7000, 9000],
+                    "mean_audio_rtf": [0.2, 0.35, 0.6, 0.85, 0.9]
                }
            },
            {
@@ -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],
-        "max_concurrency": [8],
+        "dataset_path": "linyueqian/seed-tts-eval-subset",
+        "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_rtf": [0.75],
-          "audio_throughput": [10.0]
+          "median_audio_ttfp_ms": [3500, 30000, 7000],
+          "median_audio_rtf": [0.75, 10.0, 1.5],
+          "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],
-        "max_concurrency": [8],
+        "num_prompts": [80, 128, 128],
+        "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_rtf": [0.30],
-          "audio_throughput": [30.0]
+          "median_audio_ttfp_ms": [1500, 30000, 3000],
+          "median_audio_rtf": [0.30, 10.0, 0.60],
+          "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],
-        "max_concurrency": [8],
+        "num_prompts": [80, 128, 128],
+        "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_rtf": [0.35],
-          "audio_throughput": [25.0]
+          "median_audio_ttfp_ms": [1500, 30000, 3500],
+          "median_audio_rtf": [0.35, 10.0, 0.70],
+          "audio_throughput": [25.0, 1.0, 80.0]
        }
      },
      {
@@ -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):
@@ -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"
@@ -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 tasks >= {
-        "t2t",
-        "i2t",
+    assert set(available_tasks()) == {"t2t", "i2t", "it2i", "t2i"}
+
+
+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="t2i_recaption", tokenizer=tok) == [answer_id]
+    assert resolve_stop_token_ids(task="i2t", bot_task=None, 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",
-        "i2t",
-        "it2i_think",
-        "it2i_recaption",
-        "it2i_think_recaption",
-        "t2i_think",
-        "t2i_recaption",
+        ("t2t", None),
+        ("i2t", None),
+        ("it2i", "think"),
+        ("it2i", "recaption"),
+        ("it2i", "think_recaption"),
+        ("t2i", "think"),
+        ("t2i", "recaption"),
+        ("t2i", "think_recaption"),
    ],
 )
-def test_build_prompt_string_structure_chat_template(task: str):
-    """Chat-template tasks must produce <|startoftext|>...User: ...Assistant: ...
-    with image placeholder (when applicable) and trigger tag AFTER `Assistant: `."""
-    s = build_prompt("HELLO", task=task)
-
+def test_build_prompt_string_structure_chat_template(task: str, bot_task: str | None):
+    s = build_prompt("HELLO", task=task, bot_task=bot_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")):
-        assert s.index("User: ") < s.index("<img>") < s.index("HELLO"), (
-            "<img> placeholder must sit between `User: ` and the user prompt"
-        )
+    if task in ("i2t", "it2i"):
+        assert s.index("User: ") < s.index("<img>") < s.index("HELLO")
    else:
        assert "<img>" not in s

-    # Trigger tag must be the FINAL token of the prompt (after `Assistant: `).
-    # Note: the system prompt itself mentions <think>/<recaption> as mode
-    # documentation, so substring index() catches the wrong occurrence -- use
-    # endswith() which directly captures "trigger is at the tail" (the Part A
-    # fix: trigger goes AFTER `Assistant: `, not before user_prompt).
-    if task in ("it2i_think", "t2i_think", "it2i_think_recaption"):
-        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."
+    if bot_task in ("think", "think_recaption"):
+        assert s.endswith("Assistant: <think>")
+    elif bot_task == "recaption":
+        assert s.endswith("Assistant: <recaption>")
+    elif bot_task is None:
+        assert s.endswith("Assistant: ")


 def test_build_prompt_vanilla_uses_pretrain_template():
-    """t2i_vanilla is the only task that bypasses chat structure -- direct
-    text->image generation driven by the vanilla system prompt."""
-    s = build_prompt("HELLO", task="t2i_vanilla")
+    s = build_prompt("HELLO", task="t2i", bot_task="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")
-
-    # Each canonical segment is encoded in its own call.
+    build_prompt_tokens("写诗。", tok, task="i2t", bot_task=None)
    assert "User: " 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 "写诗。" in tok.encode_calls
    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 FakeTokenizer.SPECIAL["<img>"] in ids, "<img> placeholder must be present for i2t/it2i tasks"
+    assert ids[0] == FakeTokenizer.SPECIAL["<|startoftext|>"]
+    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>"]),
-        ("t2i_think", FakeTokenizer.SPECIAL["<think>"]),
-        ("it2i_recaption", FakeTokenizer.SPECIAL["<recaption>"]),
-        ("t2i_recaption", FakeTokenizer.SPECIAL["<recaption>"]),
+        ("it2i", "think", FakeTokenizer.SPECIAL["<think>"]),
+        ("t2i", "think", FakeTokenizer.SPECIAL["<think>"]),
+        ("t2i", "think_recaption", FakeTokenizer.SPECIAL["<think>"]),
+        ("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):
-    """Trigger tag id must be the LAST token (after `Assistant: ` segment)."""
+def test_build_prompt_tokens_trigger_is_last_token(task: str, bot_task: str | None, trigger_id: int):
    tok = FakeTokenizer()
-    result = build_prompt_tokens("hi", tok, task=task)
-    ids = result.token_ids
-    assert ids[-1] == trigger_id
+    result = build_prompt_tokens("hi", tok, task=task, bot_task=bot_task)
+    assert result.token_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")
-    ids = result.token_ids
-    assert ids[-1] not in {
+    result = build_prompt_tokens("hi", tok, task="t2t", bot_task=None)
+    assert result.token_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"}
-    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."
-    )
+    assert not (local_func_names & {"build_prompt", "build_prompt_tokens"})

    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 = {
-        "_TASK_PRESETS",
-        "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 --------------------
+    expected_imports = {"build_prompt_tokens", "resolve_stop_token_ids", "resolve_sys_type"}
+    assert expected_imports <= imported_from_prompt_utils


 _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(
-    not _hf_cached(_HUNYUAN_MODEL_ID),
-    reason=f"{_HUNYUAN_MODEL_ID} tokenizer not in HF cache",
-)
+@pytest.mark.skipif(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)
-
-    assert seg_ids != 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."
-    )
+    full_ids = tok.encode(build_prompt(user_prompt, task="i2t", bot_task=None), add_special_tokens=False)
+    assert seg_ids != full_ids
+    assert len(seg_ids) >= len(full_ids)
@@ -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")
@@ -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")

@@ -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):
-    adapter, _ = build_adapter()
-
-    first: OmniPayload = {
+def test_load_poll_non_ar_merges_into_existing_additional_information(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")
+    request.additional_information = {
        "hidden_states": {"output": torch.tensor([[1.0]])},
-        "codes": {"audio": [1]},
-        "meta": {"finished": torch.tensor(False, dtype=torch.bool)},
+        "ids": {"prompt": [11, 12]},
+        "meta": {"finished": torch.tensor(False, dtype=torch.bool), "step": 1},
    }
-    adapter._update_request_payload("ext", first)
-    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)
+    request.num_computed_tokens = 9

-    assert torch.equal(merged["hidden_states"]["output"], torch.tensor([[1.0], [2.0]]))
-    assert merged["codes"]["audio"] == [1, 2]
-    assert merged["meta"]["finished"].item() is True
+    payload: OmniPayload = {
+        "hidden_states": {"output": torch.tensor([[2.0]])},
+        "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")
-
-    adapter.request_ids_mapping["req-merged"] = "ext-merged"
-    adapter.request_payload["ext-merged"] = {
+    request.additional_information = {
        "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(
-        request.additional_information["hidden_states"]["output"],
-        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
+    # AR mode now forwards the latest payload directly.
+    assert request.additional_information == payload
    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)
@@ -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
@@ -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):
-        topo = KVTPTopology(source_tp_size=2, target_tp_size=2, local_rank=0)
-        assert get_kv_target_ranks(topo) == [0]
-        assert get_kv_source_ranks(topo) == [0]
+    @pytest.mark.parametrize(
+        "tp_size,local_rank",
+        [
+            (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):
-        topo = KVTPTopology(source_tp_size=2, target_tp_size=2, local_rank=1)
-        assert get_kv_target_ranks(topo) == [1]
-        assert get_kv_source_ranks(topo) == [1]
+    @pytest.mark.parametrize(
+        "local_rank,expected_source_ranks",
+        [
+            (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):
-        topo = KVTPTopology(source_tp_size=4, target_tp_size=4, local_rank=3)
-        assert get_kv_target_ranks(topo) == [3]
-        assert get_kv_source_ranks(topo) == [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)
-        assert get_kv_source_ranks(topo) == [0, 1]
-
-    def test_sender_gt_receiver_tp4_to_tp2_rank1(self):
-        """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]
+    @pytest.mark.parametrize(
+        "local_rank,expected_target_ranks",
+        [
+            (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_lt_receiver_tp2_to_tp4(self, local_rank, expected_target_ranks):
+        """Sender with source_tp=2, target_tp=4."""
+        topo = KVTPTopology(source_tp_size=2, target_tp_size=4, local_rank=local_rank)
+        assert get_kv_target_ranks(topo) == expected_target_ranks

    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):
-        topo = KVTPTopology(source_tp_size=2, target_tp_size=2, local_rank=0)
+    @pytest.mark.parametrize(
+        "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 rank == 0
-
-    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
+        assert key == expected_key
+        assert rank == expected_rank

    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):
-        """TP=2 rank 0: metadata should point to sender rank 0's port."""
-        mgr = _make_manager(from_tp=2, to_tp=2, local_rank=0, recv_timeout=0.05)
+    @pytest.mark.parametrize(
+        "local_rank,expected_port_offset",
+        [
+            (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
-
-    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
+        assert meta["source_port"] == 50151 + expected_port_offset * 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 ───────────────────────────────────────────────────

@@ -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"),
        ),
    ],
 )
@@ -93,7 +93,13 @@ def _run(stage_config_path: str, output_path: Path) -> tuple[Image.Image, str, f
    from vllm_omni.platforms import current_omni_platform

    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,
+        task="it2i",
+        bot_task="recaption",
+        sys_type="en_unified",
+    )
    token_ids = result.token_ids
    system_prompt_type = result.system_prompt_type

@@ -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 omni_out in omni.generate(request, DEFAULT_SAMPLING):
-        mm = omni_out.multimodal_output
-        assert mm is not None, "Expected multimodal_output to be non-None"
-        audio = mm.get("audio")
-        sr = mm.get("sr")
-        assert audio is not None, "Expected 'audio' key in multimodal_output"
-        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
+    for stage_outputs in omni.generate(request, DEFAULT_SAMPLING):
+        req_output = stage_outputs.request_output
+        if req_output is not None:
+            mm = req_output.outputs[0].multimodal_output
+            assert mm is not None, "Expected multimodal_output to be non-None"
+            audio = mm.get("audio")
+            sr = mm.get("sr")
+            assert audio is not None, "Expected 'audio' key in multimodal_output"
+            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.
-    for omni_out in omni_runner.omni.generate(requests, DEFAULT_SAMPLING):
-        mm = omni_out.multimodal_output
-        assert mm is not None
-        results.append(mm["audio"].cpu())
+    # Single-stage model (num_stages=1): one sampling param for all requests.
+    for stage_outputs in omni_runner.omni.generate(requests, [DEFAULT_SAMPLING]):
+        req_output = stage_outputs.request_output
+        if req_output is not None:
+            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):
@@ -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
@@ -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",
    )
 ]
@@ -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
@@ -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):
@@ -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
-    assert captured_sampling_params.height == 16
+    # size="auto" on multi-stage pipelines deliberately leaves the diffusion
+    # 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}"
+        )
@@ -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)
@@ -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'])}"
+    )
@@ -22,6 +22,7 @@ def serving_chat():

    instance = object.__new__(OmniOpenAIServingChat)
    instance._supported_speakers = None
+    instance.has_kv_connector = False
    return instance


@@ -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
@@ -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

@@ -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
@@ -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
@@ -1504,7 +1504,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["max_num_seqs"] == 64

    def test_empty_kwargs_yaml_only(self):
        stages = self._stages({})
@@ -1598,6 +1598,41 @@ 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."""
@@ -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


@@ -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"
@@ -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)

@@ -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,14 @@ 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.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 +81,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)

@@ -201,7 +226,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 +337,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 +356,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 +384,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 +521,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:
@@ -589,6 +650,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:
@@ -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,
@@ -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:
-                waiting_queue.remove(request)
+            if to_remove:
+                # 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
@@ -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)
@@ -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
@@ -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
@@ -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_layers = config.kv_cache_skip_layer_indices
+        self._kv_cache_skip_steps = config.diffusion_kv_cache_skip_step_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
@@ -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
-    kv_cache_skip_layers: str | None = None
-    kv_cache_skip_step_indices: set[int] | None = None
-    kv_cache_skip_layer_indices: set[int] | None = None
+    diffusion_kv_cache_skip_steps: str | None = None
+    diffusion_kv_cache_skip_layers: str | None = None
+    diffusion_kv_cache_skip_step_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.kv_cache_skip_layer_indices = parse_kv_cache_skip_selector(self.kv_cache_skip_layers)
+        self.diffusion_kv_cache_skip_step_indices = parse_kv_cache_skip_selector(self.diffusion_kv_cache_skip_steps)
+        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)
@@ -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,
+    ),
 }


@@ -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(),
@@ -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.
-        # The AR-side stage input processor (``ar2diffusion``) already prepends
-        # 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_list = [
+            (p.get("extra", {}).get("ar_generated_text") if isinstance(p, dict) else None) or None for p in req.prompts
+        ]
        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
        )
@@ -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
-canonical mapping for both flows.
+module provides that builder; the (task, bot_task) -> template mapping
+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_vanilla": ("en_vanilla", "image", None),
+    "t2i": ("en_unified", None, 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`."""
-    return sorted(_TASK_PRESETS)
+    """Sorted list of `task` values accepted by the prompt builders."""
+    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",
-    bot_task: str = "think",
+    task: str = "it2i",
+    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).
-
-    NOTE: when this string is passed to the engine, the engine's tokenizer
-    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]
+    """Build a HunyuanImage-3.0 prompt as a string (legacy/compat path)."""
+    task, bot_task = _normalize_task_and_bot_task(task, bot_task)
+    preset_sys_type, trigger_tag = _resolve_preset(task, bot_task)
    effective_sys_type = sys_type or preset_sys_type

-    system_prompt = get_system_prompt(effective_sys_type, preset_bot_task, custom_system_prompt)
-    sys_text = system_prompt.strip() if system_prompt else ""
+    system_prompt = get_system_prompt(effective_sys_type, bot_task, custom_system_prompt)
+    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
-    # vanilla system prompt drives the model with no chat structure.
-    if task == "t2i_vanilla":
+    if bot_task == "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
-    system_prompt_type: str  # The effective system prompt type used
+    token_ids: list[int]
+    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.
-
-    Calling tokenizer.encode(build_prompt(...)) on the full string lets BPE
-    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]
+    """Segment-by-segment tokenization that matches HF apply_chat_template."""
+    task, bot_task = _normalize_task_and_bot_task(task, bot_task)
+    preset_sys_type, trigger_tag = _resolve_preset(task, bot_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
-    # system prompt drives the model directly. No segment boundaries to
-    # protect, fall back to whole-string encode.
-    if task == "t2i_vanilla":
-        s = build_prompt(user_prompt, task, sys_type, custom_system_prompt)
+    if bot_task == "vanilla":
+        s = build_prompt(
+            user_prompt,
+            task=task,
+            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)
-    # Do NOT strip -- HF apply_chat_template keeps the system prompt's
-    # natural trailing newline; stripping it would shift one token id.
+    system_prompt = get_system_prompt(effective_sys_type, bot_task, custom_system_prompt)
    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",
+]
@@ -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
@@ -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 {
@@ -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:
-                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
-                kv_payload = rank_payloads[0]
-                for dst_rank in range(1, cfg_size):
-                    cfg_group.send_object(rank_payloads[dst_rank], dst_rank)
-            else:
+            is_owner = cfg_rank == 0 and sp_rank == 0
+
+            # step1: only owner need to receive
+            if is_owner:
+                received = self.receive_multi_kv_cache(
+                    req,
+                    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
@@ -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 = merged_payload
+                request.additional_information = payload_data
                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,
@@ -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 {}

@@ -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,
        )

@@ -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(
-                request_id=request_id,
-                prompt=prompt,
-                params=params,
-                supported_tasks=self.supported_tasks,
-                arrival_time=arrival_time,
-                lora_request=lora_request,
-                tokenization_kwargs=tokenization_kwargs,
-                trace_headers=trace_headers,
-                priority=priority,
-                data_parallel_rank=data_parallel_rank,
-                resumable=resumable,
-            )
+            try:
+                request = self.input_processor.process_inputs(
+                    request_id=request_id,
+                    prompt=prompt,
+                    params=params,
+                    supported_tasks=self.supported_tasks,
+                    arrival_time=arrival_time,
+                    lora_request=lora_request,
+                    tokenization_kwargs=tokenization_kwargs,
+                    trace_headers=trace_headers,
+                    priority=priority,
+                    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),
-            "kv_cache_skip_steps": kwargs.get("kv_cache_skip_steps", None),
-            "kv_cache_skip_layers": kwargs.get("kv_cache_skip_layers", None),
+            "diffusion_kv_cache_dtype": kwargs.get("diffusion_kv_cache_dtype", None),
+            "diffusion_kv_cache_skip_steps": kwargs.get("diffusion_kv_cache_skip_steps", 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")
-                if 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:
+                diffusion_kv_cache_dtype = kwargs.get("diffusion_kv_cache_dtype")
+                if diffusion_kv_cache_dtype is not None:
                    if (
-                        not hasattr(cfg.engine_args, "kv_cache_skip_steps")
-                        or cfg.engine_args.kv_cache_skip_steps is None
+                        not hasattr(cfg.engine_args, "diffusion_kv_cache_dtype")
+                        or cfg.engine_args.diffusion_kv_cache_dtype is None
                    ):
-                        cfg.engine_args.kv_cache_skip_steps = kv_cache_skip_steps
-                kv_cache_skip_layers = kwargs.get("kv_cache_skip_layers")
-                if kv_cache_skip_layers is not None:
+                        cfg.engine_args.diffusion_kv_cache_dtype = diffusion_kv_cache_dtype
+                diffusion_kv_cache_skip_steps = kwargs.get("diffusion_kv_cache_skip_steps")
+                if diffusion_kv_cache_skip_steps is not None:
                    if (
-                        not hasattr(cfg.engine_args, "kv_cache_skip_layers")
-                        or cfg.engine_args.kv_cache_skip_layers is None
+                        not hasattr(cfg.engine_args, "diffusion_kv_cache_skip_steps")
+                        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)

@@ -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],
-        rpc_async_queue: janus.AsyncQueue[EngineQueueMessage],
+        output_async_queue: janus.AsyncQueue[dict[str, Any]],
+        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":
-                    await self._handle_add_request(msg)
-                elif msg.type == "streaming_update":
-                    await self._handle_streaming_update(msg)
-                else:
-                    logger.warning("[Orchestrator] Unknown stage submission type: %s", msg.type)
-            elif isinstance(msg, AddCompanionRequestMessage):
+            if msg_type == "add_request":
+                await self._handle_add_request(msg)
+            elif msg_type == "streaming_update":
+                await self._handle_streaming_update(msg)
+            elif msg_type == "add_companion_request":
                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
-        kwargs = dict(msg.kwargs)
+        args = tuple(msg.args)
+        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]
-
-        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]
+        next_client = next_pool.stage_client
        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
-            if sender_stage is None:
-                sender_stage = sender_pool.llm_stage_client
-
-            sender_info = sender_stage.get_kv_sender_info()
+            sender_info = get_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(
@@ -3,6 +3,9 @@ from typing import Any
 import numpy as np
 import torch
 from vllm.logger import init_logger
+from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
+    split_routed_experts,
+)
 from vllm.outputs import PoolingRequestOutput
 from vllm.sampling_params import RequestOutputKind
 from vllm.tokenizers import TokenizerLike
@@ -230,7 +233,17 @@ class OmniRequestState(RequestState):
                self.sent_tokens_offset = self.detokenizer.num_output_tokens()

        external_req_id = self.external_req_id
-        output = self._new_completion_output(new_token_ids, finish_reason, stop_reason, routed_experts)
+
+        # Split routing data into prompt and generation portions, matching
+        # upstream make_request_output behaviour.
+        prompt_routed_experts = None
+        gen_routed_experts = None
+        if routed_experts is not None:
+            prompt_len = len(self.prompt_token_ids) if self.prompt_token_ids else 0
+            num_gen = self.detokenizer.num_output_tokens() if self.detokenizer is not None else None
+            prompt_routed_experts, gen_routed_experts = split_routed_experts(routed_experts, prompt_len, num_gen)
+
+        output = self._new_completion_output(new_token_ids, finish_reason, stop_reason, gen_routed_experts)

        if self.parent_req is None:
            outputs = [output]
@@ -240,7 +253,13 @@ class OmniRequestState(RequestState):
                return None
            external_req_id = self.parent_req.external_req_id

-        return self._new_request_output(external_req_id, outputs, finished, kv_transfer_params)
+        return self._new_request_output(
+            external_req_id,
+            outputs,
+            finished,
+            kv_transfer_params,
+            prompt_routed_experts,
+        )

    def _new_completion_output(
        self,
@@ -80,7 +80,8 @@ class StageEngineCoreClientBase(StageClientBase):
    def make_async_mp_client(
        vllm_config: Any,
        executor_class: type,
-        metadata: StageMetadata,
+        log_stats: bool = False,
+        metadata: StageMetadata | None = None,
        client_addresses: dict[str, str] | None = None,
        proc: Any = None,
        engine_manager: Any = None,
@@ -93,6 +94,7 @@ class StageEngineCoreClientBase(StageClientBase):
        client_args = dict(
            vllm_config=vllm_config,
            executor_class=executor_class,
+            log_stats=log_stats,
            metadata=metadata,
            client_addresses=client_addresses,
            proc=proc,
@@ -441,7 +443,7 @@ class StageEngineCoreClientBase(StageClientBase):
            kwargs=kwargs,
        )

-    def shutdown(self) -> None:
+    def shutdown(self, timeout: float | None = None) -> None:
        """Shutdown managed resources and any externally spawned subprocess."""
        child_procs: list[psutil.Process] = []
        if self._proc is not None and self._proc.pid is not None:
@@ -451,7 +453,7 @@ class StageEngineCoreClientBase(StageClientBase):
                child_procs = []

        try:
-            super().shutdown()
+            super().shutdown(timeout=timeout)
        finally:
            if self._proc is not None and self._proc.is_alive():
                self._proc.terminate()
@@ -232,6 +232,7 @@ class AsyncOmni(EngineClient, OmniBase):
        priority: int = 0,
        data_parallel_rank: int | None = None,
        reasoning_ended: bool | None = None,
+        reasoning_parser_kwargs: dict[str, Any] | None = None,
    ) -> AsyncGenerator[OmniRequestOutput, None]:
        """Generate outputs for the given prompt(s) asynchronously.

@@ -938,6 +939,37 @@ class AsyncOmni(EngineClient, OmniBase):
        """Check if tracing is enabled."""
        return False

+    async def notify_kv_transfer_request_rejected(
+        self,
+        request_id: str,
+        kv_transfer_params: dict[str, Any],
+        *,
+        data_parallel_rank: int | None = None,
+    ) -> None:
+        """Notify engine that a KV-transfer request was rejected before admission.
+
+        Omni does not currently use KV-transfer pre-admission resources,
+        so this is a no-op.
+        """
+        logger.debug(
+            "KV-transfer request rejected (no-op in omni): request_id=%s",
+            request_id,
+        )
+
+    async def start_weight_update(self, is_checkpoint_format: bool = True) -> None:
+        """Start a new weight update.
+
+        Omni does not currently support weight transfer, so this is a no-op.
+        """
+        logger.debug("Weight update start requested (no-op in omni)")
+
+    async def finish_weight_update(self) -> None:
+        """Finish the current weight update.
+
+        Omni does not currently support weight transfer, so this is a no-op.
+        """
+        logger.debug("Weight update finish requested (no-op in omni)")
+
    async def do_log_stats(self) -> None:
        """Log statistics.

@@ -472,26 +472,25 @@ class OmniServeCommand(CLISubcommand):
            default=None,
            help="Scheduler flow_shift for video models (e.g., 5.0 for 720p, 12.0 for 480p).",
        )
-        # vLLM already registers --kv-cache-dtype for the serve parser. Keep
-        # this fallback only for older vLLM versions where the option is absent.
-        if "--kv-cache-dtype" not in serve_parser._option_string_actions:
-            omni_config_group.add_argument(
-                "--kv-cache-dtype",
-                type=str,
-                default=None,
-                help="Config-level KV cache dtype (e.g. fp8).",
-            )
+        # Diffusion KV-cache quantization uses dedicated flags so we do not reuse
+        # vLLM's --kv-cache-dtype (AR cache dtype, default "auto").
        omni_config_group.add_argument(
-            "--kv-cache-skip-steps",
+            "--diffusion-kv-cache-dtype",
            type=str,
            default=None,
-            help="Config-level KV-cache quantization skip-step selector, e.g. '0-9,20,25-30'.",
+            help="Diffusion attention KV cache dtype (e.g. fp8). Separate from vLLM --kv-cache-dtype.",
        )
        omni_config_group.add_argument(
-            "--kv-cache-skip-layers",
+            "--diffusion-kv-cache-skip-steps",
            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-step selector, e.g. '0-9,20,25-30'.",
+        )
+        omni_config_group.add_argument(
+            "--diffusion-kv-cache-skip-layers",
+            type=str,
+            default=None,
+            help="Diffusion KV-cache quantization skip-layer selector, e.g. '0,1,4-8'.",
        )
        omni_config_group.add_argument(
            "--cfg-parallel-size",
@@ -1700,7 +1700,10 @@ async def edit_images(
    # vllm-omni extension for layered models (e.g., Qwen-Image-Layered)
    layers: int | None = Form(None),
    resolution: int | None = Form(None),  # See SUPPORTED_LAYERED_RESOLUTIONS
+    # /v1/images/edits is always IT2I; only the prompting knobs are exposed.
    bot_task: str | None = Form(None),
+    sys_type: str | None = Form(None),
+    system_prompt: str | None = Form(None),
 ) -> ImageGenerationResponse:
    """
    OpenAI-compatible image edit endpoint.
@@ -1751,16 +1754,20 @@ async def edit_images(
                status_code=HTTPStatus.BAD_REQUEST.value,
                detail=detail,
            )
-        pil_images = await _load_input_images(input_images_list)
+        # Match the offline path: RGB normalize when the caller opts into
+        # Hunyuan-aware behavior. RGBA/P uploads otherwise diverge from offline.
+        normalize_edit_images_rgb = bot_task is not None or sys_type is not None
+        pil_images = await _load_input_images(input_images_list, normalize_rgb=normalize_edit_images_rgb)
        prompt["multi_modal_data"] = {}
        prompt["multi_modal_data"]["image"] = pil_images

        if mask_image is not None:
-            loaded = await _load_input_images([mask_image])
+            # Mask role is different (alpha channel matters); never normalize.
+            loaded = await _load_input_images([mask_image], normalize_rgb=False)
            prompt["multi_modal_data"]["mask_image"] = loaded[0]

        if reference_image is not None:
-            loaded = await _load_input_images([reference_image])
+            loaded = await _load_input_images([reference_image], normalize_rgb=normalize_edit_images_rgb)
            prompt["multi_modal_data"]["reference_image"] = loaded[0]

        # 3 Build sample params
@@ -1811,7 +1818,8 @@ async def edit_images(

        # 3.3 Parse and add size if provided
        width, height = None, None
-        if size.lower() == "auto":
+        size_was_auto = size.lower() == "auto"
+        if size_was_auto:
            if resolution is None:
                # No resolution specified, use input image size
                width, height = pil_images[0].size
@@ -1882,10 +1890,13 @@ async def edit_images(
                "seed": effective_seed,
                "num_outputs_per_prompt": n,
            }
-            if width is not None:
-                extra_body["width"] = width
-            if height is not None:
-                extra_body["height"] = height
+            # size="auto" resolves width/height from input image; forwarding
+            # those would override AR-driven `<img_ratio_*>` token selection.
+            if not size_was_auto:
+                if width is not None:
+                    extra_body["width"] = width
+                if height is not None:
+                    extra_body["height"] = height
            if negative_prompt is not None:
                extra_body["negative_prompt"] = negative_prompt
            if num_inference_steps is not None:
@@ -1907,6 +1918,10 @@ async def edit_images(
                extra_body["lora"] = lora_dict
            if bot_task is not None:
                extra_body["bot_task"] = bot_task
+            if sys_type is not None:
+                extra_body["sys_type"] = sys_type
+            if system_prompt is not None:
+                extra_body["system_prompt"] = system_prompt

            prompt_text = prompt.get("prompt", "")
            generation_result = await chat_handler.generate_diffusion_images(
@@ -2187,6 +2202,8 @@ def _extract_images_from_result(result: Any) -> list[Any]:

 async def _load_input_images(
    inputs: list[str],
+    *,
+    normalize_rgb: bool = True,
 ) -> list[Image.Image]:
    """
    convert to PIL.Image.Image list
@@ -2233,7 +2250,18 @@ async def _load_input_images(
    if not images:
        raise ValueError("No valid input images found")

-    return images
+    if not normalize_rgb:
+        return images
+
+    # Match the offline HunyuanImage3 image-edit example path, which eagerly
+    # normalizes input files with ``Image.open(...).convert("RGB")`` before
+    # they reach the AR stage. Keeping uploads as RGBA/P PIL objects makes
+    # online IT2I observe a different visual input than offline (for example
+    # transparent-logo uploads alpha-composited over white instead of black),
+    # which is enough for HunyuanImage3 AR recaption to diverge before DiT
+    # sees the request -- root cause of the "online 3 magnets vs offline 1
+    # magnet" systematic semantic mismatch.
+    return [img.convert("RGB") for img in images]


 def _choose_output_format(output_format: str | None, background: str | None) -> str:
@@ -196,7 +196,7 @@ class SpeechBatchItem(BaseModel):
    all other fields override the batch-level defaults when set."""

    input: str
-    voice: str | None = None
+    voice: str | None = Field(default=None, validation_alias=AliasChoices("voice", "speaker"))
    instructions: str | None = None
    response_format: Literal["wav", "pcm", "flac", "mp3", "aac", "opus"] | None = None
    speed: float | None = Field(default=None, ge=0.25, le=4.0)
@@ -215,7 +215,7 @@ class BatchSpeechRequest(BaseModel):

    model: str | None = None
    items: list[SpeechBatchItem] = Field(..., min_length=1)
-    voice: str | None = None
+    voice: str | None = Field(default=None, validation_alias=AliasChoices("voice", "speaker"))
    instructions: str | None = None
    response_format: Literal["wav", "pcm", "flac", "mp3", "aac", "opus"] = "wav"
    speed: float | None = Field(default=1.0, ge=0.25, le=4.0)
@@ -248,7 +248,7 @@ class StreamingSpeechSessionConfig(BaseModel):
    """Configuration sent as the first WebSocket message for streaming TTS."""

    model: str | None = None
-    voice: str | None = None
+    voice: str | None = Field(default=None, validation_alias=AliasChoices("voice", "speaker"))
    task_type: Literal["CustomVoice", "VoiceDesign", "Base"] | None = None
    language: str | None = None
    instructions: str | None = None
@@ -120,6 +120,7 @@ class RealtimeConnection(VllmRealtimeConnection):
        full_text = ""
        prompt_token_ids_len = 0
        completion_tokens_len = 0
+        last_prompt_token_ids_len = 0  # detect Stage-0 segment rollover
        self._realtime_audio_ref = None

        # Coerce cumulative outputs to delta outputs; this ensures
@@ -144,6 +145,18 @@ class RealtimeConnection(VllmRealtimeConnection):
                    new_token_ids = list(first_output.token_ids)
                    new_tokens_len = len(new_token_ids)

+                    cur_prompt_token_ids_len = len(output.prompt_token_ids or [])
+                    # Stage-0 segment rollover: buffer_realtime_audio may yield
+                    # multiple TokensPrompt segments and the second segment's
+                    # prompt_token_ids include the first segment's decoded output.
+                    # Clear accumulated full_text at the boundary so the
+                    # carried-over prefix is not duplicated in the final
+                    # TranscriptionDone.text.
+                    stage_id = getattr(output, "stage_id", None)
+                    if stage_id == 0 and cur_prompt_token_ids_len > last_prompt_token_ids_len > 0:
+                        full_text = ""
+                    last_prompt_token_ids_len = cur_prompt_token_ids_len
+
                    if not prompt_token_ids_len and output.prompt_token_ids:
                        prompt_token_ids_len = len(output.prompt_token_ids)

@@ -148,6 +148,8 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
        instance._diffusion_mode = True
        instance._diffusion_engine = diffusion_engine
        instance._diffusion_model_name = model_name
+        instance.engine_client = None
+        instance.has_kv_connector = False
        return instance

    def _get_diffusion_extra_body_params(self) -> frozenset[str]:
@@ -216,6 +218,15 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
        For diffusion models, this generates images and returns them
        in a chat completion response format.
        """
+        return await self._with_kv_transfer_rejection_cleanup(
+            self._create_chat_completion(request, raw_request), request, raw_request
+        )
+
+    async def _create_chat_completion(
+        self,
+        request: ChatCompletionRequest,
+        raw_request: Request | None = None,
+    ) -> AsyncGenerator[str, None] | ChatCompletionResponse | ErrorResponse:
        # Handle diffusion mode
        if self._diffusion_mode:
            return await self._create_diffusion_chat_completion(request, raw_request)
@@ -419,7 +430,10 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                    # consistency.  After the multimodal processor consumes
                    # the image data, the uuids remain as a stable reference.
                    tprompt["multi_modal_uuids"] = {
-                        k: [f"{request_id}-{k}-{i}"] for i, k in enumerate(engine_prompt_image)
+                        k: [f"{request_id}-{k}-{i}" for i in range(len(v))]
+                        if isinstance(v, list)
+                        else [f"{request_id}-{k}-0"]
+                        for k, v in engine_prompt_image.items()
                    }

                engine_prompts = [tprompt]
@@ -993,6 +1007,9 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                # when handling audio/image responses
                role = self.get_chat_request_role(request)

+                # Compute prompt_text once at first iteration (upstream #42052)
+                prompt_text = getattr(res, "prompt", None) if getattr(request, "return_prompt_text", None) else None
+
                # We need to do it here, because if there are exceptions in
                # the result_generator, it needs to be sent as the FIRST
                # response (by the try...catch).
@@ -1019,6 +1036,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                            choices=[choice_data],
                            model=model_name,
                            prompt_token_ids=(res.prompt_token_ids if request.return_token_ids else None),
+                            prompt_text=prompt_text,
                            modality=final_output_type,
                        )

@@ -1737,6 +1755,13 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                    response_metrics.update(extra)
            choices.extend(choices_data)

+        # Compute prompt_text for non-streaming response (upstream #42052)
+        prompt_text = (
+            getattr(final_res, "prompt", None)
+            if final_res is not None and getattr(request, "return_prompt_text", None)
+            else None
+        )
+
        response = OmniChatCompletionResponse(
            id=request_id,
            created=created_time,
@@ -1745,6 +1770,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
            usage=usage,
            prompt_logprobs=prompt_logprobs,
            prompt_token_ids=prompt_token_ids,
+            prompt_text=prompt_text,
            kv_transfer_params=kv_transfer_params,
            metrics=response_metrics,
        )
@@ -2031,13 +2057,17 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
        mm_output = final_res.outputs[0].multimodal_output
        audio_data = mm_output.get("audio")
        if isinstance(audio_data, list):
-            if stream:
+            if not audio_data:
+                audio_tensor = None
+            elif stream:
                audio_tensor = audio_data[-1]
            else:
                audio_tensor = torch.cat(audio_data, dim=-1)
        else:
            audio_tensor = audio_data
-        audio_tensor = audio_tensor.float().detach().cpu().numpy()
+        if audio_tensor is None:
+            return self._create_error_response("Audio generation completed but no audio was produced.")
+        audio_tensor = audio_tensor.detach().cpu().float().numpy()

        # Ensure audio is 1D (flatten if needed)
        if audio_tensor.ndim > 1:
@@ -2245,6 +2275,8 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
        layers = extra_body.get("layers")
        resolution = extra_body.get("resolution")
        bot_task = extra_body.get("bot_task")
+        sys_type = extra_body.get("sys_type")
+        custom_system_prompt = extra_body.get("system_prompt")

        engine_prompt_data: dict[str, Any] | None = None
        modalities = ["image"]
@@ -2255,31 +2287,48 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
            else:
                engine_prompt_data = {"image": reference_images}

-        engine_prompt: OmniTextPrompt = {"prompt": prompt}
-        if bot_task:
+        prompt_token_ids: list[int] | None = None
+        system_prompt_type: str | None = None
+        if bot_task is not None or sys_type is not None or custom_system_prompt is not None:
            from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import (
                build_prompt,
                build_prompt_tokens,
            )

-            prompt_token_ids: list[int] | None = None
-            system_prompt_type: str | None = None
+            build_kwargs: dict[str, Any] = {
+                "task": "it2i" if reference_images else "t2i",
+                "sys_type": sys_type,
+                "custom_system_prompt": custom_system_prompt,
+                "num_images": len(reference_images) if reference_images else 1,
+            }
+            if bot_task is not None:
+                build_kwargs["bot_task"] = bot_task
+            elif "bot_task" in extra_body:
+                # Explicit None from the caller is plain-mode; omitted lets
+                # each task fall back to its default trigger.
+                build_kwargs["bot_task"] = None
            if tokenizer is not None:
-                result = build_prompt_tokens(prompt, tokenizer, task=bot_task)
+                # Feed segment-tokenized prompt_token_ids so AR matches HF
+                # apply_chat_template byte-for-byte (engine BPE would merge
+                # across template boundaries, e.g. "。\n\n" -> single id).
+                result = build_prompt_tokens(prompt, tokenizer, **build_kwargs)
                prompt_token_ids = result.token_ids
                system_prompt_type = result.system_prompt_type
            else:
-                prompt = build_prompt(prompt, task=bot_task)
-                engine_prompt["prompt"] = prompt
-
+                prompt = build_prompt(prompt, **build_kwargs)
            if reference_images and len(reference_images) == 1:
                engine_prompt_data = {"image": reference_images[0]}
                modalities = ["image"]
-            if prompt_token_ids is not None:
-                engine_prompt["prompt_token_ids"] = prompt_token_ids
-            if system_prompt_type is not None:
-                engine_prompt["use_system_prompt"] = system_prompt_type

+        engine_prompt: OmniTextPrompt = {"prompt": prompt}
+        if prompt_token_ids is not None:
+            engine_prompt["prompt_token_ids"] = prompt_token_ids
+        if system_prompt_type is not None:
+            engine_prompt["use_system_prompt"] = system_prompt_type
+        # DiT's get_system_prompt(use_system_prompt, "image", system_prompt) reads
+        # this; omitting it makes sys_type=custom yield an empty DiT prefix.
+        if custom_system_prompt is not None:
+            engine_prompt["system_prompt"] = custom_system_prompt
        engine_prompt["modalities"] = modalities
        if negative_prompt is not None:
            engine_prompt["negative_prompt"] = negative_prompt
@@ -2295,7 +2344,11 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
            engine_prompt["multi_modal_data"] = engine_prompt_data
            # Provide multi_modal_uuids so that newer vLLM versions can
            # validate multi_modal_data / multi_modal_uuids consistency.
-            engine_prompt["multi_modal_uuids"] = {k: [f"img-{k}-{i}"] for i, k in enumerate(engine_prompt_data)}
+            # Generate one uuid per image when the value is a list (multi-image inputs).
+            engine_prompt["multi_modal_uuids"] = {
+                k: [f"img-{k}-{i}" for i in range(len(v))] if isinstance(v, list) else [f"img-{k}-0"]
+                for k, v in engine_prompt_data.items()
+            }

        comprehension_idx = None
        for idx, stage in enumerate(stage_configs):
@@ -2320,10 +2373,8 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
            ):
                default_stage_params.seed = seed

-            # Inject target_h/w into comprehension (AR) stage sampling params
-            # for models that need M-RoPE position pre-computation (e.g.
-            # GLM-Image).  max_tokens is handled via the deploy YAML default
-            # (upper-bound ceiling) rather than computed dynamically here.
+            # Inject target_h/w into AR stage for M-RoPE position pre-computation
+            # (e.g. GLM-Image). max_tokens comes from deploy YAML.
            if comprehension_idx is not None and idx == comprehension_idx and height is not None and width is not None:
                extra_args = getattr(default_stage_params, "extra_args", None)
                if extra_args is None:
@@ -2449,13 +2500,17 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
            diffusion_engine = cast(AsyncOmni, engine)
            stage_configs = getattr(diffusion_engine, "stage_configs", None) or []
            if len(stage_configs) > 1:
+                # Pull tokenizer from the comprehension (AR) stage so we can
+                # build HF byte-for-byte prompt_token_ids in the helper. If
+                # the engine doesn"t expose one, fall back to the legacy
+                # string-prompt path (engine re-tokenizes).
                tokenizer = None
                get_tok = getattr(diffusion_engine, "get_tokenizer", None)
                if get_tok is not None:
                    try:
                        tokenizer = await get_tok()
                    except Exception as exc:
-                        logger.warning("get_tokenizer failed: %s", exc)
+                        logger.warning("get_tokenizer failed; falling back to string prompt path: %s", exc)
                engine_prompt, sampling_params_list = self._build_multistage_generation_inputs(
                    engine=diffusion_engine,
                    prompt=prompt,
@@ -2103,6 +2103,13 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                    voice_lower = request.voice.lower()
                    tts_params["voice_name"] = [voice_lower]
                    tts_params["voice_created_at"] = [self._voice_created_at(voice_lower)]
+                # Propagate seed from sampling params for deterministic generation.
+                # MOSS-TTS-Nano uses its own internal sampling inside
+                # inference_stream(), which reads seed from additional_information
+                # (not SamplingParams). Without this the model uses ambient RNG
+                # state and produces non-deterministic output.
+                if sampling_params_list and getattr(sampling_params_list[0], "seed", None) is not None:
+                    tts_params["seed"] = [sampling_params_list[0].seed]
                prompt = tokens_input(prompt_token_ids=[1])
                prompt["additional_information"] = tts_params
            else:
@@ -351,8 +351,10 @@ class CodePredictorBaseModel(nn.Module):
        # native bf16 support (Turing, Volta).  The RMSNorm and RoPE
        # layers already upcast internally; this extends the same
        # treatment to attention and MLP.
+        # autocast to float32 is unsupported on CPU; skip fp32 upcast there
+        # (CPU uses full-precision intermediates internally).
        input_dtype = inputs_embeds.dtype
-        use_fp32 = input_dtype == torch.float16
+        use_fp32 = input_dtype == torch.float16 and inputs_embeds.device.type != "cpu"
        if use_fp32:
            inputs_embeds = inputs_embeds.float()
        hidden_states = inputs_embeds
@@ -61,8 +61,9 @@ class SnakeBeta(nn.Module):
            x = tl.load(x_ptr + bid * stride_b + cid * stride_c + t_off, mask=mask, other=0.0)
            ea = tl.load(exp_alpha_ptr + cid)
            ib = tl.load(inv_beta_ptr + cid)
-            sin_val = tl.sin(x * ea)
-            result = x + ib * sin_val * sin_val
+            x_float = x.to(tl.float32)
+            sin_val = tl.sin(x_float * ea)
+            result = x + (ib * sin_val * sin_val).to(x.dtype)

            tl.store(out_ptr + bid * stride_b + cid * stride_c + t_off, result, mask=mask)

@@ -27,17 +27,8 @@ from vllm.distributed import (
 )
 from vllm.inputs import MultiModalDataDict
 from vllm.logger import init_logger
+from vllm.model_executor.layers.fused_moe import FusedMoE as SharedFusedMoE
 from vllm.model_executor.layers.fused_moe import fused_moe_make_expert_params_mapping
-
-try:
-    from vllm.model_executor.layers.fused_moe.shared_fused_moe import SharedFusedMoE
-except ImportError:
-    # PyPI vllm 0.20.x neither exports `SharedFusedMoE` from the package top-level
-    # nor ships a `shared_fused_moe.py` submodule. The functionality lives on
-    # `FusedMoE` directly (which gained a `shared_experts` parameter), so alias
-    # the symbol — call sites only use the classmethod `make_expert_params_mapping`
-    # and `__init__(shared_experts=..., ...)` which are present on `FusedMoE`.
-    from vllm.model_executor.layers.fused_moe import FusedMoE as SharedFusedMoE
 from vllm.model_executor.layers.linear import (
    ColumnParallelLinear,
    ReplicatedLinear,
@@ -737,7 +728,7 @@ class HunyuanImage3Processor:
    class ResolutionGroup:
        """Group of resolutions for image processing."""

-        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}"
@@ -751,6 +742,11 @@ class HunyuanImage3Processor:
            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
@@ -815,7 +811,18 @@ class HunyuanImage3Processor:
    def __init__(self, tokenizer, hf_config, **kwargs: object):
        self.tokenizer = tokenizer
        self.hf_config = hf_config
-        self.reso_group = self.ResolutionGroup(base_size=hf_config.image_base_size)
+        # `HUNYUAN_IMAGE3_EXTRA_RESOLUTIONS` mirrors the official
+        # `vae_reso_group` extras (image_processor.py:147-152). Build with
+        # this processor's inner Resolution class so `data` stays
+        # type-homogeneous.
+        from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_image3_transformer import (
+            HUNYUAN_IMAGE3_EXTRA_RESOLUTIONS,
+        )
+
+        self.reso_group = self.ResolutionGroup(
+            base_size=hf_config.image_base_size,
+            extra_resolutions=[HunyuanImage3Processor.Resolution(s) for s in HUNYUAN_IMAGE3_EXTRA_RESOLUTIONS],
+        )
        self.vision_encoder_processor = Siglip2ImageProcessorFast.from_dict(hf_config.vit_processor)
        self.vae_processor = transforms.Compose(
            [
@@ -860,6 +867,13 @@ class HunyuanImage3Processor:
        else:
            raise TypeError(f"Unsupported image type: {type(image_input)}.")

+        # Each cond image keeps its own VAE bucket (mirrors official HF's
+        # ragged behavior in `_encode_cond_image`). VAE pixel tensors have
+        # different (H_i, W_i) per image, so they're flattened to 1-D and
+        # concatenated; vLLM `flat_from_sizes("image", vae_pixel_size)` slices
+        # them back per-image at consumption time. VIT (Siglip2 naflex) pads
+        # to `max_num_patches` so VIT fields keep the existing `batched`
+        # stack path.
        batch_data = []
        for image in images:
            current_info = {}
@@ -883,68 +897,80 @@ class HunyuanImage3Processor:
                _ss = torch.tensor(_ss, dtype=torch.long)
            current_info["vit_spatial_shapes"] = _ss.squeeze(0)

-            # VAE processing.
-            # The resize/crop math here mirrors HF's `resize_and_crop` with
-            # crop_type="center" (hunyuan3.0_ins/image_processor.py:61). VAE
-            # normalize uses the same transforms.Compose([ToTensor,
-            # Normalize([0.5], [0.5])]) as HF's `pil_image_to_tensor`. So
-            # numerical output of this branch should match HF up to floating-
-            # point reduction order.
+            # VAE: per-image bucket via `reso_group.get_target_size`; mirrors
+            # HF's `resize_and_crop` (crop_type="center", the official
+            # generate_image default with infer_align_image_size=False).
+            # Keep fp32 — the VAE encoder casts to model dtype at its
+            # boundary (see `_vae_encode`).
            image_width, image_height = self.reso_group.get_target_size(image.width, image.height)
            resized_image = self._resize_and_crop(image, (image_width, image_height))
-            vae_pixel_values = self.vae_processor(resized_image)
+            vae_pixel_values = self.vae_processor(resized_image).squeeze(0)
            token_height = image_height // (self.hf_config.vae_downsample_factor[0] * self.hf_config.patch_size)
            token_width = image_width // (self.hf_config.vae_downsample_factor[1] * self.hf_config.patch_size)
-            # Keep fp32 — the VAE encoder casts to model dtype at its boundary
-            # (see _vae_encode). Casting to bf16 here costs ~7e-4 mean-abs-diff
-            # bf16 quantization error on every pixel vs HF (which keeps fp32
-            # in build_cond_images), measurable as a real numerical drift in
-            # downstream image embeddings.
-            current_info["vae_pixel_values"] = vae_pixel_values.squeeze(0)
+
+            current_info["vae_pixel_values_flat"] = vae_pixel_values.reshape(-1)
+            current_info["vae_pixel_size"] = torch.tensor(vae_pixel_values.numel(), dtype=torch.long)
            current_info["vae_token_grid_hw"] = torch.tensor([token_height, token_width])

-            # size
            base_size, ratio_index = self.reso_group.get_base_size_and_ratio_index(image_width, image_height)
            current_info["base_size"] = torch.tensor(base_size)
            current_info["ratio_index"] = torch.tensor(ratio_index)

            batch_data.append(current_info)

-        # Stack the tensors in the list into a batch dimension (B, ...)
-        final_image_info = {}
-        if len(batch_data) > 0:
-            for key in batch_data[0].keys():
-                final_image_info[key] = torch.stack([d[key] for d in batch_data], dim=0)
+        final_image_info: dict[str, torch.Tensor] = {}
+        if not batch_data:
+            return final_image_info

-        if final_image_info:
-            shapes_info = {k: tuple(v.shape) for k, v in final_image_info.items()}
-            logger.info(f"Successfully processed {len(images)} image(s). Final tensor shapes: {shapes_info}")
+        # Same-shape fields: stack along a new image-batch dim as before.
+        same_shape_keys = [
+            "vit_pixel_values",
+            "vit_pixel_attention_mask",
+            "vit_spatial_shapes",
+            "vae_token_grid_hw",
+            "vae_pixel_size",
+            "base_size",
+            "ratio_index",
+        ]
+        for key in same_shape_keys:
+            final_image_info[key] = torch.stack([d[key] for d in batch_data], dim=0)
+
+        # Variable-shape VAE pixels: 1-D concat across images (paired with
+        # `vae_pixel_size` via `flat_from_sizes` in `_get_mm_fields_config`).
+        final_image_info["vae_pixel_values"] = torch.cat([d["vae_pixel_values_flat"] for d in batch_data], dim=0)
+
+        shapes_info = {k: tuple(v.shape) for k, v in final_image_info.items()}
+        logger.info(f"Successfully processed {len(images)} image(s). Final tensor shapes: {shapes_info}")

        return final_image_info

-    def _resize_and_crop(self, image: Image.Image, target_size: tuple[int, int]) -> Image.Image:
+    def _resize_and_crop(
+        self,
+        image: Image.Image,
+        target_size: tuple[int, int],
+        crop_type: str = "resize",
+    ) -> Image.Image:
+        # Default mode mirrors the official `infer_align_image_size=True`
+        # path (image_processor.py:355 → crop_type="resize") used by the
+        # IT2I demo: stretch the cond image to the bucket dims so its
+        # `<img_ratio_*>` tag and ViT/VAE features stay aligned with the
+        # bucket, instead of dropping content via center crop.
        tw, th = target_size
+        if crop_type == "resize":
+            return image.resize((tw, th), resample=Image.Resampling.LANCZOS)
        w, h = image.size
-
        tr = th / tw
        r = h / w
-
-        # resize
        if r < tr:
            resize_height = th
            resize_width = int(round(th / h * w))
        else:
            resize_width = tw
            resize_height = int(round(tw / w * h))
-
        image = image.resize((resize_width, resize_height), resample=Image.Resampling.LANCZOS)
-
-        # center crop
        crop_top = int(round((resize_height - th) / 2.0))
        crop_left = int(round((resize_width - tw) / 2.0))
-
-        image = image.crop((crop_left, crop_top, crop_left + tw, crop_top + th))
-        return image
+        return image.crop((crop_left, crop_top, crop_left + tw, crop_top + th))


 class HunyuanImage3ProcessingInfo(BaseProcessingInfo):
@@ -1030,8 +1056,13 @@ class HunyuanImage3MultiModalProcessor(BaseMultiModalProcessor[HunyuanImage3Proc
            config["vit_pixel_attention_mask"] = MultiModalFieldConfig.batched("image")
        if "vit_spatial_shapes" in hf_inputs:
            config["vit_spatial_shapes"] = MultiModalFieldConfig.batched("image")
-        if "vae_pixel_values" in hf_inputs:
-            config["vae_pixel_values"] = MultiModalFieldConfig.batched("image")
+        # `vae_pixel_values` is a 1-D concatenation of variable-shape per-image
+        # VAE tensors (see `process_image`). `vae_pixel_size` carries the
+        # per-image flat length so vLLM can split the buffer back per image.
+        if "vae_pixel_values" in hf_inputs and "vae_pixel_size" in hf_inputs:
+            config["vae_pixel_values"] = MultiModalFieldConfig.flat_from_sizes("image", hf_inputs["vae_pixel_size"])
+        if "vae_pixel_size" in hf_inputs:
+            config["vae_pixel_size"] = MultiModalFieldConfig.batched("image")
        if "vae_token_grid_hw" in hf_inputs:
            config["vae_token_grid_hw"] = MultiModalFieldConfig.batched("image")
        if "base_size" in hf_inputs:
@@ -1085,38 +1116,37 @@ class HunyuanImage3MultiModalProcessor(BaseMultiModalProcessor[HunyuanImage3Proc
            ratio_token_id = tokenizer.convert_tokens_to_ids(f"<img_ratio_{_ratio_index}>")
            if ratio_token_id is None:
                raise ValueError(f"Ratio token '<img_ratio_{_ratio_index}>' not found in tokenizer vocabulary")
+            timestep_token_id = tokenizer.convert_tokens_to_ids("<timestep>")
+            if timestep_token_id is None:
+                raise ValueError("Timestep token '<timestep>' not found in tokenizer vocabulary")

-            # NOTE on the timestep slot:
-            # HF's apply_chat_template emits the literal <timestep> token id
-            # 128017 here. HF's modeling forward (`instantiate_continuous_tokens`,
-            # see hunyuan3.0_ins/modeling_hunyuan_image_3.py:1964) then *scatter-
-            # replaces* the embedding at that position with `timestep_emb(0)`
-            # for cond images. So the wte embedding of <timestep> is irrelevant
-            # at runtime — what matters is the timestep_emb injection.
+            # Use the real <timestep> token id (HF parity). The trained wte
+            # at this slot is overwritten with timestep_emb(0) at runtime by
+            # `embed_input_ids`.
            #
-            # vllm-omni achieves the same effect via the multimodal-embedding
-            # merger: we put an <img> (128006) placeholder here and ship a
-            # `timestep_emb(0)` tensor at the head of `embed_multimodal()`'s
-            # combined_embeddings. The merger replaces this placeholder's
-            # embedding with the timestep tensor, yielding a final hidden
-            # state numerically equivalent to HF at that position.
-            #
-            # Keep this slot as <img> (NOT <timestep>): switching to <timestep>
-            # requires either (a) a second PromptReplacement targeting 128017,
-            # or (b) the merger's embed_token_id to be a list — neither is
-            # currently supported by PromptUpdateDetails.select_token_id.
+            # Mark <img>*VAE + <joint_img_sep> + <img>*ViT as one contiguous
+            # embed run so vLLM's prefix-LM mask treats it as a single
+            # bidirectional region, mirroring official `joint_image_slices`
+            # full-attention range (image_processor.py:388, with
+            # cond_token_attn_type effectively spanning VAE+sep+ViT). With the
+            # default `select_token_id(<img>)` mask, sep splits the run into
+            # two regions; that asymmetry is what biased multi-image AR
+            # ratio prediction to the first image's bucket.
            replacement = (
                [boi_token_id]
                + [base_size_token_id]
                + [ratio_token_id]
-                + [img_token_id] * timestep_token_num
+                + [timestep_token_id] * timestep_token_num
                + [img_token_id] * vae_token_num
                + [joint_img_sep_token_id]
                + [img_token_id] * vit_token_num
                + [eoi_token_id]
            )
            logger.debug(f"actual replacement token count: {timestep_token_num + vae_token_num + vit_token_num}")
-            return PromptUpdateDetails.select_token_id(replacement, embed_token_id=img_token_id)
+            return PromptUpdateDetails.select_token_ids(
+                replacement,
+                embed_token_ids=[img_token_id, joint_img_sep_token_id],
+            )

        return [
            PromptReplacement(modality="image", target=[img_token_id], replacement=get_replacement_image),
@@ -1128,7 +1158,7 @@ def _hunyuan_image3_unpack_packed_topk(
    gating_output: torch.Tensor,
    topk: int,
    renormalize: bool,
-    num_experts: int,
+    num_experts: int | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
    """Unpack pre-computed ``(topk_weights, topk_indices)`` packed by
    :class:`HunyuanImage3SparseMoeBlock` into ``gating_output``.
@@ -1502,6 +1532,7 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo
        self._end_of_answer_id = tokenizer.convert_tokens_to_ids("</answer>")
        image_base_size = getattr(config, "image_base_size", 1024)
        self._size_token_id = tokenizer.convert_tokens_to_ids(f"<img_size_{image_base_size}>")
+        self._timestep_token_id = tokenizer.convert_tokens_to_ids("<timestep>")
        self._start_ratio_id = tokenizer.convert_tokens_to_ids("<img_ratio_0>")
        self._end_ratio_id = tokenizer.convert_tokens_to_ids("<img_ratio_32>")
        ratio_33 = tokenizer.convert_tokens_to_ids("<img_ratio_33>")
@@ -1669,6 +1700,9 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo
        vit_pixel_attention_mask = kwargs.pop("vit_pixel_attention_mask", None)
        vit_spatial_shapes = kwargs.pop("vit_spatial_shapes", None)
        vae_pixel_values = kwargs.pop("vae_pixel_values", None)
+        # vae_pixel_size is only metadata for vLLM's flat_from_sizes split;
+        # we reconstruct per-image shapes from vae_token_grid_hw below.
+        kwargs.pop("vae_pixel_size", None)
        vae_token_grid_hw = kwargs.pop("vae_token_grid_hw", None)

        if vit_pixel_values is None or vae_pixel_values is None:
@@ -1678,13 +1712,36 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo
        if vit_pixel_values.numel() == 0 or vae_pixel_values.numel() == 0:
            return None

+        # `vae_pixel_values` arrives as a 1-D concatenation of per-image flat
+        # buffers (see `process_image` + `flat_from_sizes`). Reconstruct a
+        # list of per-image (3, H_i, W_i) tensors using the per-image grid
+        # dims so the downstream VAE encoder can run image-by-image.
+        vae_factor_h = self.config.vae_downsample_factor[0] * self.config.patch_size
+        vae_factor_w = self.config.vae_downsample_factor[1] * self.config.patch_size
+        num_images = vae_token_grid_hw.shape[0]
+        vae_image_list: list[torch.Tensor] = []
+        offset = 0
+        flat = vae_pixel_values.reshape(-1)
+        for i in range(num_images):
+            token_h, token_w = vae_token_grid_hw[i].tolist()
+            h_i = int(token_h) * vae_factor_h
+            w_i = int(token_w) * vae_factor_w
+            n_i = 3 * h_i * w_i
+            vae_image_list.append(flat[offset : offset + n_i].reshape(3, h_i, w_i))
+            offset += n_i
+        if offset != flat.numel():
+            raise ValueError(
+                f"vae_pixel_values size mismatch: consumed {offset} of {flat.numel()} elements "
+                f"across {num_images} images (token_grid_hw={vae_token_grid_hw.tolist()})"
+            )
+
        return HunyuanImage3PixelInputs(
            type="pixel_values",
            pixel_values={
                "vit_pixel_values": vit_pixel_values,
                "vit_pixel_attention_mask": vit_pixel_attention_mask,
                "vit_spatial_shapes": vit_spatial_shapes,
-                "vae_pixel_values": vae_pixel_values,
+                "vae_pixel_values": vae_image_list,
                "vae_token_grid_hw": vae_token_grid_hw,
            },
        )
@@ -1712,7 +1769,11 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo
            images = images.to(dtype=self.vae.dtype)

        vae_encode_result = self.vae.encode(images)
-        latents = vae_encode_result.latent_dist.sample()
+        # Match HunyuanImage-3's cond encode path: sample the posterior, but
+        # use a fixed generator so online requests do not consume the global
+        # RNG and drift across a long-running server.
+        _cond_vae_gen = torch.Generator(device=images.device).manual_seed(0)
+        latents = vae_encode_result.latent_dist.sample(_cond_vae_gen)

        # Apply shift and scaling factors if present
        if hasattr(config, "shift_factor") and config.shift_factor:
@@ -1796,22 +1857,12 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo
        # Perform ViT encoding
        vit_embeddings = self._vit_encode(vit_pixel_values, vit_pixel_attention_mask, vit_spatial_shapes)

-        # Perform VAE encoding
-        t, latents = self._vae_encode(vae_pixel_values, vae_cfg_factor)
-
-        # Process VAE latents through patch_embed to convert to token embeddings
-        # VAE latents are in (B, C, H, W) format, need to be converted to (B, seq_len, hidden_size)
+        # VAE encode + patch_embed per image — each cond image is at its own
+        # `reso_group` bucket so shapes are ragged across the image-batch dim.
        vae_token_embeddings = []
-        batch_size = latents.shape[0]
-        for i in range(batch_size):
-            t_i = t[i]
-            latents_i = latents[i : i + 1]  # Shape: (1, C, H, W)
-
-            # Time embedding for VAE processing
-            t_emb = self.time_embed(t_i)
-
-            # Process VAE latent through patch_embed
-            # Input: (1, C, H, W) -> Output: (1, seq_len, hidden_size)
+        for vae_image_i in vae_pixel_values:
+            t_i, latents_i = self._vae_encode(vae_image_i.unsqueeze(0), vae_cfg_factor)
+            t_emb = self.time_embed(t_i[0])
            vae_tokens, _, _ = self.patch_embed(latents_i, t_emb)
            vae_token_embeddings.append(vae_tokens)

@@ -1821,27 +1872,31 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo
            "Each image should have both VAE and ViT embeddings."
        )

-        # Order per image: timestep -> VAE tokens -> ViT tokens.
-        # The <img> placeholder at the timestep slot (see _get_prompt_updates)
-        # gets its embedding replaced by `timestep_emb(0)` here, which is what
-        # HF achieves via instantiate_continuous_tokens at runtime.
+        # Order per image: VAE tokens -> <joint_img_sep> wte -> ViT tokens.
+        # The <joint_img_sep> wte is included so it joins the bidirectional
+        # MM region (matching the official `joint_image_slices` full-attn
+        # range that spans VAE+sep+ViT). The merger replaces the sep slot
+        # with this wte tensor, which is numerically identical to what
+        # `model.embed_input_ids` would produce — no semantic change for
+        # single-image, but with multi-image the sep position now sits
+        # inside the bidirectional region (matching how the model was
+        # trained).
+        sep_token_id = self._mrope_joint_img_sep_token_id
+        sep_input_ids = torch.tensor([sep_token_id], device=vit_embeddings.device, dtype=torch.long)
+        sep_embed = self.model.embed_input_ids(sep_input_ids).to(vit_embeddings.dtype)
+
+        # The <timestep> slot at the head of each per-image scaffold is NOT
+        # included here — its embedding is patched in by `embed_input_ids`
+        # via a token-id mask, mirroring HF's `instantiate_continuous_tokens`
+        # scatter-replace.
        combined_embeddings: list[torch.Tensor] = []
        num_images = len(vae_token_embeddings)
        for img_idx in range(num_images):
-            # 1. Timestep embedding (cond image timestep == 0)
-            timestep = torch.zeros((1,)).to(vit_embeddings.device).to(vit_embeddings.dtype)
-            timestep_emb = self._timestep_encode(timestep)
-
-            # 2. VAE image token embeddings
            vae_token_embed = vae_token_embeddings[img_idx]
-            # Remove batch dimension if present: (B, seq_len, hidden_size) -> (seq_len, hidden_size)
            if vae_token_embed.ndim == 3:
                vae_token_embed = vae_token_embed.squeeze(0)
-
-            # 3. ViT image embeddings
            vit_embed = vit_embeddings[img_idx]
-
-            stacked_embed = torch.cat([timestep_emb, vae_token_embed, vit_embed], dim=0)
+            stacked_embed = torch.cat([vae_token_embed, sep_embed, vit_embed], dim=0)
            combined_embeddings.append(stacked_embed)

        return combined_embeddings
@@ -1854,14 +1909,23 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo
        is_multimodal: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """Embed input IDs with optional multimodal embeddings."""
-        # Get text embeddings
        inputs_embeds = self.model.embed_input_ids(input_ids)

-        # If no multimodal embeddings, return text embeddings
+        # Patch <timestep> slots with timestep_emb(0). HF parity: the trained
+        # wte at this slot is irrelevant; runtime uses
+        # `instantiate_continuous_tokens(timestep_emb(0))`. With multi-image,
+        # keeping these slots as <img> ids merged the timestep position into
+        # the bidirectional MM region and biased AR ratio prediction toward
+        # the first image's bucket.
+        timestep_mask = input_ids == self._timestep_token_id
+        n_timestep = int(timestep_mask.sum().item())
+        if n_timestep > 0:
+            timestep_input = torch.zeros((n_timestep,), device=inputs_embeds.device, dtype=inputs_embeds.dtype)
+            inputs_embeds[timestep_mask] = self._timestep_encode(timestep_input)
+
        if multimodal_embeddings is None or len(multimodal_embeddings) == 0:
            return inputs_embeds

-        # Merge multimodal embeddings with text embeddings
        merged_embeds = _merge_multimodal_embeddings(
            inputs_embeds=inputs_embeds,
            multimodal_embeddings=multimodal_embeddings,
@@ -1919,8 +1983,6 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo

        min_score = torch.finfo(logits.dtype).min

-        assert logits.shape[0] == 1, f"HunyuanImage3 sampler requires max_num_seqs=1, got batch size {logits.shape[0]}"
-
        for req_idx in range(logits.shape[0]):
            decoded_tokens: list[int] = (
                sampling_metadata.output_token_ids[req_idx] if req_idx < len(sampling_metadata.output_token_ids) else []
@@ -2077,6 +2139,7 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo
        boi_token_id = self._mrope_boi_token_id
        eoi_token_id = self._mrope_eoi_token_id
        joint_img_sep_token_id = self._mrope_joint_img_sep_token_id
+        timestep_token_id = self._timestep_token_id

        # Build position arrays
        t_pos: list[int] = []  # temporal (same as 1D for this model)
@@ -2093,7 +2156,7 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo

            if tok == boi_token_id:
                # Found start of image block.
-                # Structure: <boi> <size> <ratio> <img>*timestep <img>*vae
+                # Structure: <boi> <size> <ratio> <timestep> <img>*vae
                #            <joint_img_sep> <img>*vit <eoi>
                # <boi> token
                t_pos.append(pos)
@@ -2118,8 +2181,8 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo
                    pos += 1
                    i += 1

-                # Timestep token (1 <img> token)
-                if i < n and input_tokens[i] == img_token_id:
+                # <timestep> token (1 token)
+                if i < n and input_tokens[i] == timestep_token_id:
                    t_pos.append(pos)
                    h_pos.append(pos)
                    w_pos.append(pos)
@@ -160,7 +160,7 @@ class MossTTSNanoForGeneration(nn.Module):
    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
        with self._lock:
            if self._lm is not None:
-                return set()
+                return None
            _patch_torchaudio_load()

            try:
@@ -233,7 +233,7 @@ class MossTTSNanoForGeneration(nn.Module):
        # iterator so the underlying stream is closed cleanly.
        for _ in weights:
            pass
-        return set()
+        return None

    # ------------------------------------------------------------------
    # Dummy run support
@@ -30,7 +30,8 @@ QWEN3_OMNI_PIPELINE = PipelineConfig(
            requires_multimodal_data=True,
            hf_config_name="thinker_config",
            engine_output_type="latent",
-            custom_process_next_stage_input_func=(f"{_PROC}.thinker2talker_async_chunk"),
+            custom_process_next_stage_input_func=(f"{_PROC}.thinker2talker_full_payload"),
+            async_chunk_process_next_stage_input_func=(f"{_PROC}.thinker2talker_async_chunk"),
            sampling_constraints={"detokenize": True},
        ),
        StagePipelineConfig(
@@ -41,7 +42,9 @@ QWEN3_OMNI_PIPELINE = PipelineConfig(
            hf_config_name="talker_config",
            engine_output_type="latent",
            custom_process_input_func=f"{_PROC}.thinker2talker",
-            custom_process_next_stage_input_func=(f"{_PROC}.talker2code2wav_async_chunk"),
+            sync_process_input_func=f"{_PROC}.thinker2talker_token_only",
+            custom_process_next_stage_input_func=(f"{_PROC}.talker2code2wav_full_payload"),
+            async_chunk_process_next_stage_input_func=(f"{_PROC}.talker2code2wav_async_chunk"),
            sampling_constraints={
                "detokenize": False,
                "stop_token_ids": [2150],
@@ -237,14 +237,22 @@ class Qwen3OmniMoeForConditionalGeneration(

        prompt_token_ids = tokenizer.encode(prompt_template)

+        # In non-async-chunk (full-payload) mode the engine treats each
+        # streaming TokensPrompt as a fresh decode, so mid-stream segment
+        # yields cause the thinker to emit duplicate responses and the
+        # talker to emit duplicate audio. Defer all audio to the final
+        # flush so the thinker sees one complete prompt.
+        async_chunk = getattr(model_config, "async_chunk", False)
+
        async for audio_chunk in audio_stream:
            buffer.write_audio(audio_chunk)

-            while (segment := buffer.read_audio()) is not None:
-                yield TokensPrompt(
-                    prompt_token_ids=prompt_token_ids,
-                    multi_modal_data={"audio": segment},
-                )
+            if async_chunk:
+                while (segment := buffer.read_audio()) is not None:
+                    yield TokensPrompt(
+                        prompt_token_ids=prompt_token_ids,
+                        multi_modal_data={"audio": segment},
+                    )

        remaining = buffer.flush()
        if remaining is not None and len(remaining) > 0:
@@ -982,19 +990,10 @@ class Qwen3OmniMoeForConditionalGeneration(
        """
        embed = payload.get("embed", {})
        meta = payload.get("meta", {})
-        ids = payload.get("ids", {})

        cached_thinker_decode_embeds = embed.get("cached_decode", None)
        thinker_decode_embed = embed.get("decode", None)
        start_index = meta.get("num_processed_tokens", 0)
-        thinker_output_token_ids = ids.get("output", [])
-        if start_index >= len(thinker_output_token_ids) - 1:
-            # When the tokens output by the thinker are exhausted, an EOS token needs to be appended.
-            # Use the finished_flag to mark that all tokens output by thinker have been consumed.
-            if meta.get("eos_emitted", False):
-                return self.tts_pad_embed.to(device)
-            update_dict.setdefault("meta", {})["eos_emitted"] = True
-            return self.tts_eos_embed.to(device)

        if cached_thinker_decode_embeds is not None and start_index < cached_thinker_decode_embeds.shape[0]:
            cached_thinker_decode_embeds = cached_thinker_decode_embeds.to(device)
@@ -1003,10 +1002,20 @@ class Qwen3OmniMoeForConditionalGeneration(
                thinker_decode_embed = thinker_decode_embed.to(device)
                cached_thinker_decode_embeds = torch.cat([cached_thinker_decode_embeds, thinker_decode_embed], dim=0)
                update_dict.setdefault("embed", {})["cached_decode"] = cached_thinker_decode_embeds
-        else:
+
+        elif thinker_decode_embed is not None:
            thinker_embed = thinker_decode_embed
            if thinker_embed.device != device:
                thinker_embed = thinker_embed.to(device)
+
+        else:
+            # When the tokens output by the thinker are exhausted, an EOS token needs to be appended.
+            # Use the finished_flag to mark that all tokens output by thinker have been consumed.
+            if meta.get("eos_emitted", False):
+                return self.tts_pad_embed.to(device)
+            update_dict.setdefault("meta", {})["eos_emitted"] = True
+            return self.tts_eos_embed.to(device)
+
        update_dict.setdefault("embed", {})["decode"] = None
        return self.talker.text_projection(thinker_embed).to(device)

@@ -534,6 +534,10 @@ class Qwen3TTSTokenizerV2DecoderTransformerModel(Qwen3TTSTokenizerV2DecoderPreTr
        cache_position=None,
        **kwargs,
    ) -> BaseModelOutputWithPast:
+        r"""
+        cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
+            Indices depicting the position of the input sequence tokens in the sequence. Used to update the cache.
+        """
        if input_ids is not None:
            raise ValueError("input_ids is not expected")
        if (input_ids is None) ^ (inputs_embeds is not None):
@@ -1,10 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Stage input processor for HunyuanImage3: AR → Diffusion transition.
+"""Stage input processor for HunyuanImage3: AR to Diffusion transition.

 In IT2I (image editing) mode:
  - Stage 0 (AR) receives (image + edit instruction), generates CoT/latent tokens
-  - Stage 1 (DiT) receives the AR output + original image, denoises → edited image
+  - Stage 1 (DiT) receives the AR output + original image, denoises to edited image

 The ar2diffusion function bridges these two stages, following the same
 signature pattern as glm_image.ar2diffusion.
@@ -12,16 +12,154 @@ signature pattern as glm_image.ar2diffusion.

 from __future__ import annotations

+from functools import lru_cache
 from typing import Any

-import torch
 from vllm.inputs import TextPrompt
 from vllm.logger import init_logger

+from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import (
+    HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS,
+)
 from vllm_omni.inputs.data import OmniTokensPrompt

 logger = init_logger(__name__)

+# AR emits `<img_size_BASE><img_ratio_Y>` after `</recaption>` in IT2I/T2I
+# (see `HunyuanImage3ForCausalMM.sample` and `_stage_transitions`). The
+# ratio_index resolves to a (height, width) bucket via ResolutionGroup, which
+# is the official upstream's mechanism for AR-driven output aspect; without
+# this lookup the DiT pipeline falls back to the user-provided width/height
+# (in the `/v1/images/edits` path that defaults to `pil_images[0].size`,
+# i.e. the first reference image's bucket, usually square, see
+# api_server.py:1808-1811).
+_HUNYUAN_IMAGE3_EXTRA_RESOLUTIONS: tuple[str, ...] = (
+    "1024x768",
+    "1280x720",
+    "768x1024",
+    "720x1280",
+)
+
+
+class _Resolution:
+    def __init__(self, size: str | int | tuple[int, int], *args: int):
+        if isinstance(size, str):
+            if "x" in size:
+                h, w = size.split("x")
+                size = (int(h), int(w))
+            else:
+                size = int(size)
+        if args:
+            size = (int(size), args[0])
+        if isinstance(size, int):
+            size = (size, size)
+
+        self.height = int(size[0])
+        self.width = int(size[1])
+        self.ratio = self.height / self.width
+
+
+def _build_resolutions_by_step(base_size: int, align: int = 1) -> list[_Resolution]:
+    step = base_size // 16
+    min_height = base_size // 2
+    min_width = base_size // 2
+    max_height = base_size * 2
+    max_width = base_size * 2
+
+    resolutions = [_Resolution(base_size, base_size)]
+
+    cur_height, cur_width = base_size, base_size
+    while True:
+        if cur_height >= max_height and cur_width <= min_width:
+            break
+        cur_height = min(cur_height + step, max_height)
+        cur_width = max(cur_width - step, min_width)
+        resolutions.append(_Resolution(cur_height // align * align, cur_width // align * align))
+
+    cur_height, cur_width = base_size, base_size
+    while True:
+        if cur_height <= min_height and cur_width >= max_width:
+            break
+        cur_height = max(cur_height - step, min_height)
+        cur_width = min(cur_width + step, max_width)
+        resolutions.append(_Resolution(cur_height // align * align, cur_width // align * align))
+
+    return sorted(resolutions, key=lambda x: x.ratio)
+
+
+@lru_cache(maxsize=4)
+def _build_ratio_size_table(base_size: int) -> list[tuple[int, int]]:
+    """Return `[(height, width)]` indexed by ratio_index for HunyuanImage-3.
+
+    Mirrors `HunyuanImage3ImageProcessor.build_image_info`'s
+    `reso_group[ratio_index]` reverse lookup. Cached because the table
+    is constant per `base_size`.
+    """
+    resolutions = _build_resolutions_by_step(base_size)
+    for extra_resolution in (_Resolution(s) for s in _HUNYUAN_IMAGE3_EXTRA_RESOLUTIONS):
+        if not any(r.ratio == extra_resolution.ratio for r in resolutions):
+            resolutions.append(extra_resolution)
+    return [(r.height, r.width) for r in resolutions]
+
+
+def _truncate_at_cot_end(generated_text: str) -> str:
+    """Truncate AR output at first `</recaption>` (or `</think>` fallback).
+
+    Mirrors upstream `HunyuanImage3ForCausalMM.generate_image` which feeds
+    DiT only the cot text up to the closing tag; the trailing
+    `<answer><boi><img_size_*><img_ratio_*>` is consumed via height/width
+    extraction and must not leak into DiT's prompt builder.
+    """
+    for marker in ("</recaption>", "</think>"):
+        idx = generated_text.find(marker)
+        if idx != -1:
+            return generated_text[: idx + len(marker)]
+    return generated_text
+
+
+@lru_cache(maxsize=4)
+def _build_ratio_id_lookup() -> dict[int, int]:
+    """Return `{token_id: ratio_index}` for HunyuanImage3 ratio tokens.
+
+    The ids are fixed in tokenizer.json and already pinned in prompt_utils.
+    Avoid loading AutoTokenizer here: this bridge runs on the hot AR->DiT
+    transition path and must keep working in offline deployments where the
+    tokenizer object is not exposed to the stage-input processor.
+    """
+    ratio_0 = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<img_ratio_0>"]
+    ratio_32 = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<img_ratio_32>"]
+    ratio_33 = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<img_ratio_33>"]
+    ratio_36 = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<img_ratio_36>"]
+
+    table: dict[int, int] = {}
+    for i in range(ratio_32 - ratio_0 + 1):
+        table[ratio_0 + i] = i
+    base_idx = ratio_32 - ratio_0 + 1
+    for j in range(ratio_36 - ratio_33 + 1):
+        table[ratio_33 + j] = base_idx + j
+    return table
+
+
+def _extract_ratio_index(generated_token_ids) -> int | None:
+    """Resolve the AR-predicted ratio_index from this stage's output.
+
+    `HunyuanImage3ForCausalMM`'s `_stage_transitions` forces the AR to emit
+    exactly one `<img_ratio_*>` token after `</recaption><answer><boi>
+    <img_size_*>`, so we scan the token stream from the tail for the first
+    id that maps to a ratio. Token-ids are the source of truth; text-side
+    regex is unreliable because most deploy yamls run AR with
+    `skip_special_tokens: True` (special tokens are stripped from text but
+    still present in `cumulative_token_ids`).
+    """
+    if generated_token_ids is None:
+        return None
+    table = _build_ratio_id_lookup()
+    for tid in reversed(list(generated_token_ids)):
+        idx = table.get(int(tid))
+        if idx is not None:
+            return idx
+    return None
+

 def ar2diffusion(
    source_outputs: list[Any],
@@ -64,31 +202,65 @@ def ar2diffusion(
        width = original_prompt.get("width", 1024)
        text_prompt = original_prompt.get("prompt", "")
        use_system_prompt = original_prompt.get("use_system_prompt")
+        custom_system_prompt = original_prompt.get("system_prompt")
+
+        # Prefer the AR's predicted output aspect (`<img_size_*><img_ratio_*>`
+        # tail emitted by `HunyuanImage3ForCausalMM.sample` under the
+        # ratio-restriction logits processor) over the carried-through
+        # height/width, which the serving layer fills with the first
+        # reference image's bucket and so collapses non-square targets to
+        # square in the multi-image / mismatched-aspect case. Mirrors the
+        # official upstream where `reso_group[ratio_index]` is the
+        # canonical source of the diffusion target shape.
+        ratio_idx = _extract_ratio_index(generated_token_ids)
+        ar_predicted = False
+        if ratio_idx is not None:
+            base_size = int(original_prompt.get("image_base_size", 1024))
+            size_table = _build_ratio_size_table(base_size)
+            if 0 <= ratio_idx < len(size_table):
+                height, width = size_table[ratio_idx]
+                ar_predicted = True
+            else:
+                logger.warning(
+                    "[ar2diffusion] Request %d: ratio_index=%d out of range [0,%d), keeping prompt size %dx%d",
+                    i,
+                    ratio_idx,
+                    len(size_table),
+                    height,
+                    width,
+                )
+
+        cot_text_for_dit = _truncate_at_cot_end(generated_text)

        logger.info(
-            "[ar2diffusion] Request %d: AR generated %d tokens, text length=%d, target size=%dx%d",
+            "[ar2diffusion] Request %d: AR generated %d tokens, text length=%d, "
+            "cot_text length=%d, target size=%dx%d (%s)",
            i,
            len(generated_token_ids),
            len(generated_text),
+            len(cot_text_for_dit),
            height,
            width,
+            f"AR ratio_idx={ratio_idx}" if ar_predicted else "from prompt (no AR ratio token)",
        )

-        token_tensor = torch.tensor(generated_token_ids, dtype=torch.long)
-
        diffusion_input: dict[str, Any] = {
            "prompt": text_prompt,
            "height": height,
            "width": width,
            "extra": {
-                "ar_token_ids": token_tensor,
-                "ar_generated_text": generated_text,
+                "ar_generated_text": cot_text_for_dit,
            },
        }

-        # Forward use_system_prompt so the DiT can build the same system prefix
+        # Forward use_system_prompt so the DiT can build the same system prefix.
+        # Also forward the custom system prompt body when sys_type=custom so
+        # DiT's `get_system_prompt(use, "image", body)` doesn't fall back to
+        # an empty prefix and silently diverge from AR.
        if use_system_prompt is not None:
            diffusion_input["use_system_prompt"] = use_system_prompt
+        if custom_system_prompt is not None:
+            diffusion_input["system_prompt"] = custom_system_prompt

        # Forward multimodal data (original image for IT2I conditioning).
        # The diffusion pre_process_func reads multi_modal_data["image"], which
@@ -35,6 +35,10 @@ logger = logging.getLogger(__name__)
 # Pooling output layer keys: "0" = word embedding, "24" = accept_hidden_layer
 _EMBED_LAYER_KEY = "0"
 _HIDDEN_LAYER_KEY = "24"
+_QWEN3_CODEC_CODEBOOK_SIZE = 2048
+_QWEN3_CODEC_PAD_TOKEN_ID = 4196
+_QWEN3_CODEC_BOS_TOKEN_ID = 4197
+_QWEN3_CODEC_EOS_TOKEN_ID = 4198


 def _layer_tensor(layers: dict[Any, Any], key: str) -> torch.Tensor | None:
@@ -99,6 +103,84 @@ def _ensure_list(x):
    return list(x)


+def _as_tensor_or_none(value: Any) -> torch.Tensor | None:
+    if isinstance(value, torch.Tensor):
+        return value.detach().cpu()
+    if isinstance(value, list) and value and isinstance(value[0], torch.Tensor):
+        return value[0].detach().cpu()
+    return None
+
+
+def _is_valid_qwen3_codec_token_id(token_id: Any) -> bool:
+    try:
+        token_id = int(token_id)
+    except (TypeError, ValueError):
+        return False
+    return 0 <= token_id < _QWEN3_CODEC_CODEBOOK_SIZE
+
+
+def should_accumulate_qwen3_omni_full_payload_output(
+    model_config: Any,
+    custom_process_func: Any,
+) -> bool:
+    """Return whether Qwen3-Omni should accumulate full-payload outputs."""
+    return (
+        custom_process_func is not None
+        and not getattr(model_config, "async_chunk", False)
+        and getattr(model_config, "model_arch", None) == "Qwen3OmniMoeForConditionalGeneration"
+        and getattr(model_config, "model_stage", None) in {"thinker", "talker"}
+    )
+
+
+def _extract_qwen3_full_payload_codec_rows(
+    code_predictor_codes: torch.Tensor,
+    output_token_ids: list[int],
+) -> tuple[torch.Tensor, dict[str, int]]:
+    """Filter full-payload codec rows by the authoritative output ids."""
+    if code_predictor_codes.ndim != 2 or code_predictor_codes.numel() == 0:
+        return code_predictor_codes, {
+            "raw_rows": int(code_predictor_codes.shape[0]) if code_predictor_codes.ndim > 0 else 0,
+            "aligned_rows": 0,
+            "valid_rows": 0,
+            "trailing_placeholder_count": 0,
+        }
+
+    trailing_placeholder_count = 0
+    while (
+        trailing_placeholder_count < len(output_token_ids) and output_token_ids[-1 - trailing_placeholder_count] == -1
+    ):
+        trailing_placeholder_count += 1
+
+    aligned_len = min(int(code_predictor_codes.shape[0]), len(output_token_ids))
+    if aligned_len <= 0:
+        return code_predictor_codes[:0], {
+            "raw_rows": int(code_predictor_codes.shape[0]),
+            "aligned_rows": 0,
+            "valid_rows": 0,
+            "trailing_placeholder_count": trailing_placeholder_count,
+        }
+
+    aligned_rows = code_predictor_codes[-aligned_len:]
+    aligned_token_ids = output_token_ids[-aligned_len:]
+    aligned_token_mask = torch.tensor(
+        [_is_valid_qwen3_codec_token_id(token_id) for token_id in aligned_token_ids],
+        dtype=torch.bool,
+        device=aligned_rows.device,
+    )
+    row_valid_mask = (aligned_rows.max(dim=1).values < _QWEN3_CODEC_CODEBOOK_SIZE) & (
+        aligned_rows.min(dim=1).values >= 0
+    )
+    filtered_rows = aligned_rows[aligned_token_mask & row_valid_mask]
+    if filtered_rows.numel() == 0:
+        filtered_rows = aligned_rows[:0]
+    return filtered_rows, {
+        "raw_rows": int(code_predictor_codes.shape[0]),
+        "aligned_rows": aligned_len,
+        "valid_rows": int(filtered_rows.shape[0]) if filtered_rows.ndim > 0 else 0,
+        "trailing_placeholder_count": trailing_placeholder_count,
+    }
+
+
 # =========================
 # PD disaggregation helpers
 # =========================
@@ -349,27 +431,101 @@ def thinker2talker_async_chunk(
            payload.hidden_states.output = torch.cat(
                (save_payload.get("hidden_states", {}).get("output"), payload.hidden_states.output), dim=0
            )
+            prefill_shape = payload.embed.prefill.shape[0]
+            if not is_finished and prefill_shape <= len(prompt_token_ids):
+                transfer_manager.request_payload[request_id] = to_dict(payload)
+                return None
    else:
-        output_token_ids = _ensure_list(request.output_token_ids)
+        if thinker_emb.shape[0] > 1:
+            logger.warning(
+                "Unexpected multiple embeddings in thinker2talker_async_chunk for chunk_id %d: "
+                "request_id %s, num_computed_tokens%d %s. Expected shape [1, D].",
+                chunk_id,
+                request_id,
+                request.num_computed_tokens,
+                thinker_emb.shape,
+            )
+            return None
        meta = MetaStruct(finished=torch.tensor(is_finished, dtype=torch.bool))
-        if output_token_ids:
-            meta.override_keys = [("embed", "decode"), ("ids", "output")]
-            payload = OmniPayloadStruct(
-                meta=meta,
-                embed=EmbeddingsStruct(decode=thinker_emb.detach().cpu()),
-                ids=IdsStruct(output=output_token_ids),
-                speaker=speaker,
-                language=language,
-            )
-        else:
-            # When prefilling a chunked thinker, thinker_hidden_states needs to be updated.
-            payload = OmniPayloadStruct(
-                meta=meta,
-                embed=EmbeddingsStruct(prefill=thinker_emb.detach().cpu()),
-                hidden_states=HiddenStatesStruct(output=thinker_hid.detach().cpu()),
-                speaker=speaker,
-                language=language,
-            )
+        payload = OmniPayloadStruct(
+            meta=meta,
+            embed=EmbeddingsStruct(decode=thinker_emb.detach().cpu()),
+            speaker=speaker,
+            language=language,
+        )
+    return payload
+
+
+def thinker2talker_full_payload(
+    transfer_manager: Any,
+    pooling_output: dict[str, Any],
+    request: OmniEngineCoreRequest,
+) -> dict[str, Any] | None:
+    """Pack complete thinker output for the non-async connector path."""
+    if not isinstance(pooling_output, dict):
+        return None
+
+    layers = {
+        0: pooling_output.get("hidden_states.layer_0"),
+        24: pooling_output.get("hidden_states.layer_24"),
+    }
+    thinker_emb = _layer_tensor(layers, _EMBED_LAYER_KEY)
+    thinker_hid = _layer_tensor(layers, _HIDDEN_LAYER_KEY)
+    if thinker_emb is None:
+        hidden = pooling_output.get("hidden")
+        thinker_emb = hidden if isinstance(hidden, torch.Tensor) else None
+    if thinker_emb is None or thinker_hid is None:
+        logger.debug(
+            "thinker2talker_full_payload: missing thinker tensors for req=%s (embed=%s hidden=%s)",
+            getattr(request, "request_id", None),
+            thinker_emb is not None,
+            thinker_hid is not None,
+        )
+        return None
+
+    prompt_token_ids = _ensure_list(getattr(request, "prompt_token_ids", []) or [])
+    all_token_ids = _ensure_list(getattr(request, "all_token_ids", None) or [])
+    if not all_token_ids:
+        output_token_ids = _ensure_list(getattr(request, "output_token_ids", []) or [])
+        all_token_ids = list(prompt_token_ids) + list(output_token_ids)
+
+    # Trim the trailing stop-token row from the accumulated thinker output.
+    # The accumulator captures one hidden-state row per executed thinker
+    # forward (prefill + every decode step including the one that emitted
+    # the stop_token), so for a finished request thinker_emb has exactly one
+    # row more than the rows the talker should consume.  async_chunk's
+    # chunk-0 path naturally captures only the prefill / non-stop portion,
+    # which is why the [async_chunk] parametrization passes while [default]
+    # over-generates one codec frame on short outputs (e.g.
+    # test_one_word_prompt_001[default]: audio extends "London" with
+    # spurious phonemes).
+    if isinstance(thinker_emb, torch.Tensor) and thinker_emb.shape[0] > 0:
+        thinker_emb_prefill = thinker_emb[:-1]
+    else:
+        thinker_emb_prefill = thinker_emb
+    if isinstance(thinker_hid, torch.Tensor) and thinker_hid.shape[0] > 0:
+        thinker_hid_prefill = thinker_hid[:-1]
+    else:
+        thinker_hid_prefill = thinker_hid
+
+    payload: OmniPayload = {
+        "embed": {
+            "prefill": thinker_emb_prefill.detach().cpu(),
+            "tts_bos": _as_tensor_or_none(pooling_output.get("embed.tts_bos")),
+            "tts_eos": _as_tensor_or_none(pooling_output.get("embed.tts_eos")),
+            "tts_pad": _as_tensor_or_none(pooling_output.get("embed.tts_pad")),
+        },
+        "hidden_states": {"output": thinker_hid_prefill.detach().cpu()},
+        "ids": {"all": list(all_token_ids), "prompt": list(prompt_token_ids)},
+        "meta": {"finished": torch.tensor(True, dtype=torch.bool)},
+    }
+    payload["next_stage_prompt_len"] = _compute_talker_prompt_ids_length(payload, device="cpu")
+    speaker = extract_speaker_from_request(request)
+    if speaker is not None:
+        payload["speaker"] = speaker
+    language = extract_language_from_request(request)
+    if language is not None:
+        payload["language"] = language
    return payload


@@ -482,6 +638,85 @@ def thinker2talker(
    return talker_inputs


+def thinker2talker_token_only(
+    source_outputs: list[Any],
+    prompt: OmniTokensPrompt | TextPrompt | None = None,
+    requires_multimodal_data: bool = False,
+    streaming_context: Any | None = None,
+) -> list[OmniTokensPrompt]:
+    """Non-async-chunk Stage-1 input builder for the connector data plane.
+
+    The worker connector (Stage-0 ``thinker2talker_full_payload`` →
+    ``_sync_local_stage_payloads``) supplies the bulk talker conditioning
+    tensors (embed / hidden_states / ids) via ``model_intermediate_buffer``.
+    The orchestrator only needs to ship a placeholder prefill prompt of the
+    correct length so the scheduler can allocate KV-cache slots.
+
+    Small per-request voice metadata (``speaker`` / ``language``) is forwarded
+    here from the user prompt so the worker's line-408 buffer seed picks it
+    up. The connector-side ``extract_speaker_from_request`` reads the
+    strongly-typed ``request.additional_information.entries["speaker"]`` which
+    currently does not always round-trip the user-supplied voice; until that
+    plumbing is normalized, providing the small fields directly preserves
+    voice selection (regression discovered on Buildkite 9668:
+    ``test_speaker_002[default]`` lost the preset voice).
+    """
+    talker_inputs: list[OmniTokensPrompt] = []
+    for i, thinker_output in enumerate(source_outputs):
+        output = thinker_output.outputs[0]
+        req_id = str(getattr(thinker_output, "request_id", f"idx-{i}"))
+        # Skip-on-missing parity with thinker2talker_full_payload: if the
+        # connector builder would drop this request (no MM dict or missing
+        # hidden-state layers), do the same here so the worker buffer
+        # presence agrees with the orchestrator's scheduling decision.
+        thinker_mm_raw = getattr(output, "multimodal_output", None)
+        if not isinstance(thinker_mm_raw, dict):
+            logger.debug("thinker2talker_token_only: skip req=%s due to empty multimodal_output", req_id)
+            continue
+        mm_hs = thinker_mm_raw.get("hidden_states", {})
+        mm_layers = mm_hs.get("layers", {}) if isinstance(mm_hs, dict) else {}
+        if _layer_tensor(mm_layers, _EMBED_LAYER_KEY) is None or _layer_tensor(mm_layers, _HIDDEN_LAYER_KEY) is None:
+            logger.debug("thinker2talker_token_only: skip req=%s due to missing hidden-state layers", req_id)
+            continue
+        prompt_token_ids = _ensure_list(thinker_output.prompt_token_ids)
+        output_ids = _ensure_list(output.cumulative_token_ids)
+        is_streaming_session = bool(getattr(streaming_context, "enabled", False))
+        if is_streaming_session:
+            prompt_token_ids, output_ids, thinker_sequences, thinker_input_ids = _get_streaming_talker_tokens(
+                req_id,
+                prompt_token_ids,
+                output_ids,
+                getattr(streaming_context, "new_prompt_len_snapshot", None),
+                streaming_context,
+                clear_state=bool(getattr(thinker_output, "finished", False)),
+            )
+        else:
+            thinker_sequences = prompt_token_ids + output_ids
+            thinker_input_ids = prompt_token_ids
+        info_for_len = {"ids": {"all": thinker_sequences, "prompt": thinker_input_ids}}
+        prompt_len = _compute_talker_prompt_ids_length(info_for_len, device="cpu")
+
+        # Forward only small voice metadata; bulk tensors come from the
+        # connector path via _sync_local_stage_payloads.
+        small_info: dict[str, Any] = {}
+        speaker = extract_speaker_from_prompt(prompt, index=i)
+        if speaker is not None:
+            small_info["speaker"] = speaker
+        language = extract_language_from_prompt(prompt, index=i)
+        if language is not None:
+            small_info["language"] = language
+
+        talker_inputs.append(
+            OmniTokensPrompt(
+                prompt_token_ids=[0] * prompt_len,
+                additional_information=(small_info if small_info else None),
+                multi_modal_data=None,
+                mm_processor_kwargs=None,
+            )
+        )
+    return talker_inputs
+
+
 # =========================
 # Talker -> Code2Wav
 # =========================
@@ -556,6 +791,56 @@ def talker2code2wav_async_chunk(
    )


+def talker2code2wav_full_payload(
+    transfer_manager: Any,
+    pooling_output: dict[str, Any],
+    request: OmniEngineCoreRequest,
+) -> dict[str, Any] | None:
+    """Pack complete talker codec output for the non-async connector path."""
+    if not isinstance(pooling_output, dict):
+        return None
+    code_predictor_codes = pooling_output.get("codes.audio")
+    if code_predictor_codes is None:
+        codes = pooling_output.get("codes")
+        if isinstance(codes, dict):
+            code_predictor_codes = codes.get("audio")
+    if code_predictor_codes is None:
+        return None
+    if not isinstance(code_predictor_codes, torch.Tensor):
+        code_predictor_codes = torch.as_tensor(code_predictor_codes)
+    if code_predictor_codes.numel() == 0:
+        return None
+
+    output_token_ids = _ensure_list(getattr(request, "output_token_ids", []) or [])
+    raw_shape = tuple(code_predictor_codes.shape)
+    code_predictor_codes, codec_stats = _extract_qwen3_full_payload_codec_rows(
+        code_predictor_codes.to(torch.long),
+        list(output_token_ids),
+    )
+    if code_predictor_codes.numel() == 0:
+        return None
+
+    codec_codes = code_predictor_codes.transpose(0, 1).cpu().reshape(-1).tolist()
+    logger.debug(
+        "talker2code2wav_full_payload: raw_shape=%s output_ids_len=%s aligned_rows=%s "
+        "valid_rows=%s placeholders=%s flattened_len=%s pad4196=%s bos4197=%s eos4198=%s",
+        raw_shape,
+        len(output_token_ids),
+        codec_stats["aligned_rows"],
+        codec_stats["valid_rows"],
+        codec_stats["trailing_placeholder_count"],
+        len(codec_codes),
+        sum(1 for tid in output_token_ids if tid == _QWEN3_CODEC_PAD_TOKEN_ID),
+        sum(1 for tid in output_token_ids if tid == _QWEN3_CODEC_BOS_TOKEN_ID),
+        sum(1 for tid in output_token_ids if tid == _QWEN3_CODEC_EOS_TOKEN_ID),
+    )
+    return {
+        "codes": {"audio": codec_codes},
+        "code_predictor_codes": codec_codes,
+        "meta": {"finished": torch.tensor(True, dtype=torch.bool)},
+    }
+
+
 def talker2code2wav(
    source_outputs: list[Any],
    _prompt: OmniTokensPrompt | TextPrompt | None = None,
@@ -102,6 +102,11 @@ if not hasattr(RequestStatus, "WAITING_FOR_CHUNK"):
    # as a non-finished state and remains compatible with existing comparisons.
    extend_enum(RequestStatus, "WAITING_FOR_CHUNK", -1)

+if not hasattr(RequestStatus, "WAITING_FOR_INPUT"):
+    # Full-payload stage handoff uses a distinct waiting state so the
+    # scheduler can restore the request once non-stage-0 inputs arrive.
+    extend_enum(RequestStatus, "WAITING_FOR_INPUT", -2)
+
 # Snapshot sys.modules: `hasattr` below can trigger lazy submodule imports
 # (e.g. transformers' `_LazyModule.__getattr__`), which mutate sys.modules
 # during iteration and raise `dictionary changed size during iteration`.
@@ -204,3 +204,7 @@ class OmniPlatform(Platform):
 class UnspecifiedOmniPlatform(OmniPlatform):
    _omni_enum = OmniPlatformEnum.UNSPECIFIED
    device_type = ""
+
+    @classmethod
+    def get_device_count(cls) -> int:
+        return 0
@@ -16,7 +16,11 @@ from vllm.distributed.kv_transfer import get_kv_transfer_group, has_kv_transfer_
 from vllm.distributed.parallel_state import get_pp_group, get_tp_group
 from vllm.forward_context import BatchDescriptor
 from vllm.logger import logger
-from vllm.model_executor.layers.fused_moe.routed_experts_capturer import RoutedExpertsCapturer
+from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
+    extract_routed_experts_for_current_batch,
+    get_global_experts_capturer,
+    issue_routing_d2h_copy,
+)
 from vllm.sequence import IntermediateTensors
 from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
 from vllm.v1.outputs import (
@@ -28,8 +32,8 @@ from vllm.v1.outputs import (
 from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
 from vllm.v1.structured_output.utils import apply_grammar_bitmask
 from vllm.v1.utils import record_function_or_nullcontext
-from vllm.v1.worker import mamba_utils
 from vllm.v1.worker.gpu_model_runner import AsyncGPUModelRunnerOutput, PerLayerAttnMetadata
+from vllm.v1.worker.mamba_utils import preprocess_mamba
 from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices
 from vllm_ascend.ascend_forward_context import set_ascend_forward_context
 from vllm_ascend.attention.utils import AscendCommonAttentionMetadata
@@ -246,11 +250,9 @@ class NPUARModelRunner(OmniNPUModelRunner):
        intermediate_tensors: IntermediateTensors | None = None,
    ) -> OmniModelRunnerOutput | IntermediateTensors | None:
        if self.vllm_config.model_config.enable_return_routed_experts:
-            capturer = RoutedExpertsCapturer.get_instance()
-            if capturer is not None:
-                capturer.clear_buffer()
-            else:
-                logger.warning("RoutedExpertsCapturer is not initialized.")
+            capturer = get_global_experts_capturer()
+            if capturer is not None and hasattr(capturer, "finalize_pending_copy"):
+                capturer.finalize_pending_copy()
        if self.ascend_config.profiling_chunk_config.enabled:
            self._sync_device()
            self._execution_start_time = time.perf_counter()
@@ -299,7 +301,9 @@ class NPUARModelRunner(OmniNPUModelRunner):
        # TODO(Ronald1995): deepcopy is expensive when there is a large
        # number of requests, optimize it later.
        if ((
-            self.use_async_scheduling and self.num_spec_tokens and self._draft_token_ids is None  # type: ignore[has-type]
+            self.use_async_scheduling
+            and self.num_spec_tokens
+            and self._draft_token_ids is None  # type: ignore[has-type]
        ) or (
            # NOTE: This branch specifically triggers a deepcopy during the prefill phase
            # only for PCP (Parallel Context Processing) + Multi-Modal (MM) scenarios.
@@ -455,7 +459,7 @@ class NPUARModelRunner(OmniNPUModelRunner):
                    if deferred_state_corrections_fn:
                        deferred_state_corrections_fn()
                        deferred_state_corrections_fn = None
-                    mamba_utils.preprocess_mamba(
+                    preprocess_mamba(
                        scheduler_output,
                        self.kv_cache_config,
                        self.cache_config,
@@ -715,6 +719,15 @@ class NPUARModelRunner(OmniNPUModelRunner):
        # previous model forward without breaking async scheduling.
        if deferred_state_corrections_fn:
            deferred_state_corrections_fn()
+
+        if self.vllm_config.model_config.enable_return_routed_experts and hasattr(self, "_positions_cpu"):
+            issue_routing_d2h_copy(
+                input_batch_req_ids=self.input_batch.req_ids,
+                num_scheduled_tokens=scheduler_output.num_scheduled_tokens,
+                positions=self.positions,
+                positions_cpu=self._positions_cpu,
+            )
+
        return None

    @torch.inference_mode()
@@ -858,12 +871,17 @@ class NPUARModelRunner(OmniNPUModelRunner):
            if self.speculative_config is not None:
                self.finalize_kv_connector()

+        routed_experts_dict = None
        if self.model_config.enable_return_routed_experts:
-            capturer = RoutedExpertsCapturer.get_instance()
-            if capturer is not None:
-                capturer.save_captured_experts(indices=self.cpu_slot_mapping)
-            else:
-                logger.warning("RoutedExpertsCapturer is not initialized.")
+            capturer = get_global_experts_capturer()
+            if capturer is not None and hasattr(self.input_batch, "num_tokens_no_spec"):
+                routed_experts_dict = extract_routed_experts_for_current_batch(
+                    req_ids=req_ids_output_copy,
+                    requests=self.requests,
+                    req_id_to_index=self.input_batch.req_id_to_index,
+                    num_tokens_no_spec=self.input_batch.num_tokens_no_spec,
+                    max_model_len=self.max_model_len,
+                )

        #  -------------------------------------- Omni-new -------------------------------------------------
        engine_output_type, downstream_req_ids = self._resolve_pooler_payload_req_ids(req_ids_output_copy)
@@ -984,6 +1002,7 @@ class NPUARModelRunner(OmniNPUModelRunner):
            kv_connector_output=kv_connector_output,
            ec_connector_output=ec_connector_output if self.supports_mm_inputs else None,
            cudagraph_stats=cudagraph_stats,
+            routed_experts_dict=routed_experts_dict,
        )
        model_runner_output.kv_extracted_req_ids = kv_extracted_req_ids
        #  -------------------------------------- Omni-new -------------------------------------------------
@@ -15,14 +15,18 @@ from vllm.distributed.ec_transfer import get_ec_transfer, has_ec_transfer
 from vllm.distributed.kv_transfer import get_kv_transfer_group, has_kv_transfer_group
 from vllm.distributed.parallel_state import get_pp_group
 from vllm.logger import logger
-from vllm.model_executor.layers.fused_moe.routed_experts_capturer import RoutedExpertsCapturer
+from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
+    extract_routed_experts_for_current_batch,
+    get_global_experts_capturer,
+    issue_routing_d2h_copy,
+)
 from vllm.sequence import IntermediateTensors
 from vllm.utils.math_utils import cdiv
 from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
 from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT, AsyncModelRunnerOutput, make_empty_encoder_model_runner_output
 from vllm.v1.utils import record_function_or_nullcontext
-from vllm.v1.worker import mamba_utils
 from vllm.v1.worker.gpu_model_runner import AsyncGPUModelRunnerOutput, PerLayerAttnMetadata
+from vllm.v1.worker.mamba_utils import preprocess_mamba
 from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices
 from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs
 from vllm_ascend.ascend_forward_context import set_ascend_forward_context
@@ -74,11 +78,9 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
        intermediate_tensors: IntermediateTensors | None = None,
    ) -> OmniModelRunnerOutput | IntermediateTensors | None:
        if self.vllm_config.model_config.enable_return_routed_experts:
-            capturer = RoutedExpertsCapturer.get_instance()
-            if capturer is not None:
-                capturer.clear_buffer()
-            else:
-                logger.warning("RoutedExpertsCapturer is not initialized.")
+            capturer = get_global_experts_capturer()
+            if capturer is not None and hasattr(capturer, "finalize_pending_copy"):
+                capturer.finalize_pending_copy()
        if self.ascend_config.profiling_chunk_config.enabled:
            self._sync_device()
            self._execution_start_time = time.perf_counter()
@@ -227,7 +229,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
                    if deferred_state_corrections_fn:
                        deferred_state_corrections_fn()
                        deferred_state_corrections_fn = None
-                    mamba_utils.preprocess_mamba(
+                    preprocess_mamba(
                        scheduler_output,
                        self.kv_cache_config,
                        self.cache_config,
@@ -383,6 +385,15 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
        # previous model forward without breaking async scheduling.
        if deferred_state_corrections_fn:
            deferred_state_corrections_fn()
+
+        if self.vllm_config.model_config.enable_return_routed_experts and hasattr(self, "_positions_cpu"):
+            issue_routing_d2h_copy(
+                input_batch_req_ids=self.input_batch.req_ids,
+                num_scheduled_tokens=scheduler_output.num_scheduled_tokens,
+                positions=self.positions,
+                positions_cpu=self._positions_cpu,
+            )
+
        return None

    @torch.inference_mode()
@@ -468,6 +479,17 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
        # [Omni] Copy req_id mappings to avoid async scheduling mutation.
        req_ids_output_copy = self.input_batch.req_ids.copy()
        req_id_to_index_output_copy = self.input_batch.req_id_to_index.copy()
+        routed_experts_dict = None
+        if self.vllm_config.model_config.enable_return_routed_experts and hasattr(
+            self.input_batch, "num_tokens_no_spec"
+        ):
+            routed_experts_dict = extract_routed_experts_for_current_batch(
+                req_ids=req_ids_output_copy,
+                requests=self.requests,
+                req_id_to_index=self.input_batch.req_id_to_index,
+                num_tokens_no_spec=self.input_batch.num_tokens_no_spec,
+                max_model_len=self.max_model_len,
+            )
        output = OmniModelRunnerOutput(
            req_ids=req_ids_output_copy,
            req_id_to_index=req_id_to_index_output_copy,
@@ -478,6 +500,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
            kv_connector_output=kv_connector_output,
            ec_connector_output=ec_connector_output if self.supports_mm_inputs else None,
            cudagraph_stats=cudagraph_stats,
+            routed_experts_dict=routed_experts_dict,
        )
        #  -------------------------------------- Omni-new -------------------------------------------------
Author	SHA1	Message	Date
khluuandClaude Opus 4.6	04c4c6742a	Enable twine upload to PyPI in release pipeline Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>	2026-05-17 03:59:46 -07:00
Kevin H. Luu	67d672db95	[CI/Build] Unify release pipeline with NIGHTLY=1 option, add x86_64/aarch64 image builds (#3428 ) Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com> (cherry picked from commit `bc26cad19a`)	2026-05-17 03:42:31 -07:00
+20 Zhou Taichang Claude Opus 4.7 NumberWan dsinghvi Dnoob Yueqian Lin Hongsheng Liu knlnguyen1802 Samit rein yang Nick Cao zhumingjue138 Ricardo Noriega lyj-jjj Cursor gcanlin wangyu weizhoublue weizhoublue dengyunyang 汪志鹏 WeiQing Chen Jie Liu Yueqian Lin NATURE bjf-frz amy-why-3459 Y. Fisher Zeyu Huang \| 黃澤宇 Copilot Autofix powered by AI	eb4e60ee64	[Rebase] Rebase to vllm v0.21.0 (#3530 ) Signed-off-by: tzhouam <tzhouam@connect.ust.hk> Signed-off-by: Zhou Taichang <tzhouam@connect.ust.hk> Signed-off-by: NumberWan <wantszkin2003@gmail.com> Signed-off-by: Divyansh Singhvi <divyanshsinghvi@gmail.com> Signed-off-by: dsinghvi <divyanshsinghvi@gmail.com> Signed-off-by: Dnoob <dxpouo@gmail.com> Signed-off-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com> Signed-off-by: knlnguyen1802 <knlnguyen1802@gmail.com> Signed-off-by: rein yang <ruiruyang2@gmail.com> Signed-off-by: Nick Cao <ncao@redhat.com> Signed-off-by: zhumingjue <zhumingjue@huawei.com> Signed-off-by: Ricardo Noriega De Soto <rnoriega@redhat.com> Signed-off-by: lyj-jjj <liuyingjun5@huawei.com> Signed-off-by: gcanlin <canlinguosdu@gmail.com> Signed-off-by: wangyu <410167048@qq.com> Signed-off-by: weizhoublue <weizhoublue@github.com> Signed-off-by: weizhou.lan@daocloud.io <weizhou.lan@daocloud.io> Signed-off-by: dengyunyang <584797741@qq.com> Signed-off-by: 汪志鹏 <wangzhipeng628@gmail.com> Signed-off-by: David Chen <530634352@qq.com> Signed-off-by: Jie Liu <33612777+keeper-jie@users.noreply.github.com> Signed-off-by: Yueqian Lin <linyueqian@outlook.com> Signed-off-by: princepride <wangzhipeng628@gmail.com> Signed-off-by: natureofnature <wzliu@connect.hku.hk> Signed-off-by: bjf-frz <frz123db@gmail.com> Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com> Signed-off-by: KexiongYu <yukexiong1@huawei.com> Signed-off-by: Huang, Zeyu <11222265+fhfuih@users.noreply.github.com> Signed-off-by: Zeyu Huang \| 黃澤宇 <11222265+fhfuih@users.noreply.github.com> Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com> Co-authored-by: NumberWan <wantszkin2003@gmail.com> Co-authored-by: dsinghvi <divyanshsinghvi@gmail.com> Co-authored-by: Dnoob <dxpouo@gmail.com> Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com> Co-authored-by: knlnguyen1802 <knlnguyen1802@gmail.com> Co-authored-by: Samit <285365963@qq.com> Co-authored-by: rein yang <73573651+R2-Y@users.noreply.github.com> Co-authored-by: Nick Cao <ncao@redhat.com> Co-authored-by: zhumingjue138 <zhumingjue@huawei.com> Co-authored-by: Ricardo Noriega <rnoriega@redhat.com> Co-authored-by: lyj-jjj <liuyingjun5@huawei.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: gcanlin <canlinguosdu@gmail.com> Co-authored-by: wangyu <53896905+yenuo26@users.noreply.github.com> Co-authored-by: weizhoublue <45163302+weizhoublue@users.noreply.github.com> Co-authored-by: weizhoublue <weizhoublue@github.com> Co-authored-by: dengyunyang <584797741@qq.com> Co-authored-by: 汪志鹏 <wangzhipeng628@gmail.com> Co-authored-by: WeiQing Chen <40507679+david6666666@users.noreply.github.com> Co-authored-by: Jie Liu <33612777+keeper-jie@users.noreply.github.com> Co-authored-by: Yueqian Lin <linyueqian@outlook.com> Co-authored-by: NATURE <wzliu@connect.hku.hk> Co-authored-by: bjf-frz <frz123db@gmail.com> Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com> Co-authored-by: Y. Fisher <yukexiong1@huawei.com> Co-authored-by: Zeyu Huang \| 黃澤宇 <11222265+fhfuih@users.noreply.github.com> Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>	2026-05-16 18:18:43 +08:00
Yueqian LinandHongsheng Liu	5e1986206f	[CI] Replace c=128 perf cell with c=16; loosen new-cell baselines (#3637 ) Signed-off-by: Yueqian Lin <linyueqian@outlook.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-05-15 23:35:44 +08:00
amy-why-3459	d1c65bdffa	[BugFix] fix ci (#3650 ) Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com>	2026-05-15 22:50:03 +08:00
Zeyu Huang \| 黃澤宇andCopilot Autofix powered by AI	d18168ccbb	[bugfix] Fix diffusers backend input bug after #2913 (#3644 ) Signed-off-by: Huang, Zeyu <11222265+fhfuih@users.noreply.github.com> Signed-off-by: Zeyu Huang \| 黃澤宇 <11222265+fhfuih@users.noreply.github.com> Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>	2026-05-15 09:55:09 +00:00
Y. Fisher	779bf3118b	[Bugfix] fix compatibility of _hunyuan_image3_unpack_packed_topk between vllm / vllm ascend (#3640 ) Signed-off-by: KexiongYu <yukexiong1@huawei.com>	2026-05-15 17:35:22 +08:00
amy-why-3459	e7ee5de09f	[BugFix] Fix the issue of thinker requests being preempted, causing shape mismatch. (#3147 ) Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com>	2026-05-15 17:07:42 +08:00
bjf-frz	440c718d2a	[Bugfix]Fix multimodal cache routing for AR replicas (#3605 ) Signed-off-by: bjf-frz <frz123db@gmail.com>	2026-05-15 16:42:18 +08:00
NATUREandHongsheng Liu	82a0b3a467	[2/5] [core]refactor communication layer: PR 2 of 5 Qwen3 Omni non async (#2677 ) Signed-off-by: natureofnature <wzliu@connect.hku.hk> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-05-15 14:14:08 +08:00
汪志鹏	c7178d89bb	[Bugfix] UnspecifiedOmniPlatform.get_device_count returns 0 instead o… (#3636 ) Signed-off-by: princepride <wangzhipeng628@gmail.com>	2026-05-15 05:28:40 +00:00
wangyu	fdb0efea94	[CI] add cuda marker to Diffusion X2V function pytest (#3625 ) Signed-off-by: wangyu <410167048@qq.com>	2026-05-15 11:50:22 +08:00
Jie LiuandYueqian Lin	90f5b3c3a1	Update streaming_speech_client.py to solve Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice voice problem (#3380 ) Signed-off-by: Jie Liu <33612777+keeper-jie@users.noreply.github.com> Signed-off-by: Yueqian Lin <linyueqian@outlook.com> Co-authored-by: Yueqian Lin <linyueqian@outlook.com>	2026-05-14 20:43:05 -07:00
lyj-jjjandCursor	bbc00f9f86	[BugFix] fix(omni): isolate diffusion KV-cache dtype from vLLM --kv-cache-dtype #3585 (#3596 ) Signed-off-by: lyj-jjj <liuyingjun5@huawei.com> Co-authored-by: Cursor <cursoragent@cursor.com>	2026-05-15 11:22:59 +08:00
WeiQing Chen	adb2291c27	Update WeChat group QR code (#3624 ) Signed-off-by: David Chen <530634352@qq.com>	2026-05-15 09:38:47 +08:00
Canlin Guo	4f13b871f9	[CI] Add Qwen3-TTS tests for ready tag (#3600 ) Signed-off-by: gcanlin <canlinguosdu@gmail.com>	2026-05-14 20:37:51 +08:00
dengyunyangand汪志鹏	94254e015f	[BugFix] fix shm connector (#3583 ) Signed-off-by: dengyunyang <584797741@qq.com> Signed-off-by: 汪志鹏 <wangzhipeng628@gmail.com> Co-authored-by: 汪志鹏 <wangzhipeng628@gmail.com>	2026-05-14 20:07:12 +08:00
bjf-frzandCanlin Guo	f7161b07d0	[Bugfix]Allow HunyuanImage3 AR sampler batching (#3590 ) Signed-off-by: bjf-frz <frz123db@gmail.com> Co-authored-by: Canlin Guo <canlinguosdu@gmail.com>	2026-05-14 09:27:59 +00:00
Hongsheng Liuandhsliu_ustc	c0b7509f07	update v0.20.0 readme (#3594 ) Signed-off-by: hsliu_ustc <hsliu_ustc@noreply.gitcode.com> Co-authored-by: hsliu_ustc <hsliu_ustc@noreply.gitcode.com>	2026-05-14 17:23:55 +08:00
TaffyOfficial TaffyOfficial TaffyOfficial skf1999	3f63aaf982	[Feature] HunyuanImage-3.0 IT2I: multi-image input + prompt API cleanup (#3444 ) Signed-off-by: TaffyOfficial <2324465096@qq.com> Signed-off-by: TaffyOfficial <wu15922848573@outlook.com> Signed-off-by: skf1999 <13234016272@163.com> Signed-off-by: zuiho <2324465096@qq.com> Signed-off-by: Claude Code <noreply@anthropic.com> Signed-off-by: zuiho <wu15922848573@outlook.com> Signed-off-by: TaffyOfficial <2587297563@qq.com> Co-authored-by: TaffyOfficial <2324465096@qq.com> Co-authored-by: TaffyOfficial <wu15922848573@outlook.com> Co-authored-by: skf1999 <13234016272@163.com>	2026-05-14 16:58:40 +08:00
Yueqian Lin	c4f859bf56	[CI] Harden Qwen3-TTS perf nightly: enable Base voice_clone, add c=64/128, 2-GPU split (#3491 ) Signed-off-by: Yueqian Lin <linyueqian@outlook.com>	2026-05-14 14:26:59 +08:00