[Rebase] Rebase to vllm 0.23.0 (#4286 )

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
[XPU] update Dreamzero to support any HW (#4399 )
2026-06-13 16:43:02 +00:00 · 2026-06-13 15:32:39 +00:00 · 2026-06-13 18:12:37 +08:00 · 2026-06-13 15:42:16 +08:00 · 2026-06-13 04:22:54 +00:00 · 2026-06-13 10:59:15 +08:00
157 changed files with 9452 additions and 935 deletions
@@ -10,7 +10,7 @@ steps:
          DOCKER_BUILDKIT: "1"
          # Buildkite will automatically replace this with the actual commit hash
          VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}"
-          VLLM_VERSION: "v0.22.0"
+          VLLM_VERSION: "v0.23.0"
        priority: 100
        timeout_in_minutes: 60
        soft_fail: false
@@ -149,7 +149,7 @@ steps:
        commands:
          - pytest -s -v tests/entrypoints/ -m "advanced_model and cuda and L4" --run-level "advanced_model"
        agents:
-          queue: "gpu_4_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
+          queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
        plugins:
          - docker#v5.2.0:
              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
@@ -164,7 +164,35 @@ steps:
  - group: ":card_index_dividers: E2E Test"
    depends_on: upload-merge-pipeline
    steps:
+      - label: "TTS · Step-Audio2 E2E Test"
+        timeout_in_minutes: 20
+        commands:
+          - |
+            timeout 20m bash -c "
+              export VLLM_LOGGING_LEVEL=DEBUG
+              export VLLM_WORKER_MULTIPROC_METHOD=spawn
+              pytest -s -v tests/e2e/offline_inference/test_step_audio2.py -m 'advanced_model' --run-level 'advanced_model'
+            "
+        agents:
+          queue: "gpu_4_queue"
+        plugins:
+          - docker#v5.2.0:
+              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+              always-pull: true
+              propagate-environment: true
+              shm-size: "8gb"
+              environment:
+                - "HF_HOME=/fsx/hf_cache"
+              volumes:
+                - "/fsx/hf_cache:/fsx/hf_cache"
+
      - label: "TTS · Qwen3-TTS CustomVoice Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_qwen3_tts_customvoice.py
+          - tests/e2e/offline_inference/test_qwen3_tts_customvoice.py
+          - vllm_omni/model_executor/models/qwen3_tts/
+          - vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
+          - vllm_omni/deploy/qwen3_tts.yaml
        commands:
          - |
            timeout 20m bash -c "
@@ -173,7 +201,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_4_queue"
+          queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
        plugins:
          - docker#v5.2.0:
              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
@@ -186,6 +214,12 @@ steps:
                - "/fsx/hf_cache:/fsx/hf_cache"

      - label: "TTS · Qwen3-TTS Base Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_qwen3_tts_base.py
+          - tests/e2e/offline_inference/test_qwen3_tts_base.py
+          - vllm_omni/model_executor/models/qwen3_tts/
+          - vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
+          - vllm_omni/deploy/qwen3_tts.yaml
        commands:
          - |
            timeout 20m bash -c "
@@ -194,7 +228,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_4_queue"
+          queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
        plugins:
          - docker#v5.2.0:
              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
@@ -207,6 +241,12 @@ steps:
                - "/fsx/hf_cache:/fsx/hf_cache"

      - label: "TTS · Ming-Omni-TTS Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_ming_tts.py
+          - tests/e2e/offline_inference/test_ming_tts.py
+          - vllm_omni/model_executor/models/ming_tts/
+          - vllm_omni/model_executor/stage_input_processors/ming_tts.py
+          - vllm_omni/deploy/ming_tts.yaml
        timeout_in_minutes: 20
        commands:
          - |
@@ -411,6 +451,11 @@ steps:
                      type: DirectoryOrCreate

      - label: "TTS · VoxCPM2 Test"
+        source_file_dependencies:
+          - tests/e2e/offline_inference/test_voxcpm2_tts.py
+          - tests/e2e/online_serving/test_voxcpm2_tts.py
+          - vllm_omni/model_executor/models/voxcpm2/
+          - vllm_omni/deploy/voxcpm2.yaml
        timeout_in_minutes: 30
        commands:
          - |
@@ -434,6 +479,12 @@ steps:
                - "/fsx/hf_cache:/fsx/hf_cache"

      - label: "Omni · Higgs-Audio V2 Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_higgs_audio_v2.py
+          - tests/e2e/offline_inference/test_higgs_audio_v2.py
+          - vllm_omni/model_executor/models/higgs_audio_v2/
+          - vllm_omni/model_executor/stage_input_processors/higgs_audio_v2.py
+          - vllm_omni/deploy/higgs_audio_v2.yaml
        timeout_in_minutes: 30
        commands:
          - |
@@ -805,6 +856,10 @@ steps:
                      type: DirectoryOrCreate

      - label: "Diffusion · Z Image Test"
+        source_file_dependencies:
+          - tests/e2e/offline_inference/test_z_image.py
+          - tests/e2e/online_serving/test_z_image.py
+          - vllm_omni/diffusion/models/z_image/
        timeout_in_minutes: 30
        commands:
          - pytest -s -v tests/e2e/offline_inference/test_z_image.py tests/e2e/online_serving/test_z_image.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
@@ -822,6 +877,10 @@ steps:
                - "/fsx/hf_cache:/fsx/hf_cache"

      - label: "Diffusion · FLUX.2 Klein Test"
+        source_file_dependencies:
+          - tests/e2e/offline_inference/test_flux2_klein.py
+          - tests/e2e/online_serving/test_flux2_klein.py
+          - vllm_omni/diffusion/models/flux2_klein/
        timeout_in_minutes: 30
        commands:
          - pytest -s -v tests/e2e/offline_inference/test_flux2_klein.py tests/e2e/online_serving/test_flux2_klein.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
@@ -215,47 +215,78 @@ steps:
                  path: /mnt/hf-cache
                  type: DirectoryOrCreate

-  - label: "Entrypoints Test"
+  - group: ":card_index_dividers: Entrypoints Test"
    depends_on: upload-ready-pipeline
-    commands:
-      - |
-        timeout 30m bash -c "
-          export VLLM_WORKER_MULTIPROC_METHOD=spawn
-          pytest -sv tests/entrypoints/ -m 'core_model and cuda' --run-level 'core_model'
-        "
-    agents:
-      queue: "mithril-h100-pool"
-    plugins:
-      - kubernetes:
-          podSpec:
-            containers:
-              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-                resources:
-                  limits:
-                    nvidia.com/gpu: 2
-                volumeMounts:
+    steps:
+      - label: "Entrypoints Test with H100"
+        commands:
+          - |
+            timeout 30m bash -c "
+              export VLLM_WORKER_MULTIPROC_METHOD=spawn
+              pytest -sv tests/entrypoints/ -m 'core_model and cuda and H100' --run-level 'core_model'
+            "
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 2
+                    volumeMounts:
+                      - name: devshm
+                        mountPath: /dev/shm
+                      - name: hf-cache
+                        mountPath: /root/.cache/huggingface
+                    env:
+                      - name: HF_HOME
+                        value: /root/.cache/huggingface
+                nodeSelector:
+                  node.kubernetes.io/instance-type: gpu-h100-sxm
+                volumes:
                  - name: devshm
-                    mountPath: /dev/shm
+                    emptyDir:
+                      medium: Memory
                  - name: hf-cache
-                    mountPath: /root/.cache/huggingface
-                env:
-                  - name: HF_HOME
-                    value: /root/.cache/huggingface
-            nodeSelector:
-              node.kubernetes.io/instance-type: gpu-h100-sxm
-            volumes:
-              - name: devshm
-                emptyDir:
-                  medium: Memory
-              - name: hf-cache
-                hostPath:
-                  path: /mnt/hf-cache
-                  type: DirectoryOrCreate
+                    hostPath:
+                      path: /mnt/hf-cache
+                      type: DirectoryOrCreate
+
+      - label: "Entrypoints Test with L4"
+        commands:
+          - |
+            timeout 30m bash -c "
+              export VLLM_WORKER_MULTIPROC_METHOD=spawn
+              pytest -sv tests/entrypoints/ -m 'core_model and cuda and L4' --run-level 'core_model'
+            "
+        agents:
+          queue: "gpu_1_queue"
+        plugins:
+          - docker#v5.2.0:
+              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+              always-pull: true
+              propagate-environment: true
+              shm-size: "8gb"
+              environment:
+                - "HF_HOME=/fsx/hf_cache"
+                - "HF_TOKEN"
+              volumes:
+                - "/fsx/hf_cache:/fsx/hf_cache"

  - group: ":card_index_dividers: E2E Test"
    depends_on: upload-ready-pipeline
    steps:
      - label: "RLHF · VeRL-Omni E2E Test"
+        source_file_dependencies:
+          - tests/e2e/offline_inference/rlhf_test/test_verl_omni_e2e.py
+          - tests/e2e/offline_inference/custom_pipeline/
+          - vllm_omni/diffusion/
+          - vllm_omni/engine/
+          - vllm_omni/entrypoints/
+          - vllm_omni/inputs/
+          - vllm_omni/outputs/
        commands:
          - |
            timeout 20m bash -c "
@@ -316,13 +347,17 @@ steps:
                      type: DirectoryOrCreate

      - label: "TTS · VoxCPM2 Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_voxcpm2_tts.py
+          - vllm_omni/model_executor/models/voxcpm2/
+          - vllm_omni/deploy/voxcpm2.yaml
        timeout_in_minutes: 30
        commands:
          - |
            timeout 30m bash -c "
              export VLLM_LOGGING_LEVEL=DEBUG
              export VLLM_WORKER_MULTIPROC_METHOD=spawn
-              pytest -s -v tests/e2e/offline_inference/test_voxcpm2_tts.py tests/e2e/online_serving/test_voxcpm2_tts.py -m 'core_model' --run-level 'core_model'
+              pytest -s -v tests/e2e/online_serving/test_voxcpm2_tts.py -m 'core_model' --run-level 'core_model'
            "
        agents:
          queue: "gpu_1_queue"
@@ -339,6 +374,11 @@ steps:
                - "/fsx/hf_cache:/fsx/hf_cache"

      - label: "TTS · Qwen3-TTS CustomVoice Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_qwen3_tts_customvoice.py
+          - vllm_omni/model_executor/models/qwen3_tts/
+          - vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
+          - vllm_omni/deploy/qwen3_tts.yaml
        commands:
          - |
            timeout 20m bash -c "
@@ -346,7 +386,7 @@ steps:
              pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m 'core_model' --run-level 'core_model'
            "
        agents:
-          queue: "gpu_4_queue"
+          queue: "gpu_1_queue"
        plugins:
          - docker#v5.2.0:
              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
@@ -360,6 +400,11 @@ steps:
                - "/fsx/hf_cache:/fsx/hf_cache"

      - label: "TTS · Qwen3-TTS Base Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_qwen3_tts_base.py
+          - vllm_omni/model_executor/models/qwen3_tts/
+          - vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
+          - vllm_omni/deploy/qwen3_tts.yaml
        commands:
          - |
            timeout 20m bash -c "
@@ -369,7 +414,7 @@ steps:
              pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py -m 'core_model' --run-level 'core_model'
            "
        agents:
-          queue: "gpu_4_queue"
+          queue: "gpu_1_queue"
        plugins:
          - docker#v5.2.0:
              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
@@ -383,6 +428,11 @@ steps:
                - "/fsx/hf_cache:/fsx/hf_cache"

      - label: "TTS · Higgs-Audio-V3 Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_higgs_audio_v3.py
+          - vllm_omni/model_executor/models/higgs_audio_v3/
+          - vllm_omni/model_executor/stage_input_processors/higgs_audio_v3.py
+          - vllm_omni/deploy/higgs_multimodal_qwen3.yaml
        commands:
          - |
            timeout 20m bash -c "
@@ -392,7 +442,7 @@ steps:
              pytest -s -v tests/e2e/online_serving/test_higgs_audio_v3.py -m 'core_model' --run-level 'core_model'
            "
        agents:
-          queue: "gpu_4_queue"
+          queue: "gpu_1_queue"
        plugins:
          - docker#v5.2.0:
              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
@@ -2,7 +2,7 @@
 # The pinned VLLM_PRECOMPILED_WHEEL_COMMIT and VLLM_BASE_TAG below must
 # be updated on every rebase to match the new upstream vllm commit/tag.
 ARG VLLM_BASE_IMAGE=vllm/vllm-openai
-ARG VLLM_BASE_TAG=v0.22.0
+ARG VLLM_BASE_TAG=v0.23.0
 FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG}
 ARG APP_DIR=/workspace/vllm-omni
 WORKDIR ${APP_DIR}
@@ -16,6 +16,7 @@ RUN apt-get update && \

 RUN uv pip install --system ".[dev]"

+
 RUN ln -sf /usr/bin/python3 /usr/bin/python

 ENTRYPOINT []
@@ -1,4 +1,4 @@
-ARG BASE_IMAGE=vllm/vllm-openai:v0.22.0
+ARG BASE_IMAGE=vllm/vllm-openai:v0.23.0
 FROM ${BASE_IMAGE}

 ARG COMMON_WORKDIR=/app
@@ -1,4 +1,4 @@
-ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.22.0
+ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.23.0
 FROM ${BASE_IMAGE} AS base

 # Declare a variable to know if we want to use the nightly build or the stable build.
@@ -88,7 +88,7 @@ ENV UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
 ENV UV_INDEX_STRATEGY="unsafe-best-match"
 ENV UV_LINK_MODE="copy"

-ARG VLLM_VERSION=v0.22.0
+ARG VLLM_VERSION=v0.23.0
 RUN git clone -b ${VLLM_VERSION} https://github.com/vllm-project/vllm
 WORKDIR /workspace/vllm

@@ -142,15 +142,12 @@ WORKDIR /workspace/vllm-omni
 COPY . .

 ENV VLLM_OMNI_TARGET_DEVICE=xpu
-RUN uv pip install --no-cache-dir ".[dev]" --no-build-isolation
-
-# FIX triton
-RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip uninstall triton triton-xpu && \
-    uv pip install triton-xpu==3.7.0
-
-# remove torch bundled oneccl to avoid conflicts
 RUN --mount=type=cache,target=/root/.cache/uv \
+    echo "triton @ https://example.com/blocked" > /tmp/vllm_omni_constraints.txt && \
+    PIP_CONSTRAINT=/tmp/vllm_omni_constraints.txt \
+    uv pip install --no-build-isolation \
+    --extra-index-url https://mirrors.ustc.edu.cn/pypi/simple/ \
+    ".[dev]" && \
    uv pip uninstall oneccl oneccl-devel

 FROM vllm-omni AS vllm-omni-openai
@@ -268,8 +268,8 @@ Before entering specific testing levels, the project establishes two common spec
 1.  ***PR Checklist ([Tests Style](../ci/tests_style.md))***: This template defines the self-check items that must be completed before submitting a code review (Pull Request). It ensures that each code change meets basic requirements such as code style, dependency updates, and documentation synchronization before entering the automated testing pipeline, serving as the first manual line of defense for quality assurance.
 2.  ***CI Failure Explanation ([CI Failures](../ci/failures.md))***: This document archives and explains common failure patterns in the Continuous Integration (CI) pipeline, error log interpretation, and preliminary troubleshooting steps. It helps developers and testers quickly diagnose the causes of automated test failures, improving problem-solving efficiency.

-<<<<<<< ready
 ## Notes
+
 ### Diff-aware Buildkite uploads (`source_file_dependencies`)

 L2 (`.buildkite/test-ready.yml`) and L3 (`.buildkite/test-merge.yml`) pipelines can **skip unrelated GPU jobs at upload time** based on the PR diff. This is implemented by `.buildkite/scripts/upload_pipeline.py`, which filters steps before calling `buildkite-agent pipeline upload`.
@@ -302,12 +302,17 @@ Steps **without** `source_file_dependencies` are always uploaded (subject to the

 To balance CI cost and coverage:

- **Always run** (no `source_file_dependencies`): baseline groups such as Simple Test, Diffusion unit tests, Engine/Model Executor, Distributed, Custom Pipeline, Entrypoints (L2), LoRA / Entrypoints (L3), and **all E2E jobs on L4 docker queues** (`gpu_1_queue`, `gpu_4_queue`).
- **Diff-gated** (`source_file_dependencies` set): **H100 E2E jobs only** (`queue: mithril-h100-pool` under the E2E Test group)—typically the model’s pytest file(s) plus related `vllm_omni/` model, deploy, and stage paths.
+- **Always run** (no `source_file_dependencies`): baseline groups outside E2E Test—e.g. Simple Test, Diffusion unit tests, Engine/Model Executor, Distributed, Custom Pipeline, Entrypoints (L2), LoRA / Entrypoints (L3).
+- **Diff-gated** (`source_file_dependencies` set): **every leaf job under the E2E Test group** in `test-ready.yml` and `test-merge.yml`, regardless of queue (`mithril-h100-pool`, `gpu_1_queue`, or `gpu_4_queue`). Each step lists the smallest set of prefixes that should trigger it—typically:
+  - pytest file(s) exercised by the job (online and/or offline);
+  - model code under `vllm_omni/model_executor/models/` or `vllm_omni/diffusion/models/`;
+  - related `vllm_omni/model_executor/stage_input_processors/` and `vllm_omni/deploy/*.yaml` when applicable.

-Adding a new expensive H100 E2E step: list the test file(s) and the smallest set of source prefixes that should trigger it. Prefer **per-step** deps on leaf jobs rather than a broad group-level list unless every child shares the same paths.
+Adding a new E2E step: add `source_file_dependencies` on the leaf job with those prefixes. Prefer **per-step** deps rather than a broad group-level list unless every child shares the same paths.

-#### YAML example
+#### YAML examples
+
+H100 E2E (kubernetes / `mithril-h100-pool`):

 ```yaml
      - label: "Diffusion · Qwen Image Test"
@@ -320,6 +325,21 @@ Adding a new expensive H100 E2E step: list the test file(s) and the smallest set
          queue: "mithril-h100-pool"
 ```

+Docker E2E (`gpu_1_queue` / `gpu_4_queue`)—same key, same upload-time filtering:
+
+```yaml
+      - label: "TTS · Qwen3-TTS CustomVoice Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_qwen3_tts_customvoice.py
+          - vllm_omni/model_executor/models/qwen3_tts/
+          - vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
+          - vllm_omni/deploy/qwen3_tts.yaml
+        commands:
+          - pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py ...
+        agents:
+          queue: "gpu_4_queue"
+```
+
 A **group** may also define `source_file_dependencies`; nested steps inherit filtering as a unit—the whole group is dropped if no prefix matches.

 #### Local dry-run
@@ -339,7 +359,7 @@ On a PR build, Buildkite logs from `upload_pipeline.py` include lines such as `s

 - Implementation: `.buildkite/scripts/upload_pipeline.py`
 - L2/L3 diff skip does **not** replace label-based triggers (`ready`, `merge-test`); it only reduces which steps appear **after** the pipeline is already scheduled.
-=======
+
 ### Test helper environment variables

 Some shared helpers under `tests/helpers/` honor optional environment variables for local debugging. These are **not** set in CI by default.
@@ -361,7 +381,6 @@ Example (Windows PowerShell):
 $env:VLLM_OMNI_KEEP_REQUEST_MEDIA = "1"
 pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -k test_mix_to_text_audio
 ```
->>>>>>> main

 ## Chapter 1: L1 & L2 Level Testing - Unit Testing and Basic End-to-End Verification

@@ -397,7 +416,7 @@ A clear directory structure is key to managing test cases efficiently.
 ### 1.4 Execution Method and Example

 -   ***Trigger Timing***: **`PR with ready label`**. That is, when a developer adds a "ready for review" or similar label to a PR on platforms like GitHub, L1 and L2 tests are automatically triggered.
-   ***Diff-aware step skipping***: On L2, expensive **H100 E2E** jobs may be omitted at pipeline upload when the PR diff does not touch their [`source_file_dependencies`](#diff-aware-buildkite-uploads-source_file_dependencies) paths; L4 docker jobs and non-E2E steps still run. See [Diff-aware Buildkite uploads](#diff-aware-buildkite-uploads-source_file_dependencies).
+-   ***Diff-aware step skipping***: On L2, **E2E Test** jobs may be omitted at pipeline upload when the PR diff does not touch their [`source_file_dependencies`](#diff-aware-buildkite-uploads-source_file_dependencies) prefixes (H100 and docker queues alike); non-E2E groups still always upload. See [Diff-aware Buildkite uploads](#diff-aware-buildkite-uploads-source_file_dependencies).
 -   ***Execution Environment***: L1 uses ***CPU*** environment; L2 requires ***GPU*** environment.
 -   ***Script Example***:

@@ -501,7 +520,7 @@ L3 level testing executes after code is merged into the main branch. Its core pu
 ### 2.4 Execution Method and Example

 -   ***Trigger Timing***: **`PR Merged`**. Automatically triggered after code review is approved and merged into the main branch (typically via `merge-test` label on the PR before merge).
-   ***Diff-aware step skipping***: Same [`source_file_dependencies`](#diff-aware-buildkite-uploads-source_file_dependencies) mechanism as L2—only **H100 E2E** steps are diff-gated; L4 docker E2E and other L3 groups always upload. See [Diff-aware Buildkite uploads](#diff-aware-buildkite-uploads-source_file_dependencies).
+-   ***Diff-aware step skipping***: Same [`source_file_dependencies`](#diff-aware-buildkite-uploads-source_file_dependencies) mechanism as L2—**all E2E Test** leaf steps are diff-gated; other L3 groups always upload. See [Diff-aware Buildkite uploads](#diff-aware-buildkite-uploads-source_file_dependencies).
 -   ***Execution Environment***: ***GPU*** servers.
 -   ***Script Example***:

@@ -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, 0.18.0, 0.20.0, 0.21.0, and 0.22.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, 0.21.0, 0.22.0, and 0.23.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==0.22.0 --torch-backend=auto
+uv pip install vllm==0.23.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.22.0 release version of the library
+If you do not need to modify source code of vLLM, you can directly install the stable 0.23.0 release version of the library

 ```bash
-uv pip install vllm==0.22.0 --torch-backend=auto
+uv pip install vllm==0.23.0 --torch-backend=auto
 ```

-The 0.22.0 release of vLLM ships CUDA 12.9-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.23.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.22.0
+git checkout v0.23.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 12.9 (the default for v0.22.0; the wheel filename has no `+cu129` suffix)
-export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.22.0/vllm-0.22.0-cp38-abi3-manylinux_2_28_x86_64.whl
+#For CUDA 13.0 (the default for v0.23.0; the wheel filename has no `+cu130` suffix)
+export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.23.0/vllm-0.23.0-cp38-abi3-manylinux_2_28_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.22.0 \
+    vllm/vllm-omni:v0.23.0 \
    vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --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.22.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.22.0/rocm722
+uv pip install vllm==0.23.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.23.0/rocm722
 ```

 #### Installation of vLLM-Omni
@@ -34,13 +34,13 @@ uv pip install onnxruntime-rocm
 # --8<-- [start:build-wheel-from-source]

 #### Installation of vLLM
-If you do not need to modify source code of vLLM, you can directly install the stable 0.22.0 release version of the library
+If you do not need to modify source code of vLLM, you can directly install the stable 0.23.0 release version of the library

 ```bash
-uv pip install vllm==0.22.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.22.0/rocm722
+uv pip install vllm==0.23.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.23.0/rocm722
 ```

-The pre-built 0.22.0 vLLM wheel targets ROCm 7.2.2. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
+The pre-built 0.23.0 vLLM wheel targets ROCm 7.2.2. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.

 #### Installation of vLLM-Omni
 Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -58,7 +58,7 @@ If you want to check, modify or debug with source code of vLLM, install the libr
 ```bash
 git clone https://github.com/vllm-project/vllm.git
 cd vllm
-git checkout v0.22.0
+git checkout v0.23.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.22.0 \
+  vllm/vllm-omni-rocm:v0.23.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.22.0
+  vllm/vllm-omni-rocm:v0.23.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.22.0 --torch-backend=auto
+uv pip install vllm==0.23.0 --torch-backend=auto

 # On ROCm
-uv pip install vllm==0.22.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.22.0/rocm722
+uv pip install vllm==0.23.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.23.0/rocm722

 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.22.0` and vLLM Omni `0.22.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.23.0` and vLLM Omni `0.23.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue.

 ## Offline Inference

@@ -42,7 +42,7 @@ th {
 | `LTX2ImageToVideoTwoStagesPipeline` | LTX-2-I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |
 | `LTX23Pipeline` | LTX-2.3-T2V | `dg845/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
 | `LTX23ImageToVideoPipeline` | LTX-2.3-I2V | `dg845/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
-| `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | | | |
+| `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | |
 | `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | |
 | `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | ✅︎ | ✅︎ | | ✅︎ |
@@ -0,0 +1,375 @@
+# Step-Audio2 Offline Inference Examples
+
+This directory contains examples for running offline inference with Step-Audio2 using vLLM-Omni.
+
+## Model Overview
+
+Step-Audio2 is a two-stage audio model:
+
+- **Stage 0 (Thinker)**: Audio understanding → Text + Audio tokens
+  - Input: Audio (16kHz)
+  - Output: Text transcription + Audio tokens for synthesis
+
+- **Stage 1 (Token2Wav)**: Audio synthesis
+  - Input: Audio tokens + Speaker prompt wav
+  - Output: Synthesized audio waveform (24kHz)
+
+## Hardware Requirements
+
+| Mode | GPU Configuration | VRAM Required |
+|------|-------------------|---------------|
+| ASR (S2T) | 1x GPU | ~20-25GB |
+| TTS/S2ST (single GPU) | 1x GPU | ~40-50GB |
+| TTS/S2ST (multi GPU) | 2x GPU | GPU0: ~28GB, GPU1: ~22GB |
+
+**Tested on:**
+- 1x NVIDIA H100 80GB (single-card S2ST)
+- 2x NVIDIA A10 40GB (multi-card S2ST)
+
+**Notes:**
+- Single GPU mode requires high VRAM due to both stages sharing memory
+- Multi GPU mode separates Stage 0 (Thinker) and Stage 1 (Token2Wav) across GPUs
+- VRAM usage can be adjusted via `gpu_memory_utilization` in stage config
+
+## Performance Benchmark
+
+### vLLM-Omni vs Official Step-Audio2
+
+Single request latency comparison between vLLM-Omni and official Step-Audio2 implementation.
+
+| Task | Tokens | vllm-omni | Step-Audio2 | Speedup |
+|------|--------|-----------|-------------|---------|
+| S2ST | ~85 | 5.45s | 7.36s | **1.35x** |
+| S2ST | ~160 | 6.67s | 13.92s | **2.09x** |
+| S2ST | ~315 | 9.42s | 31.50s | **3.34x** |
+| TTS | ~1024 | 16.65s | ~87s | **~5.2x** |
+
+**Key observations:**
+- Speedup increases with sequence length due to vLLM's efficient KV cache management
+- TTS (pure generation) shows the largest speedup (~5x)
+- S2ST benefits from optimized multi-stage pipeline
+
+**Benchmark environment:**
+- GPU: NVIDIA H100 80GB (single card)
+- Model: Step-Audio2-mini
+- Warmup: 1 run, Measured: 3 runs (averaged)
+
+### Async Chunk Streaming Performance
+
+Comparison between sequential (non-async) and async chunk modes via `/v1/audio/speech` TTS endpoint.
+
+| Mode | Mean TTFP | Mean E2E | Mean RTF | Audio Throughput |
+|------|-----------|----------|----------|-----------------|
+| Sequential | 4316ms | 4316ms | 0.938 | 1.07x realtime |
+| **Async Chunk** | **1437ms** | 4362ms | 0.949 | 1.06x realtime |
+| **Improvement** | **-67% (3x faster)** | ~same | ~same | ~same |
+
+**Key observations:**
+- Async chunk reduces **time-to-first-audio (TTFP) by 67%** (4.3s → 1.4s)
+- E2E latency remains comparable — async chunk overlaps Thinker decode with Token2Wav synthesis
+- RTF < 1 in both modes (real-time capable)
+- Sequential mode: TTFP ≈ E2E (must wait for all audio tokens before synthesis starts)
+- Async chunk mode: Token2Wav starts after first 28 tokens (chunk_size=25 + lookahead=3)
+
+**Benchmark environment:**
+- GPU: 4x NVIDIA RTX 3090 24GB (TP=2 for Thinker, 1 GPU for Token2Wav)
+- Model: Step-Audio2-mini
+- Endpoint: `/v1/audio/speech` (10 prompts, concurrency=1)
+- Measured via `bench_tts_serve.py`
+
+## Installation
+
+Make sure you have installed vLLM-Omni and all required dependencies:
+
+```bash
+# Install vLLM-Omni
+pip install vllm-omni
+
+# Install Step-Audio2 (REQUIRED for Token2Wav stage)
+pip install step-audio2
+```
+
+## Model Setup
+
+### Option 1: Auto-download from HuggingFace (Recommended)
+
+The script will **automatically download** the model on first run:
+
+```bash
+# Just run without specifying --model, it will auto-download stepfun-ai/Step-Audio2-mini
+python end2end.py --query-type audio_to_text
+
+# Or explicitly specify the HuggingFace model
+python end2end.py --query-type audio_to_text --model stepfun-ai/Step-Audio2-mini
+```
+
+Models will be cached in `~/.cache/huggingface/hub/` for future use.
+
+**Available models**:
+- `stepfun-ai/Step-Audio2-mini` (smaller, faster)
+- `stepfun-ai/Step-Audio2-7B` (larger, better quality)
+
+### Option 2: Manual Download (for offline use)
+
+Download and use locally:
+
+```bash
+# Download from HuggingFace
+huggingface-cli download stepfun-ai/Step-Audio2-mini --local-dir ./models/Step-Audio2-mini
+
+# Then use the local path
+python end2end.py --query-type audio_to_text --model ./models/Step-Audio2-mini
+```
+
+Ensure the model directory contains:
+```
+Step-Audio2-mini/
+├── config.json
+├── model.safetensors (or pytorch_model.bin)
+├── tokenizer.json
+├── tokenizer_config.json
+└── token2wav/                           # Token2Wav models (REQUIRED)
+    ├── speech_tokenizer_v2_25hz.onnx   # Audio tokenizer
+    ├── campplus.onnx                    # Speaker encoder
+    ├── flow.yaml                        # Flow model config
+    ├── flow.pt                          # Flow model weights
+    └── hift.pt                          # HiFT vocoder weights
+```
+
+## Usage Examples
+
+### 1. Audio to Text (ASR - Speech Recognition)
+
+Transcribe audio to text:
+
+```bash
+# Quick start - Using default model and test audio
+python end2end.py --query-type audio_to_text
+
+# Using your own audio file (model will auto-download)
+python end2end.py --query-type audio_to_text \
+    --audio-path /path/to/input.wav
+
+# With specific model
+python end2end.py --query-type audio_to_text \
+    --audio-path /path/to/input.wav \
+    --model stepfun-ai/Step-Audio2-7B
+
+# With custom question
+python end2end.py --query-type audio_to_text \
+    --audio-path input.wav \
+    --question "What is the speaker saying?"
+```
+
+**Output**: Text transcription saved to `output_step_audio2/00000_text.txt`
+
+### 2. Text to Audio (TTS - Speech Synthesis)
+
+Convert text to speech:
+
+```bash
+# Basic TTS (model auto-downloads)
+python end2end.py --query-type text_to_audio \
+    --text "Hello, this is a test of Step Audio 2 synthesis."
+
+# With specific model
+python end2end.py --query-type text_to_audio \
+    --text "Hello, this is a test." \
+    --model stepfun-ai/Step-Audio2-7B
+```
+
+**Note**: Speaker voice is controlled by the `STEP_AUDIO2_DEFAULT_PROMPT_WAV` environment variable or the default prompt wav bundled with the model.
+
+**Output**:
+- Text: `output_step_audio2/00000_text.txt`
+- Audio: `output_step_audio2/00000_output.wav` (24kHz)
+
+### 3. Audio to Audio (Voice Conversion)
+
+Process input audio and generate output audio:
+
+```bash
+# Basic voice conversion (model auto-downloads)
+python end2end.py --query-type audio_to_audio \
+    --audio-path /path/to/source_audio.wav
+
+# With specific model
+python end2end.py --query-type audio_to_audio \
+    --audio-path source.wav \
+    --model stepfun-ai/Step-Audio2-7B
+```
+
+This mode:
+1. Understands the content in `--audio-path` (source)
+2. Generates audio output with the default voice
+
+**Note**: To use a custom speaker voice, set the `STEP_AUDIO2_DEFAULT_PROMPT_WAV` environment variable.
+
+### Advanced Options
+
+```bash
+# Use custom stage configuration
+python end2end.py --query-type audio_to_text \
+    --stage-configs-path /path/to/custom_config.yaml
+
+# Multiple prompts (for batch testing)
+python end2end.py --query-type audio_to_text \
+    --audio-path input.wav \
+    --num-prompts 5
+
+# Custom output directory
+python end2end.py --query-type text_to_audio \
+    --text "Test synthesis" \
+    --output-dir ./my_outputs
+
+# Enable detailed logging
+python end2end.py --query-type audio_to_text \
+    --audio-path input.wav \
+    --enable-stats
+
+# Adjust generation parameters
+python end2end.py --query-type audio_to_text \
+    --audio-path input.wav \
+    --max-tokens 2048
+
+# Use custom speaker voice via environment variable
+STEP_AUDIO2_DEFAULT_PROMPT_WAV=/path/to/speaker.wav python end2end.py \
+    --query-type text_to_audio \
+    --text "Hello world"
+
+# Use Ray backend for distributed processing
+python end2end.py --query-type text_to_audio \
+    --text "Hello world" \
+    --worker-backend ray \
+    --ray-address "auto"
+```
+
+## Configuration
+
+### Stage Configuration
+
+The default configuration (`step_audio_2.yaml`) uses:
+
+- **Stage 0 (Thinker)**: GPU 0, 80% memory
+- **Stage 1 (Token2Wav)**: GPU 1, 30% memory
+
+For **single GPU** setup, edit the config to use `devices: "0"` for both stages.
+
+### Sampling Parameters
+
+- **Thinker (Stage 0)**:
+  - Temperature: 0.7 (balanced creativity)
+  - Top-p: 0.9
+  - Max tokens: 1024 (configurable)
+
+- **Token2Wav (Stage 1)**:
+  - Temperature: 0.0 (deterministic)
+  - Operates in generation mode (not sampling)
+
+## Common Issues
+
+### 1. ImportError: No module named 's3tokenizer'
+
+**Solution**: Install Step-Audio2 package:
+```bash
+pip install step-audio2
+```
+
+### 2. FileNotFoundError: prompt_wav file not found
+
+**Solution**: Set the `STEP_AUDIO2_DEFAULT_PROMPT_WAV` environment variable to a valid audio file:
+```bash
+export STEP_AUDIO2_DEFAULT_PROMPT_WAV=/path/to/speaker.wav
+python end2end.py --query-type text_to_audio --text "Hello"
+```
+Or ensure the default prompt wav (`default_female.wav`) exists in your model directory.
+
+### 3. FileNotFoundError: token2wav models not found
+
+**Solution**: Ensure your model directory has the complete `token2wav/` subdirectory with all ONNX and PyTorch models.
+
+### 4. CUDA Out of Memory
+
+**Solutions**:
+- Use single GPU mode (set both stages to `devices: "0"`)
+- Reduce `gpu_memory_utilization` in config
+- Reduce `max_num_batched_tokens`
+- Process fewer prompts at once
+
+### 5. Model not found in registry
+
+**Solution**: Ensure you're using vLLM-Omni's entry point with `--omni` flag or install vllm-omni properly:
+```bash
+pip install vllm-omni
+```
+
+## Output Files
+
+The script generates files in the output directory (default: `output_step_audio2/`):
+
+```
+output_step_audio2/
+├── 00000_text.txt        # Text output from Thinker stage
+├── 00000_output.wav      # Audio output from Token2Wav stage (24kHz)
+├── 00001_text.txt        # (if multiple prompts)
+└── 00001_output.wav
+```
+
+## Performance Tips
+
+1. **First run is slow**: Stage initialization takes 20-60 seconds
+2. **Single GPU**: Set both stages to `devices: "0"` in config
+3. **Multiple prompts**: Use `--num-prompts N` for batch testing
+4. **Ray backend**: For multi-node or advanced scheduling
+5. **Logging**: Use `--enable-stats` to debug performance issues
+
+## Speaker Voice Configuration
+
+The Token2Wav stage requires a speaker prompt wav for voice conditioning. It is automatically resolved in this order:
+
+1. `STEP_AUDIO2_DEFAULT_PROMPT_WAV` environment variable (if set)
+2. `{model_dir}/assets/default_female.wav`
+3. `{model_dir}/default_female.wav`
+
+If none are found, set the environment variable explicitly:
+```bash
+export STEP_AUDIO2_DEFAULT_PROMPT_WAV=/path/to/speaker.wav
+```
+
+**Guidelines for custom speaker prompt:**
+
+- **Duration**: 3-10 seconds recommended
+- **Quality**: Clean audio, minimal background noise
+- **Format**: WAV, MP3, FLAC (will be resampled internally)
+- **Content**: Clear speech, representative of target voice
+
+## Example Workflow
+
+Complete example from audio to final output:
+
+```bash
+# 1. ASR: Transcribe audio
+python end2end.py --query-type audio_to_text \
+    --audio-path interview.wav \
+    --model ./models/Step-Audio2-7B \
+    --output-dir ./outputs
+
+# 2. Check the transcription
+cat ./outputs/00000_text.txt
+
+# 3. TTS: Synthesize new speech (with custom voice)
+STEP_AUDIO2_DEFAULT_PROMPT_WAV=./speaker_samples/female_voice.wav \
+python end2end.py --query-type text_to_audio \
+    --text "The quick brown fox jumps over the lazy dog" \
+    --model ./models/Step-Audio2-7B \
+    --output-dir ./outputs
+
+# 4. Listen to the result
+# Audio saved to: ./outputs/00000_output.wav
+```
+
+## References
+
+- [Step-Audio2 Paper](https://arxiv.org/abs/2507.16632)
+- [vLLM-Omni Documentation](https://vllm-omni.readthedocs.io)
+- [Model on HuggingFace](https://huggingface.co/stepfun-ai/Step-Audio-2-mini)
@@ -0,0 +1,527 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+This example shows how to use vLLM-Omni for running offline inference
+with Step-Audio2 model for audio understanding and generation.
+
+Step-Audio2 is a two-stage model:
+  - Stage 0 (Thinker): Audio understanding → Text + Audio tokens
+  - Stage 1 (Token2Wav): Audio tokens → Waveform (24kHz)
+
+Usage examples:
+  # Audio to text (ASR)
+  python end2end.py --query-type audio_to_text --audio-path input.wav --model-path "$MODEL_PATH"
+
+  # Text to audio (TTS with default voice)
+  python end2end.py --query-type text_to_audio --model-path "$MODEL_PATH"
+
+  # Audio to audio (Voice conversion with default voice)
+  python end2end.py --query-type audio_to_audio --audio-path input.wav --model-path "$MODEL_PATH"
+"""
+
+import os
+import sys
+from pathlib import Path
+from typing import NamedTuple
+
+import librosa  # noqa: TID251
+import numpy as np
+import soundfile as sf
+import torch
+from vllm.assets.audio import AudioAsset
+from vllm.sampling_params import SamplingParams
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+from vllm_omni.entrypoints.omni import Omni
+
+SEED = 42
+
+
+class QueryResult(NamedTuple):
+    inputs: dict
+    limit_mm_per_prompt: dict[str, int]
+
+
+def get_audio_to_text_query(
+    audio_path: str | None = None,
+    question: str | None = None,
+    sampling_rate: int = 16000,
+) -> QueryResult:
+    """
+    Audio Speech Recognition (ASR) - Convert audio to text.
+
+    Args:
+        audio_path: Path to audio file
+        question: Question about the audio (optional)
+        sampling_rate: Target sampling rate (16kHz for Step-Audio2)
+
+    Returns:
+        QueryResult with prompt and audio data
+    """
+    if question is None:
+        question = "请记录下你所听到的语音内容。"
+
+    prompt = (
+        f"<|im_start|>system\n{question}<|im_end|>\n<|im_start|>user\n<audio_patch><|im_end|>\n<|im_start|>assistant\n"
+    )
+
+    # Load audio
+    if audio_path:
+        if not os.path.exists(audio_path):
+            raise FileNotFoundError(f"Audio file not found: {audio_path}")
+        audio_signal, sr = librosa.load(audio_path, sr=sampling_rate)
+        audio_data = (audio_signal.astype(np.float32), sr)
+    else:
+        # Use default test audio
+        audio_data = AudioAsset("mary_had_lamb").audio_and_sample_rate
+        # Resample to 16kHz if needed
+        if audio_data[1] != sampling_rate:
+            audio_signal = librosa.resample(audio_data[0], orig_sr=audio_data[1], target_sr=sampling_rate)
+            audio_data = (audio_signal.astype(np.float32), sampling_rate)
+
+    return QueryResult(
+        inputs={
+            "prompt": prompt,
+            "multi_modal_data": {
+                "audio": audio_data,
+            },
+        },
+        limit_mm_per_prompt={"audio": 1},
+    )
+
+
+def get_text_to_audio_query(
+    text: str | None = None,
+    sampling_rate: int = 16000,
+) -> QueryResult:
+    """
+    Text-to-Speech (TTS) - Convert text to audio.
+
+    Args:
+        text: Text to synthesize
+        sampling_rate: Target sampling rate (16kHz for Step-Audio2)
+
+    Returns:
+        QueryResult with text prompt and speaker wav path
+
+    Note:
+        For audio generation, the prompt must end with "<tts_start>"
+        to signal the model to generate audio tokens instead of just text.
+    """
+    if text is None:
+        text = "Hello, this is a test of Step Audio 2 text to speech synthesis."
+
+    prompt = (
+        "<|im_start|>system\n"
+        "请逐字朗读用户提供的文本，不要改写或补充。<|im_end|>\n"
+        "<|im_start|>user\n"
+        f"{text}<|im_end|>\n"
+        "<|im_start|>assistant\n"
+        "<tts_start>"
+    )
+
+    inputs = {
+        "prompt": prompt,
+    }
+
+    return QueryResult(
+        inputs=inputs,
+        limit_mm_per_prompt={},
+    )
+
+
+def get_audio_to_audio_query(
+    audio_path: str | None = None,
+    question: str | None = None,
+    sampling_rate: int = 16000,
+) -> QueryResult:
+    """
+    Audio-to-Audio - Voice conversion/cloning.
+
+    Args:
+        audio_path: Path to source audio file
+        question: Question/instruction about the audio
+        sampling_rate: Target sampling rate (16kHz for Step-Audio2)
+
+    Returns:
+        QueryResult with audio input
+
+    Note:
+        This mode processes input audio and generates output audio.
+        Speaker voice is controlled by STEP_AUDIO2_DEFAULT_PROMPT_WAV env var.
+    """
+    if question is None:
+        question = "请仔细聆听这段语音，然后复述其内容。"
+
+    prompt = (
+        "<|im_start|>system\n"
+        f"{question}<|im_end|>\n"
+        "<|im_start|>user\n"
+        "<audio_patch><|im_end|>\n"
+        "<|im_start|>assistant\n"
+        "<tts_start>"
+    )
+
+    # Load source audio
+    if audio_path:
+        if not os.path.exists(audio_path):
+            raise FileNotFoundError(f"Audio file not found: {audio_path}")
+        audio_signal, sr = librosa.load(audio_path, sr=sampling_rate)
+        audio_data = (audio_signal.astype(np.float32), sr)
+
+    else:
+        audio_data = AudioAsset("mary_had_lamb").audio_and_sample_rate
+        if audio_data[1] != sampling_rate:
+            audio_signal = librosa.resample(audio_data[0], orig_sr=audio_data[1], target_sr=sampling_rate)
+            audio_data = (audio_signal.astype(np.float32), sampling_rate)
+
+    return QueryResult(
+        inputs={
+            "prompt": prompt,
+            "multi_modal_data": {
+                "audio": audio_data,
+            },
+        },
+        limit_mm_per_prompt={"audio": 1},
+    )
+
+
+query_map = {
+    "audio_to_text": get_audio_to_text_query,
+    "text_to_audio": get_text_to_audio_query,
+    "audio_to_audio": get_audio_to_audio_query,
+}
+
+
+def main(args):
+    # Use local model path or HuggingFace model name
+    if not args.model:
+        sys.exit(1)
+
+    model_name = args.model
+
+    # Check if it's a local path
+    if os.path.isdir(model_name):
+        # Verify essential files exist
+        config_path = os.path.join(model_name, "config.json")
+        if not os.path.exists(config_path):
+            sys.exit(1)
+
+    # Get query configuration
+    query_func = query_map[args.query_type]
+
+    # Build query based on type
+    if args.query_type == "audio_to_text":
+        query_result = query_func(
+            audio_path=args.audio_path,
+            question=args.question,
+            sampling_rate=args.sampling_rate,
+        )
+    elif args.query_type == "text_to_audio":
+        query_result = query_func(
+            text=args.text,
+            sampling_rate=args.sampling_rate,
+        )
+    elif args.query_type == "audio_to_audio":
+        query_result = query_func(
+            audio_path=args.audio_path,
+            question=args.question,
+            sampling_rate=args.sampling_rate,
+        )
+    else:
+        raise ValueError(f"Unknown query type: {args.query_type}")
+
+    # Initialize vLLM-Omni with Step-Audio2
+    # Resolve stage config path
+    if args.stage_configs_path:
+        stage_config_path = args.stage_configs_path
+    else:
+        configs_dir = Path(__file__).parent.parent.parent.parent / "vllm_omni/model_executor/stage_configs"
+        if args.query_type == "audio_to_text":
+            # ASR only needs Thinker (Stage 0), no Token2Wav
+            stage_config_path = str(configs_dir / "step_audio_2_asr.yaml")
+            if not os.path.exists(stage_config_path):
+                stage_config_path = str(configs_dir / "step_audio_2.yaml")
+        else:
+            # TTS/S2ST need both stages
+            stage_config_path = str(configs_dir / "step_audio_2.yaml")
+
+    omni_llm = Omni(
+        model=model_name,
+        log_stats=args.enable_stats,
+        log_file=("step_audio2_pipeline.log" if args.enable_stats else None),
+        init_sleep_seconds=args.init_sleep_seconds,
+        batch_timeout=args.batch_timeout,
+        init_timeout=args.init_timeout,
+        shm_threshold_bytes=args.shm_threshold_bytes,
+        stage_configs_path=stage_config_path,
+        worker_backend=args.worker_backend,
+        ray_address=args.ray_address,
+    )
+
+    # Configure sampling parameters for each stage
+    # Stage 0 (Thinker): Generate text and audio tokens
+    # Note: For audio generation tasks, we need to allow the model to generate audio tokens
+    # Audio tokens are in range [151696, 158257], so we should not stop too early
+    # Qwen EOS token is typically 151645 (<|im_end|>), but we want to allow audio generation
+    # So we don't set stop_token_ids for audio tasks, but we do set it for text-only tasks
+    stop_token_ids = None
+    if args.query_type == "audio_to_text":
+        # For ASR, stop at EOS token
+        stop_token_ids = [151645]  # Qwen EOS token <|im_end|>
+
+    # Adjust sampling parameters based on query type
+    # For audio generation tasks, we need higher repetition_penalty to avoid text loops
+    # and allow the model to transition to audio token generation
+    repetition_penalty = 1.1 if args.query_type in ["text_to_audio", "audio_to_audio"] else 1.05
+
+    thinker_sampling_params = SamplingParams(
+        temperature=0.7,
+        top_p=0.9,
+        top_k=-1,
+        max_tokens=args.max_tokens,
+        seed=SEED,
+        detokenize=True,
+        repetition_penalty=repetition_penalty,
+        stop_token_ids=stop_token_ids,
+    )
+
+    # Stage 1 (Token2Wav): Convert audio tokens to waveform
+    # Note: This stage doesn't actually sample, it's deterministic generation
+    token2wav_sampling_params = SamplingParams(
+        temperature=0.0,
+        top_p=1.0,
+        top_k=-1,
+        max_tokens=1,
+        seed=SEED,
+        detokenize=False,
+    )
+
+    if args.query_type == "audio_to_text":
+        # ASR: single stage, only Thinker
+        sampling_params_list = [thinker_sampling_params]
+    else:
+        # TTS/S2ST: two stages
+        sampling_params_list = [thinker_sampling_params, token2wav_sampling_params]
+
+    # Prepare prompts
+    if args.num_prompts > 1:
+        prompts = [query_result.inputs for _ in range(args.num_prompts)]
+    else:
+        prompts = [query_result.inputs]
+
+    # Run inference
+    omni_outputs = omni_llm.generate(prompts, sampling_params_list)
+
+    # Create output directory
+    output_dir = args.output_dir
+    os.makedirs(output_dir, exist_ok=True)
+
+    def _get_request_index(req_id: str, fallback: int) -> int:
+        """Parse request index from vLLM request_id.
+
+        vLLM may return IDs like "0_<uuid>". Use the leading integer if present,
+        otherwise fall back to the loop index.
+        """
+        try:
+            return int(req_id)
+        except (TypeError, ValueError):
+            if isinstance(req_id, str):
+                head = req_id.split("_", 1)[0]
+                if head.isdigit():
+                    return int(head)
+        return fallback
+
+    # Process outputs from each stage
+    for stage_idx, stage_outputs in enumerate(omni_outputs):
+        # request_output may be a single object or a list
+        raw_output = stage_outputs.request_output
+        outputs_list = raw_output if isinstance(raw_output, list) else [raw_output]
+
+        if stage_outputs.final_output_type == "text":
+            # Stage 0 (Thinker) text output
+            for i, output in enumerate(outputs_list):
+                request_id = _get_request_index(output.request_id, i)
+                text_output = output.outputs[0].text
+
+                # Save to file
+                out_txt = os.path.join(output_dir, f"{request_id:05d}_text.txt")
+                with open(out_txt, "w", encoding="utf-8") as f:
+                    f.write(f"Prompt:\n{prompts[request_id]['prompt']}\n\n")
+                    f.write(f"Output:\n{text_output}\n")
+                print(f"Text output saved to: {out_txt}")
+                print(f"Text: {text_output[:200]}")
+
+        elif stage_outputs.final_output_type == "audio":
+            # Stage 1 (Token2Wav) audio output
+            for i, output in enumerate(outputs_list):
+                request_id = _get_request_index(output.request_id, i)
+
+                # Get audio tensor (24kHz)
+                mm_out = {}
+                if hasattr(output, "multimodal_output") and output.multimodal_output:
+                    mm_out = output.multimodal_output
+                elif output.outputs and hasattr(output.outputs[0], "multimodal_output"):
+                    mm_out = output.outputs[0].multimodal_output or {}
+                audio_tensor = None
+                for key in ("audio", "wav", "waveform", "audio_pcm", "pcm", "model_outputs"):
+                    if key in mm_out and mm_out[key] is not None:
+                        audio_tensor = mm_out[key]
+                        break
+
+                if audio_tensor is not None:
+                    # Save audio file
+                    output_wav = os.path.join(output_dir, f"{request_id:05d}_output.wav")
+                    if isinstance(audio_tensor, list):
+                        chunks = []
+                        for chunk in audio_tensor:
+                            if chunk is None:
+                                continue
+                            if isinstance(chunk, torch.Tensor):
+                                chunk_np = chunk.detach().cpu().numpy()
+                            else:
+                                chunk_np = np.asarray(chunk)
+                            if chunk_np.size > 0:
+                                chunks.append(chunk_np)
+                        if not chunks:
+                            continue
+                        audio_numpy = np.concatenate(chunks, axis=-1)
+                    elif isinstance(audio_tensor, torch.Tensor):
+                        audio_numpy = audio_tensor.detach().cpu().numpy()
+                    else:
+                        audio_numpy = np.asarray(audio_tensor)
+                    sf.write(output_wav, audio_numpy, samplerate=24000)
+
+
+def parse_args():
+    parser = FlexibleArgumentParser(description="Demo on using vLLM-Omni for offline inference with Step-Audio2")
+
+    # Model and config
+    parser.add_argument(
+        "--model",
+        "-m",
+        type=str,
+        required=True,
+        help=(
+            "Model path (required). Can be:\n"
+            "  - Local path: /path/to/step-audio-2\n"
+            "  - HuggingFace ID: stepfun-ai/Step-Audio2-mini"
+        ),
+    )
+    parser.add_argument(
+        "--stage-configs-path",
+        type=str,
+        default=None,
+        help="Path to stage config YAML file (default: use built-in step_audio_2.yaml)",
+    )
+
+    # Query configuration
+    parser.add_argument(
+        "--query-type",
+        "-q",
+        type=str,
+        default="audio_to_text",
+        choices=query_map.keys(),
+        help="Query type: audio_to_text (ASR), text_to_audio (TTS), audio_to_audio (voice conversion)",
+    )
+    parser.add_argument(
+        "--audio-path",
+        "-a",
+        type=str,
+        default=None,
+        help="Path to input audio file (for audio_to_text and audio_to_audio)",
+    )
+    parser.add_argument(
+        "--text",
+        "-t",
+        type=str,
+        default=None,
+        help="Text to synthesize (for text_to_audio mode)",
+    )
+    parser.add_argument(
+        "--question",
+        type=str,
+        default=None,
+        help="Question/instruction about the audio",
+    )
+    parser.add_argument(
+        "--sampling-rate",
+        type=int,
+        default=16000,
+        help="Audio sampling rate for input (default: 16000 Hz)",
+    )
+
+    # Generation parameters
+    parser.add_argument(
+        "--max-tokens",
+        type=int,
+        default=1024,
+        help="Maximum tokens to generate in Thinker stage (default: 1024)",
+    )
+    parser.add_argument(
+        "--num-prompts",
+        type=int,
+        default=1,
+        help="Number of times to repeat the prompt (for testing)",
+    )
+
+    # System configuration
+    parser.add_argument(
+        "--enable-stats",
+        action="store_true",
+        default=False,
+        help="Enable detailed statistics logging",
+    )
+    parser.add_argument(
+        "--init-sleep-seconds",
+        type=int,
+        default=20,
+        help="Sleep seconds after starting each stage (default: 20)",
+    )
+    parser.add_argument(
+        "--batch-timeout",
+        type=int,
+        default=5,
+        help="Batch timeout in seconds (default: 5)",
+    )
+    parser.add_argument(
+        "--init-timeout",
+        type=int,
+        default=300,
+        help="Initialization timeout in seconds (default: 300)",
+    )
+    parser.add_argument(
+        "--shm-threshold-bytes",
+        type=int,
+        default=65536,
+        help="Shared memory threshold in bytes (default: 65536)",
+    )
+    parser.add_argument(
+        "--worker-backend",
+        type=str,
+        default="multi_process",
+        choices=["multi_process", "ray"],
+        help="Worker backend (default: multi_process)",
+    )
+    parser.add_argument(
+        "--ray-address",
+        type=str,
+        default=None,
+        help="Ray cluster address (for ray backend)",
+    )
+
+    # Output
+    parser.add_argument(
+        "--output-dir",
+        "-o",
+        type=str,
+        default="output_step_audio2",
+        help="Output directory for generated files (default: output_step_audio2)",
+    )
+
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    main(args)
@@ -0,0 +1,193 @@
+# Step-Audio2 Online Serving
+
+This directory contains examples for running Step-Audio2 with vLLM-Omni's online serving API.
+
+## Installation
+
+Please refer to [README.md](../../../README.md)
+
+## Launch the Server
+
+```bash
+# Async chunk mode (recommended — lower first-packet latency for TTS)
+vllm serve stepfun-ai/Step-Audio2-mini --omni --port 8092 \
+    --stage-configs-path vllm_omni/model_executor/stage_configs/step_audio_2_async_chunk.yaml \
+    --trust-remote-code --enforce-eager
+```
+
+Sequential mode:
+```bash
+vllm serve stepfun-ai/Step-Audio2-mini --omni --port 8092 \
+    --stage-configs-path vllm_omni/model_executor/stage_configs/step_audio_2.yaml \
+    --trust-remote-code --enforce-eager
+```
+
+With local model:
+```bash
+vllm serve /path/to/Step-Audio-2-mini --omni --port 8092 \
+    --trust-remote-code --enforce-eager
+```
+
+## Send Requests
+
+### TTS via `/v1/audio/speech` (Recommended)
+
+```bash
+cd examples/online_serving/step_audio2
+
+# Python client
+python openai_speech_client.py --text "你好世界"
+
+# With custom system prompt
+python openai_speech_client.py --text "Hello, how are you?" \
+    --instructions "You are a friendly assistant."
+
+# Save to specific file
+python openai_speech_client.py --text "你好世界" -o output.wav
+```
+
+Or via curl:
+
+```bash
+curl -X POST http://localhost:8092/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d '{"model":"stepfun-ai/Step-Audio2-mini","input":"你好世界","voice":"default"}' \
+    --output output.wav
+```
+
+This endpoint bypasses the chat template and directly triggers TTS mode. It supports async chunk streaming for low first-packet latency.
+
+**Note**: Speaker voice is controlled by `STEP_AUDIO2_DEFAULT_PROMPT_WAV` env var on the server side.
+
+### Chat Completions (ASR / S2ST)
+
+```bash
+# Audio to Text (ASR)
+python openai_chat_completion_client.py --query-type audio_to_text
+
+# Audio to Audio (S2ST)
+python openai_chat_completion_client.py --query-type audio_to_audio --audio-path /path/to/input.wav
+```
+
+| Argument | Description |
+|----------|-------------|
+| `--query-type`, `-q` | Query type: `audio_to_text`, `text_to_audio`, `audio_to_audio` |
+| `--audio-path`, `-a` | Path to input audio file (local or URL) |
+| `--text`, `-t` | Text to synthesize (for TTS mode) |
+| `--prompt`, `-p` | Custom prompt/question |
+| `--output-dir`, `-o` | Output directory for audio files (default: `output_online`) |
+| `--api-base` | API base URL (default: `http://localhost:8092/v1`) |
+
+### Curl (Chat Completions)
+
+```bash
+# Audio to Text
+bash run_curl.sh audio_to_text
+
+# Text to Audio
+bash run_curl.sh text_to_audio
+
+# Audio to Audio
+bash run_curl.sh audio_to_audio
+```
+
+## Query Types
+
+### 1. Audio to Text (ASR)
+
+Transcribe audio to text.
+
+```bash
+python openai_chat_completion_client.py \
+    --query-type audio_to_text \
+    --audio-path /path/to/speech.wav \
+    --prompt "Transcribe this audio."
+```
+
+### 2. Text to Audio (TTS)
+
+Convert text to speech:
+
+```bash
+# Via speech endpoint (recommended, returns WAV directly)
+python openai_speech_client.py --text "Hello, welcome to Step-Audio2."
+
+# Via chat completions
+python openai_chat_completion_client.py \
+    --query-type text_to_audio \
+    --text "Hello, welcome to Step-Audio2."
+```
+
+### 3. Audio to Audio (S2ST)
+
+Process input audio and generate text transcription + audio output.
+
+```bash
+python openai_chat_completion_client.py \
+    --query-type audio_to_audio \
+    --audio-path /path/to/source.wav
+```
+
+## Output
+
+- **Text output**: Printed to console
+- **Audio output**: Saved to `output_online/audio_0.wav` (24kHz WAV)
+
+## API Format
+
+Step-Audio2 uses the OpenAI-compatible chat completions API:
+
+```json
+{
+  "model": "stepfun-ai/Step-Audio2-mini",
+  "messages": [
+    {
+      "role": "system",
+      "content": [{"type": "text", "text": "Transcribe the audio."}]
+    },
+    {
+      "role": "user",
+      "content": [
+        {"type": "audio_url", "audio_url": {"url": "..."}},
+        {"type": "text", "text": "Please transcribe."}
+      ]
+    }
+  ],
+  "sampling_params_list": [
+    {"temperature": 0.7, "max_tokens": 1024},
+    {"temperature": 0.0, "max_tokens": 1}
+  ]
+}
+```
+
+## Performance
+
+### Async Chunk vs Sequential
+
+Benchmark via `/v1/audio/speech` (4x RTX 3090, 10 prompts, concurrency=1):
+
+| Mode | Mean TTFP | Mean E2E | Mean RTF |
+|------|-----------|----------|----------|
+| Sequential | 4316ms | 4316ms | 0.938 |
+| **Async Chunk** | **1437ms** | 4362ms | 0.949 |
+
+Async chunk reduces TTFP by **67%** by streaming audio token chunks from Thinker to Token2Wav as they are generated. RTF < 1 in both modes (real-time capable).
+
+## Troubleshooting
+
+### Server not responding
+- Check if the server is running: `curl http://localhost:8092/v1/models`
+- Verify the port number matches
+
+### FileNotFoundError: prompt_wav file not found
+- Ensure `default_female.wav` exists at `{model_dir}/assets/default_female.wav`
+- Or set `STEP_AUDIO2_DEFAULT_PROMPT_WAV` environment variable when launching the server
+
+### Audio not generated
+- For TTS, use the `/v1/audio/speech` endpoint (recommended) or `openai_speech_client.py`
+- For chat completions TTS, ensure the prompt ends with `<tts_start>`
+- Check server logs for errors
+
+### Out of memory
+- Reduce `gpu_memory_utilization` in stage configs
+- Use a smaller batch size
@@ -0,0 +1,342 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+OpenAI-compatible client for Step-Audio2 online serving.
+
+Step-Audio2 is a two-stage model:
+  - Stage 0 (Thinker): Audio understanding -> Text + Audio tokens
+  - Stage 1 (Token2Wav): Audio tokens -> Waveform (24kHz)
+
+Usage examples:
+  # Start server first:
+  vllm serve stepfun-ai/Step-Audio2-mini --omni --port 8092
+
+  # Audio to text (ASR)
+  python openai_chat_completion_client.py --query-type audio_to_text
+
+  # Text to audio (TTS)
+  python openai_chat_completion_client.py --query-type text_to_audio --text "Hello world"
+
+  # Audio to audio (Voice conversion)
+  python openai_chat_completion_client.py --query-type audio_to_audio --audio-path input.wav
+"""
+
+import base64
+import logging
+import os
+
+import requests
+from openai import OpenAI
+from vllm.assets.audio import AudioAsset
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
+logger = logging.getLogger(__name__)
+
+SEED = 42
+
+
+def encode_base64_content_from_file(file_path: str) -> str:
+    """Encode a local file to base64 format."""
+    with open(file_path, "rb") as f:
+        return base64.b64encode(f.read()).decode("utf-8")
+
+
+def encode_base64_content_from_url(url: str) -> str:
+    """Encode content from URL to base64 format."""
+    with requests.get(url) as response:
+        response.raise_for_status()
+        return base64.b64encode(response.content).decode("utf-8")
+
+
+def get_audio_url(audio_path: str | None = None) -> str:
+    """Convert audio path to URL format for the API.
+
+    Args:
+        audio_path: Local file path or URL. If None, uses default test audio.
+
+    Returns:
+        Audio URL (either original URL or base64 data URL)
+    """
+    if not audio_path:
+        return AudioAsset("mary_had_lamb").url
+
+    if audio_path.startswith(("http://", "https://")):
+        return audio_path
+
+    if not os.path.exists(audio_path):
+        raise FileNotFoundError(f"Audio file not found: {audio_path}")
+
+    # Detect MIME type
+    audio_path_lower = audio_path.lower()
+    if audio_path_lower.endswith(".wav"):
+        mime_type = "audio/wav"
+    elif audio_path_lower.endswith((".mp3", ".mpeg")):
+        mime_type = "audio/mpeg"
+    elif audio_path_lower.endswith(".ogg"):
+        mime_type = "audio/ogg"
+    elif audio_path_lower.endswith(".flac"):
+        mime_type = "audio/flac"
+    else:
+        mime_type = "audio/wav"
+
+    audio_base64 = encode_base64_content_from_file(audio_path)
+    return f"data:{mime_type};base64,{audio_base64}"
+
+
+def get_audio_to_text_query(
+    audio_path: str | None = None,
+    custom_prompt: str | None = None,
+) -> dict:
+    """Build query for Audio-to-Text (ASR) mode."""
+    question = custom_prompt or "Please transcribe the audio content."
+    audio_url = get_audio_url(audio_path)
+
+    return {
+        "role": "user",
+        "content": [
+            {"type": "audio_url", "audio_url": {"url": audio_url}},
+            {"type": "text", "text": question},
+        ],
+    }
+
+
+def get_text_to_audio_query(
+    text: str | None = None,
+    custom_prompt: str | None = None,
+) -> dict:
+    """Build query for Text-to-Audio (TTS) mode.
+
+    Returns a list of messages: user message + assistant message with <tts_start>.
+    The assistant message triggers TTS mode and must be used with
+    continue_final_message=True so the chat template does not append <|im_end|>.
+    """
+    text_to_speak = text or "Hello, this is a test of Step Audio 2 text to speech."
+
+    return [
+        {
+            "role": "user",
+            "content": [{"type": "text", "text": text_to_speak}],
+        },
+        {
+            "role": "assistant",
+            "content": "<tts_start>",
+        },
+    ]
+
+
+def get_audio_to_audio_query(
+    audio_path: str | None = None,
+    custom_prompt: str | None = None,
+) -> dict:
+    """Build query for Audio-to-Audio (S2ST) mode.
+
+    Returns a list of messages: user message (with audio) + assistant message
+    with <tts_start>. Must be used with continue_final_message=True.
+    """
+    question = custom_prompt or "Please listen to this audio and repeat its content."
+    audio_url = get_audio_url(audio_path)
+
+    return [
+        {
+            "role": "user",
+            "content": [
+                {"type": "audio_url", "audio_url": {"url": audio_url}},
+                {"type": "text", "text": question},
+            ],
+        },
+        {
+            "role": "assistant",
+            "content": "<tts_start>",
+        },
+    ]
+
+
+def get_system_prompt(query_type: str) -> dict:
+    """Get system prompt based on query type."""
+    if query_type == "audio_to_text":
+        system_text = "You are a speech recognition assistant. Transcribe the audio accurately."
+    elif query_type == "text_to_audio":
+        system_text = "You are a text-to-speech assistant. Read the text aloud exactly as provided."
+    else:  # audio_to_audio
+        system_text = "You are an audio processing assistant. Listen and repeat the audio content."
+
+    return {
+        "role": "system",
+        "content": [{"type": "text", "text": system_text}],
+    }
+
+
+QUERY_MAP = {
+    "audio_to_text": get_audio_to_text_query,
+    "text_to_audio": get_text_to_audio_query,
+    "audio_to_audio": get_audio_to_audio_query,
+}
+
+
+def run_step_audio2(args) -> None:
+    """Run Step-Audio2 inference via OpenAI-compatible API."""
+
+    # Initialize OpenAI client
+    client = OpenAI(
+        api_key="EMPTY",
+        base_url=args.api_base,
+    )
+
+    model_name = args.model
+
+    # Sampling parameters for each stage
+    thinker_sampling_params = {
+        "temperature": 0.7,
+        "top_p": 0.9,
+        "top_k": -1,
+        "max_tokens": args.max_tokens,
+        "seed": SEED,
+        "detokenize": True,
+        "repetition_penalty": 1.1 if args.query_type != "audio_to_text" else 1.05,
+    }
+
+    # For ASR, stop at EOS token
+    if args.query_type == "audio_to_text":
+        thinker_sampling_params["stop_token_ids"] = [151645]
+
+    token2wav_sampling_params = {
+        "temperature": 0.0,
+        "top_p": 1.0,
+        "top_k": -1,
+        "max_tokens": 1,
+        "seed": SEED,
+        "detokenize": False,
+    }
+
+    sampling_params_list = [
+        thinker_sampling_params,
+        token2wav_sampling_params,
+    ]
+
+    # Build query based on type
+    query_func = QUERY_MAP[args.query_type]
+    if args.query_type == "audio_to_text":
+        user_message = query_func(audio_path=args.audio_path, custom_prompt=args.prompt)
+    elif args.query_type == "text_to_audio":
+        user_message = query_func(text=args.text, custom_prompt=args.prompt)
+    else:  # audio_to_audio
+        user_message = query_func(audio_path=args.audio_path, custom_prompt=args.prompt)
+
+    system_message = get_system_prompt(args.query_type)
+
+    logger.info(f"Query type: {args.query_type}")
+    logger.info(f"Model: {model_name}")
+    logger.info("Sending request to server...")
+
+    # Build messages list.
+    # TTS and S2ST return a list of messages (user + assistant with <tts_start>);
+    # ASR returns a single user message dict.
+    if isinstance(user_message, list):
+        messages = [system_message] + user_message
+    else:
+        messages = [system_message, user_message]
+
+    # For TTS/S2ST, use continue_final_message so <tts_start> in the
+    # assistant turn is not followed by <|im_end|>.
+    extra_body: dict = {"sampling_params_list": sampling_params_list}
+    if args.query_type in ("text_to_audio", "audio_to_audio"):
+        extra_body["continue_final_message"] = True
+        extra_body["add_generation_prompt"] = False
+
+    # Send request
+    chat_completion = client.chat.completions.create(
+        messages=messages,
+        model=model_name,
+        extra_body=extra_body,
+    )
+
+    # Process response
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    audio_count = 0
+    for choice in chat_completion.choices:
+        if choice.message.audio:
+            # Audio response
+            audio_data = base64.b64decode(choice.message.audio.data)
+            audio_file = os.path.join(args.output_dir, f"audio_{audio_count}.wav")
+            with open(audio_file, "wb") as f:
+                f.write(audio_data)
+            logger.info(f"Audio saved to: {audio_file}")
+            audio_count += 1
+
+        if choice.message.content:
+            # Text response
+            logger.info(f"Text output: {choice.message.content}")
+
+    if audio_count == 0 and args.query_type != "audio_to_text":
+        logger.warning("No audio output received. Check server logs for details.")
+
+    logger.info("Done!")
+
+
+def parse_args():
+    parser = FlexibleArgumentParser(description="OpenAI-compatible client for Step-Audio2 online serving")
+
+    parser.add_argument(
+        "--query-type",
+        "-q",
+        type=str,
+        default="audio_to_text",
+        choices=QUERY_MAP.keys(),
+        help="Query type: audio_to_text (ASR), text_to_audio (TTS), audio_to_audio",
+    )
+    parser.add_argument(
+        "--audio-path",
+        "-a",
+        type=str,
+        default=None,
+        help="Path to input audio file (local path or URL)",
+    )
+    parser.add_argument(
+        "--text",
+        "-t",
+        type=str,
+        default=None,
+        help="Text to synthesize (for text_to_audio mode)",
+    )
+    parser.add_argument(
+        "--prompt",
+        "-p",
+        type=str,
+        default=None,
+        help="Custom prompt/question",
+    )
+    parser.add_argument(
+        "--model",
+        "-m",
+        type=str,
+        default="stepfun-ai/Step-Audio2-mini",
+        help="Model name",
+    )
+    parser.add_argument(
+        "--api-base",
+        type=str,
+        default="http://localhost:8092/v1",
+        help="API base URL",
+    )
+    parser.add_argument(
+        "--max-tokens",
+        type=int,
+        default=1024,
+        help="Maximum tokens for Thinker stage",
+    )
+    parser.add_argument(
+        "--output-dir",
+        "-o",
+        type=str,
+        default="output_online",
+        help="Output directory for audio files",
+    )
+
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    run_step_audio2(args)
@@ -0,0 +1,130 @@
+"""OpenAI-compatible client for Step-Audio2 TTS via /v1/audio/speech endpoint.
+
+Examples:
+    # Basic TTS
+    python openai_speech_client.py --text "你好世界"
+
+    # With custom system prompt
+    python openai_speech_client.py --text "Hello, how are you?" \
+        --instructions "You are a friendly assistant."
+
+    # Save to specific file
+    python openai_speech_client.py --text "你好世界" -o output.wav
+"""
+
+import argparse
+
+import httpx
+
+DEFAULT_API_BASE = "http://localhost:8092"
+DEFAULT_API_KEY = "EMPTY"
+
+
+def run_tts_generation(args) -> None:
+    """Run TTS generation via /v1/audio/speech API."""
+    payload = {
+        "model": args.model,
+        "input": args.text,
+        "voice": args.voice,
+        "response_format": args.response_format,
+    }
+
+    if args.instructions:
+        payload["instructions"] = args.instructions
+
+    print(f"Model: {args.model}")
+    print(f"Text: {args.text}")
+    print(f"Voice: {args.voice}")
+    print("Generating audio...")
+
+    api_url = f"{args.api_base}/v1/audio/speech"
+    headers = {
+        "Content-Type": "application/json",
+        "Authorization": f"Bearer {args.api_key}",
+    }
+
+    with httpx.Client(timeout=300.0) as client:
+        response = client.post(api_url, json=payload, headers=headers)
+
+    if response.status_code != 200:
+        print(f"Error: {response.status_code}")
+        print(response.text)
+        return
+
+    try:
+        text = response.content.decode("utf-8")
+        if text.startswith('{"error"'):
+            print(f"Error: {text}")
+            return
+    except UnicodeDecodeError:
+        pass
+
+    output_path = args.output or "tts_output.wav"
+    with open(output_path, "wb") as f:
+        f.write(response.content)
+    print(f"Audio saved to: {output_path}")
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="Step-Audio2 TTS client via /v1/audio/speech",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog=__doc__,
+    )
+    parser.add_argument(
+        "--api-base",
+        type=str,
+        default=DEFAULT_API_BASE,
+        help=f"API base URL (default: {DEFAULT_API_BASE})",
+    )
+    parser.add_argument(
+        "--api-key",
+        type=str,
+        default=DEFAULT_API_KEY,
+        help="API key (default: EMPTY)",
+    )
+    parser.add_argument(
+        "--model",
+        "-m",
+        type=str,
+        default="stepfun-ai/Step-Audio2-mini",
+        help="Model name/path",
+    )
+    parser.add_argument(
+        "--text",
+        type=str,
+        required=True,
+        help="Text to synthesize",
+    )
+    parser.add_argument(
+        "--voice",
+        type=str,
+        default="default",
+        help="Voice name (default: default)",
+    )
+    parser.add_argument(
+        "--instructions",
+        type=str,
+        default=None,
+        help="System prompt for the Thinker stage",
+    )
+    parser.add_argument(
+        "--response-format",
+        type=str,
+        default="wav",
+        choices=["wav", "pcm"],
+        help="Audio output format (default: wav)",
+    )
+    parser.add_argument(
+        "--output",
+        "-o",
+        type=str,
+        default=None,
+        help="Output audio file path (default: tts_output.wav)",
+    )
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    run_tts_generation(args)
@@ -0,0 +1,149 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+# Step-Audio2 curl client for online serving
+# Usage: bash run_curl.sh [audio_to_text|text_to_audio|audio_to_audio]
+
+QUERY_TYPE="${1:-audio_to_text}"
+API_BASE="${API_BASE:-http://localhost:8092}"
+
+# Validate query type
+if [[ ! "$QUERY_TYPE" =~ ^(audio_to_text|text_to_audio|audio_to_audio)$ ]]; then
+    echo "Error: Invalid query type '$QUERY_TYPE'"
+    echo "Usage: $0 [audio_to_text|text_to_audio|audio_to_audio]"
+    echo "  audio_to_text: Speech recognition (ASR)"
+    echo "  text_to_audio: Text-to-speech (TTS)"
+    echo "  audio_to_audio: Voice conversion"
+    exit 1
+fi
+
+SEED=42
+
+# Default test audio URL
+MARY_HAD_LAMB_URL="https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/mary_had_lamb.ogg"
+
+# Sampling parameters for Thinker stage
+thinker_sampling_params='{
+  "temperature": 0.7,
+  "top_p": 0.9,
+  "top_k": -1,
+  "max_tokens": 1024,
+  "seed": 42,
+  "detokenize": true,
+  "repetition_penalty": 1.05
+}'
+
+# Sampling parameters for Token2Wav stage
+token2wav_sampling_params='{
+  "temperature": 0.0,
+  "top_p": 1.0,
+  "top_k": -1,
+  "max_tokens": 1,
+  "seed": 42,
+  "detokenize": false
+}'
+
+# Build request based on query type
+case "$QUERY_TYPE" in
+  audio_to_text)
+    system_content='[{"type": "text", "text": "You are a speech recognition assistant. Transcribe the audio accurately."}]'
+    user_content='[
+      {"type": "audio_url", "audio_url": {"url": "'"$MARY_HAD_LAMB_URL"'"}},
+      {"type": "text", "text": "Please transcribe this audio."}
+    ]'
+    # Add stop token for ASR
+    thinker_sampling_params='{
+      "temperature": 0.7,
+      "top_p": 0.9,
+      "top_k": -1,
+      "max_tokens": 1024,
+      "seed": 42,
+      "detokenize": true,
+      "repetition_penalty": 1.05,
+      "stop_token_ids": [151645]
+    }'
+    ;;
+  text_to_audio)
+    system_content='[{"type": "text", "text": "You are a text-to-speech assistant. Read the text aloud exactly as provided."}]'
+    user_content='[
+      {"type": "text", "text": "Hello, this is a test of Step Audio 2 text to speech synthesis.<tts_start>"}
+    ]'
+    thinker_sampling_params='{
+      "temperature": 0.7,
+      "top_p": 0.9,
+      "top_k": -1,
+      "max_tokens": 1024,
+      "seed": 42,
+      "detokenize": true,
+      "repetition_penalty": 1.1
+    }'
+    ;;
+  audio_to_audio)
+    system_content='[{"type": "text", "text": "You are an audio processing assistant. Listen and repeat the audio content."}]'
+    user_content='[
+      {"type": "audio_url", "audio_url": {"url": "'"$MARY_HAD_LAMB_URL"'"}},
+      {"type": "text", "text": "Please listen to this audio and repeat its content.<tts_start>"}
+    ]'
+    thinker_sampling_params='{
+      "temperature": 0.7,
+      "top_p": 0.9,
+      "top_k": -1,
+      "max_tokens": 1024,
+      "seed": 42,
+      "detokenize": true,
+      "repetition_penalty": 1.1
+    }'
+    ;;
+esac
+
+sampling_params_list='[
+  '"$thinker_sampling_params"',
+  '"$token2wav_sampling_params"'
+]'
+
+echo "Query type: $QUERY_TYPE"
+echo "API base: $API_BASE"
+echo "Sending request..."
+echo ""
+
+output=$(curl -sS -X POST "${API_BASE}/v1/chat/completions" \
+    -H "Content-Type: application/json" \
+    -d @- <<EOF
+{
+  "model": "stepfun-ai/Step-Audio2-mini",
+  "sampling_params_list": $sampling_params_list,
+  "messages": [
+    {
+      "role": "system",
+      "content": $system_content
+    },
+    {
+      "role": "user",
+      "content": $user_content
+    }
+  ]
+}
+EOF
+)
+
+# Extract and display text content
+text_content=$(echo "$output" | jq -r '.choices[0].message.content // empty')
+if [[ -n "$text_content" ]]; then
+    echo "Text output: $text_content"
+fi
+
+# Check for audio content
+audio_data=$(echo "$output" | jq -r '.choices[0].message.audio.data // empty')
+if [[ -n "$audio_data" ]]; then
+    echo "Audio output received (base64 encoded)"
+    echo "To save audio, use the Python client or decode the base64 data"
+fi
+
+# Check for errors
+error=$(echo "$output" | jq -r '.error // empty')
+if [[ -n "$error" ]]; then
+    echo "Error: $error"
+fi
+
+echo ""
+echo "Done!"
@@ -281,6 +281,7 @@ extend-ignore-identifiers-re = [
    ".*[Oo]no_[Aa]nna.*",
    ".*cann.*",
    ".*[Nn][Dd]3.*",
+    ".*updator.*",
    "byt5.*",
 ]
 extend-ignore-re = [
@@ -54,7 +54,6 @@ recipes/
 | [`StabilityAI/Stable-Diffusion-3.5.md`](./StabilityAI/Stable-Diffusion-3.5.md) | Text-to-image serving (SD 3.5-medium and SD 3.5-large) | 1x RTX A6000 48GB |
 | [`zai-org/GLM-TTS.md`](./zai-org/GLM-TTS.md) | Online serving for Chinese/English zero-shot voice-cloned TTS | 1x A40 48GB |
 | [`GLM/GLM-Image.md`](./GLM/GLM-Image.md) | Online serving for image generation | 1x A800 80GB / 2x A800 80GB |
-|

 Within a single recipe file, include different hardware support sections such
 as `GPU`, `ROCm`, and `NPU`, and add concrete tested configurations like
@@ -1,82 +0,0 @@
-# Ming-omni-tts 0.5B
-
-> Offline and online dense Ming text-to-speech on ROCm
-
-## Summary
-
- Vendor: inclusionAI
- Model: `inclusionAI/Ming-omni-tts-0.5B`
- Task: Text-to-speech with style, dialect, cloning, and multi-speaker controls
- Mode: Online serving via the OpenAI-compatible `/v1/audio/speech` API; offline inference
- Maintainer: Community
-
-## References
-
- [Model card](https://huggingface.co/inclusionAI/Ming-omni-tts-0.5B)
- [Upstream repository](https://github.com/inclusionAI/Ming-omni-tts)
- [Offline example](../../examples/offline_inference/text_to_speech/ming_tts/)
- [Online example](../../examples/online_serving/text_to_speech/ming_tts/)
-
-## Hardware Support
-
-## ROCm
-
-### 1x AMD `gfx942`
-
-#### Environment
-
- OS: Ubuntu 22.04.5 LTS, x86_64
- Python: 3.12.13
- ROCm / HIP: 7.2.53211
- vLLM version: 0.22.0
- vLLM-Omni version or commit: 0.1.dev1873 / `99c35c410`
- Docker image: `vllm/vllm-omni-rocm:v0.22.0`
-
-#### Command
-
-From the vLLM-Omni repository root:
-
-```bash
-docker run --rm \
-    --group-add=video \
-    --ipc=host \
-    --cap-add=SYS_PTRACE \
-    --security-opt seccomp=unconfined \
-    --device /dev/kfd \
-    --device /dev/dri \
-    -v ~/.cache/huggingface:/root/.cache/huggingface \
-    -v "$PWD":/app/vllm-omni \
-    -w /app/vllm-omni \
-    -e VLLM_ROCM_USE_AITER=0 \
-    -p 8091:8091 \
-    vllm/vllm-omni-rocm:v0.22.0 \
-    --model inclusionAI/Ming-omni-tts-0.5B \
-    --deploy-config vllm_omni/deploy/ming_tts.yaml \
-    --omni \
-    --port 8091 \
-    --enforce-eager
-```
-
-#### Verification
-
-```bash
-python examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py \
-    --text "我觉得社会企业同个人都有责任" \
-    --instruction-json '{"方言":"广粤话"}' \
-    --ref-audio /path/to/yue_prompt.wav \
-    --max-new-tokens 200 \
-    --output dialect.wav
-```
-
-`--ref-audio` matches upstream `use_spk_emb=True`; do not add `--ref-text`
-for the dialect case.
-
-## Notes
-
- The official ROCm image includes the platform dependencies.
- See the [ROCm installation guide](../../docs/getting_started/installation/gpu.md)
-  for interactive and source-build workflows.
- The [offline example](../../examples/offline_inference/text_to_speech/ming_tts/)
-  can be run from an interactive container.
- The tested environment uses `--enforce-eager`.
- Output is mono 44.1 kHz audio.
@@ -0,0 +1,181 @@
+# Ming-omni-tts
+
+> Online serving for voice synthesis, cloning, and design
+
+## Summary
+
+- Vendor: inclusionAI
+- Model: `inclusionAI/Ming-omni-tts-0.5B`
+- Task: Text-to-speech with style, dialect, cloning, and multi-speaker controls
+- Mode: Online serving via the OpenAI-compatible `/v1/audio/speech` API
+- Maintainer: Community
+
+## References
+
+- [Huggingface Model card](https://huggingface.co/inclusionAI/Ming-omni-tts-0.5B)
+- Upstream repository [inclusionAI/Ming-omni-tts](https://github.com/inclusionAI/Ming-omni-tts)
+
+
+## Hardware Support
+
+This recipe documents a validated ROCm configuration and a CUDA configuration for the dense 0.5B two-stage TTS pipeline deployment.
+Other hardware is welcome as community validation lands.
+
+## CUDA
+
+### 1x H100 80GB
+
+#### Environment
+
+- OS: Linux
+- Python: 3.10+
+- CUDA Driver Version: 590.48.01
+- CUDA: 13.0
+- vLLM version: 0.22.0
+- vLLM-Omni version or commit: 0342827d
+
+#### Command
+
+Launch the two-stage talker:
+
+```bash
+vllm-omni serve inclusionAI/Ming-omni-tts-0.5B \
+    --deploy-config vllm_omni/deploy/ming_tts.yaml \
+    --omni \
+    --port 8091 \
+    --enforce-eager
+```
+
+#### Verification
+
+Basic synthesis (save the WAV bytes):
+
+```bash
+curl -X POST http://localhost:8091/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d '{
+      "model": "inclusionAI/Ming-omni-tts-0.5B",
+      "input": "你好，这是 Ming 在线语音合成测试。",
+      "response_format": "wav"
+    }' --output ming_tts_basic.wav
+```
+
+The following requests condition on reference clips from upstream audio assets. The server resolves `ref_audio` as an `http(s)://` URL, a `file://` path, or an inline base64 `data:` URL, so the examples point it straight at the public cookbook fixtures in the vendor repo's [`inclusionAI/Ming-omni-tts/tree/main/data/wavs`](https://github.com/inclusionAI/Ming-omni-tts/tree/main/data/wavs) directory:
+
+```bash
+BASE=https://raw.githubusercontent.com/inclusionAI/Ming-omni-tts/main/data/wavs
+```
+
+And note that the server latency will include automatically resolving and downloading input files.
+
+Reference-audio zero-shot cloning:
+
+```bash
+REF_AUDIO="$BASE/10002287-00000094.wav"
+jq -n --arg ref_audio "$REF_AUDIO" \
+    '{
+       model: "inclusionAI/Ming-omni-tts-0.5B",
+       input: "我们的愿景是构建未来服务业的数字化基础设施，为世界带来更多微小而美好的改变。",
+       ref_audio: $ref_audio,
+       ref_text: "在此奉劝大家别乱打美白针。",
+       max_new_tokens: 200,
+       response_format: "wav"
+     }' \
+| curl -X POST http://localhost:8091/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d @- --output ming_tts_zero_shot.wav
+```
+
+Podcast-style multi-speaker generation uses one reference clip plus one transcript per speaker, with the `speaker_N` labels in `input`/`ref_text` lined up with the reference order:
+
+```bash
+REF_A="$BASE/CTS-CN-F2F-2019-11-11-423-012-A.wav"
+REF_B="$BASE/CTS-CN-F2F-2019-11-11-423-012-B.wav"
+
+INPUT=' speaker_1:你可以说一下，就大概说一下，可能虽然我也不知道，我看过那部电影没有。
+ speaker_2:就是那个叫什么，变相一节课的嘛。
+ speaker_1:嗯。
+ speaker_2:一部搞笑的电影。
+ speaker_1:一部搞笑的。'
+
+REF_TEXT=' speaker_1:并且我们还要进行每个月还要考核 笔试的话还要进行笔试，做个，当服务员还要去笔试了
+ speaker_2:对啊，这真的很奇怪，就是 单纯的因，单纯自己工资不高，只是因为可能人家那个店比较出名一点，就对你苛刻要求'
+
+jq -n \
+    --arg input "$INPUT" \
+    --arg ref_text "$REF_TEXT" \
+    --arg a "$REF_A" \
+    --arg b "$REF_B" \
+    '{
+       model: "inclusionAI/Ming-omni-tts-0.5B",
+       input: $input,
+       ref_audio: [$a, $b],
+       ref_text: $ref_text,
+       response_format: "wav"
+     }' \
+| curl -X POST http://localhost:8091/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d @- --output ming_tts_podcast_style.wav
+```
+
+## ROCm
+
+### 1x AMD MI300X
+
+#### Environment
+
+- OS: Ubuntu 22.04.5 LTS, x86_64
+- Python: 3.12.13
+- ROCm / HIP: 7.2.53211
+- vLLM version: 0.22.0
+- vLLM-Omni version or commit: 0.1.dev1873 / `99c35c410`
+- Docker image: `vllm/vllm-omni-rocm:v0.22.0`
+
+#### Command
+
+From the vLLM-Omni repository root:
+
+```bash
+docker run --rm \
+    --group-add=video \
+    --ipc=host \
+    --cap-add=SYS_PTRACE \
+    --security-opt seccomp=unconfined \
+    --device /dev/kfd \
+    --device /dev/dri \
+    -v ~/.cache/huggingface:/root/.cache/huggingface \
+    -v "$PWD":/app/vllm-omni \
+    -w /app/vllm-omni \
+    -e VLLM_ROCM_USE_AITER=0 \
+    -p 8091:8091 \
+    vllm/vllm-omni-rocm:v0.22.0 \
+    --model inclusionAI/Ming-omni-tts-0.5B \
+    --deploy-config vllm_omni/deploy/ming_tts.yaml \
+    --omni \
+    --port 8091 \
+    --enforce-eager
+```
+
+#### Verification
+
+```bash
+python examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py \
+    --text "我觉得社会企业同个人都有责任" \
+    --instruction-json '{"方言":"广粤话"}' \
+    --ref-audio /path/to/yue_prompt.wav \
+    --max-new-tokens 200 \
+    --output dialect.wav
+```
+
+`--ref-audio` matches upstream `use_spk_emb=True`; do not add `--ref-text`
+for the dialect case.
+
+## Notes
+
+- The official ROCm image includes the platform dependencies.
+- See the [ROCm installation guide](../../docs/getting_started/installation/gpu.md) for interactive and source-build workflows.
+- The reference clips above are a subset of the upstream [`inclusionAI/Ming-omni-tts/tree/main/data/wavs`](https://github.com/inclusionAI/Ming-omni-tts/tree/main/data/wavs) cookbook fixtures.
+- To clone from local audio instead of a URL, pass a `file://` URI (launch the server with `--allowed-local-media-path <dir>`), or base64-encode the clip into a `data:` URL written to a payload file and send it with `curl -d @payload.json`. Putting the base64 directly on the `curl` command line overflows the OS argument limit (`bash: /usr/bin/curl: Argument list too long`).
+- The tested environment uses `--enforce-eager`.
+- Non-streaming responses return WAV bytes; streaming responses return PCM.
+- Output is mono 44.1 kHz audio.
@@ -1,10 +1,15 @@
 # Common dependencies for all platforms
-# transformers: keep aligned with upstream vLLM constraint.
+# transformers: kept exactly aligned with the upstream vLLM constraint.
 # vLLM pins ">= 4.56.0, != 5.0.*, != 5.1.*, != 5.2.*, != 5.3.*, != 5.4.*, != 5.5.0".
-# 5.9.0 breaks models that define _keys_to_ignore_on_load_unexpected as a list
-# (e.g. OpenMOSS-Team/MOSS-TTS-Nano) because it changed
-# _adjust_missing_and_unexpected_keys from list+list to set|set.
-transformers >= 4.56.0, != 5.0.*, != 5.1.*, != 5.2.*, != 5.3.*, != 5.4.*, != 5.5.0, < 5.9.0
+# The previous extra "< 5.9.0" exclusion was removed when re-aligning with
+# upstream: transformers 5.9 changed _adjust_missing_and_unexpected_keys from
+# list+list to set|set, breaking models that declare
+# _keys_to_ignore_on_load_unexpected as a list (e.g. OpenMOSS-Team/MOSS-TTS-Nano)
+# with "TypeError: unsupported operand type(s) for |: 'list' and 'set'".
+# That breakage is now handled in code instead of by pinning:
+# MossTTSNanoForGeneration wraps weight loading in
+# transformers_keys_to_ignore_compat (vllm_omni/model_executor/models/utils.py).
+transformers >= 4.56.0, != 5.0.*, != 5.1.*, != 5.2.*, != 5.3.*, != 5.4.*, != 5.5.0
 av>=14.0.0
 omegaconf>=2.3.0
 diffusers==0.38.0
@@ -1,3 +1,3 @@
 -r common.txt
 onnxruntime-cann>=1.23.2
-torchaudio==2.9.0
+torchaudio==2.10.0
@@ -42,10 +42,12 @@ def test_bench_serve_cli_mocks_http_request(tmp_path: Path):
                        yield chunk

            class MockResponse:
-                def __init__(self):
-                    self.status = 200
-                    self.reason = "OK"
-                    self.content = _Content(SSE_CHUNKS)
+                def __init__(self, status=200, reason="OK", chunks=None):
+                    self.status = status
+                    self.reason = reason
+                    self.content = _Content(chunks or SSE_CHUNKS)
+                    self._json_data = None
+                    self._text_data = None

                async def __aenter__(self):
                    return self
@@ -53,20 +55,49 @@ def test_bench_serve_cli_mocks_http_request(tmp_path: Path):
                async def __aexit__(self, exc_type, exc, tb):
                    return False

+                async def text(self):
+                    if self._text_data is not None:
+                        return self._text_data
+                    return json.dumps({"tokens": [1, 2, 3, 4, 5]})
+
+                async def json(self):
+                    if self._json_data is not None:
+                        return self._json_data
+                    return {"tokens": [1, 2, 3, 4, 5]}
+
+                def raise_for_status(self):
+                    pass
+
            class MockClientSession:
                def __init__(self, *args, **kwargs):
                    pass

+                async def __aenter__(self):
+                    return self
+
+                async def __aexit__(self, exc_type, exc, tb):
+                    return False
+
                def post(self, url=None, *args, **kwargs):
                    if url is not None:
                        POST_CALLS.append(url)
-                    return MockResponse()
+                    return MockResponse(status=200)
+
+                async def get(self, url=None, *args, **kwargs):
+                    if url is not None:
+                        POST_CALLS.append(url)
+                    return MockResponse(status=200)

                async def close(self):
                    return None

-            import vllm_omni.benchmarks.patch.patch as patch_mod
+            # Patch globally so modules loaded after this also see the mock
+            import aiohttp
+            aiohttp.ClientSession = MockClientSession
+            aiohttp.TCPConnector = lambda *args, **kwargs: object()

+            # Also patch in the patch module
+            import vllm_omni.benchmarks.patch.patch as patch_mod
            patch_mod.aiohttp.ClientSession = MockClientSession
            patch_mod.aiohttp.TCPConnector = lambda *args, **kwargs: object()

@@ -136,8 +167,13 @@ def test_bench_serve_cli_mocks_http_request(tmp_path: Path):

    expected_url = f"http://127.0.0.1:{port}/v1/chat/completions"
    requested_urls = calls["requested_urls"]
-    sent_requests = len(requested_urls)
+    # Count only benchmark requests (to the chat completions endpoint),
+    # excluding tokenize/detokenize alignment calls from the upstream.
+    bench_requests = [url for url in requested_urls if url == expected_url]
+    sent_bench_requests = len(bench_requests)

-    assert result["completed"] == sent_requests == num_prompts
-    assert requested_urls
-    assert all(url == expected_url for url in requested_urls), f"Unexpected target URLs: {requested_urls}"
+    assert result["completed"] == sent_bench_requests == num_prompts, (
+        f"completed={result['completed']}, bench_requests={sent_bench_requests}, all_requested_urls={requested_urls}"
+    )
+    assert bench_requests
+    assert all(url == expected_url for url in bench_requests), f"Unexpected target URLs: {bench_requests}"
@@ -578,12 +578,27 @@ def test_rmsnorm_fused_native_parity_2d_input():


 def test_rmsnorm_fused_native_parity_custom_weight():
-    """Verify fused and native RMSNorm parity with non-default weights."""
+    """Verify fused and native RMSNorm parity with non-default weights.
+
+    Note: bfloat16 requires wider tolerance because the fused CUDA kernel
+    (vllm._custom_ops.rms_norm) computes in bf16 precision internally,
+    while forward_native uses float32 intermediates. With random custom
+    weights (not the default ones-weight), this precision gap exceeds 1e-3.
+
+    The gap surfaced after upstream vLLM routed RMSNorm through the vLLM IR
+    op (ir.ops.rms_norm) — see upstream commit 40bb17502
+    "[vLLM IR] 1/N Implement IR skeleton and rms_norm op (#33825)" — which
+    changed which fused kernel/precision path the layer dispatches to.
+    """
    from vllm_omni.diffusion.layers.norm import RMSNorm

    hidden_size = 64
    eps = 1e-6

+    # Per-dtype tolerances: bf16 fused kernel has inherent precision limit
+    atol_map = {torch.float32: 1e-3, torch.float16: 1e-3, torch.bfloat16: 2e-2}
+    rtol_map = {torch.float32: 1e-3, torch.float16: 1e-3, torch.bfloat16: 1e-2}
+
    for dtype in [torch.float32, torch.float16, torch.bfloat16]:
        torch.manual_seed(789)
        norm = RMSNorm(hidden_size, eps=eps).cuda()
@@ -594,4 +609,4 @@ def test_rmsnorm_fused_native_parity_custom_weight():
        out_fused = norm._forward_fused(x)
        out_native = norm.forward_native(x)

-        torch.testing.assert_close(out_fused, out_native, atol=1e-3, rtol=1e-3)
+        torch.testing.assert_close(out_fused, out_native, atol=atol_map[dtype], rtol=rtol_map[dtype])
@@ -9,6 +9,7 @@ import torch
 from vllm_omni.diffusion.data import DiffusionOutput
 from vllm_omni.diffusion.ipc import (
    _SHM_TENSOR_THRESHOLD,
+    DIFFUSION_RPC_RESULT_ENVELOPE,
    _pack_value_if_large,
    _unpack_if_shm_handle,
    pack_diffusion_output_shm,
@@ -81,6 +82,31 @@ def test_diffusion_output_list_tensors_round_trip_through_shm() -> None:
    torch.testing.assert_close(output.output[1], frames[1])


+def test_rpc_result_envelope_diffusion_output_round_trips_through_shm() -> None:
+    tensor = torch.arange(300_000, dtype=torch.float32)
+    envelope = {
+        "type": DIFFUSION_RPC_RESULT_ENVELOPE,
+        "result": DiffusionOutput(output=tensor),
+        "rank_statuses": [{"rank": 0, "ok": True}],
+    }
+
+    packed = pack_diffusion_output_shm(envelope)
+
+    assert packed is envelope
+    result = packed["result"]
+    assert isinstance(result, DiffusionOutput)
+    assert result.output["__tensor_shm__"] is True
+    assert packed["rank_statuses"] == [{"rank": 0, "ok": True}]
+
+    unpacked = unpack_diffusion_output_shm(packed)
+
+    assert unpacked is envelope
+    result = unpacked["result"]
+    assert isinstance(result, DiffusionOutput)
+    torch.testing.assert_close(result.output, tensor)
+    assert unpacked["rank_statuses"] == [{"rank": 0, "ok": True}]
+
+
 def test_pack_value_keeps_tensor_at_threshold_inline() -> None:
    tensor = torch.arange(
        _SHM_TENSOR_THRESHOLD // torch.empty((), dtype=torch.float32).element_size(),
@@ -17,8 +17,10 @@ from vllm.v1.engine.exceptions import EngineDeadError
 from vllm_omni.diffusion.data import DiffusionOutput
 from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
 from vllm_omni.diffusion.executor.multiproc_executor import MultiprocDiffusionExecutor
+from vllm_omni.diffusion.ipc import DIFFUSION_RPC_RESULT_ENVELOPE
 from vllm_omni.diffusion.sched import RequestScheduler
 from vllm_omni.diffusion.stage_diffusion_proc import StageDiffusionProc
+from vllm_omni.diffusion.worker.diffusion_worker import WorkerProc
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput

@@ -305,6 +307,80 @@ class TestSerialEngineOperations:
        assert len(results) == 1
        assert results[0].error == "rank0"

+    def test_collective_rpc_all_rank_status_error_propagation(self):
+        engine, _, _, res_q = _make_engine(num_gpus=2)
+
+        res_q.put(
+            {
+                "type": DIFFUSION_RPC_RESULT_ENVELOPE,
+                "method": "add_lora",
+                "result": True,
+                "rank_statuses": [
+                    {"rank": 0, "ok": True, "bool_result": True},
+                    {
+                        "rank": 1,
+                        "ok": False,
+                        "error": "rank1 boom",
+                        "error_type": "RuntimeError",
+                        "traceback": "rank1 traceback",
+                    },
+                ],
+            }
+        )
+
+        with pytest.raises(RuntimeError) as excinfo:
+            engine.collective_rpc("add_lora")
+        error = str(excinfo.value)
+        assert "rank 1" in error
+        assert "rank1 boom" in error
+        assert "rank1 traceback" in error
+
+    def test_collective_rpc_all_rank_bool_false_is_aggregated(self):
+        engine, _, _, res_q = _make_engine(num_gpus=2)
+
+        res_q.put(
+            {
+                "type": DIFFUSION_RPC_RESULT_ENVELOPE,
+                "method": "remove_lora",
+                "result": True,
+                "rank_statuses": [
+                    {"rank": 0, "ok": True, "bool_result": True},
+                    {"rank": 1, "ok": True, "bool_result": False},
+                ],
+            }
+        )
+
+        assert engine.collective_rpc("remove_lora") == [False]
+
+    def test_collective_rpc_collects_rank_status_only_for_control_plane_all_rank_rpc(self):
+        executor, req_q, res_q = _make_executor(num_gpus=2)
+
+        res_q.put(_tagged_output("forward"))
+        result = executor.collective_rpc(
+            "execute_stepwise",
+            unique_reply_rank=0,
+            exec_all_ranks=True,
+        )
+        forward_rpc = req_q.get_nowait()
+
+        assert result.error == "forward"
+        assert forward_rpc["exec_all_ranks"] is True
+        assert forward_rpc["collect_rank_status"] is False
+
+        res_q.put(
+            {
+                "type": DIFFUSION_RPC_RESULT_ENVELOPE,
+                "method": "remove_lora",
+                "result": True,
+                "rank_statuses": [{"rank": 0, "ok": True, "bool_result": True}],
+            }
+        )
+        assert executor.collective_rpc("remove_lora") == [True]
+        control_rpc = req_q.get_nowait()
+
+        assert control_rpc["exec_all_ranks"] is True
+        assert control_rpc["collect_rank_status"] is True
+
    def test_serial_add_req_then_collective_rpc(self):
        engine, _, req_q, res_q = _make_engine()
        wt = _start_worker(req_q, res_q, count=2)
@@ -348,6 +424,115 @@ class TestSerialEngineOperations:
            engine.collective_rpc("anything")


+class TestWorkerProcRpcRankStatus:
+    def _make_worker_proc(self, has_result_mq: bool = True):
+        proc = object.__new__(WorkerProc)
+        proc.gpu_id = 0
+        proc.result_mq = object() if has_result_mq else None
+        proc.worker = SimpleNamespace(execute_method=Mock(return_value=True))
+        return proc
+
+    def test_execute_rpc_returns_rank_status_envelope(self, monkeypatch):
+        proc = self._make_worker_proc()
+
+        monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True)
+        monkeypatch.setattr(torch.distributed, "get_world_size", lambda: 2)
+
+        def _all_gather_object(out, local):
+            out[0] = local
+            out[1] = {
+                "rank": 1,
+                "ok": False,
+                "error": "rank1 boom",
+                "error_type": "RuntimeError",
+                "traceback": "trace",
+                "bool_result": None,
+            }
+
+        monkeypatch.setattr(torch.distributed, "all_gather_object", _all_gather_object)
+
+        result, should_reply = proc.execute_rpc(
+            {
+                "method": "remove_lora",
+                "args": (),
+                "kwargs": {},
+                "output_rank": 0,
+                "exec_all_ranks": True,
+                "collect_rank_status": True,
+            }
+        )
+
+        assert should_reply is True
+        assert result["type"] == DIFFUSION_RPC_RESULT_ENVELOPE
+        assert result["result"] is True
+        assert result["rank_statuses"][0]["rank"] == 0
+        assert result["rank_statuses"][1]["rank"] == 1
+        assert result["rank_statuses"][1]["ok"] is False
+
+    def test_execute_rpc_local_exception_is_reported_in_envelope(self, monkeypatch):
+        proc = self._make_worker_proc()
+        proc.worker.execute_method = Mock(side_effect=RuntimeError("local boom"))
+        monkeypatch.setattr(torch.distributed, "is_initialized", lambda: False)
+
+        result, should_reply = proc.execute_rpc(
+            {
+                "method": "add_lora",
+                "args": (),
+                "kwargs": {},
+                "output_rank": 0,
+                "exec_all_ranks": True,
+                "collect_rank_status": True,
+            }
+        )
+
+        assert should_reply is True
+        assert result["type"] == DIFFUSION_RPC_RESULT_ENVELOPE
+        assert len(result["rank_statuses"]) == 1
+        status = result["rank_statuses"][0]
+        assert status["rank"] == 0
+        assert status["ok"] is False
+        assert status["error"] == "local boom"
+        assert status["error_type"] == "RuntimeError"
+        assert status["bool_result"] is None
+        assert "local boom" in status["traceback"]
+
+    def test_execute_rpc_rejects_collect_rank_status_without_all_ranks(self):
+        proc = self._make_worker_proc()
+
+        with pytest.raises(ValueError, match="collect_rank_status requires exec_all_ranks=True"):
+            proc.execute_rpc(
+                {
+                    "method": "ping",
+                    "args": (),
+                    "kwargs": {},
+                    "output_rank": 0,
+                    "exec_all_ranks": False,
+                    "collect_rank_status": True,
+                }
+            )
+
+        proc.worker.execute_method.assert_not_called()
+
+    def test_execute_rpc_non_collect_exception_preserves_original_type(self):
+        proc = self._make_worker_proc()
+        original = ValueError("local boom")
+        proc.worker.execute_method = Mock(side_effect=original)
+
+        with pytest.raises(ValueError) as excinfo:
+            proc.execute_rpc(
+                {
+                    "method": "bad",
+                    "args": (),
+                    "kwargs": {},
+                    "output_rank": 0,
+                    "exec_all_ranks": False,
+                    "collect_rank_status": False,
+                }
+            )
+
+        assert excinfo.value is original
+
+
 # ───────── error handling: EngineDeadError propagation through layers ─────


@@ -200,9 +200,9 @@ def test_save_async(build_adapter):
    request = _req("req-1", RequestStatus.WAITING, external_req_id="external-1")

    adapter.custom_process_next_stage_input_func = lambda **kwargs: {"x": [1], "finished": False}
-    adapter.save_async(pooling_output=None, request=request)
+    adapter.save_async(multimodal_output=None, request=request)
    adapter.custom_process_next_stage_input_func = lambda **kwargs: {}
-    adapter.save_async(pooling_output=None, request=request)
+    adapter.save_async(multimodal_output=None, request=request)

    task = adapter._pending_save_reqs.popleft()
    assert task["is_finished"] is False
@@ -214,7 +214,7 @@ def test_save_async_uses_confirmed_tokens_for_async_scheduler_watermark(build_ad
    request.num_computed_tokens = 10
    request.num_output_placeholders = 2

-    adapter.save_async(pooling_output=None, request=request)
+    adapter.save_async(multimodal_output=None, request=request)

    assert adapter.requests_num_chunks_sent["external-async"] == 8
    assert len(adapter._pending_save_reqs) == 1
@@ -236,7 +236,7 @@ def test_send_single_request_struct_without_meta_does_not_crash(build_adapter, m
    monkeypatch.setattr(adapter, "cleanup", lambda *a, **kw: cleanup_calls.append((a, kw)))

    adapter._send_single_request(
-        {"pooling_output": None, "request": request, "is_finished": False, "is_segment_finished": False}
+        {"multimodal_output": None, "request": request, "is_finished": False, "is_segment_finished": False}
    )

    assert cleanup_calls == []  # no terminal cleanup; meta.finished is false
@@ -256,7 +256,7 @@ def test_send_single_request_empty_struct_goes_on_wire(build_adapter, monkeypatc
    monkeypatch.setattr(adapter, "cleanup", lambda *a, **kw: None)

    adapter._send_single_request(
-        {"pooling_output": None, "request": request, "is_finished": False, "is_segment_finished": False}
+        {"multimodal_output": None, "request": request, "is_finished": False, "is_segment_finished": False}
    )

    assert connector.put.called
@@ -274,7 +274,7 @@ def test_send_single_request_struct_preserves_segment_finished(build_adapter, mo
    monkeypatch.setattr(adapter, "cleanup", lambda *a, **kw: None)

    adapter._send_single_request(
-        {"pooling_output": None, "request": request, "is_finished": False, "is_segment_finished": True}
+        {"multimodal_output": None, "request": request, "is_finished": False, "is_segment_finished": True}
    )

    sent_payload = connector.put.call_args.kwargs["data"]
@@ -289,13 +289,13 @@ def test_save_async_skips_stale_resumable_chunk_until_dedup_is_reset(build_adapt
    request.num_computed_tokens = 0
    adapter.requests_num_chunks_sent["ext-stream"] = 111

-    adapter.save_async(pooling_output=None, request=request, is_segment_finished=False)
+    adapter.save_async(multimodal_output=None, request=request, is_segment_finished=False)

    assert len(adapter._pending_save_reqs) == 0
    assert adapter.requests_num_chunks_sent["ext-stream"] == 111

    adapter.requests_num_chunks_sent.pop("ext-stream")
-    adapter.save_async(pooling_output=None, request=request, is_segment_finished=False)
+    adapter.save_async(multimodal_output=None, request=request, is_segment_finished=False)

    assert len(adapter._pending_save_reqs) == 1
    assert adapter.requests_num_chunks_sent["ext-stream"] == 0
@@ -312,7 +312,7 @@ def test_send_single_request_cleans_up_after_finished_payload(build_adapter, mon
    monkeypatch.setattr(adapter, "cleanup", lambda *a, **kw: cleanup_calls.append((a, kw)))

    adapter._send_single_request(
-        {"pooling_output": None, "request": request, "is_finished": True, "is_segment_finished": True}
+        {"multimodal_output": None, "request": request, "is_finished": True, "is_segment_finished": True}
    )

    assert len(cleanup_calls) == 1
@@ -0,0 +1,37 @@
+# Stage config for Step-Audio2 CI tests
+# Single Thinker stage (audio/text understanding → text). ASR/text-output
+# tests don't need Token2Wav; a two-stage config would deadlock on stage-1 skip.
+
+# The following config has been verified on 2x 24GB GPU (L4/RTX3090/RTX4090).
+stage_args:
+  - stage_id: 0
+    runtime:
+      process: true
+      devices: "0"
+      max_batch_size: 1
+    engine_args:
+      model_stage: thinker
+      model_arch: StepAudio2ThinkerForConditionalGeneration
+      tensor_parallel_size: 1
+      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      max_model_len: 4096
+      max_num_batched_tokens: 4096
+      max_num_seqs: 1
+      gpu_memory_utilization: 0.8
+      skip_mm_profiling: true
+      enforce_eager: true
+      trust_remote_code: true
+      enable_prefix_caching: false
+      engine_output_type: text  # Output text + audio tokens
+    is_comprehension: true
+    final_output: true
+    final_output_type: text
+    default_sampling_params:
+      temperature: 0.7
+      top_p: 0.9
+      top_k: -1
+      max_tokens: 256
+      seed: 42
+      detokenize: true
+      repetition_penalty: 1.05
@@ -5,6 +5,8 @@

 import asyncio
 import uuid
+from collections.abc import Mapping
+from typing import Any

 import numpy as np
 import pytest
@@ -102,8 +104,8 @@ def _flatten_audio(audio) -> torch.Tensor:
    return torch.as_tensor(audio, dtype=torch.float32).reshape(-1).cpu()


-def _extract_audio(multimodal_output: dict) -> torch.Tensor:
-    assert isinstance(multimodal_output, dict), f"Expected dict, got {type(multimodal_output)}"
+def _extract_audio(multimodal_output: Mapping[str, Any]) -> torch.Tensor:
+    assert isinstance(multimodal_output, Mapping), f"Expected Mapping, got {type(multimodal_output)}"
    audio = multimodal_output.get("audio", multimodal_output.get("model_outputs"))
    assert audio is not None, f"No audio output found, keys={list(multimodal_output.keys())}"
    waveform = _flatten_audio(audio)
@@ -112,7 +114,7 @@ def _extract_audio(multimodal_output: dict) -> torch.Tensor:
    return waveform


-def _extract_sample_rate(multimodal_output: dict) -> int:
+def _extract_sample_rate(multimodal_output: Mapping[str, Any]) -> int:
    sample_rate = multimodal_output.get("sr")
    if sample_rate is None:
        raise RuntimeError("Expected multimodal_output['sr']")
@@ -133,29 +135,29 @@ def _extract_final_audio_outputs(outputs):
        if request_output is None:
            continue
        multimodal_output = getattr(request_output, "multimodal_output", None)
-        if isinstance(multimodal_output, dict):
+        if isinstance(multimodal_output, Mapping):
            final_outputs.append(item)
            continue
        completions = getattr(request_output, "outputs", None) or []
-        if any(isinstance(getattr(completion, "multimodal_output", None), dict) for completion in completions):
+        if any(isinstance(getattr(completion, "multimodal_output", None), Mapping) for completion in completions):
            final_outputs.append(item)
    return final_outputs


-def _extract_multimodal_output(output) -> dict:
+def _extract_multimodal_output(output) -> Mapping[str, Any]:
    multimodal_output = getattr(output, "multimodal_output", None)
-    if isinstance(multimodal_output, dict):
+    if isinstance(multimodal_output, Mapping):
        return multimodal_output

    request_output = getattr(output, "request_output", None)
    if request_output is not None:
        multimodal_output = getattr(request_output, "multimodal_output", None)
-        if isinstance(multimodal_output, dict):
+        if isinstance(multimodal_output, Mapping):
            return multimodal_output
        completions = getattr(request_output, "outputs", None) or []
        for completion in completions:
            multimodal_output = getattr(completion, "multimodal_output", None)
-            if isinstance(multimodal_output, dict):
+            if isinstance(multimodal_output, Mapping):
                return multimodal_output

    raise AssertionError("No multimodal audio output found in Ming generate results")
@@ -32,6 +32,7 @@ _OMNI_RUNNER_PARAM = (
 )

 pytestmark = [
+    pytest.mark.skip(reason="issue#4361"),
    pytest.mark.full_model,
    pytest.mark.tts,
    pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True),
@@ -0,0 +1,303 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+E2E tests for Step-Audio2 model with audio input and audio output.
+
+These tests verify:
+1. Processor loads correctly and handles audio input
+2. Field configuration for vLLM batching (flat_from_sizes)
+3. Tensor shapes are correct (3D not 4D for profiling)
+4. Basic E2E inference works with audio input
+"""
+
+from pathlib import Path
+
+import numpy as np
+import pytest
+import torch
+
+from tests.helpers.runtime import OmniRunner
+
+models = ["stepfun-ai/Step-Audio-2-mini"]
+
+stage_configs = [str(Path(__file__).parent / "stage_configs" / "step_audio2_ci.yaml")]
+
+test_params = [(model, stage_config) for model in models for stage_config in stage_configs]
+
+
+def create_dummy_audio(sample_rate: int = 16000, duration_sec: float = 1.0) -> tuple[np.ndarray, int]:
+    """Create a dummy audio signal (sine wave) for testing."""
+    t = np.linspace(0, duration_sec, int(sample_rate * duration_sec), dtype=np.float32)
+    audio = 0.5 * np.sin(2 * np.pi * 440 * t)
+    return audio, sample_rate
+
+
+class StepAudio2OmniRunner(OmniRunner):
+    """Test runner specialized for Step-Audio2 model."""
+
+    def get_omni_inputs(
+        self,
+        prompts: list[str] | str,
+        system_prompt: str | None = None,
+        audios=None,
+        **kwargs,
+    ) -> list[dict]:
+        """Construct Omni input format for Step-Audio2."""
+        if system_prompt is None:
+            system_prompt = "你是一个语音对话助手，能够理解音频输入并生成语音回复。"
+
+        audio_placeholder = "<audio_patch>"
+
+        if isinstance(prompts, str):
+            prompts = [prompts]
+
+        num_prompts = len(prompts)
+
+        if audios is None:
+            audios_list = [None] * num_prompts
+        elif isinstance(audios, list):
+            audios_list = audios if len(audios) == num_prompts else [audios[0]] * num_prompts
+        else:
+            audios_list = [audios] * num_prompts
+
+        omni_inputs = []
+        for i, prompt_text in enumerate(prompts):
+            user_content = ""
+            multi_modal_data = {}
+
+            audio = audios_list[i]
+            if audio is not None:
+                user_content += audio_placeholder
+                multi_modal_data["audio"] = audio
+
+            user_content += prompt_text
+
+        full_prompt = (
+            f"<|im_start|>system\n{system_prompt}<|im_end|>\n"
+            f"<|im_start|>user\n{user_content}<|im_end|>\n"
+            f"<|im_start|>assistant\n"
+        )
+
+        input_dict = {"prompt": full_prompt}
+        if multi_modal_data:
+            input_dict["multi_modal_data"] = multi_modal_data
+
+        omni_inputs.append(input_dict)
+
+        return omni_inputs
+
+
+@pytest.fixture(scope="session")
+def step_audio2_runner():
+    return StepAudio2OmniRunner
+
+
+@pytest.mark.core_model
+def test_processor_loads_correctly():
+    """Test that Step-Audio2 processor can be imported and initialized."""
+    from vllm_omni.model_executor.models.step_audio2.step_audio2_thinker import (
+        StepAudio2MultiModalProcessor,
+        StepAudio2ProcessingInfo,
+        StepAudio2Processor,
+    )
+
+    assert StepAudio2MultiModalProcessor is not None
+    assert StepAudio2ProcessingInfo is not None
+    assert StepAudio2Processor is not None
+
+
+@pytest.mark.core_model
+def test_audio_preprocessing():
+    """Test that audio preprocessing produces correct tensor shapes."""
+    from vllm_omni.model_executor.models.step_audio2.step_audio2_thinker import (
+        log_mel_spectrogram,
+        padding_mels,
+    )
+
+    audio, sample_rate = create_dummy_audio(sample_rate=16000, duration_sec=1.0)
+
+    mel = log_mel_spectrogram(audio)
+
+    assert mel.ndim == 2, f"Expected 2D tensor, got {mel.ndim}D"
+    assert mel.shape[0] == 128, f"Expected 128 mel bins, got {mel.shape[0]}"
+
+    mels = [mel, mel]
+    padded_mels, mel_lens = padding_mels(mels)
+
+    assert padded_mels.ndim == 3, f"Expected 3D tensor, got {padded_mels.ndim}D"
+    assert padded_mels.shape[0] == 2, f"Expected batch size 2, got {padded_mels.shape[0]}"
+    assert padded_mels.shape[1] == 128, f"Expected 128 mel bins, got {padded_mels.shape[1]}"
+
+
+@pytest.mark.core_model
+def test_feature_length_calculation():
+    """Test that audio feature length calculation is correct."""
+    from vllm_omni.model_executor.models.step_audio2.step_audio2_thinker import (
+        calculate_audio_feature_length,
+    )
+
+    assert calculate_audio_feature_length(1000) == 125
+
+    assert calculate_audio_feature_length(100) > 0
+
+    assert calculate_audio_feature_length(1) >= 1
+
+
+@pytest.mark.core_model
+def test_token_config_constants():
+    """Test that token configuration constants are set correctly."""
+    from vllm_omni.model_executor.models.step_audio2.step_audio2_constants import (
+        DEFAULT_TOKEN_CONFIG,
+        STEP_AUDIO2_AUDIO_END,
+        STEP_AUDIO2_AUDIO_PATCH_TOKEN_ID,
+        STEP_AUDIO2_AUDIO_START,
+        STEP_AUDIO2_AUDIO_VOCAB_SIZE,
+        STEP_AUDIO2_TEXT_MAX,
+    )
+
+    assert STEP_AUDIO2_TEXT_MAX < STEP_AUDIO2_AUDIO_START, "Text tokens should come before audio tokens"
+    assert STEP_AUDIO2_AUDIO_END == STEP_AUDIO2_AUDIO_START + STEP_AUDIO2_AUDIO_VOCAB_SIZE - 1
+    assert STEP_AUDIO2_AUDIO_PATCH_TOKEN_ID > STEP_AUDIO2_TEXT_MAX
+    assert STEP_AUDIO2_AUDIO_PATCH_TOKEN_ID < STEP_AUDIO2_AUDIO_START
+
+    assert DEFAULT_TOKEN_CONFIG.text_max == STEP_AUDIO2_TEXT_MAX
+    assert DEFAULT_TOKEN_CONFIG.audio_start == STEP_AUDIO2_AUDIO_START
+
+
+@pytest.mark.core_model
+def test_mm_field_config_structure():
+    """Test that multimodal field config is properly structured for vLLM batching."""
+    from vllm.multimodal.inputs import MultiModalFieldConfig
+
+    audio_lens = torch.tensor([100, 200, 150], dtype=torch.int32)
+
+    field_config = MultiModalFieldConfig.flat_from_sizes("audio", audio_lens)
+
+    assert field_config is not None
+
+
+@pytest.mark.core_model
+def test_audio_encoder_output_shape():
+    """Test that audio encoder produces correct output shapes."""
+    from vllm_omni.model_executor.models.step_audio2.step_audio2_thinker import (
+        Adaptor,
+        AudioEncoder,
+    )
+
+    encoder = AudioEncoder(n_mels=128, n_ctx=1500, n_state=512, n_head=8, n_layer=6)
+
+    batch_size = 2
+    n_mels = 128
+    time_steps = 400
+    x = torch.randn(batch_size, n_mels, time_steps)
+    x_len = torch.tensor([time_steps, time_steps // 2], dtype=torch.int32)
+
+    encoded, encoded_lens = encoder(x, x_len)
+
+    assert encoded.ndim == 3, f"Expected 3D tensor, got {encoded.ndim}D"
+    assert encoded.shape[0] == batch_size
+    assert encoded.shape[2] == 512
+
+    adapter = Adaptor(n_state=512, n_hidden=4096, kernel_size=3, stride=2)
+    adapted = adapter(encoded)
+
+    assert adapted.ndim == 3
+    assert adapted.shape[0] == batch_size
+    assert adapted.shape[2] == 4096
+
+
+@pytest.mark.core_model
+def test_token_separation():
+    """Test that token separation works correctly."""
+    from vllm_omni.model_executor.models.step_audio2.step_audio2_thinker import (
+        StepAudio2ThinkerForConditionalGeneration,
+    )
+
+    token_ids = [100, 200, 151700, 300, 151800, 400]
+
+    text_tokens, audio_tokens = StepAudio2ThinkerForConditionalGeneration.separate_tokens(token_ids)
+
+    assert text_tokens == [100, 200, 300, 400]
+    assert audio_tokens == [4, 104]
+
+
+@pytest.mark.core_model
+def test_has_audio_output():
+    """Test detection of audio tokens in output."""
+    from vllm_omni.model_executor.models.step_audio2.step_audio2_thinker import (
+        StepAudio2ThinkerForConditionalGeneration,
+    )
+
+    text_only = [100, 200, 300]
+    assert not StepAudio2ThinkerForConditionalGeneration.has_audio_output(text_only)
+
+    with_audio = [100, 200, 151700, 300]
+    assert StepAudio2ThinkerForConditionalGeneration.has_audio_output(with_audio)
+
+
+@pytest.mark.advanced_model
+@pytest.mark.parametrize("test_config", test_params)
+def test_audio_to_text_and_audio(step_audio2_runner: type[StepAudio2OmniRunner], test_config: tuple[str, str]) -> None:
+    """Test processing audio input and generating text + audio output."""
+    model, stage_config_path = test_config
+
+    with step_audio2_runner(model, seed=42, stage_configs_path=stage_config_path) as runner:
+        audio = create_dummy_audio(sample_rate=16000, duration_sec=2.0)
+
+        outputs = runner.generate_multimodal(
+            prompts="请复述这段音频的内容。",
+            audios=audio,
+        )
+
+        assert len(outputs) > 0
+
+        text_output = None
+        for stage_output in outputs:
+            if stage_output.final_output_type == "text":
+                text_output = stage_output
+                break
+
+        assert text_output is not None
+        assert text_output.request_output is not None
+        text_content = text_output.request_output.outputs[0].text
+        assert text_content is not None
+        assert len(text_content.strip()) > 0
+
+        audio_output = None
+        for stage_output in outputs:
+            if stage_output.final_output_type == "audio":
+                audio_output = stage_output
+                break
+
+        if audio_output is not None:
+            assert audio_output.request_output is not None
+            audio_tensor = audio_output.request_output.outputs[0].multimodal_output.get("audio")
+            if audio_tensor is not None:
+                assert audio_tensor.numel() > 0
+
+
+@pytest.mark.advanced_model
+@pytest.mark.parametrize("test_config", test_params)
+def test_text_only_input(step_audio2_runner: type[StepAudio2OmniRunner], test_config: tuple[str, str]) -> None:
+    """Test processing text-only input (no audio)."""
+    model, stage_config_path = test_config
+
+    with step_audio2_runner(model, seed=42, stage_configs_path=stage_config_path) as runner:
+        outputs = runner.generate_multimodal(
+            prompts="你好，请用中文回答。",
+            audios=None,
+        )
+
+        assert len(outputs) > 0
+
+        text_output = None
+        for stage_output in outputs:
+            if stage_output.final_output_type == "text":
+                text_output = stage_output
+                break
+
+        assert text_output is not None
+        assert text_output.request_output is not None
+        text_content = text_output.request_output.outputs[0].text
+        assert text_content is not None
+        assert len(text_content.strip()) > 0
@@ -1,6 +1,7 @@
 """E2E test for VoxCPM2 native AR offline inference."""

 import os
+from collections.abc import Mapping

 import pytest
 import torch
@@ -23,7 +24,7 @@ def voxcpm2_engine():

 def _extract_audio(multimodal_output: dict) -> torch.Tensor:
    """Extract the final complete audio tensor from multimodal output."""
-    assert isinstance(multimodal_output, dict), f"Expected dict, got {type(multimodal_output)}"
+    assert isinstance(multimodal_output, (dict, Mapping)), f"Expected dict/Mapping, got {type(multimodal_output)}"

    # Output processor accumulates per-step audio chunks under "audio".
    audio = multimodal_output.get("audio")
@@ -0,0 +1,81 @@
+# Stage config for Step-Audio2 CI tests
+# Stage 0: Thinker (audio understanding → text + audio tokens)
+# Stage 1: Token2Wav (audio tokens → waveform)
+
+# The following config has been verified on 2x 24GB GPU (L4/RTX3090/RTX4090).
+stage_args:
+  - stage_id: 0
+    runtime:
+      process: true
+      devices: "0"
+      max_batch_size: 1
+    engine_args:
+      model_stage: thinker
+      model_arch: StepAudio2ThinkerForConditionalGeneration
+      tensor_parallel_size: 1
+      hf_overrides:
+          architectures: ["StepAudio2ThinkerForConditionalGeneration"]
+
+      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      max_model_len: 1024
+      max_num_batched_tokens: 1024
+      max_num_seqs: 1
+      gpu_memory_utilization: 0.7
+      skip_mm_profiling: true
+      enforce_eager: true
+      trust_remote_code: true
+      enable_prefix_caching: false
+      engine_output_type: text  # Output text + audio tokens
+    is_comprehension: true
+    final_output: true
+    final_output_type: text
+    default_sampling_params:
+      temperature: 0.7
+      top_p: 0.9
+      top_k: -1
+      max_tokens: 256
+      seed: 42
+      detokenize: true
+      repetition_penalty: 1.05
+
+  - stage_id: 1
+    runtime:
+      process: true
+      devices: "0"
+      max_batch_size: 1
+    engine_args:
+
+      model_stage: token2wav
+      model_arch: StepAudio2Token2WavModel
+      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
+      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
+      hf_overrides:
+          architectures: ["StepAudio2Token2WavModel"]
+      gpu_memory_utilization: 0.2
+      enforce_eager: true
+      trust_remote_code: true
+      enable_prefix_caching: false
+      max_num_batched_tokens: 1024
+      engine_output_type: audio  # Output waveform
+    engine_input_source: [0]
+    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.step_audio2.thinker2token2wav
+    final_output: true
+    final_output_type: audio
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 1
+      seed: 42
+      detokenize: false
+
+runtime:
+  enabled: true
+  defaults:
+    window_size: -1
+    max_inflight: 1
+  edges:
+    - from: 0
+      to: 1
+      window_size: -1
@@ -140,6 +140,7 @@ class TestHiggsAudioV3OnlineHappyPath:
            }
        )

+    @pytest.mark.skip(reason="issue#4411")
    @pytest.mark.core_model
    @pytest.mark.tts
    @hardware_test(res={"cuda": "H100"}, num_cards=1)
@@ -25,11 +25,11 @@ from tests.helpers.runtime import OmniServerParams
 from tests.helpers.stage_config import get_deploy_config_path

 # ``omni`` for all tests; 001/002 also get ``full_model`` via ``TestMossTtsNanoFull`` (003 is core+advanced only, no full_model).
-pytestmark = [pytest.mark.full_model, pytest.mark.tts]
-
-_SKIP_ISSUE_3168 = pytest.mark.skip(
-    reason="https://github.com/vllm-project/vllm-omni/issues/3168",
-)
+pytestmark = [
+    pytest.mark.skip(reason="issue#4361"),
+    pytest.mark.full_model,
+    pytest.mark.tts,
+]

 MODEL = "OpenMOSS-Team/MOSS-TTS-Nano"
 REF_AUDIO_URL = "https://raw.githubusercontent.com/OpenMOSS/MOSS-TTS-Nano/main/assets/audio/zh_1.wav"
@@ -0,0 +1,368 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+E2E tests for Step-Audio2 online serving with concurrent requests.
+
+Tests:
+1. Single request via API
+2. Concurrent throughput (1/2/4/8 concurrent requests)
+3. Latency percentiles (p50/p95/p99)
+"""
+
+import asyncio
+import os
+import socket
+import subprocess
+import sys
+import time
+from pathlib import Path
+
+import aiohttp
+import numpy as np
+import pytest
+from vllm.utils.network_utils import get_open_port
+
+MODEL = "stepfun-ai/Step-Audio-2-mini"
+STAGE_CONFIG = str(Path(__file__).parent / "stage_configs" / "step_audio2_ci.yaml")
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+
+class OmniServer:
+    """Omniserver for vLLM-Omni tests."""
+
+    def __init__(
+        self,
+        model: str,
+        serve_args: list[str],
+        *,
+        env_dict: dict[str, str] | None = None,
+    ) -> None:
+        self.model = model
+        self.serve_args = serve_args
+        self.env_dict = env_dict
+        self.proc: subprocess.Popen | None = None
+        self.host = "127.0.0.1"
+        self.port = get_open_port()
+
+    def _start_server(self) -> None:
+        """Start the vLLM-Omni server subprocess."""
+        env = os.environ.copy()
+        env["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+        if self.env_dict is not None:
+            env.update(self.env_dict)
+
+        cmd = [
+            sys.executable,
+            "-m",
+            "vllm_omni.entrypoints.cli.main",
+            "serve",
+            self.model,
+            "--omni",
+            "--host",
+            self.host,
+            "--port",
+            str(self.port),
+        ] + self.serve_args
+
+        print(f"Launching OmniServer with: {' '.join(cmd)}")
+        self.proc = subprocess.Popen(
+            cmd,
+            env=env,
+            cwd=os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
+        )
+
+        max_wait = 600  # 10 minutes
+        start_time = time.time()
+        while time.time() - start_time < max_wait:
+            try:
+                with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+                    sock.settimeout(1)
+                    result = sock.connect_ex((self.host, self.port))
+                    if result == 0:
+                        print(f"Server ready on {self.host}:{self.port}")
+                        return
+            except Exception:
+                pass
+            time.sleep(2)
+
+        raise RuntimeError(f"Server failed to start within {max_wait} seconds")
+
+    def __enter__(self):
+        self._start_server()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if self.proc:
+            self.proc.terminate()
+            try:
+                self.proc.wait(timeout=30)
+            except subprocess.TimeoutExpired:
+                self.proc.kill()
+                self.proc.wait()
+
+
+def create_dummy_audio_base64(duration_sec: float = 5.0, sample_rate: int = 16000) -> str:
+    """Create dummy audio as base64."""
+    import base64
+    import io
+    import wave
+
+    num_samples = int(duration_sec * sample_rate)
+    t = np.linspace(0, duration_sec, num_samples, dtype=np.float32)
+    audio = (0.5 * np.sin(2 * np.pi * 440 * t) * 32767).astype(np.int16)
+
+    buffer = io.BytesIO()
+    with wave.open(buffer, "wb") as wav_file:
+        wav_file.setnchannels(1)
+        wav_file.setsampwidth(2)
+        wav_file.setframerate(sample_rate)
+        wav_file.writeframes(audio.tobytes())
+
+    return base64.b64encode(buffer.getvalue()).decode("utf-8")
+
+
+def create_s2st_request(audio_base64: str) -> dict:
+    """Create S2ST request payload."""
+    # Use audio_url format with base64 data URL
+    audio_url = f"data:audio/wav;base64,{audio_base64}"
+    return {
+        "model": MODEL,
+        "messages": [
+            {"role": "system", "content": "请仔细聆听这段语音，然后复述其内容。"},
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "audio_url",
+                        "audio_url": {"url": audio_url},
+                    },
+                    {
+                        "type": "text",
+                        "text": "<tts_start>",
+                    },
+                ],
+            },
+        ],
+        "max_tokens": 512,
+        "temperature": 0.7,
+        "top_p": 0.9,
+    }
+
+
+async def send_request(session: aiohttp.ClientSession, url: str, payload: dict) -> tuple[float, int, str]:
+    """Send single request and return (latency, num_tokens, response_text)."""
+    start = time.perf_counter()
+    async with session.post(url, json=payload) as response:
+        result = await response.json()
+    latency = time.perf_counter() - start
+
+    if "choices" in result and len(result["choices"]) > 0:
+        text = result["choices"][0].get("message", {}).get("content", "")
+        tokens = result.get("usage", {}).get("completion_tokens", len(text.split()))
+    else:
+        text = str(result)
+        tokens = 0
+
+    return latency, tokens, text
+
+
+async def benchmark_concurrent(
+    url: str,
+    payload: dict,
+    num_concurrent: int,
+    num_requests: int,
+    warmup: int = 2,
+) -> dict:
+    """Run concurrent benchmark with true sustained concurrency.
+
+    Uses semaphore to maintain exactly num_concurrent requests in flight at all times,
+    rather than batch-and-wait approach. Measures real wall clock time.
+    """
+    semaphore = asyncio.Semaphore(num_concurrent)
+
+    async def bounded_request(session: aiohttp.ClientSession):
+        async with semaphore:
+            return await send_request(session, url, payload)
+
+    async with aiohttp.ClientSession() as session:
+        if warmup > 0:
+            print(f"    Warming up ({warmup} requests)...", end=" ", flush=True)
+            for _ in range(warmup):
+                try:
+                    await send_request(session, url, payload)
+                except Exception as e:
+                    print(f"Warmup failed: {e}")
+            print("done")
+
+        print(f"    Running {num_requests} requests with concurrency={num_concurrent}...", end=" ", flush=True)
+        wall_start = time.perf_counter()
+
+        tasks = [bounded_request(session) for _ in range(num_requests)]
+        results = await asyncio.gather(*tasks, return_exceptions=True)
+
+        wall_time = time.perf_counter() - wall_start
+        print(f"done ({wall_time:.2f}s)")
+
+    all_latencies = []
+    all_tokens = []
+    errors = 0
+
+    for result in results:
+        if isinstance(result, Exception):
+            print(f"    Request failed: {result}")
+            errors += 1
+            continue
+        latency, tokens, _ = result
+        all_latencies.append(latency)
+        all_tokens.append(tokens)
+
+    if not all_latencies:
+        return {"error": "All requests failed"}
+
+    latencies = np.array(all_latencies)
+    total_tokens = sum(all_tokens)
+    successful_requests = len(all_latencies)
+
+    return {
+        "num_concurrent": num_concurrent,
+        "num_requests": num_requests,
+        "successful_requests": successful_requests,
+        "errors": errors,
+        "wall_time": wall_time,
+        "p50_latency": float(np.percentile(latencies, 50)),
+        "p95_latency": float(np.percentile(latencies, 95)),
+        "p99_latency": float(np.percentile(latencies, 99)),
+        "mean_latency": float(np.mean(latencies)),
+        "min_latency": float(np.min(latencies)),
+        "max_latency": float(np.max(latencies)),
+        "throughput_req_per_sec": successful_requests / wall_time,
+        "throughput_tokens_per_sec": total_tokens / wall_time,
+    }
+
+
+class TestStepAudio2OnlineServing:
+    """Online serving tests for Step-Audio2."""
+
+    @pytest.fixture(scope="class")
+    def server(self):
+        """Start vLLM-Omni server."""
+        with OmniServer(MODEL, ["--stage-configs-path", STAGE_CONFIG]) as omni_server:
+            yield f"http://{omni_server.host}:{omni_server.port}/v1/chat/completions"
+
+    @pytest.mark.asyncio
+    async def test_single_request(self, server):
+        """Test single request works."""
+        audio_base64 = create_dummy_audio_base64(duration_sec=2.0)
+        payload = create_s2st_request(audio_base64)
+
+        async with aiohttp.ClientSession() as session:
+            latency, tokens, text = await send_request(session, server, payload)
+
+        assert latency > 0
+        assert tokens >= 0
+        print(f"Single request: latency={latency:.2f}s, tokens={tokens}")
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("num_concurrent", [1, 2, 4])
+    async def test_concurrent_throughput(self, server, num_concurrent):
+        """Test concurrent throughput at different concurrency levels."""
+        audio_base64 = create_dummy_audio_base64(duration_sec=5.0)
+        payload = create_s2st_request(audio_base64)
+
+        num_requests = num_concurrent * 3
+
+        stats = await benchmark_concurrent(server, payload, num_concurrent, num_requests)
+
+        assert "error" not in stats
+        print(f"\nConcurrency={num_concurrent}:")
+        print(f"  p50={stats['p50_latency']:.2f}s, p95={stats['p95_latency']:.2f}s")
+        print(f"  throughput={stats['throughput_req_per_sec']:.2f} req/s")
+
+
+async def run_throughput_benchmark(
+    server_url: str,
+    audio_duration: float = 5.0,
+    concurrency_levels: list[int] = [1, 2, 4, 8],
+    requests_per_level: int = 10,
+    warmup: int = 2,
+):
+    """Run full throughput benchmark."""
+    audio_base64 = create_dummy_audio_base64(duration_sec=audio_duration)
+    payload = create_s2st_request(audio_base64)
+
+    print(f"{'=' * 70}")
+    print("Step-Audio2 Online Serving Throughput Benchmark")
+    print(f"{'=' * 70}")
+    print(f"Server: {server_url}")
+    print(f"Audio duration: {audio_duration}s")
+    print(f"Requests per level: {requests_per_level}")
+    print(f"Warmup requests: {warmup}")
+    print(f"Concurrency levels: {concurrency_levels}")
+    print(f"{'=' * 70}")
+
+    results = []
+    for num_concurrent in concurrency_levels:
+        print(f"\n[Concurrency={num_concurrent}]")
+        stats = await benchmark_concurrent(server_url, payload, num_concurrent, requests_per_level, warmup)
+        results.append(stats)
+
+        if "error" not in stats:
+            print(
+                f"    Result: p50={stats['p50_latency']:.2f}s, "
+                f"p95={stats['p95_latency']:.2f}s, "
+                f"throughput={stats['throughput_req_per_sec']:.2f} req/s, "
+                f"{stats['throughput_tokens_per_sec']:.1f} tok/s"
+            )
+
+    print(f"\n{'=' * 80}")
+    print("SUMMARY")
+    print(f"{'=' * 80}")
+    print(
+        f"{'Conc':<6} {'Reqs':<6} {'Err':<5} {'Wall(s)':<9} {'p50':<8} {'p95':<8} {'p99':<8} {'Req/s':<8} {'Tok/s':<8}"
+    )
+    print("-" * 80)
+
+    for stats in results:
+        if "error" in stats:
+            print(f"{stats.get('num_concurrent', '?'):<6} ERROR: {stats['error']}")
+            continue
+        print(
+            f"{stats['num_concurrent']:<6} "
+            f"{stats['successful_requests']:<6} "
+            f"{stats['errors']:<5} "
+            f"{stats['wall_time']:<9.2f} "
+            f"{stats['p50_latency']:<8.2f} "
+            f"{stats['p95_latency']:<8.2f} "
+            f"{stats['p99_latency']:<8.2f} "
+            f"{stats['throughput_req_per_sec']:<8.2f} "
+            f"{stats['throughput_tokens_per_sec']:<8.1f}"
+        )
+
+    print("=" * 80)
+    return results
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Step-Audio2 Online Serving Benchmark")
+    parser.add_argument(
+        "--server-url", type=str, default="http://localhost:8000/v1/chat/completions", help="vLLM server URL"
+    )
+    parser.add_argument("--audio-duration", type=float, default=5.0, help="Input audio duration in seconds")
+    parser.add_argument("--concurrency", type=str, default="1,2,4,8", help="Comma-separated concurrency levels to test")
+    parser.add_argument("--requests-per-level", type=int, default=10, help="Number of requests per concurrency level")
+    parser.add_argument("--warmup", type=int, default=2, help="Number of warmup requests before each benchmark")
+
+    args = parser.parse_args()
+    concurrency_levels = [int(x) for x in args.concurrency.split(",")]
+
+    asyncio.run(
+        run_throughput_benchmark(
+            args.server_url,
+            args.audio_duration,
+            concurrency_levels,
+            args.requests_per_level,
+            args.warmup,
+        )
+    )
@@ -73,6 +73,7 @@ def test_voice_clone_streaming_001(omni_server, openai_client) -> None:

@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
+@pytest.mark.skip(reason="#4335")
 def test_response_format_001(omni_server, openai_client) -> None:
    """
    Test voice-clone non-stream PCM output via OpenAI API.
@@ -0,0 +1,59 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for _ensure_tensor_values boundary enforcement."""
+
+import pytest
+import torch
+
+from vllm_omni.worker.gpu_ar_model_runner import _ensure_tensor_values
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_tensors_pass_through():
+    """Tensor values pass through unchanged."""
+    t = torch.randn(10)
+    result = _ensure_tensor_values({"audio": t})
+    assert torch.equal(result["audio"], t)
+
+
+def test_scalar_int_wrapped():
+    """Integer scalars are wrapped in torch.tensor."""
+    result = _ensure_tensor_values({"sr": 24000})
+    assert isinstance(result["sr"], torch.Tensor)
+    assert result["sr"].item() == 24000
+
+
+def test_scalar_float_wrapped():
+    """Float scalars are wrapped in torch.tensor."""
+    result = _ensure_tensor_values({"gain": 0.5})
+    assert isinstance(result["gain"], torch.Tensor)
+    assert result["gain"].item() == pytest.approx(0.5)
+
+
+def test_list_converted():
+    """Lists of numbers are converted to tensors."""
+    result = _ensure_tensor_values({"ids": [1, 2, 3]})
+    assert isinstance(result["ids"], torch.Tensor)
+    assert list(result["ids"]) == [1, 2, 3]
+
+
+def test_non_tensorizable_dropped():
+    """Values that cannot be tensorized are dropped with warning."""
+    result = _ensure_tensor_values({"ok": torch.randn(5), "bad": {"nested": "dict"}})
+    assert "ok" in result
+    assert "bad" not in result
+
+
+def test_mixed_payload():
+    """Mixed tensor + scalar + list payload is properly sanitized."""
+    payload = {
+        "audio": torch.randn(100),
+        "sr": 24000,
+        "codes": [1, 2, 3],
+        "hidden": torch.randn(10, 64),
+    }
+    result = _ensure_tensor_values(payload)
+    assert len(result) == 4
+    for v in result.values():
+        assert isinstance(v, torch.Tensor)
@@ -81,7 +81,7 @@ def test_multimodal_payload_and_completion_output():
    assert p is not None
    assert "waveform" in p.tensors and torch.equal(p.primary_tensor, data["waveform"])
    assert p.metadata["sample_rate"] == 16000
-    assert not p.is_empty and len(p) == 1
+    assert not p.is_empty and len(p) == 2  # 1 tensor + 1 metadata

    # None/empty returns None
    assert MultimodalPayload.from_dict(None) is None
@@ -125,7 +125,7 @@ def test_output_modality_printed_examples(capsys):
    print(f"    tensors keys : {list(p.tensors.keys())}")
    print(f"    primary_tensor: shape={p.primary_tensor.shape}, dtype={p.primary_tensor.dtype}")
    print(f"    metadata      : {p.metadata}")
-    print(f"    is_empty={p.is_empty}, len={len(p)}")
+    print(f"    is_empty={p.is_empty}, len={len(p)}")  # len = tensors + metadata
    print(f"  from_dict(None) -> {MultimodalPayload.from_dict(None)}")
    print(f"  from_dict({{}})   -> {MultimodalPayload.from_dict({})}")

@@ -209,3 +209,64 @@ def test_cumulative_token_ids_is_a_copy():
    s = _make_state(RequestOutputKind.DELTA)
    out = s._new_completion_output([42], None, None)
    assert out.cumulative_token_ids is not _DETOK.output_token_ids
+
+
+# ---------------------------------------------------------------------------
+# No-detokenizer regression tests (Concern 2)
+# ---------------------------------------------------------------------------
+
+_NO_DETOK_STATE_KWARGS = {
+    **_DEFAULT_STATE_KWARGS,
+    "detokenizer": None,
+    "logprobs_processor": None,
+}
+
+
+def _make_no_detok_state(output_kind: RequestOutputKind):
+    return OmniRequestState(**_NO_DETOK_STATE_KWARGS, output_kind=output_kind)
+
+
+def test_no_detokenizer_completion_output():
+    """_new_completion_output works when detokenizer is None (generation stages)."""
+    s = _make_no_detok_state(RequestOutputKind.CUMULATIVE)
+    s.add_multimodal_tensor(torch.randn(10), mm_type=AUDIO)
+    out = s._new_completion_output([1, 2], None, None)
+    # Should still produce output with empty text and the token_ids as-is
+    assert out.text == ""
+    assert list(out.token_ids) == [1, 2]
+    assert hasattr(out, "cumulative_token_ids")
+    assert list(out.cumulative_token_ids) == [1, 2]
+    # Multimodal data should be attached
+    assert AUDIO in out.multimodal_output
+
+
+def test_no_detokenizer_make_request_output():
+    """make_request_output works without detokenizer when multimodal data is present."""
+    s = _make_no_detok_state(RequestOutputKind.CUMULATIVE)
+    s.add_multimodal_tensor(torch.randn(10), mm_type=AUDIO)
+    result = s.make_request_output([], None, FinishReason.STOP, None)
+    assert result is not None
+    assert not isinstance(result, PoolingRequestOutput)
+    assert AUDIO in result.outputs[0].multimodal_output
+
+
+def test_no_detokenizer_stream_interval_skipped():
+    """stream_interval > 1 does not crash when detokenizer is None."""
+    kwargs = {**_NO_DETOK_STATE_KWARGS, "stream_interval": 5}
+    s = OmniRequestState(**kwargs, output_kind=RequestOutputKind.DELTA)
+    s.add_multimodal_tensor(torch.randn(10), mm_type=AUDIO)
+    # Should not assert; stream_interval logic is simply skipped
+    result = s.make_request_output([], None, FinishReason.STOP, None)
+    assert result is not None
+
+
+def test_no_detokenizer_final_only():
+    """FINAL_ONLY with no detokenizer returns None until finish."""
+    s = _make_no_detok_state(RequestOutputKind.FINAL_ONLY)
+    s.add_multimodal_tensor(torch.randn(10), mm_type=AUDIO)
+    # Non-finish step should return None
+    assert s.make_request_output([], None, None, None) is None
+    # Finish step should return output
+    result = s.make_request_output([], None, FinishReason.STOP, None)
+    assert result is not None
+    assert AUDIO in result.outputs[0].multimodal_output
@@ -0,0 +1,88 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Msgpack roundtrip tests for OmniEngineCoreOutputs with multimodal_output.
+
+Validates that tensor-only payloads survive msgspec encode/decode and that
+non-tensor values fail decoding (enforcing the wire invariant).
+"""
+
+import pytest
+import torch
+from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder
+
+from vllm_omni.engine import OmniEngineCoreOutput, OmniEngineCoreOutputs
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _roundtrip(outputs: OmniEngineCoreOutputs) -> OmniEngineCoreOutputs:
+    encoder = MsgpackEncoder()
+    encoded = encoder.encode(outputs)
+    decoder = MsgpackDecoder(OmniEngineCoreOutputs)
+    return decoder.decode(encoded)
+
+
+def test_tensor_only_roundtrip():
+    """Tensor-only multimodal_output survives msgpack roundtrip."""
+    audio = torch.randn(16000)
+    sr = torch.tensor(24000)
+    eco = OmniEngineCoreOutput(
+        request_id="req-1",
+        new_token_ids=[1, 2, 3],
+        finish_reason=None,
+        multimodal_output={"audio": audio, "sr": sr},
+    )
+    decoded = _roundtrip(OmniEngineCoreOutputs(outputs=[eco]))
+    assert len(decoded.outputs) == 1
+    out = decoded.outputs[0]
+    assert out.request_id == "req-1"
+    assert isinstance(out.multimodal_output, dict)
+    assert torch.allclose(out.multimodal_output["audio"], audio)
+    assert out.multimodal_output["sr"].item() == 24000
+
+
+def test_empty_multimodal_roundtrip():
+    """None multimodal_output roundtrips correctly."""
+    eco = OmniEngineCoreOutput(
+        request_id="req-2",
+        new_token_ids=[4],
+        finish_reason=None,
+        multimodal_output=None,
+    )
+    decoded = _roundtrip(OmniEngineCoreOutputs(outputs=[eco]))
+    assert decoded.outputs[0].multimodal_output is None
+
+
+def test_multiple_tensor_keys_roundtrip():
+    """Multiple tensor keys survive roundtrip."""
+    hidden = torch.randn(10, 4096)
+    codes = torch.randint(0, 1024, (1, 100))
+    eco = OmniEngineCoreOutput(
+        request_id="req-3",
+        new_token_ids=[],
+        finish_reason=None,
+        multimodal_output={"hidden": hidden, "codes": codes},
+    )
+    decoded = _roundtrip(OmniEngineCoreOutputs(outputs=[eco]))
+    out = decoded.outputs[0]
+    assert torch.allclose(out.multimodal_output["hidden"], hidden)
+    assert torch.equal(out.multimodal_output["codes"], codes)
+
+
+def test_non_tensor_value_fails_decode():
+    """Non-tensor values in multimodal_output cause decode failure.
+
+    This validates that the wire field must be tensor-only and that
+    the _ensure_tensor_values boundary in model runners is necessary.
+    """
+    eco = OmniEngineCoreOutput(
+        request_id="req-bad",
+        new_token_ids=[],
+        finish_reason=None,
+        multimodal_output={"audio": torch.randn(10), "sr": 24000},  # scalar!
+    )
+    encoder = MsgpackEncoder()
+    encoded = encoder.encode(OmniEngineCoreOutputs(outputs=[eco]))
+    decoder = MsgpackDecoder(OmniEngineCoreOutputs)
+    with pytest.raises(Exception):
+        decoder.decode(encoded)
@@ -0,0 +1,101 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Unit tests for the TTS serving adapter registry (RFC #4327).
+
+Pure-Python registry/resolution logic; no model or GPU resources are loaded.
+"""
+
+import pytest
+
+from vllm_omni.entrypoints.openai.tts_adapters import (
+    TTS_ADAPTER_REGISTRY,
+    ARTTSAdapter,
+    DiffusionTTSAdapter,
+    all_tts_model_types,
+    resolve_adapter,
+)
+from vllm_omni.entrypoints.openai.tts_adapters.qwen3_tts import Qwen3TTSAdapter
+
+# Every dedicated TTS model-type must have an adapter so the orchestrator's
+# uniform ``self._adapter.build(...)`` dispatch covers it.
+EXPECTED_MODEL_TYPES = {
+    "qwen3_tts",
+    "voxcpm2",
+    "voxtral_tts",
+    "fish_tts",
+    "cosyvoice3",
+    "omnivoice",
+    "covo_audio",
+    "ming_tts",
+    "moss_tts_nano",
+    "moss_tts",
+    "higgs_audio_v2",
+    "higgs_audio_v3",
+    "glm_tts",
+    "step_audio2",
+}
+
+
+def test_all_model_types_registered():
+    assert EXPECTED_MODEL_TYPES <= all_tts_model_types()
+
+
+def test_registry_keyed_by_name():
+    for name, cls in TTS_ADAPTER_REGISTRY.items():
+        assert cls.name == name
+
+
+def test_resolve_each_model_type():
+    for model_type in EXPECTED_MODEL_TYPES:
+        cls = resolve_adapter(model_type)
+        assert cls is not None, model_type
+        assert cls.name == model_type
+
+
+def test_resolve_qwen3_tts_class():
+    assert resolve_adapter("qwen3_tts") is Qwen3TTSAdapter
+
+
+def test_resolve_unknown_returns_none():
+    assert resolve_adapter("not_a_real_model") is None
+    assert resolve_adapter(None) is None
+
+
+def test_ming_flash_omni_not_migrated():
+    """ming_flash_omni is intentionally excluded from the adapter migration in
+    this PR; it stays on the legacy inline dispatch in serving_speech.py."""
+    assert resolve_adapter("ming_flash_omni_tts") is None
+
+
+def test_voxcpm2_resolves():
+    """VoxCPM2 (the served ``latent_generator`` model) resolves cleanly.
+
+    Detection never returns the legacy ``voxcpm`` type, so there is no shared
+    stage-key ambiguity to resolve.
+    """
+    assert resolve_adapter("voxcpm2") is not None
+    assert resolve_adapter("voxcpm") is None
+
+
+def test_all_adapters_are_ar_or_diffusion():
+    for cls in TTS_ADAPTER_REGISTRY.values():
+        assert issubclass(cls, (ARTTSAdapter, DiffusionTTSAdapter))
+        assert cls.backend in ("ar", "diffusion")
+
+
+def test_qwen3_tts_metadata():
+    assert Qwen3TTSAdapter.backend == "ar"
+    assert issubclass(Qwen3TTSAdapter, ARTTSAdapter)
+
+
+def test_diffusion_adapter_extra_body_params_fallback():
+    class _DiffAdapter(DiffusionTTSAdapter):
+        name = "diff_probe"
+
+        async def build(self, request, sampling_params_list):  # pragma: no cover
+            raise NotImplementedError
+
+    assert _DiffAdapter.extra_body_params() == frozenset()
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])
@@ -481,12 +481,32 @@ def assert_omni_response(response: Any, request_config: dict[str, Any], run_leve
            if "text" in modalities:
                transcript = (response.audio_content or "").strip()
                text_output = (response.text_content or "").strip()
-                similarity = cosine_similarity_text(
-                    transcript.lower(),
-                    text_output.lower(),
-                )
-                print(f"similarity is: {similarity}")
-                assert similarity > similarity_threshold, "The audio content is not same as the text"
+                # For very short outputs (e.g. one-word answers), n-gram cosine
+                # similarity with length penalty is unreliable because Whisper
+                # may hallucinate extra context around the short utterance.  Use
+                # a containment check instead: the shorter text must appear in
+                # the longer one (after preprocessing removes punctuation).
+                _SHORT_TEXT_THRESHOLD = 15
+                if len(text_output) <= _SHORT_TEXT_THRESHOLD or len(transcript) <= _SHORT_TEXT_THRESHOLD:
+                    shorter = text_output.lower() if len(text_output) <= len(transcript) else transcript.lower()
+                    longer = transcript.lower() if len(text_output) <= len(transcript) else text_output.lower()
+                    import re as _re
+
+                    shorter_clean = _re.sub(r"[^\w\s]", "", shorter).strip()
+                    longer_clean = _re.sub(r"[^\w\s]", "", longer).strip()
+                    assert shorter_clean and (shorter_clean in longer_clean), (
+                        f"The audio content is not same as the text "
+                        f"(short-text containment check failed: "
+                        f"text={text_output!r}, transcript={transcript!r})"
+                    )
+                    print(f"short-text containment check passed: {shorter_clean!r} in {longer_clean!r}")
+                else:
+                    similarity = cosine_similarity_text(
+                        transcript.lower(),
+                        text_output.lower(),
+                    )
+                    print(f"similarity is: {similarity}")
+                    assert similarity > similarity_threshold, "The audio content is not same as the text"
            if audio_ref_text:
                audio_similarity = cosine_similarity_text(
                    response.audio_content.lower(),
@@ -500,9 +520,25 @@ def assert_omni_response(response: Any, request_config: dict[str, Any], run_leve
 def assert_audio_speech_response(response: Any, request_config: dict[str, Any], run_level: str) -> None:
    """Validate speech API results from :class:`~tests.helpers.runtime.OmniResponse`.

-    Success path only. Use :meth:`OpenAIClientHandler.send_audio_speech_http_request` with ``err_code`` /
-    ``err_message`` for expected HTTP errors.
+    When ``request_config`` carries ``status_code`` and/or ``err_message``, the
+    request is expected to be rejected: assert it failed and that the HTTP status
+    / error text match. Otherwise the normal success-path checks run.
    """
+    expected_status = request_config.get("status_code")
+    expected_err = request_config.get("err_message")
+    if expected_status is not None or expected_err is not None:
+        assert not response.success, "Expected an error response, but the request succeeded."
+        if expected_status is not None:
+            allowed = expected_status if isinstance(expected_status, (list, tuple)) else (expected_status,)
+            assert response.status_code in allowed, f"Expected HTTP status in {allowed}, got {response.status_code}"
+        if expected_err is not None:
+            alternatives = expected_err if isinstance(expected_err, (list, tuple)) else (expected_err,)
+            error_text = response.error_message or ""
+            assert any(alt in error_text for alt in alternatives), (
+                f"Expected one of {alternatives} in error text, got: {error_text!r}"
+            )
+        return
+
    assert response.success, "The request failed."

    # Optional floor on decoded audio size (models with very short clips may use a lower value).
@@ -762,6 +762,10 @@ class OmniResponse:
    prompt_tokens: int | None = None
    cached_tokens: int | None = None
    logprobs: list | None = None
+    #: HTTP status + error text for the error-handling path (e.g. validator
+    #: rejections); populated when the OpenAI client raises an APIError.
+    status_code: int | None = None
+    error_message: str | None = None


@dataclass
@@ -0,0 +1,76 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Regression tests for cross-request ``codes.ref`` leakage (#4370).
+
+``make_omni_output`` used to collapse every request's reference codec
+frames into a single last-writer-wins slot and emit it as a length-1
+list. ``to_payload_element`` indexes list payloads with
+``element[idx] if idx < len(element) else element[0]``, so that length-1
+list was broadcast to every request in the batch: each utterance's first
+vocoder chunks then decoded with another request's reference voice as
+context (audible onset timbre deformation at any concurrency > 1).
+"""
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import (
+    Qwen3TTSTalkerForConditionalGeneration,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_NUM_CODE_GROUPS = 16
+
+
+def _make_talker() -> Qwen3TTSTalkerForConditionalGeneration:
+    # make_omni_output reads no instance state, so skip __init__ (no real
+    # config / checkpoint needed), same as the preprocess tests.
+    return Qwen3TTSTalkerForConditionalGeneration.__new__(Qwen3TTSTalkerForConditionalGeneration)
+
+
+def _info(span_frames: int, ref_code: torch.Tensor | None) -> dict:
+    codes: dict = {"audio": torch.zeros((span_frames, _NUM_CODE_GROUPS), dtype=torch.long)}
+    meta: dict = {}
+    if ref_code is not None:
+        codes["ref"] = ref_code
+        meta["ref_code_len"] = int(ref_code.shape[0])
+    return {"codes": codes, "meta": meta}
+
+
+def test_make_omni_output_keeps_ref_codes_per_request() -> None:
+    ref_a = torch.ones((3, _NUM_CODE_GROUPS), dtype=torch.long)
+    ref_b = torch.full((5, _NUM_CODE_GROUPS), 2, dtype=torch.long)
+
+    out = _make_talker().make_omni_output(
+        torch.zeros((4, 8)),
+        model_intermediate_buffer=[_info(2, ref_a), _info(2, ref_b)],
+    )
+
+    ref_list = out.multimodal_outputs["codes"]["ref"]
+    assert len(ref_list) == 2, "codes.ref must stay batch-aligned (one entry per request)"
+    assert torch.equal(ref_list[0], ref_a)
+    assert torch.equal(ref_list[1], ref_b)
+
+
+def test_make_omni_output_pads_requests_without_ref_code() -> None:
+    ref_b = torch.full((5, _NUM_CODE_GROUPS), 2, dtype=torch.long)
+
+    out = _make_talker().make_omni_output(
+        torch.zeros((4, 8)),
+        model_intermediate_buffer=[_info(2, None), _info(2, ref_b)],
+    )
+
+    ref_list = out.multimodal_outputs["codes"]["ref"]
+    assert len(ref_list) == 2
+    assert ref_list[0].numel() == 0, "no-ref request must get an empty placeholder, not a neighbor's ref"
+    assert torch.equal(ref_list[1], ref_b)
+
+
+def test_make_omni_output_omits_ref_key_when_no_request_has_one() -> None:
+    out = _make_talker().make_omni_output(
+        torch.zeros((4, 8)),
+        model_intermediate_buffer=[_info(2, None), _info(2, None)],
+    )
+
+    assert "ref" not in out.multimodal_outputs["codes"]
@@ -0,0 +1,142 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+# Token2Wav imports the optional `step-audio2` extras (flashcosyvoice / s3tokenizer
+# / hyperpyyaml) at module load. Skip instead of erroring at collection time on CI
+# images that do not install them (e.g. the CPU core_model runner); a genuine
+# import error (e.g. vLLM API drift) still surfaces. See
+# examples/offline_inference/step_audio2/README.md.
+try:
+    from vllm_omni.model_executor.models.step_audio2.step_audio2_token2wav import (
+        StepAudio2Token2WavForConditionalGeneration,
+    )
+except ModuleNotFoundError as exc:
+    pytest.skip(
+        f"Step-Audio2 Token2Wav extras not installed ({exc.name})",
+        allow_module_level=True,
+    )
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _DummyModelConfig:
+    def __init__(self) -> None:
+        self.hf_config = SimpleNamespace(
+            token2wav_path="/tmp/fake-token2wav",
+            token2wav_float16=False,
+        )
+        self.model = "/tmp/fake-model"
+
+    def get_hidden_size(self) -> int:
+        return 8
+
+
+def _build_model() -> StepAudio2Token2WavForConditionalGeneration:
+    vllm_config = SimpleNamespace(
+        model_config=_DummyModelConfig(),
+        device_config=SimpleNamespace(device="cpu"),
+    )
+    return StepAudio2Token2WavForConditionalGeneration(vllm_config=vllm_config)
+
+
+def test_step_audio2_token2wav_sync_path_not_misdetected_by_empty_runtime_info(monkeypatch):
+    model = _build_model()
+
+    monkeypatch.setattr("os.path.exists", lambda _: True)
+    monkeypatch.setattr(
+        model.token2wav,
+        "forward",
+        lambda *args, **kwargs: torch.zeros(1, 16, dtype=torch.float32),
+    )
+
+    out = model.forward(
+        input_ids=torch.tensor([1, 2, 3]),
+        positions=torch.tensor([0, 1, 2]),
+        runtime_additional_information=[{}],
+    )
+
+    audio = out.multimodal_outputs["model_outputs"]
+    # sync path returns the full waveform wrapped in a single-element list,
+    # same contract as the async-chunk path (see tests below)
+    assert isinstance(audio, list)
+    assert len(audio) == 1
+    assert isinstance(audio[0], torch.Tensor)
+    assert model._stream_states == []
+
+
+def test_step_audio2_token2wav_async_chunk_batch_guard():
+    model = _build_model()
+
+    with pytest.raises(RuntimeError, match="only supports batch=1"):
+        model._forward_async_chunk(
+            input_ids=torch.tensor([1, 2, 3]),
+            runtime_additional_information=[
+                {"left_context_size": 0},
+                {"left_context_size": 1},
+            ],
+        )
+
+
+def test_step_audio2_token2wav_async_chunk_last_chunk_resets_state(monkeypatch):
+    model = _build_model()
+
+    monkeypatch.setattr("os.path.exists", lambda _: True)
+    calls = {"setup": 0, "stream": 0, "reset": 0}
+
+    def _setup_stream_for(prompt_wav, state):
+        calls["setup"] += 1
+        state.setup_done = True
+        state.stream_cache = {}
+        state.hift_cache_dict = {
+            "mel": torch.zeros(1, 1, 0),
+            "source": torch.zeros(1, 1, 0),
+            "speech": torch.zeros(1, 0),
+        }
+
+    def _stream_chunk_for(audio_tokens, prompt_wav, last_chunk, state):
+        calls["stream"] += 1
+        return torch.tensor([0.1, 0.2, 0.3], dtype=torch.float32)
+
+    def _reset_stream_for(state):
+        calls["reset"] += 1
+        state.setup_done = False
+        state.finished = True
+        state.stream_cache = None
+        state.hift_cache_dict = {}
+
+    monkeypatch.setattr(model.token2wav, "setup_stream_for", _setup_stream_for)
+    monkeypatch.setattr(model.token2wav, "stream_chunk_for", _stream_chunk_for)
+    monkeypatch.setattr(model.token2wav, "reset_stream_for", _reset_stream_for)
+
+    out = model.forward(
+        input_ids=torch.tensor([10, 11, 12]),
+        positions=torch.tensor([0, 1, 2]),
+        runtime_additional_information=[{"left_context_size": 1}],
+    )
+
+    audio_list = out.multimodal_outputs["model_outputs"]
+    assert isinstance(audio_list, list)
+    assert len(audio_list) == 1
+    assert torch.allclose(audio_list[0], torch.tensor([0.1, 0.2, 0.3], dtype=torch.float32))
+    assert calls == {"setup": 1, "stream": 1, "reset": 1}
+
+
+def test_step_audio2_token2wav_async_chunk_empty_eof_returns_zero_chunk():
+    model = _build_model()
+
+    out = model._forward_async_chunk(
+        input_ids=torch.tensor([], dtype=torch.int64),
+        runtime_additional_information=[{"left_context_size": 1}],
+    )
+
+    audio_list = out.multimodal_outputs["model_outputs"]
+    assert isinstance(audio_list, list)
+    assert len(audio_list) == 1
+    assert torch.equal(audio_list[0], torch.zeros(1, dtype=torch.float32))
+    assert model._stream_states
+    assert model._stream_states[0].finished is True
@@ -140,7 +140,7 @@ def test_talker2code2wav_async_chunk_final_payload_uses_absolute_token_offset():

    payload = talker2code2wav_async_chunk(
        transfer_manager=transfer_manager,
-        pooling_output=None,
+        multimodal_output=None,
        request=request,
        is_finished=True,
    )
@@ -167,7 +167,7 @@ def test_talker2code2wav_async_chunk_emits_eof_when_finished_without_valid_codes

    payload = talker2code2wav_async_chunk(
        transfer_manager=transfer_manager,
-        pooling_output=None,
+        multimodal_output=None,
        request=request,
        is_finished=True,
    )
@@ -188,13 +188,13 @@ def test_talker2code2wav_async_chunk_does_not_reemit_without_new_tokens():

    payload1 = talker2code2wav_async_chunk(
        transfer_manager=transfer_manager,
-        pooling_output=None,
+        multimodal_output=None,
        request=request,
        is_finished=False,
    )
    payload2 = talker2code2wav_async_chunk(
        transfer_manager=transfer_manager,
-        pooling_output=None,
+        multimodal_output=None,
        request=request,
        is_finished=False,
    )
@@ -216,14 +216,14 @@ def test_talker2code2wav_async_chunk_waits_for_prelookahead_and_emits_cumulative

    payload_pending = talker2code2wav_async_chunk(
        transfer_manager=transfer_manager,
-        pooling_output=None,
+        multimodal_output=None,
        request=request,
        is_finished=False,
    )
    request.output_token_ids = [1, 2, 3]
    payload_ready = talker2code2wav_async_chunk(
        transfer_manager=transfer_manager,
-        pooling_output=None,
+        multimodal_output=None,
        request=request,
        is_finished=False,
    )
@@ -246,14 +246,14 @@ def test_talker2code2wav_async_chunk_final_flush_uses_previous_token_offset():

    payload_stream = talker2code2wav_async_chunk(
        transfer_manager=transfer_manager,
-        pooling_output=None,
+        multimodal_output=None,
        request=request,
        is_finished=False,
    )
    request.output_token_ids = [3, 4, 5, 6]
    payload_final = talker2code2wav_async_chunk(
        transfer_manager=transfer_manager,
-        pooling_output=None,
+        multimodal_output=None,
        request=request,
        is_finished=True,
    )
@@ -281,14 +281,14 @@ def test_talker2code2wav_async_chunk_respects_prompt_token_pad_on_first_chunk():

    payload_pending = talker2code2wav_async_chunk(
        transfer_manager=transfer_manager,
-        pooling_output=None,
+        multimodal_output=None,
        request=request,
        is_finished=False,
    )
    request.output_token_ids = [8, 9, 10, 11]
    payload_ready = talker2code2wav_async_chunk(
        transfer_manager=transfer_manager,
-        pooling_output=None,
+        multimodal_output=None,
        request=request,
        is_finished=False,
    )
@@ -310,13 +310,13 @@ def test_talker2code2wav_async_chunk_emits_terminal_eof_without_duplicate_audio(

    payload_stream = talker2code2wav_async_chunk(
        transfer_manager=transfer_manager,
-        pooling_output=None,
+        multimodal_output=None,
        request=request,
        is_finished=False,
    )
    payload_final = talker2code2wav_async_chunk(
        transfer_manager=transfer_manager,
-        pooling_output=None,
+        multimodal_output=None,
        request=request,
        is_finished=True,
    )
@@ -56,7 +56,7 @@ def _call(tm, rid, *, n_frames, finished=False, req_ic=None):
    tm.code_prompt_token_ids[rid] = [_FRAME[:] for _ in range(n_frames)]
    return talker2code2wav_async_chunk(
        transfer_manager=tm,
-        pooling_output={"codes": {"audio": torch.zeros((0,))}},
+        multimodal_output={"codes": {"audio": torch.zeros((0,))}},
        request=_req(rid, finished=finished, initial_codec_chunk_frames=req_ic),
        is_finished=finished,
    )
@@ -66,7 +66,7 @@ def test_empty_returns_none():
    tm = _tm()
    p = talker2code2wav_async_chunk(
        transfer_manager=tm,
-        pooling_output={"codes": {"audio": torch.zeros((0,))}},
+        multimodal_output={"codes": {"audio": torch.zeros((0,))}},
        request=_req("r", finished=False),
    )
    assert p is None
@@ -76,7 +76,7 @@ def test_eof_marker_when_finished_empty():
    tm = _tm()
    p = talker2code2wav_async_chunk(
        transfer_manager=tm,
-        pooling_output=None,
+        multimodal_output=None,
        request=_req("r", finished=True),
        is_finished=True,
    )
@@ -89,7 +89,7 @@ def test_flush_on_finish():
    tm.code_prompt_token_ids["r"] = [_FRAME[:] for _ in range(24)]
    p = talker2code2wav_async_chunk(
        transfer_manager=tm,
-        pooling_output=None,
+        multimodal_output=None,
        request=_req("r", finished=True),
        is_finished=True,
    )
@@ -262,7 +262,7 @@ def test_first_streaming_chunk_prepends_ref_code_context():

    payload = talker2code2wav_async_chunk(
        transfer_manager=tm,
-        pooling_output={"codes": {"audio": torch.zeros((0,)), "ref": ref_code}},
+        multimodal_output={"codes": {"audio": torch.zeros((0,)), "ref": ref_code}},
        request=_req(rid, finished=False, initial_codec_chunk_frames=10),
        is_finished=False,
    )
@@ -286,7 +286,7 @@ def test_followup_ref_code_context_is_sent_as_metadata_handle():

    payload = talker2code2wav_async_chunk(
        transfer_manager=tm,
-        pooling_output={"codes": {"audio": torch.zeros((0,)), "ref": ref_code}},
+        multimodal_output={"codes": {"audio": torch.zeros((0,)), "ref": ref_code}},
        request=_req(rid, finished=False, initial_codec_chunk_frames=10),
        is_finished=False,
    )
@@ -319,7 +319,7 @@ def test_streaming_ref_code_context_is_bounded_for_batchable_shapes():

    payload = talker2code2wav_async_chunk(
        transfer_manager=tm,
-        pooling_output={"codes": {"audio": torch.zeros((0,)), "ref": ref_code}},
+        multimodal_output={"codes": {"audio": torch.zeros((0,)), "ref": ref_code}},
        request=_req(rid, finished=False),
        is_finished=False,
    )
@@ -338,7 +338,7 @@ def test_ref_code_context_can_be_buffered_before_first_emit():

    first_payload = talker2code2wav_async_chunk(
        transfer_manager=tm,
-        pooling_output={"codes": {"audio": torch.tensor([[1, 2, 3, 4]]), "ref": ref_code}},
+        multimodal_output={"codes": {"audio": torch.tensor([[1, 2, 3, 4]]), "ref": ref_code}},
        request=_req(rid, finished=False, initial_codec_chunk_frames=10),
        is_finished=False,
    )
@@ -348,14 +348,14 @@ def test_ref_code_context_can_be_buffered_before_first_emit():
    for _ in range(8):
        talker2code2wav_async_chunk(
            transfer_manager=tm,
-            pooling_output={"codes": {"audio": torch.tensor([[1, 2, 3, 4]])}},
+            multimodal_output={"codes": {"audio": torch.tensor([[1, 2, 3, 4]])}},
            request=_req(rid, finished=False, initial_codec_chunk_frames=10),
            is_finished=False,
        )

    payload = talker2code2wav_async_chunk(
        transfer_manager=tm,
-        pooling_output={"codes": {"audio": torch.tensor([[1, 2, 3, 4]])}},
+        multimodal_output={"codes": {"audio": torch.tensor([[1, 2, 3, 4]])}},
        request=_req(rid, finished=False, initial_codec_chunk_frames=10),
        is_finished=False,
    )
@@ -0,0 +1,132 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from collections import defaultdict
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.step_audio2.step_audio2_constants import (
+    DEFAULT_STREAM_CONFIG,
+    DEFAULT_TOKEN_CONFIG,
+)
+from vllm_omni.model_executor.stage_input_processors.step_audio2 import (
+    thinker2token2wav_async_chunk,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _req(external_req_id: str, *, prompt_token_ids: list[int], all_token_ids: list[int], finished: bool):
+    return SimpleNamespace(
+        external_req_id=external_req_id,
+        prompt_token_ids=prompt_token_ids,
+        all_token_ids=all_token_ids,
+        is_finished=lambda: finished,
+    )
+
+
+def test_step_audio2_async_chunk_uses_decode_only_tokens_not_prompt_history():
+    audio_start = DEFAULT_TOKEN_CONFIG.audio_start
+    audio_eos = DEFAULT_TOKEN_CONFIG.audio_eos
+    transfer_manager = SimpleNamespace(code_prompt_token_ids=defaultdict(list))
+
+    prompt = [11, 12, audio_start + 5]  # historical prompt audio token
+    generated = [99, audio_start + 1, audio_start + 2, audio_start + audio_eos]
+    request = _req(
+        "rid-decode-only",
+        prompt_token_ids=prompt,
+        all_token_ids=prompt + generated,
+        finished=True,
+    )
+
+    payload = thinker2token2wav_async_chunk(
+        transfer_manager=transfer_manager,
+        pooling_output=None,
+        request=request,
+    )
+
+    assert payload is not None
+    assert payload["code_predictor_codes"] == [1, 2]
+    assert payload["left_context_size"] == 1
+    assert payload["finished"].item() is True
+
+
+def test_step_audio2_async_chunk_returns_none_when_not_enough_tokens():
+    audio_start = DEFAULT_TOKEN_CONFIG.audio_start
+    required = DEFAULT_STREAM_CONFIG.chunk_size + DEFAULT_STREAM_CONFIG.pre_lookahead_len
+    transfer_manager = SimpleNamespace(code_prompt_token_ids=defaultdict(list))
+
+    generated = [audio_start + i for i in range(required - 1)]
+    request = _req(
+        "rid-not-ready",
+        prompt_token_ids=[1, 2, 3],
+        all_token_ids=[1, 2, 3] + generated,
+        finished=False,
+    )
+
+    payload = thinker2token2wav_async_chunk(
+        transfer_manager=transfer_manager,
+        pooling_output=None,
+        request=request,
+    )
+
+    assert payload is None
+
+
+def test_step_audio2_async_chunk_emits_non_last_chunk_and_advances_consumed_by_chunk_size():
+    audio_start = DEFAULT_TOKEN_CONFIG.audio_start
+    chunk_size = DEFAULT_STREAM_CONFIG.chunk_size
+    required = chunk_size + DEFAULT_STREAM_CONFIG.pre_lookahead_len
+    transfer_manager = SimpleNamespace(code_prompt_token_ids=defaultdict(list))
+
+    generated = [audio_start + i for i in range(required + 10)]
+    request = _req(
+        "rid-ready",
+        prompt_token_ids=[7, 8],
+        all_token_ids=[7, 8] + generated,
+        finished=False,
+    )
+
+    payload = thinker2token2wav_async_chunk(
+        transfer_manager=transfer_manager,
+        pooling_output=None,
+        request=request,
+    )
+
+    assert payload is not None
+    assert payload["left_context_size"] == 0
+    assert payload["finished"].item() is False
+    assert payload["code_predictor_codes"] == list(range(required))
+    assert transfer_manager.code_prompt_token_ids["rid-ready"] == list(range(chunk_size))
+
+
+def test_step_audio2_async_chunk_emits_eof_when_finished_with_no_remaining_audio():
+    transfer_manager = SimpleNamespace(code_prompt_token_ids=defaultdict(list))
+    transfer_manager.code_prompt_token_ids["rid-eof"] = [1, 2, 3]
+
+    request = _req(
+        "rid-eof",
+        prompt_token_ids=[10, 11],
+        all_token_ids=[
+            10,
+            11,
+            DEFAULT_TOKEN_CONFIG.audio_start + 1,
+            DEFAULT_TOKEN_CONFIG.audio_start + 2,
+            DEFAULT_TOKEN_CONFIG.audio_start + 3,
+        ],
+        finished=True,
+    )
+
+    payload = thinker2token2wav_async_chunk(
+        transfer_manager=transfer_manager,
+        pooling_output=None,
+        request=request,
+    )
+
+    assert payload == {
+        "code_predictor_codes": [],
+        "left_context_size": 1,
+        "finished": torch.tensor(True, dtype=torch.bool),
+    }
@@ -92,7 +92,7 @@ def test_empty_chunk_when_not_finished():

    payload = generator2tokenizer_async_chunk(
        transfer_manager=transfer_manager,
-        pooling_output={"audio": torch.zeros((0,))},
+        multimodal_output={"audio": torch.zeros((0,))},
        request=request,
    )

@@ -109,7 +109,7 @@ def test_flush_tail_when_finished():

    payload = generator2tokenizer_async_chunk(
        transfer_manager=transfer_manager,
-        pooling_output=None,  # e.g. EOS step with no audio
+        multimodal_output=None,  # e.g. EOS step with no audio
        request=request,
    )

@@ -133,7 +133,7 @@ def test_eof_marker_when_finished_with_no_frames():

    payload = generator2tokenizer_async_chunk(
        transfer_manager=transfer_manager,
-        pooling_output=None,
+        multimodal_output=None,
        request=request,
    )

@@ -153,10 +153,10 @@ def test_normal_chunk_emission():

    # Feed 25 frames one by one
    for i in range(25):
-        pooling_output = {"audio": torch.tensor([float(i)] * 4)}
+        multimodal_output = {"audio": torch.tensor([float(i)] * 4)}
        payload = generator2tokenizer_async_chunk(
            transfer_manager=transfer_manager,
-            pooling_output=pooling_output,
+            multimodal_output=multimodal_output,
            request=request,
        )

@@ -181,10 +181,10 @@ def test_small_initial_chunks():

    # Feed 5 frames (should trigger emission because codec_chunk_frames_at_begin=5)
    for i in range(5):
-        pooling_output = {"audio": torch.tensor([float(i + 1)] * 3)}
+        multimodal_output = {"audio": torch.tensor([float(i + 1)] * 3)}
        payload = generator2tokenizer_async_chunk(
            transfer_manager=transfer_manager,
-            pooling_output=pooling_output,
+            multimodal_output=multimodal_output,
            request=request,
        )

@@ -208,10 +208,10 @@ def test_no_emission_between_boundaries():

    # Feed 3 frames (3 % 5 != 0, should not emit)
    for i in range(3):
-        pooling_output = {"audio": torch.tensor([float(i)] * 4)}
+        multimodal_output = {"audio": torch.tensor([float(i)] * 4)}
        payload = generator2tokenizer_async_chunk(
            transfer_manager=transfer_manager,
-            pooling_output=pooling_output,
+            multimodal_output=multimodal_output,
            request=request,
        )

@@ -230,10 +230,10 @@ def test_context_handling_format():

    # Feed 5 frames to trigger a chunk (chunk_size_at_begin=5)
    for i in range(5):
-        pooling_output = {"audio": torch.tensor([float(i + 10)] * 2)}  # codebook_dim=2
+        multimodal_output = {"audio": torch.tensor([float(i + 10)] * 2)}  # codebook_dim=2
        payload = generator2tokenizer_async_chunk(
            transfer_manager=transfer_manager,
-            pooling_output=pooling_output,
+            multimodal_output=multimodal_output,
            request=request,
        )

@@ -250,14 +250,14 @@ def test_context_handling_format():
    assert flat_codes.numel() == total_window_frames * 2  # codebook_dim=2


-def test_none_pooling_output_not_finished_returns_none():
-    """None pooling_output when not finished returns None."""
+def test_none_multimodal_output_not_finished_returns_none():
+    """None multimodal_output when not finished returns None."""
    transfer_manager = _make_transfer_manager()
    request = _req("rid-none", finished=False)

    payload = generator2tokenizer_async_chunk(
        transfer_manager=transfer_manager,
-        pooling_output=None,
+        multimodal_output=None,
        request=request,
    )

@@ -0,0 +1,338 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for 310P patch wiring.
+
+The tests load patch modules from source with fake Qwen3-TTS dependencies, so
+they validate the patch contract without loading real model or NPU kernels.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import sys
+import types
+from pathlib import Path
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _repo_root() -> Path:
+    marker = Path("vllm_omni") / "platforms" / "npu" / "_310p" / "patch"
+    for parent in Path(__file__).resolve().parents:
+        if (parent / marker).is_dir():
+            return parent
+    raise FileNotFoundError(f"could not locate repo root containing {marker}")
+
+
+def _load_source_module(name: str, path: Path):
+    spec = importlib.util.spec_from_file_location(name, path)
+    if spec is None or spec.loader is None:
+        raise RuntimeError(f"cannot load import spec for {path}")
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+def _install_fake_module(monkeypatch: pytest.MonkeyPatch, name: str, **attrs):
+    module = types.ModuleType(name)
+    for key, value in attrs.items():
+        setattr(module, key, value)
+    monkeypatch.setitem(sys.modules, name, module)
+    return module
+
+
+def _install_qwen3_tts_patch_fakes(monkeypatch: pytest.MonkeyPatch):
+    class FakeAudioResampler:
+        def __init__(self, *, target_sr: int):
+            self.target_sr = target_sr
+
+        def resample(self, wav, *, orig_sr: int):
+            del orig_sr
+            return wav
+
+    class FakeCodePredictorAttention(torch.nn.Module):
+        def __init__(self, *args, **kwargs):
+            del args, kwargs
+            super().__init__()
+            self.register_buffer("_fusion_causal_mask", torch.ones(1), persistent=False)
+
+    class FakeCodePredictorDecoderLayer(torch.nn.Module):
+        pass
+
+    class FakeCodePredictorBaseModel(torch.nn.Module):
+        pass
+
+    class FakeMimiEuclideanCodebook(torch.nn.Module):
+        @property
+        def embed(self):
+            return self._embed
+
+    class FakeEncoder:
+        def __init__(self):
+            self.to_calls: list[dict] = []
+            self.last_input_dtype = None
+
+        def to(self, **kwargs):
+            self.to_calls.append(kwargs)
+            return self
+
+        def encode(self, *, input_values, return_dict: bool):
+            assert return_dict
+            self.last_input_dtype = input_values.dtype
+            return SimpleNamespace(audio_codes=torch.arange(12, dtype=torch.long).reshape(1, 3, 4))
+
+    class FakeFeatureBatch(dict):
+        def to(self, target):
+            if isinstance(target, torch.dtype):
+                for key, value in list(self.items()):
+                    if torch.is_floating_point(value):
+                        self[key] = value.to(dtype=target)
+                self.dtype = target
+            else:
+                self.device = torch.device(target)
+                for key, value in list(self.items()):
+                    self[key] = value.to(device=self.device)
+            return self
+
+    class FakeFeatureExtractor:
+        sampling_rate = 24000
+
+        def __call__(self, *, raw_audio, sampling_rate: int, return_tensors: str):
+            assert len(raw_audio) == 1
+            assert sampling_rate == self.sampling_rate
+            assert return_tensors == "pt"
+            return FakeFeatureBatch(
+                input_values=torch.ones(1, 1, 8, dtype=torch.float32),
+                padding_mask=torch.ones(1, 1, 8, dtype=torch.long),
+            )
+
+    class FakeTalkerBase(torch.nn.Module):
+        def __init__(self, *, vllm_config, prefix: str = ""):
+            del vllm_config, prefix
+            super().__init__()
+            self._embedding_dtype = torch.bfloat16
+            self._prompt_builder = SimpleNamespace(_embedding_dtype=torch.bfloat16)
+            self.encoder = FakeEncoder()
+            self._encoder_feature_extractor = FakeFeatureExtractor()
+            self._encoder_valid_num_quantizers = 2
+            self._encoder_downsample_rate = 2
+
+        def load_weights(self, weights):
+            del weights
+            self.encoder.to(dtype=torch.bfloat16)
+            return {"loaded"}
+
+    class FakePromptEmbedsBuilder:
+        pass
+
+    fake_qwen3_code_predictor = _install_fake_module(
+        monkeypatch,
+        "vllm_omni.model_executor.models.common.qwen3_code_predictor",
+        CodePredictorAttention=FakeCodePredictorAttention,
+        CodePredictorDecoderLayer=FakeCodePredictorDecoderLayer,
+        CodePredictorBaseModel=FakeCodePredictorBaseModel,
+        _rotate_half=lambda x: x,
+    )
+    fake_prompt_builder = _install_fake_module(
+        monkeypatch,
+        "vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder",
+        Qwen3TTSPromptEmbedsBuilder=FakePromptEmbedsBuilder,
+        mel_spectrogram=lambda *_args, **_kwargs: torch.empty(0),
+    )
+    fake_talker = _install_fake_module(
+        monkeypatch,
+        "vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker",
+        Qwen3TTSTalkerForConditionalGeneration=FakeTalkerBase,
+        Qwen3TTSPromptEmbedsBuilder=FakePromptEmbedsBuilder,
+    )
+    fake_modeling_mimi = _install_fake_module(
+        monkeypatch,
+        "transformers.models.mimi.modeling_mimi",
+        MimiEuclideanCodebook=FakeMimiEuclideanCodebook,
+    )
+    fake_mimi = _install_fake_module(
+        monkeypatch,
+        "transformers.models.mimi",
+        modeling_mimi=fake_modeling_mimi,
+    )
+
+    _install_fake_module(monkeypatch, "vllm")
+    _install_fake_module(monkeypatch, "vllm.multimodal")
+    _install_fake_module(monkeypatch, "vllm.multimodal.audio", AudioResampler=FakeAudioResampler)
+    _install_fake_module(monkeypatch, "transformers")
+    _install_fake_module(monkeypatch, "transformers.models", mimi=fake_mimi)
+    _install_fake_module(monkeypatch, "vllm_omni")
+    _install_fake_module(monkeypatch, "vllm_omni.model_executor")
+    _install_fake_module(monkeypatch, "vllm_omni.model_executor.models")
+    _install_fake_module(
+        monkeypatch,
+        "vllm_omni.model_executor.models.common",
+        qwen3_code_predictor=fake_qwen3_code_predictor,
+    )
+    _install_fake_module(
+        monkeypatch,
+        "vllm_omni.model_executor.models.qwen3_tts",
+        prompt_embeds_builder=fake_prompt_builder,
+        qwen3_tts_talker=fake_talker,
+    )
+    return fake_qwen3_code_predictor, fake_prompt_builder, fake_talker
+
+
+def _load_qwen3_tts_patch(monkeypatch: pytest.MonkeyPatch):
+    fakes = _install_qwen3_tts_patch_fakes(monkeypatch)
+    path = _repo_root() / "vllm_omni" / "platforms" / "npu" / "_310p" / "patch" / "qwen3_tts.py"
+    module = _load_source_module("vllm_omni_test_310p_qwen3_tts_patch", path)
+    return module, fakes
+
+
+def test_registry_applies_worker_once_and_model_patch_lazily(monkeypatch: pytest.MonkeyPatch) -> None:
+    registry_path = _repo_root() / "vllm_omni" / "platforms" / "npu" / "_310p" / "patch" / "__init__.py"
+    registry = _load_source_module("vllm_omni_test_310p_patch_registry", registry_path)
+    calls = {"worker": 0, "talker": 0}
+
+    _install_fake_module(
+        monkeypatch,
+        "vllm_omni.platforms.npu._310p.patch.worker",
+        apply_patch=lambda: calls.__setitem__("worker", calls["worker"] + 1),
+    )
+    _install_fake_module(
+        monkeypatch,
+        "vllm_omni.platforms.npu._310p.patch.qwen3_tts",
+        apply_talker_patches=lambda: calls.__setitem__("talker", calls["talker"] + 1),
+    )
+
+    registry.apply_patches()
+    registry.apply_patches()
+    registry.apply_model_patches(SimpleNamespace(model_arch="OtherModel"))
+    registry.apply_model_patches(SimpleNamespace(model_arch="Qwen3TTSTalkerForConditionalGeneration"))
+
+    assert calls == {"worker": 1, "talker": 1}
+
+
+def test_worker_patch_replaces_base_and_runs_disable_jit(monkeypatch: pytest.MonkeyPatch) -> None:
+    calls: list[str] = []
+
+    class FakeOmniNPUWorkerBase:
+        def _init_device(self):
+            calls.append("parent")
+            return "npu:0"
+
+    fake_worker_base = _install_fake_module(
+        monkeypatch,
+        "vllm_omni.platforms.npu.worker.base",
+        OmniNPUWorkerBase=FakeOmniNPUWorkerBase,
+    )
+    _install_fake_module(monkeypatch, "vllm_omni")
+    _install_fake_module(monkeypatch, "vllm_omni.platforms")
+    _install_fake_module(monkeypatch, "vllm_omni.platforms.npu")
+    _install_fake_module(
+        monkeypatch,
+        "vllm_omni.platforms.npu._310p",
+        disable_jit_compile=lambda: calls.append("disable_jit"),
+    )
+    _install_fake_module(monkeypatch, "vllm_omni.platforms.npu.worker", base=fake_worker_base)
+
+    path = _repo_root() / "vllm_omni" / "platforms" / "npu" / "_310p" / "patch" / "worker.py"
+    module = _load_source_module("vllm_omni_test_310p_worker_patch", path)
+    module.apply_patch()
+
+    assert fake_worker_base.OmniNPUWorkerBase is module._OmniNPUWorkerBase310P
+    assert fake_worker_base.OmniNPUWorkerBase()._init_device() == "npu:0"
+    assert calls == ["parent", "disable_jit"]
+
+
+def test_qwen3_tts_patch_replaces_target_classes(monkeypatch: pytest.MonkeyPatch) -> None:
+    module, (fake_code_predictor, fake_prompt_builder, fake_talker) = _load_qwen3_tts_patch(monkeypatch)
+
+    module.apply_talker_patches()
+
+    assert fake_talker.Qwen3TTSTalkerForConditionalGeneration is module._Qwen3TTSTalker310P
+    assert fake_talker.Qwen3TTSPromptEmbedsBuilder is module._Qwen3TTSPromptEmbedsBuilder310P
+    assert fake_prompt_builder.Qwen3TTSPromptEmbedsBuilder is module._Qwen3TTSPromptEmbedsBuilder310P
+    assert fake_code_predictor.CodePredictorAttention is module._Qwen3CodePredictorAttention310P
+    assert fake_code_predictor.CodePredictorDecoderLayer is module._Qwen3CodePredictorDecoderLayer310P
+    assert fake_code_predictor.CodePredictorBaseModel is module._Qwen3CodePredictorBaseModel310P
+
+
+def test_qwen3_tts_talker_patch_uses_fp16_runtime_dtype(monkeypatch: pytest.MonkeyPatch) -> None:
+    module, _ = _load_qwen3_tts_patch(monkeypatch)
+    talker = module._Qwen3TTSTalker310P(vllm_config=object())
+
+    assert talker._embedding_dtype is torch.float16
+    assert talker._prompt_builder._embedding_dtype is torch.float16
+    assert talker.load_weights([]) == {"loaded"}
+    assert talker.encoder.to_calls[-1] == {"dtype": torch.float16}
+
+    codes = talker._encode_ref_audio_batch([np.zeros(8, dtype=np.float32)], 24000, device=torch.device("cpu"))
+
+    assert talker.encoder.last_input_dtype is torch.float16
+    assert len(codes) == 1
+    assert codes[0].dtype is torch.long
+    assert codes[0].shape == (4, 2)
+
+
+def test_qwen3_tts_prompt_patch_runs_stft_frontend_on_cpu(monkeypatch: pytest.MonkeyPatch) -> None:
+    module, _ = _load_qwen3_tts_patch(monkeypatch)
+    captured = {}
+
+    def fake_mel_spectrogram(wav_tensor, **kwargs):
+        captured["wav_device"] = wav_tensor.device
+        captured["wav_dtype"] = wav_tensor.dtype
+        captured["kwargs"] = kwargs
+        return torch.ones(1, 128, 3, dtype=torch.float32)
+
+    class FakeSpeakerEncoder(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.param = torch.nn.Parameter(torch.zeros(1, dtype=torch.float16))
+
+        def forward(self, mels):
+            captured["speaker_input_dtype"] = mels.dtype
+            return (torch.ones(4, dtype=mels.dtype),)
+
+    monkeypatch.setattr(module.prompt_embeds_builder, "mel_spectrogram", fake_mel_spectrogram)
+    builder = object.__new__(module._Qwen3TTSPromptEmbedsBuilder310P)
+    builder._device = lambda: torch.device("cpu")
+    builder._embedding_dtype = torch.float16
+    builder._speaker_encoder = FakeSpeakerEncoder()
+    builder._config = SimpleNamespace(speaker_encoder_config=SimpleNamespace(sample_rate=24000))
+
+    speaker = builder.extract_speaker_embedding(np.zeros(16, dtype=np.float32), 24000)
+
+    assert captured["wav_device"] == torch.device("cpu")
+    assert captured["wav_dtype"] is torch.float32
+    assert captured["kwargs"]["sampling_rate"] == 24000
+    assert captured["speaker_input_dtype"] is torch.float16
+    assert speaker.dtype is torch.float16
+
+
+def test_qwen3_tts_mimi_codebook_quantize_uses_cpu_fp32_cdist(monkeypatch: pytest.MonkeyPatch) -> None:
+    module, _ = _load_qwen3_tts_patch(monkeypatch)
+    real_cdist = torch.cdist
+    captured = {}
+
+    def fake_cdist(x1, x2, p=2):
+        captured["x1_device"] = x1.device
+        captured["x2_device"] = x2.device
+        captured["x1_dtype"] = x1.dtype
+        captured["x2_dtype"] = x2.dtype
+        return real_cdist(x1, x2, p=p)
+
+    monkeypatch.setattr(torch, "cdist", fake_cdist)
+    codebook = object.__new__(module._MimiEuclideanCodebook310P)
+    codebook._embed = torch.tensor([[0.0, 0.0], [2.0, 0.0]], dtype=torch.float16)
+
+    indices = codebook.quantize(torch.tensor([[0.1, 0.0], [1.8, 0.0]], dtype=torch.float16))
+
+    assert captured == {
+        "x1_device": torch.device("cpu"),
+        "x2_device": torch.device("cpu"),
+        "x1_dtype": torch.float32,
+        "x2_dtype": torch.float32,
+    }
+    assert indices.tolist() == [0, 1]
@@ -62,6 +62,16 @@ class TestParallelConfigPropagation:
        assert od.parallel_config.ulysses_degree == 2
        assert od.parallel_config.ring_degree == 1

+    def test_mask_sp_padding_roundtrip(self):
+        pc = DiffusionParallelConfig(ulysses_degree=2, mask_sp_padding=True)
+        od = _roundtrip_diffusion_config(model="x", parallel_config=pc)
+        assert od.parallel_config.mask_sp_padding is True
+
+    def test_mask_sp_padding_defaults_false(self):
+        pc = DiffusionParallelConfig(ulysses_degree=2)
+        od = _roundtrip_diffusion_config(model="x", parallel_config=pc)
+        assert od.parallel_config.mask_sp_padding is False
+
    def test_cfg_parallel_roundtrip(self):
        pc = DiffusionParallelConfig(cfg_parallel_size=2)
        od = _roundtrip_diffusion_config(model="x", parallel_config=pc)
@@ -8,7 +8,8 @@ without requiring the actual checkpoint or GPU.

 import asyncio
 import time
-from collections import OrderedDict
+from collections import OrderedDict, defaultdict
+from types import SimpleNamespace

 import pytest
 import torch
@@ -1156,6 +1157,31 @@ class TestStageInputProcessor:
        assert result[0, 1].item() == 200
        assert result[0, 2].item() == 300

+    def test_async_chunk_accepts_multimodal_output_keyword(self):
+        from vllm_omni.model_executor.stage_input_processors.higgs_audio_v3 import (
+            talker2code2wav_async_chunk,
+        )
+
+        request_id = "req-higgs"
+        transfer_manager = SimpleNamespace(
+            code_prompt_token_ids=defaultdict(list),
+            connector=SimpleNamespace(config={"extra": {"codec_chunk_frames": 1, "codec_right_holdback_frames": 0}}),
+        )
+        request = SimpleNamespace(external_req_id=request_id, is_finished=lambda: False)
+
+        result = None
+        for step in range(8):
+            result = talker2code2wav_async_chunk(
+                transfer_manager=transfer_manager,
+                multimodal_output={"codes": {"audio": torch.full((1, 8), step, dtype=torch.long)}},
+                request=request,
+                is_finished=False,
+            )
+
+        assert result is not None
+        assert result.codes is not None
+        assert result.codes.audio.numel() > 0
+

 # ---- AC-10: Registry ----

@@ -92,6 +92,9 @@ def test_sample_tokens_tail_only_prefix_cache_uses_staged_cpu_hidden_states(monk
    runner.supports_mm_inputs = False
    runner.use_async_scheduling = False
    runner._omni_num_scheduled_tokens_np = None
+    runner.vllm_config = SimpleNamespace(
+        model_config=SimpleNamespace(engine_output_type="audio"),
+    )

    monkeypatch.setattr(
        GPUARModelRunner, "_sample", lambda self, logits, spec_decode_metadata: SimpleNamespace(sampled_token_ids=[])
@@ -125,8 +128,8 @@ def test_sample_tokens_tail_only_prefix_cache_uses_staged_cpu_hidden_states(monk

    output = GPUARModelRunner.sample_tokens(runner, grammar_output=None)

-    assert torch.equal(output.pooler_output[0]["hidden"], torch.tensor([[1.0, 10.0]]))
+    assert torch.equal(output.multimodal_outputs[0]["hidden"], torch.tensor([[1.0, 10.0]]))
    assert torch.equal(
-        output.pooler_output[1]["hidden"],
+        output.multimodal_outputs[1]["hidden"],
        torch.tensor([[2.0, 20.0], [3.0, 30.0]]),
    )
@@ -49,8 +49,8 @@ def test_sample_tokens_tensor_output():

    output = GPUGenerationModelRunner.sample_tokens(runner)

-    assert len(output.pooler_output) == 1
-    assert output.pooler_output[0]["model_outputs"].shape == (2, 3)
+    assert len(output.multimodal_outputs) == 1
+    assert output.multimodal_outputs[0]["model_outputs"].shape == (2, 3)


 def test_sample_tokens_list_output():
@@ -59,8 +59,8 @@ def test_sample_tokens_list_output():

    output = GPUGenerationModelRunner.sample_tokens(runner)

-    assert len(output.pooler_output) == 1
-    assert output.pooler_output[0]["model_outputs"].shape == (2, 1)
+    assert len(output.multimodal_outputs) == 1
+    assert output.multimodal_outputs[0]["model_outputs"].shape == (2, 1)


 def test_sample_tokens_list_allows_none_output():
@@ -69,8 +69,8 @@ def test_sample_tokens_list_allows_none_output():

    output = GPUGenerationModelRunner.sample_tokens(runner)

-    assert len(output.pooler_output) == 1
-    assert output.pooler_output[0]["model_outputs"] is None
+    assert len(output.multimodal_outputs) == 1
+    assert output.multimodal_outputs[0]["model_outputs"] is None


 def test_sample_tokens_dict_output():
@@ -79,7 +79,7 @@ def test_sample_tokens_dict_output():

    output = GPUGenerationModelRunner.sample_tokens(runner)

-    assert len(output.pooler_output) == 1
-    assert "audio" in output.pooler_output[0]
-    assert "unused" not in output.pooler_output[0]
-    assert output.pooler_output[0]["audio"].shape == (1, 4)
+    assert len(output.multimodal_outputs) == 1
+    assert "audio" in output.multimodal_outputs[0]
+    assert "unused" not in output.multimodal_outputs[0]
+    assert output.multimodal_outputs[0]["audio"].shape == (1, 4)
@@ -1193,6 +1193,7 @@ async def benchmark(
    ramp_up_end_rps: int | None = None,
    ready_check_timeout_sec: int = 600,
    ssl_context: ssl.SSLContext | bool | None = None,
+    self_timed: bool = False,
 ):
    try:
        request_func = ASYNC_REQUEST_FUNCS[endpoint_type]
@@ -1355,6 +1356,7 @@ async def benchmark(
        ramp_up_strategy,
        ramp_up_start_rps,
        ramp_up_end_rps,
+        self_timed,
    ):
        if ramp_up_strategy is not None:
            current_int_rps = int(current_request_rate)
@@ -1443,6 +1445,7 @@ async def benchmark(
            defs.AVERAGE_PIXELS_PER_IMAGE: getattr(metrics, defs.AVERAGE_PIXELS_PER_IMAGE),
            defs.MEAN_DENOISE_STEP_LATENCY_MS: getattr(metrics, defs.MEAN_DENOISE_STEP_LATENCY_MS),
            "input_lens": [output.prompt_len for output in outputs],
+            "start_times": [output.start_time for output in outputs],
            "output_lens": actual_output_lens,
            "ttfts": [output.ttft for output in outputs],
            "itls": [output.itl for output in outputs],
@@ -288,6 +288,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
        prompt_logprobs_dict = model_runner_output.prompt_logprobs_dict
        num_scheduled_tokens = scheduler_output.num_scheduled_tokens
        pooler_outputs = model_runner_output.pooler_output
+        mm_outputs = getattr(model_runner_output, "multimodal_outputs", None)
        num_nans_in_logits = model_runner_output.num_nans_in_logits
        kv_connector_output = model_runner_output.kv_connector_output
        cudagraph_stats: CUDAGraphStat | None = model_runner_output.cudagraph_stats
@@ -389,6 +390,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
            new_logprobs = None
            new_token_ids = generated_token_ids
            pooler_output = pooler_outputs[req_index] if pooler_outputs else None
+            mm_output = mm_outputs[req_index] if mm_outputs else None
            kv_transfer_params = None
            status_before_stop = request.status
            finish_reason = None
@@ -463,7 +465,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):

            # Get prompt logprobs for this request.
            prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id)
-            if new_token_ids or pooler_output is not None or kv_transfer_params or stopped:
+            if new_token_ids or mm_output is not None or pooler_output is not None or kv_transfer_params or stopped:
                # Add EngineCoreOutput for this Request.
                outputs[request.client_index].append(
                    OmniEngineCoreOutput(
@@ -473,6 +475,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                        new_logprobs=new_logprobs,
                        new_prompt_logprobs_tensors=prompt_logprobs_tensors,
                        pooling_output=pooler_output,
+                        multimodal_output=mm_output,
                        stop_reason=request.stop_reason,
                        events=request.take_events(),
                        prefill_stats=request.take_prefill_stats(),
@@ -485,7 +488,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                    )
                )
                if self.chunk_transfer_adapter is not None:
-                    self.chunk_transfer_adapter.save_async(pooler_output, request, is_segment_finished)
+                    self.chunk_transfer_adapter.save_async(mm_output, request, is_segment_finished)
            else:
                # Invariant: EngineCore returns no partial prefill outputs.
                assert not prompt_logprobs_tensors
@@ -709,6 +712,11 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):

        self._omits_kv_transfer_cache.pop(request.request_id, None)

+        # [Upstream compat] Discard request from in-flight prefills set added
+        # upstream for routed-experts in-flight reservation tracking.
+        # Use getattr for safety with test __new__ code paths.
+        getattr(self, "_inflight_prefills", set()).discard(request)
+
        # 1. Standard cleanup parts from base _free_request
        connector_delay_free_blocks, kv_xfer_params = self._connector_finished(request)

@@ -420,6 +420,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
        prompt_logprobs_dict = model_runner_output.prompt_logprobs_dict
        num_scheduled_tokens = scheduler_output.num_scheduled_tokens
        pooler_outputs = model_runner_output.pooler_output
+        mm_outputs = getattr(model_runner_output, "multimodal_outputs", None)
        num_nans_in_logits = model_runner_output.num_nans_in_logits
        kv_connector_output = model_runner_output.kv_connector_output

@@ -498,6 +499,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
            new_token_ids = generated_token_ids
            kv_transfer_params = None
            pooler_output = pooler_outputs[req_index] if pooler_outputs else None
+            mm_output = mm_outputs[req_index] if mm_outputs else None
            status_before_stop = request.status
            finish_reason = None
            routed_experts = None
@@ -558,7 +560,8 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):

            # Get prompt logprobs for this request.
            prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id)
-            if new_token_ids or pooler_output is not None or kv_transfer_params or stopped:
+            if new_token_ids or mm_output is not None or pooler_output is not None or kv_transfer_params or stopped:
+                # Add EngineCoreOutput for this Request.
                outputs[request.client_index].append(
                    OmniEngineCoreOutput(
                        request_id=req_id,
@@ -567,6 +570,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
                        new_logprobs=new_logprobs,
                        new_prompt_logprobs_tensors=prompt_logprobs_tensors,
                        pooling_output=pooler_output,
+                        multimodal_output=mm_output,
                        stop_reason=request.stop_reason,
                        events=request.take_events(),
                        prefill_stats=request.take_prefill_stats(),
@@ -130,6 +130,14 @@ class DiffusionParallelConfig:
    use_hsdp: bool = False
    """Enable Hybrid Sharded Data Parallel (HSDP) for model weight sharding."""

+    mask_sp_padding: bool = False
+    """If True, generate a boolean attention mask for zero-padded SP tokens
+    when sequence length is not divisible by the SP world size. The mask
+    routes attention through the varlen path (unpad→kernel→repad), which is
+    correct but carries additional overhead. When False (default), padding
+    tokens are left unmasked; since _shard_with_auto_pad always pads with
+    zeros, their contribution to attention output is negligible."""
+
    hsdp_shard_size: int = -1
    """Number of GPUs to shard weights across within each replica group. -1 means auto-calculate."""

@@ -6,8 +6,7 @@ from typing import TYPE_CHECKING, Any

 from vllm.utils.import_utils import resolve_obj_by_qualname

-from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
-from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.data import OmniDiffusionConfig

 if TYPE_CHECKING:
    from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput
@@ -69,11 +68,6 @@ class DiffusionExecutor(ABC):
        """Initialize the executor (e.g., launch workers, setup IPC)."""
        pass

-    @abstractmethod
-    def add_req(self, requests: OmniDiffusionRequest) -> DiffusionOutput:
-        """Add requests to the execution queue."""
-        pass
-
    @abstractmethod
    def execute_request(self, scheduler_output: DiffusionSchedulerOutput) -> BaseRunnerOutput:
        """Execute request-mode work from a scheduler output."""
@@ -16,8 +16,7 @@ from vllm.v1.engine.exceptions import EngineDeadError

 from vllm_omni.diffusion.data import SHUTDOWN_MESSAGE, DiffusionOutput
 from vllm_omni.diffusion.executor.abstract import DiffusionExecutor
-from vllm_omni.diffusion.ipc import unpack_diffusion_output_shm
-from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.ipc import DIFFUSION_RPC_RESULT_ENVELOPE, unpack_diffusion_output_shm
 from vllm_omni.diffusion.worker import WorkerProc

 if TYPE_CHECKING:
@@ -133,6 +132,57 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
                    raise EngineDeadError()
                continue

+    @staticmethod
+    def _raise_for_rpc_error_dict(response: Any) -> None:
+        if isinstance(response, dict) and response.get("status") == "error":
+            raise RuntimeError(
+                f"Worker failed with error '{response.get('error')}', "
+                "please check the stack trace above for the root cause"
+            )
+
+    @staticmethod
+    def _unwrap_rpc_result_envelope(response: Any) -> Any:
+        if not (isinstance(response, dict) and response.get("type") == DIFFUSION_RPC_RESULT_ENVELOPE):
+            return response
+
+        rank_statuses = response.get("rank_statuses") or []
+        failed = [status for status in rank_statuses if not status.get("ok", False)]
+        if failed:
+            details = "; ".join(
+                f"rank {status.get('rank')}: {status.get('error_type') or 'Error'}: {status.get('error')}"
+                for status in failed
+            )
+            tracebacks = "\n\n".join(
+                f"rank {status.get('rank')} traceback:\n{status['traceback']}"
+                for status in failed
+                if status.get("traceback")
+            )
+            if tracebacks:
+                details = f"{details}\n\n{tracebacks}"
+            method = response.get("method", "<unknown>")
+            raise RuntimeError(f"RPC '{method}' failed on worker rank(s): {details}")
+
+        result = response.get("result")
+        if isinstance(result, bool):
+            # Only bool-returning RPCs participate in the all-rank AND.
+            # Non-bool results leave bool_result unset and are ignored here.
+            bool_results = [
+                status.get("bool_result") for status in rank_statuses if status.get("bool_result") is not None
+            ]
+            if bool_results and not all(bool_results):
+                return False
+        return result
+
+    @staticmethod
+    def _handle_rpc_response(response: Any) -> Any:
+        MultiprocDiffusionExecutor._raise_for_rpc_error_dict(response)
+        response = MultiprocDiffusionExecutor._unwrap_rpc_result_envelope(response)
+        # After unwrapping, a worker method result may itself be the same
+        # {"status": "error"} shape produced by worker_busy_loop transport
+        # failures. Preserve the pre-envelope error handling for that case.
+        MultiprocDiffusionExecutor._raise_for_rpc_error_dict(response)
+        return response
+
    def _launch_workers(self, broadcast_handle, wake_events):
        od_config = self.od_config
        logger.info("Starting server...")
@@ -257,38 +307,6 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
        """Register a callback invoked when a worker process dies."""
        self._failure_callbacks.append(callback)

-    def add_req(self, request: OmniDiffusionRequest) -> DiffusionOutput:
-        self._ensure_open()
-        rpc_request = {
-            "type": "rpc",
-            "method": "generate",
-            "args": (request,),
-            "kwargs": {},
-            "output_rank": 0,
-            "exec_all_ranks": True,
-        }
-
-        try:
-            self._broadcast_mq.enqueue(rpc_request)
-            response = self._result_mq.dequeue()
-
-            try:
-                unpack_diffusion_output_shm(response)
-            except Exception as e:
-                logger.warning("SHM unpack failed (data may already be inline): %s", e)
-
-            if isinstance(response, dict) and response.get("status") == "error":
-                raise RuntimeError(
-                    f"Worker failed with error '{response.get('error')}', "
-                    "please check the stack trace above for the root cause"
-                )
-            if not isinstance(response, DiffusionOutput):
-                raise RuntimeError(f"Unexpected response type for generate: {type(response)!r}")
-            return response
-        except Exception as e:
-            logger.error(f"Generate call failed: {e}")
-            raise
-
    def execute_request(self, scheduler_output: DiffusionSchedulerOutput) -> BaseRunnerOutput:
        """Adapt request-mode scheduler output to worker execute_model RPC."""
        from vllm_omni.diffusion.worker.utils import RunnerOutput
@@ -358,16 +376,21 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
        deadline = None if timeout is None else time.monotonic() + timeout
        kwargs = kwargs or {}

-        # Prepare RPC request message
-        # When unique_reply_rank is None, all workers must execute the RPC
-        # but only rank 0 can reply (it's the only one with a result_mq).
+        # Prepare RPC request message. When unique_reply_rank is None, all
+        # workers must execute the RPC but only rank 0 can reply (it's the
+        # only one with a result_mq). Collect detailed rank statuses only for
+        # this control-plane all-rank path; forward-path exec_all_ranks RPCs
+        # avoid the per-step host object gather.
+        execute_all_ranks = unique_reply_rank is None or exec_all_ranks
+        collect_rank_status = unique_reply_rank is None
        rpc_request = {
            "type": "rpc",
            "method": method,
            "args": args,
            "kwargs": kwargs,
            "output_rank": unique_reply_rank if unique_reply_rank is not None else 0,
-            "exec_all_ranks": unique_reply_rank is None or exec_all_ranks,
+            "exec_all_ranks": execute_all_ranks,
+            "collect_rank_status": collect_rank_status,
        }

        try:
@@ -386,12 +409,7 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
                except Exception as e:
                    logger.warning("SHM unpack failed (data may already be inline): %s", e)

-                # Check if response indicates an error
-                if isinstance(response, dict) and response.get("status") == "error":
-                    raise RuntimeError(
-                        f"Worker failed with error '{response.get('error')}', "
-                        "please check the stack trace above for the root cause"
-                    )
+                response = MultiprocDiffusionExecutor._handle_rpc_response(response)

                responses.append(response)

@@ -18,6 +18,7 @@ import torch
 from vllm_omni.diffusion.data import DiffusionOutput

 _SHM_TENSOR_THRESHOLD = 1_000_000  # 1 MB
+DIFFUSION_RPC_RESULT_ENVELOPE = "diffusion_rpc_result"


 def _tensor_to_shm(tensor: torch.Tensor) -> dict[str, Any]:
@@ -130,15 +131,26 @@ def _pack_diffusion_fields(output: DiffusionOutput) -> DiffusionOutput:
    return output


+def _is_rpc_result_envelope(output: object) -> bool:
+    return isinstance(output, dict) and output.get("type") == DIFFUSION_RPC_RESULT_ENVELOPE
+
+
 def pack_diffusion_output_shm(output: object) -> object:
    """Replace large tensors in diffusion worker outputs with SHM handles.

    Supports either a bare ``DiffusionOutput`` or a wrapper object carrying one
-    in ``.result`` (for example ``RunnerOutput``).
+    in ``.result`` (for example ``RunnerOutput``), or an RPC result envelope
+    carrying the diffusion output in ``["result"]``.
    """
    if isinstance(output, DiffusionOutput):
        return _pack_diffusion_fields(output)

+    if _is_rpc_result_envelope(output):
+        result = output.get("result")
+        if isinstance(result, DiffusionOutput):
+            output["result"] = _pack_diffusion_fields(result)
+        return output
+
    result = getattr(output, "result", None)
    if isinstance(result, DiffusionOutput):
        output.result = _pack_diffusion_fields(result)
@@ -158,6 +170,12 @@ def unpack_diffusion_output_shm(output: object) -> object:
    if isinstance(output, DiffusionOutput):
        return _unpack_diffusion_fields(output)

+    if _is_rpc_result_envelope(output):
+        result = output.get("result")
+        if isinstance(result, DiffusionOutput):
+            output["result"] = _unpack_diffusion_fields(result)
+        return output
+
    result = getattr(output, "result", None)
    if isinstance(result, DiffusionOutput):
        output.result = _unpack_diffusion_fields(result)
@@ -617,6 +617,8 @@ class CausalWanModel(nn.Module):
    Architecture (14B): 40 layers, dim=5120, heads=40, ffn=13824
    """

+    _layerwise_offload_blocks_attrs = ["blocks"]
+
    def __init__(
        self,
        model_type: str = "t2v",
@@ -49,6 +49,7 @@ from vllm_omni.diffusion.models.dreamzero.utils import (
 )
 from vllm_omni.diffusion.models.schedulers.scheduling_flow_unipc_multistep import FlowUniPCMultistepScheduler
 from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.platforms import current_omni_platform

 logger = logging.getLogger(__name__)
 MAX_DREAMZERO_SESSIONS = 64
@@ -363,12 +364,10 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
        latents: tuple[torch.Tensor, torch.Tensor],
        do_true_cfg: bool,
    ) -> tuple[torch.Tensor, torch.Tensor]:
-        """Post-step sync: .contiguous() + cuda.synchronize()"""
+        """Post-step sync: .contiguous() + platform synchronize."""
        latents = tuple(t.contiguous() for t in latents)
        if do_true_cfg and get_classifier_free_guidance_world_size() > 1:
-            device = next((t.device for t in latents if t.is_cuda), None)
-            if device is not None:
-                torch.cuda.current_stream(device).synchronize()
+            current_omni_platform.synchronize()
        return latents

    # -----------------------------------------------------------------------
@@ -4,7 +4,6 @@
 from typing import Any

 import vllm.forward_context as _vllm_fc
-from vllm.model_executor.layers.fused_moe import FusedMoE
 from vllm.utils.import_utils import resolve_obj_by_qualname

 from vllm_omni.platforms import current_omni_platform
@@ -25,22 +24,93 @@ def _set_forward_context_num_tokens(num_tokens: int) -> None:
        forward_context.in_profile_run = False


-class HunyuanFusedMoEDefault(FusedMoE):
-    def __init__(self, *, prefix: str = "", **kwargs: Any) -> None:
+class HunyuanFusedMoEDefault:
+    """Adapter that configures the upstream FusedMoE ``MoERunner`` for HunyuanImage3.
+
+    Upstream commit dc68bd8c41 refactored FusedMoE from a class (``nn.Module``)
+    into a factory function that returns a ``MoERunner`` instance, whose expert
+    weights live in a ``routed_experts`` submodule
+    (``...experts.routed_experts.w13_weight`` / ``...w2_weight``).
+
+    This adapter builds that runner, installs the omni-specific forward-context
+    setup and one-shot kernel-initialisation hook, and returns the runner
+    *directly* from ``__new__`` so the parent MoE block registers it as a real
+    ``nn.Module`` submodule.
+
+    Returning the runner (rather than wrapping it in a plain object that holds it
+    in an attribute) is required for correctness: a non-Module wrapper hides the
+    runner's parameters from ``named_parameters()``, so ``load_weights`` cannot
+    find ``...experts.routed_experts.w13_weight`` and raises ``KeyError`` during
+    weight loading.
+    """
+
+    def __new__(cls, *, prefix: str = "", **kwargs: Any) -> Any:
        # Current vLLM FusedMoE handles output reduction internally.
        kwargs.pop("reduce_results", None)
-        super().__init__(prefix=prefix, **kwargs)
-        self._prefix = prefix
-        self._init_hook_handle = self.register_forward_pre_hook(self._initialize_kernel_hook, with_kwargs=True)
+        # FusedMoE is now a factory function — call it to get a MoERunner.
+        from vllm.model_executor.layers.fused_moe import FusedMoE as _FusedMoE

-    def _initialize_kernel_hook(self, module: Any, args: Any, kwargs: Any) -> None:
-        if self.quant_method and getattr(self.quant_method, "moe_kernel", None) is None:
-            self.quant_method.process_weights_after_loading(self)
-        self._init_hook_handle.remove()
+        moe_runner = _FusedMoE(prefix=prefix, **kwargs)

-    def forward(self, hidden_states: Any, router_logits: Any) -> Any:
-        _set_forward_context_num_tokens(hidden_states.shape[0])
-        return super().forward(hidden_states, router_logits)
+        # Set ForwardContext.num_tokens before each forward. After the rebase to
+        # vLLM 0.18.0 FusedMoE requires this; without it MoE routing is silently
+        # incorrect. Previously done in the wrapper's forward(); now a pre-hook
+        # on the runner so we can return the runner directly.
+        def _num_tokens_pre_hook(module: Any, args: Any, kwargs: Any) -> None:
+            hidden_states = kwargs.get("hidden_states")
+            if hidden_states is None and args:
+                hidden_states = args[0]
+            if hidden_states is not None:
+                _set_forward_context_num_tokens(hidden_states.shape[0])
+
+        moe_runner.register_forward_pre_hook(_num_tokens_pre_hook, with_kwargs=True)
+
+        # One-shot lazy kernel initialisation on the first forward (no-op unless
+        # the runner exposes an uninitialised quant_method). Mirrors the prior
+        # wrapper behaviour exactly, just bound to the runner module.
+        init_handle: Any = None
+
+        def _kernel_init_pre_hook(module: Any, args: Any, kwargs: Any) -> None:
+            nonlocal init_handle
+            quant_method = getattr(module, "quant_method", None)
+            if quant_method is not None and getattr(quant_method, "moe_kernel", None) is None:
+                quant_method.process_weights_after_loading(module)
+            if init_handle is not None:
+                init_handle.remove()
+
+        init_handle = moe_runner.register_forward_pre_hook(_kernel_init_pre_hook, with_kwargs=True)
+
+        return moe_runner
+
+    @staticmethod
+    def make_expert_params_mapping(
+        model: Any,
+        ckpt_gate_proj_name: str,
+        ckpt_down_proj_name: str,
+        ckpt_up_proj_name: str,
+        num_experts: int,
+        num_redundant_experts: int = 0,
+    ) -> list[tuple[str, str, int, str]]:
+        """Delegate to the upstream standalone function.
+
+        Upstream vLLM refactored ``FusedMoE`` from a class (which had
+        ``make_expert_params_mapping`` as a classmethod) to a factory
+        function.  The method was moved to a standalone function
+        ``fused_moe_make_expert_params_mapping`` in
+        ``vllm.model_executor.layers.fused_moe``.
+        """
+        from vllm.model_executor.layers.fused_moe import (
+            fused_moe_make_expert_params_mapping,
+        )
+
+        return fused_moe_make_expert_params_mapping(
+            model,
+            ckpt_gate_proj_name=ckpt_gate_proj_name,
+            ckpt_down_proj_name=ckpt_down_proj_name,
+            ckpt_up_proj_name=ckpt_up_proj_name,
+            num_experts=num_experts,
+            num_redundant_experts=num_redundant_experts,
+        )


 class HunyuanFusedMoE:
@@ -1009,28 +1009,43 @@ class WanTransformer3DModel(nn.Module):
        if encoder_hidden_states_image is not None:
            encoder_hidden_states = torch.concat([encoder_hidden_states_image, encoder_hidden_states], dim=1)

-        # Check for SP auto_pad: create attention mask dynamically if padding was applied
-        hidden_states_mask = None  # default
-        config = get_forward_context().omni_diffusion_config
-        parallel_config = config.parallel_config
-        if parallel_config is not None and parallel_config.sequence_parallel_size > 1:
-            ctx = get_forward_context()
-            if ctx.sp_original_seq_len is not None and ctx.sp_padding_size > 0:
-                # Create mask for the full (padded) sequence
-                # valid positions = True, padding positions = False
-                batch_size = hidden_states.shape[0]
-                padded_seq_len = ctx.sp_original_seq_len + ctx.sp_padding_size
-                hidden_states_mask = torch.ones(
-                    batch_size,
-                    padded_seq_len,
-                    dtype=torch.bool,
-                    device=hidden_states.device,
-                )
-                hidden_states_mask[:, ctx.sp_original_seq_len :] = False
-
-        # if mask is all true, set it to None
-        if hidden_states_mask is not None and hidden_states_mask.all():
-            hidden_states_mask = None
+        hidden_states_mask = None
+        ctx = get_forward_context()
+        parallel_config = ctx.omni_diffusion_config.parallel_config
+        if (
+            parallel_config is not None
+            and parallel_config.sequence_parallel_size > 1
+            and parallel_config.mask_sp_padding
+            and ctx.sp_original_seq_len is not None
+            and ctx.sp_padding_size > 0
+        ):
+            batch_size = hidden_states.shape[0]
+            padded_seq_len = ctx.sp_original_seq_len + ctx.sp_padding_size
+            hidden_states_mask = torch.ones(
+                batch_size,
+                padded_seq_len,
+                dtype=torch.bool,
+                device=hidden_states.device,
+            )
+            hidden_states_mask[:, ctx.sp_original_seq_len :] = False
+            if hidden_states_mask.all():
+                hidden_states_mask = None
+        elif (
+            parallel_config is not None
+            and parallel_config.sequence_parallel_size > 1
+            and not parallel_config.mask_sp_padding
+            and ctx.sp_original_seq_len is not None
+            and ctx.sp_padding_size > 0
+        ):
+            logger.warning_once(
+                "SP auto-padding applied %d token(s) (seq_len=%d, ulysses_degree=%d). "
+                "Padding tokens are not masked from attention (mask_sp_padding=False), "
+                "which avoids the varlen attention path but may produce minor numerical differences. "
+                "Set parallel_config.mask_sp_padding=True to restore strict masking.",
+                ctx.sp_padding_size,
+                ctx.sp_original_seq_len,
+                parallel_config.sequence_parallel_size,
+            )

        # Transformer blocks
        for block in self.blocks[self.start_layer : self.end_layer]:
@@ -10,6 +10,7 @@ from typing import Any
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from vllm.logger import init_logger
 from vllm.model_executor.layers.quantization.base_config import QuantizationConfig

 from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelInput
@@ -21,6 +22,8 @@ from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import (
    WanTransformerBlock,
 )

+logger = init_logger(__name__)
+

 class VaceWanTransformerBlock(WanTransformerBlock):
    """VACE variant of WanTransformerBlock with proj_in/proj_out for skip connections."""
@@ -214,12 +217,16 @@ class WanVACETransformer3DModel(WanTransformer3DModel):
        # Shard hidden_states via _sp_plan hook (before VACE, not at blocks.0)
        hidden_states = self._sp_shard_point(hidden_states)

-        # SP state and attention mask for padding
        hidden_states_mask = None
        ctx = get_forward_context()
        parallel_config = ctx.omni_diffusion_config.parallel_config
        sp_size = parallel_config.sequence_parallel_size if parallel_config is not None else 1
-        if ctx.sp_original_seq_len is not None and ctx.sp_padding_size > 0:
+        if (
+            parallel_config is not None
+            and parallel_config.mask_sp_padding
+            and ctx.sp_original_seq_len is not None
+            and ctx.sp_padding_size > 0
+        ):
            padded_seq_len = ctx.sp_original_seq_len + ctx.sp_padding_size
            hidden_states_mask = torch.ones(
                batch_size,
@@ -228,6 +235,21 @@ class WanVACETransformer3DModel(WanTransformer3DModel):
                device=hidden_states.device,
            )
            hidden_states_mask[:, ctx.sp_original_seq_len :] = False
+        elif (
+            parallel_config is not None
+            and not parallel_config.mask_sp_padding
+            and ctx.sp_original_seq_len is not None
+            and ctx.sp_padding_size > 0
+        ):
+            logger.warning_once(
+                "SP auto-padding applied %d token(s) (seq_len=%d, ulysses_degree=%d). "
+                "Padding tokens are not masked from attention (mask_sp_padding=False), "
+                "which avoids the varlen attention path but may produce minor numerical differences. "
+                "Set parallel_config.mask_sp_padding=True to restore strict masking.",
+                ctx.sp_padding_size,
+                ctx.sp_original_seq_len,
+                sp_size,
+            )

        # VACE: embed context and run conditioning blocks
        vace_hints = None
@@ -12,6 +12,7 @@ import gc
 import multiprocessing as mp
 import os
 import signal
+import traceback
 from collections.abc import Iterable, Iterator
 from contextlib import AbstractContextManager, contextmanager, nullcontext
 from dataclasses import dataclass
@@ -41,7 +42,7 @@ from vllm_omni.diffusion.distributed.parallel_state import (
    initialize_model_parallel,
 )
 from vllm_omni.diffusion.forward_context import set_forward_context
-from vllm_omni.diffusion.ipc import pack_diffusion_output_shm
+from vllm_omni.diffusion.ipc import DIFFUSION_RPC_RESULT_ENVELOPE, pack_diffusion_output_shm
 from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput
@@ -740,6 +741,18 @@ class WorkerProc:
        """Receive messages from broadcast queue."""
        return self.mq.dequeue(indefinite=True)

+    def _gather_rpc_rank_statuses(self, status: dict[str, Any]) -> list[dict[str, Any]]:
+        if not torch.distributed.is_initialized():
+            return [status]
+
+        world_size = torch.distributed.get_world_size()
+        statuses: list[dict[str, Any] | None] = [None] * world_size
+        torch.distributed.all_gather_object(statuses, status)
+        missing_ranks = [rank for rank, rank_status in enumerate(statuses) if rank_status is None]
+        if missing_ranks:
+            logger.warning("RPC rank status gather returned missing entries for ranks: %s", missing_ranks)
+        return [s for s in statuses if s is not None]
+
    def execute_rpc(self, rpc_request: dict) -> tuple[object | None, bool]:
        """Execute an RPC request and indicate whether to reply."""
        method = rpc_request["method"]
@@ -747,6 +760,10 @@ class WorkerProc:
        kwargs = rpc_request.get("kwargs", {})
        output_rank = rpc_request.get("output_rank")
        exec_all_ranks = rpc_request.get("exec_all_ranks", False)
+        collect_rank_status = rpc_request.get("collect_rank_status", False)
+
+        if collect_rank_status and not exec_all_ranks:
+            raise ValueError("collect_rank_status requires exec_all_ranks=True so all ranks enter the status gather")

        should_execute = exec_all_ranks or output_rank is None or output_rank == self.gpu_id
        should_reply = (output_rank is None or output_rank == self.gpu_id) and self.result_mq is not None
@@ -754,13 +771,52 @@ class WorkerProc:
        if not should_execute:
            return None, False

+        result = None
+        rpc_exception: Exception | None = None
+        status: dict[str, Any] = {
+            "rank": self.gpu_id,
+            "ok": True,
+            "error": None,
+            "error_type": None,
+            "traceback": None,
+            "bool_result": None,
+        }
+
        try:
            # Use execute_method from WorkerWrapperBase for consistent method resolution
            result = self.worker.execute_method(method, *args, **kwargs)
-            return result, should_reply
        except Exception as e:
            logger.error(f"Error executing RPC: {e}", exc_info=True)
-            raise e
+            rpc_exception = e
+            status.update(
+                {
+                    "ok": False,
+                    "error": str(e),
+                    "error_type": type(e).__name__,
+                    "traceback": traceback.format_exc(),
+                }
+            )
+
+        if isinstance(result, bool):
+            status["bool_result"] = result
+
+        if collect_rank_status:
+            rank_statuses = self._gather_rpc_rank_statuses(status)
+            if should_reply:
+                return (
+                    {
+                        "type": DIFFUSION_RPC_RESULT_ENVELOPE,
+                        "method": method,
+                        "result": result,
+                        "rank_statuses": rank_statuses,
+                    },
+                    True,
+                )
+            return None, False
+
+        if rpc_exception is not None:
+            raise rpc_exception
+        return result, should_reply

    def worker_busy_loop(self) -> None:
        """Main busy loop for Multiprocessing Workers."""
@@ -801,7 +857,7 @@ class WorkerProc:
                except Exception as e:
                    logger.error(f"Error processing RPC: {e}", exc_info=True)
                    if self.result_mq is not None:
-                        self.return_result(DiffusionOutput.from_exception(e))
+                        self.return_result({"status": "error", "error": str(e)})

            elif isinstance(msg, dict) and msg.get("type") == "shutdown":
                logger.info("Worker %s: Received shutdown message", self.gpu_id)
@@ -114,6 +114,7 @@ class MooncakeTransferEngineConnector(OmniConnectorBase):
        self._listener_ready = threading.Event()
        self._local_buffers: dict[str, Any] = {}
        self._local_buffers_lock = threading.Lock()
+        self._put_lock = threading.Lock()
        self._req_local = threading.local()
        self._worker_local = threading.local()
        self._last_ttl_check: float = _time_mod.monotonic()
@@ -382,6 +383,14 @@ class MooncakeTransferEngineConnector(OmniConnectorBase):

        put_key = self._make_key(put_key, from_stage, to_stage)

+        # Serialize concurrent put() calls on the same connector to prevent
+        # races in the alloc -> copy -> store-metadata flow at the Mooncake
+        # C++ engine level.
+        with self._put_lock:
+            return self._put_impl(put_key, data)
+
+    def _put_impl(self, put_key: str, data: Any) -> tuple[bool, int, dict[str, Any] | None]:
+        """Internal put implementation, called under _put_lock."""
        try:
            src_addr = 0
            size = 0
@@ -3,7 +3,7 @@

 import importlib
 from collections import defaultdict, deque
-from collections.abc import Callable
+from collections.abc import Callable, Mapping
 from typing import Any

 import torch
@@ -146,7 +146,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):

    def save_async(
        self,
-        pooling_output: torch.Tensor | None = None,
+        multimodal_output: dict[str, Any] | None = None,
        request: Request | None = None,
        is_segment_finished: bool = False,
    ):
@@ -163,7 +163,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
        as ``request.is_finished()``.

        Args:
-            pooling_output: Partial pooling output dictionary
+            multimodal_output: Per-request multimodal output dictionary
            request: Request object
            is_segment_finished: whether the segment of request is finished
        """
@@ -183,7 +183,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):

        self.requests_num_chunks_sent[request.external_req_id] = confirmed_num_computed_tokens
        task = {
-            "pooling_output": pooling_output,
+            "multimodal_output": multimodal_output,
            "request": request,
            "is_finished": is_finished,
            "is_segment_finished": is_segment_finished,
@@ -292,8 +292,8 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
        return False

    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
+        raw_mm = task["multimodal_output"]
+        multimodal_output = unflatten_payload(raw_mm) if isinstance(raw_mm, Mapping) else raw_mm
        request = task["request"]
        is_finished = task["is_finished"]
        is_segment_finished = task["is_segment_finished"]
@@ -308,7 +308,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
            try:
                payload_data = self.custom_process_next_stage_input_func(
                    transfer_manager=self,
-                    pooling_output=pooling_output,
+                    multimodal_output=multimodal_output,
                    request=request,
                    # Existing processors use is_finished as a flush signal.
                    is_finished=is_segment_finished,
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

+from collections.abc import Mapping
 from dataclasses import asdict, is_dataclass
 from typing import Any

@@ -150,10 +151,13 @@ class OmniMsgpackEncoder:
            "multi_modal_placeholders": getattr(obj, "multi_modal_placeholders", None),
            "kv_transfer_params": obj.kv_transfer_params,
        }
-        # Handle dynamically added multimodal_output attribute
+        # Handle multimodal_output attribute (MultimodalPayload or dict)
        mm_output = getattr(obj, "multimodal_output", None)
        if mm_output is not None:
-            result["multimodal_output"] = mm_output
+            if isinstance(mm_output, Mapping):
+                result["multimodal_output"] = dict(mm_output)
+            else:
+                result["multimodal_output"] = mm_output
        return result

    def _encode_completion_output(self, obj: CompletionOutput) -> dict[str, Any]:
@@ -161,7 +165,11 @@ class OmniMsgpackEncoder:
        result = asdict(obj)
        mm_output = getattr(obj, "multimodal_output", None)
        if mm_output is not None:
-            result["multimodal_output"] = mm_output
+            # Convert MultimodalPayload to plain dict for wire format
+            if isinstance(mm_output, Mapping):
+                result["multimodal_output"] = dict(mm_output)
+            else:
+                result["multimodal_output"] = mm_output
        return result


@@ -117,7 +117,10 @@ class OmniEngineCoreRequest(EngineCoreRequest):


 class OmniEngineCoreOutput(EngineCoreOutput):
-    pooling_output: dict[str, torch.Tensor] | None = None
+    # Dedicated channel for multimodal outputs (image/audio/latent).
+    # pooling_output is inherited from EngineCoreOutput as torch.Tensor | None
+    # and retains its original vLLM semantics for pooling/embedding tasks.
+    multimodal_output: dict[str, torch.Tensor] | None = None
    # Finished flag for streaming input segment
    is_segment_finished: bool | None = False
    # Streaming update prompt length
@@ -1480,7 +1480,7 @@ class AsyncOmniEngine:
            return

        mm_data = prompt.get("multi_modal_data")
-        if not isinstance(mm_data, dict) or not mm_data:
+        if not isinstance(mm_data, Mapping) or not mm_data:
            return

        from vllm.multimodal.hasher import MultiModalHasher
@@ -2480,6 +2480,28 @@ class AsyncOmniEngine:
                logger.exception("[AsyncOmniEngine] Failed to close OmniCoordinatorRuntime during shutdown")
            self._coordinator_runtime = None

+        # ── Release CuMem allocator memory pool ──────────────────────────────
+        # When enable_sleep_mode is in use, the CuMem (CUDA Virtual Memory
+        # Management) allocator holds model weights in a singleton memory pool
+        # that lives in the parent process.  Killing the engine-core subprocess
+        # does NOT release this pool — the weights stay resident on the GPU
+        # and can cause CUDA OOM for subsequent engine instances (especially
+        # large models like BAGEL-7B-MoT whose weights alone consume ~134 GiB).
+        #
+        # Discard mode (level=2) is correct at shutdown: there is no benefit to
+        # keeping a CPU backup when the engine is being torn down.
+        try:
+            from vllm.device_allocator.cumem import CuMemAllocator, cumem_available
+
+            if cumem_available:
+                allocator = CuMemAllocator.get_instance()
+                # Sleep at level 2 discards all pool memory from the GPU
+                # without creating CPU backups — cheapest and fastest.
+                allocator.sleep()
+                logger.debug("[AsyncOmniEngine] Released CuMem memory pool during shutdown")
+        except Exception:
+            pass
+
    def _try_shutdown(self, *args, **kwargs) -> None:
        try:
            self.shutdown()
@@ -6,6 +6,7 @@ This module defines structured types for multimodal outputs.

 from __future__ import annotations

+from collections.abc import Iterator, Mapping
 from dataclasses import dataclass, field
 from typing import Any

@@ -13,10 +14,15 @@ import torch
 from vllm.outputs import CompletionOutput


-@dataclass
-class MultimodalPayload:
+@dataclass(eq=False)
+class MultimodalPayload(Mapping):
    """Structured multimodal output payload.

+    Implements ``collections.abc.Mapping`` so that ``isinstance(payload, dict)``
+    style checks in downstream code can be replaced with duck-typing, and
+    ``payload.get(key)``, ``payload[key]``, ``key in payload``, ``len(payload)``
+    all work seamlessly for both tensors and metadata.
+
    Attributes:
        tensors: Dictionary mapping modality/key names to their tensors.
        metadata: Optional dictionary for non-tensor metadata
@@ -35,18 +41,49 @@ class MultimodalPayload:

    @property
    def is_empty(self) -> bool:
-        """Return True if the payload has no tensors."""
-        return len(self.tensors) == 0
+        """Return True if the payload has no tensors and no metadata."""
+        return len(self.tensors) == 0 and len(self.metadata) == 0

-    def get(self, key: str) -> torch.Tensor | None:
-        """Get a tensor by key, returning None if not found."""
-        return self.tensors.get(key)
+    def get(self, key: str, default: Any = None) -> Any:
+        """Get a value by key, searching tensors first then metadata."""
+        if key in self.tensors:
+            return self.tensors[key]
+        return self.metadata.get(key, default)

    def __contains__(self, key: str) -> bool:
-        return key in self.tensors
+        return key in self.tensors or key in self.metadata
+
+    def __getitem__(self, key: str) -> Any:
+        """Dict-like indexing: search tensors first, then metadata."""
+        if key in self.tensors:
+            return self.tensors[key]
+        if key in self.metadata:
+            return self.metadata[key]
+        raise KeyError(key)

    def __len__(self) -> int:
-        return len(self.tensors)
+        return len(self.tensors) + len(self.metadata)
+
+    def __iter__(self) -> Iterator[str]:
+        yield from self.tensors
+        yield from self.metadata
+
+    def __bool__(self) -> bool:
+        return bool(self.tensors) or bool(self.metadata)
+
+    def __eq__(self, other: object) -> bool:
+        """Support equality with plain dicts and other Mappings."""
+        if isinstance(other, MultimodalPayload):
+            return self.tensors == other.tensors and self.metadata == other.metadata
+        if isinstance(other, Mapping):
+            return dict(self) == dict(other)
+        return NotImplemented
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert back to a plain dict (tensors + metadata merged)."""
+        result: dict[str, Any] = dict(self.tensors)
+        result.update(self.metadata)
+        return result

    @classmethod
    def from_dict(cls, data: dict[str, Any] | None) -> MultimodalPayload | None:
@@ -21,6 +21,8 @@ _MODALITY_ALIASES: dict[str, str] = {
    "waveform": "audio",
    "pixel_values": "image",
    "pixels": "image",
+    "token_ids": "text",
+    "tokens": "text",
 }


@@ -1,8 +1,9 @@
+from dataclasses import fields as dataclass_fields
 from typing import Any

 import torch
 from vllm.logger import init_logger
-from vllm.outputs import PoolingRequestOutput
+from vllm.outputs import CompletionOutput, PoolingRequestOutput, RequestOutput
 from vllm.sampling_params import RequestOutputKind
 from vllm.tokenizers import TokenizerLike
 from vllm.v1.engine import EngineCoreOutput, EngineCoreRequest, FinishReason
@@ -16,12 +17,98 @@ from vllm.v1.engine.parallel_sampling import ParentRequest
 from vllm.v1.metrics.stats import IterationStats, RequestStateStats

 from vllm_omni.data_entry_keys import unflatten_payload
-from vllm_omni.engine.output_modality import DRAINABLE_MODALITIES
+from vllm_omni.engine.mm_outputs import MultimodalCompletionOutput, MultimodalPayload
+from vllm_omni.engine.output_modality import (
+    DRAINABLE_MODALITIES,
+    OutputModality,
+    TensorAccumulationStrategy,
+    get_accumulation_strategy,
+)
 from vllm_omni.outputs import OmniRequestOutput

 logger = init_logger(__name__)


+# ---------------------------------------------------------------------------
+# Module-level helpers
+# ---------------------------------------------------------------------------
+
+
+def _to_cpu(x: Any) -> Any:
+    """Move a tensor to CPU; pass through non-tensors unchanged."""
+    if isinstance(x, torch.Tensor):
+        try:
+            return x.detach().to("cpu", non_blocking=True).contiguous()
+        except (RuntimeError, AttributeError):
+            return x
+    if isinstance(x, dict):
+        return {str(sk): _to_cpu(sv) for sk, sv in x.items()}
+    return x
+
+
+def _merge_payload(target: MultimodalPayload, incoming: MultimodalPayload) -> None:
+    """Merge *incoming* into *target* using list-based deferred concatenation."""
+    for k, v in incoming.tensors.items():
+        if k not in target.tensors:
+            target.tensors[k] = v
+        else:
+            existing = target.tensors[k]
+            if isinstance(existing, list):
+                existing.append(v)
+            elif isinstance(existing, torch.Tensor) and isinstance(v, torch.Tensor):
+                target.tensors[k] = [existing, v]
+            else:
+                target.tensors[k] = v
+
+    for k, v in incoming.metadata.items():
+        if k not in target.metadata:
+            target.metadata[k] = v
+        else:
+            existing = target.metadata[k]
+            if isinstance(v, torch.Tensor) and isinstance(existing, torch.Tensor):
+                target.metadata[k] = [existing, v]
+            elif isinstance(v, torch.Tensor) and isinstance(existing, list):
+                existing.append(v)
+            else:
+                # Metadata: replace with latest value
+                target.metadata[k] = v
+
+
+def _cat_tensors(
+    tensors: list[torch.Tensor],
+    strategy: TensorAccumulationStrategy,
+) -> torch.Tensor:
+    """Concatenate a list of tensors according to *strategy*."""
+    if strategy == TensorAccumulationStrategy.CONCAT_LAST:
+        return torch.cat(tensors, dim=-1)
+    if strategy == TensorAccumulationStrategy.REPLACE:
+        return tensors[-1]
+    # CONCAT_DIM0 / APPEND_LIST / default
+    return torch.cat(tensors, dim=0)
+
+
+def _modality_to_type_string(modality: OutputModality) -> str:
+    """Convert an OutputModality flag to a lowercase type string."""
+    try:
+        if OutputModality.AUDIO in modality:
+            return "audio"
+        if OutputModality.IMAGE in modality:
+            return "image"
+        if OutputModality.LATENT in modality:
+            return "latent"
+    except TypeError:
+        # Flag identity mismatch (e.g. after module reload in tests).
+        name = getattr(modality, "name", "") or ""
+        lowered = name.lower()
+        if "audio" in lowered:
+            return "audio"
+        if "image" in lowered:
+            return "image"
+        if "latent" in lowered:
+            return "latent"
+    return "text"
+
+
 class OmniRequestState(RequestState):
    """Request state for omni models, tracking multimodal outputs.

@@ -35,16 +122,20 @@ class OmniRequestState(RequestState):
        *args,
        **kwargs,
    ):
+        # arrival_time is always passed as a keyword argument by the sole caller
+        # (from_new_request).  The *args fallback via positional index was removed
+        # because upstream RequestState.__init__ reorders parameters between
+        # releases, making args[12] fragile.
        arrival_time = kwargs.get("arrival_time")
-        if arrival_time is None and len(args) > 12:
-            arrival_time = args[12]
        super().__init__(*args, **kwargs)
        self.native_text_stats = RequestStateStats(arrival_time=float(arrival_time or 0.0))
        # Omni-specific: multimodal output accumulation
-        # NOTE: Keys in mm_accumulated matter, because they dictate which
-        # outputs are drained (i.e., only drain modality keys, don't drain
-        # hidden states).
-        self.mm_accumulated: dict[str, Any] = {}
+        # TODO: mm_type is per-request, not per-key. If a model ever produces
+        # both audio and latent outputs, the modality type would flip on each
+        # add_multimodal_tensor() call. Consider tracking per-key modality
+        # types (e.g. dict[str, str]) for future multi-output models.
+        self.mm_type: str | None = None
+        self.mm_accumulated: MultimodalPayload = MultimodalPayload()

    def apply_streaming_update(self, update) -> None:
        super().apply_streaming_update(update)
@@ -53,8 +144,6 @@ class OmniRequestState(RequestState):
    @staticmethod
    def _to_cpu(x):
        """Try to convert to CPU tensor if needed."""
-        # TODO: Make this more robust and unify with other payload
-        # building utils, we do this in multiple places.
        if isinstance(x, torch.Tensor):
            try:
                return x.detach().to("cpu", non_blocking=True).contiguous()
@@ -63,109 +152,94 @@ class OmniRequestState(RequestState):
        return x

    def add_multimodal_tensor(self, payload: Any | None, mm_type: str | None) -> None:
+        """Accumulate a multimodal tensor payload into the request state.
+
+        Normalizes incoming payloads (dict or raw tensor) into a
+        MultimodalPayload and merges with any previously accumulated data.
+        Uses list-based deferred concatenation to avoid O(n²) repeated
+        torch.cat calls.
+        """
        if payload is None:
            return
-
-        mm_type = (mm_type or "").lower()
        try:
+            if mm_type:
+                self.mm_type = (mm_type or "").lower()
+
+            target_key = self.mm_type or "hidden"
+
+            # Build incoming MultimodalPayload
            if isinstance(payload, dict):
-                # Keep payload flat (dotted keys like "hidden_states.layer_0")
-                # during accumulation so that all values are tensors/scalars and
-                # the merge logic below works correctly.  Unflatten happens
-                # later in _consolidate_multimodal_tensors after concatenation.
-
-                incoming: dict[str, Any] = {}
-                # TODO (Alex): Clean up and simplify key management
-                target_key = mm_type or "hidden"
-
+                remapped: dict[str, Any] = {}
                for k, v in payload.items():
+                    # Normalize producer keys to the modality name.
+                    # AR runners produce {"hidden": ...} and generation
+                    # runners produce {"model_outputs": ...}; remap both
+                    # to the semantic modality key (e.g. "audio", "latent").
                    if k == "model_outputs":
                        k = target_key
                    elif k == "hidden" and target_key != "hidden":
                        k = target_key
-
-                    incoming[k] = self._to_cpu(v)
+                    remapped[k] = _to_cpu(v)
+                incoming = MultimodalPayload.from_dict(remapped)
+            elif isinstance(payload, MultimodalPayload):
+                incoming = payload
            else:
-                key = mm_type or "hidden"
-                incoming = {key: self._to_cpu(payload)}
+                incoming = MultimodalPayload.from_dict({target_key: _to_cpu(payload)})

-            if not self.mm_accumulated:
+            if incoming is None:
+                return
+
+            if self.mm_accumulated.is_empty:
                self.mm_accumulated = incoming
            else:
-                for k, v in incoming.items():
-                    if k not in self.mm_accumulated:
-                        self.mm_accumulated[k] = v
-                    else:
-                        existing = self.mm_accumulated[k]
-                        if isinstance(v, torch.Tensor) and isinstance(existing, torch.Tensor):
-                            self.mm_accumulated[k] = [existing, v]
-                        elif isinstance(v, torch.Tensor) and isinstance(existing, list):
-                            existing.append(v)
-                        elif isinstance(v, dict) and isinstance(existing, dict):
-                            for sk, sv in v.items():
-                                if sk not in existing:
-                                    existing[sk] = sv
-                                elif isinstance(sv, torch.Tensor) and isinstance(existing[sk], torch.Tensor):
-                                    existing[sk] = [existing[sk], sv]
-                                elif isinstance(sv, torch.Tensor) and isinstance(existing[sk], list):
-                                    existing[sk].append(sv)
-                                else:
-                                    existing[sk] = sv
-                        else:
-                            self.mm_accumulated[k] = v
-        except Exception:
-            # Log and continue without crashing the output pipeline
+                _merge_payload(self.mm_accumulated, incoming)
+        except (ValueError, TypeError, RuntimeError):
            logger.exception("Error accumulating multimodal tensor")

    def _consolidate_multimodal_tensors(self) -> None:
-        """Consolidate accumulated tensor lists into single tensors via concatenation.
+        """Consolidate accumulated tensor lists into single tensors.

-        Only DELTA drains modality keys per-step, so they will never be lists here and
-        can be skipped.  For CUMULATIVE and FINAL_ONLY, modality keys accumulate across
-        steps and need consolidation.
+        Uses TensorAccumulationStrategy derived from the output modality
+        to determine concatenation behavior. Metadata values always use
+        REPLACE (keep latest).
        """
-        if not self.mm_accumulated:
+        if self.mm_accumulated.is_empty:
            return

-        skip_modality = self.output_kind == RequestOutputKind.DELTA
        try:
-            for k, v in self.mm_accumulated.items():
-                if skip_modality and k in DRAINABLE_MODALITIES:
-                    continue
+            modality = OutputModality.from_string(self.mm_type)
+        except (ValueError, KeyError):
+            modality = OutputModality.TEXT
+        strategy = get_accumulation_strategy(modality)
+
+        try:
+            for k, v in list(self.mm_accumulated.tensors.items()):
                if isinstance(v, list) and v and isinstance(v[0], torch.Tensor):
                    try:
+                        self.mm_accumulated.tensors[k] = _cat_tensors(v, strategy)
+                    except RuntimeError:
                        if k == "audio":
-                            # Preserve channel dimension when chunks are compatible;
-                            # fall back to 1-D waveform concatenation for uneven chunks.
                            try:
-                                self.mm_accumulated[k] = torch.cat(v, dim=-1)
+                                self.mm_accumulated.tensors[k] = torch.cat(v, dim=-1)
                            except RuntimeError:
-                                self.mm_accumulated[k] = torch.cat([t.reshape(-1) for t in v], dim=0)
-                        elif k == "sr":
-                            # Sample rate is a constant scalar, keep last value.
-                            self.mm_accumulated[k] = v[-1]
+                                self.mm_accumulated.tensors[k] = torch.cat([t.reshape(-1) for t in v], dim=0)
                        else:
-                            self.mm_accumulated[k] = torch.cat(v, dim=0)
-                    except Exception:
-                        # Keep last tensor on failure
-                        logger.warning(f"Error concatenating tensor for key {k}; keeping last tensor")
-                        self.mm_accumulated[k] = v[-1]
-                elif isinstance(v, dict):
-                    for sk, sv in v.items():
-                        if isinstance(sv, list) and sv and isinstance(sv[0], torch.Tensor):
-                            try:
-                                v[sk] = torch.cat(sv, dim=0)
-                            except Exception:
-                                v[sk] = sv[-1]
-        except Exception:
-            logger.exception("Error consolidating multimodal tensors")
+                            logger.warning(
+                                "Error concatenating tensor for key %s; keeping last tensor",
+                                k,
+                            )
+                            self.mm_accumulated.tensors[k] = v[-1]

-        # Restore nested structure from flat dotted keys now that all tensor
-        # lists have been concatenated into single tensors.
-        try:
-            self.mm_accumulated = unflatten_payload(self.mm_accumulated)
-        except Exception:
-            logger.exception("Error unflattening consolidated multimodal tensors")
+            # Metadata: consolidate any deferred tensor lists (REPLACE semantics)
+            for k, v in list(self.mm_accumulated.metadata.items()):
+                if isinstance(v, list) and v and isinstance(v[0], torch.Tensor):
+                    self.mm_accumulated.metadata[k] = v[-1]
+                elif isinstance(v, dict):
+                    for sk, sv in list(v.items()):
+                        if isinstance(sv, list) and sv and isinstance(sv[0], torch.Tensor):
+                            v[sk] = sv[-1]
+        except (RuntimeError, TypeError, KeyError):
+            logger.exception("Error consolidating multimodal tensors")

    # Override: do not route to pooling-only path; always create completion
    # outputs, and attach pooling_result into the CompletionOutput.
@@ -205,18 +279,18 @@ class OmniRequestState(RequestState):
            )

        finished = finish_reason is not None
-        is_final_only = self.output_kind == RequestOutputKind.FINAL_ONLY
+        final_only = self.output_kind == RequestOutputKind.FINAL_ONLY
        is_delta = self.output_kind == RequestOutputKind.DELTA

-        if not finished and is_final_only:
+        if not finished and final_only:
            return None

+        # Consolidate accumulated tensors when finishing, or for every
+        # CUMULATIVE step so the consumer always sees a single tensor.
        if finished or not is_delta:
            self._consolidate_multimodal_tensors()

-        if self.stream_interval > 1:
-            assert self.detokenizer is not None
-
+        if self.stream_interval > 1 and self.detokenizer is not None:
            # Send output request only when
            # 1. It has finished, or
            # 2. It is the first token, or
@@ -228,7 +302,7 @@ class OmniRequestState(RequestState):
            ):
                return None

-            if is_delta:
+            if self.output_kind == RequestOutputKind.DELTA:
                # Send tokens from the offset in DELTA mode, otherwise all
                # tokens are sent.
                new_token_ids = self.detokenizer.output_token_ids[self.sent_tokens_offset :]
@@ -258,50 +332,127 @@ class OmniRequestState(RequestState):
        token_ids: list[int],
        finish_reason: FinishReason | None,
        stop_reason: int | str | None,
-    ) -> Any:
-        # Reuse base text/logprobs logic, then annotate with pooling_result.
-        # Note: upstream _new_completion_output no longer accepts routed_experts
-        # as a parameter; it reads from ``self.routed_experts_chunks`` internally.
-        base_output = super()._new_completion_output(token_ids, finish_reason, stop_reason)
+    ) -> MultimodalCompletionOutput | CompletionOutput:
+        """Create a completion output with multimodal data attached.

-        # Inter-stage processors need the full cumulative token sequence.
-        # In DELTA mode, base_output.token_ids only has the latest step's
-        # tokens, so we always store a cumulative copy here.
-        base_output.cumulative_token_ids = list(self.detokenizer.output_token_ids)
+        Returns a MultimodalCompletionOutput when multimodal data has been
+        accumulated, otherwise returns the base CompletionOutput.  Snapshots
+        the accumulated payload before draining DELTA-mode modality keys so
+        that callers receive an immutable view.
+        """
+        # When there is no detokenizer (generation stages), build a minimal
+        # CompletionOutput directly since upstream asserts detokenizer != None.
+        if self.detokenizer is None:
+            finished = finish_reason is not None
+            base_output = CompletionOutput(
+                index=self.request_index,
+                text="",
+                token_ids=token_ids,
+                logprobs=None,
+                cumulative_logprob=None,
+                finish_reason=str(finish_reason) if finished else None,
+                stop_reason=stop_reason if finished else None,
+            )
+        else:
+            base_output = super()._new_completion_output(token_ids, finish_reason, stop_reason)
+
+        # Always provide cumulative token IDs for inter-stage processors.
+        if self.detokenizer is not None:
+            base_output.cumulative_token_ids = list(self.detokenizer.output_token_ids)
+        else:
+            base_output.cumulative_token_ids = list(token_ids)

        # Attach cumulative_text only at the final step for inter-stage use.
        if finish_reason is not None and hasattr(self.detokenizer, "output_text"):
            base_output.cumulative_text = self.detokenizer.output_text

-        if not hasattr(base_output, "multimodal_output"):
-            setattr(base_output, "multimodal_output", {})
-        if self.mm_accumulated:
-            mm_out = getattr(base_output, "multimodal_output")
-            if isinstance(mm_out, dict):
-                for k, v in self.mm_accumulated.items():
-                    mm_out[k] = v
-            else:
-                setattr(base_output, "multimodal_output", self.mm_accumulated)
+        try:
+            if self.mm_accumulated and not self.mm_accumulated.is_empty:
+                # Snapshot: copy current tensors/metadata so drain doesn't
+                # mutate the payload already handed to the caller.
+                # Unflatten dotted keys (e.g. "hidden_states.layer_0") back
+                # to nested dicts so downstream consumers (thinker2talker etc.)
+                # can access with .get("hidden_states", {}).get("layers", {}).
+                merged = {**self.mm_accumulated.tensors, **self.mm_accumulated.metadata}
+                nested = unflatten_payload(merged)
+                snapshot = MultimodalPayload.from_dict(nested) or MultimodalPayload()
+                kwargs = {f.name: getattr(base_output, f.name) for f in dataclass_fields(CompletionOutput)}
+                output = MultimodalCompletionOutput(
+                    multimodal_output=snapshot,
+                    **kwargs,
+                )
+                output.cumulative_token_ids = base_output.cumulative_token_ids
+                if hasattr(base_output, "cumulative_text"):
+                    output.cumulative_text = base_output.cumulative_text

-        if self.output_kind == RequestOutputKind.DELTA:
-            for modality_key in DRAINABLE_MODALITIES:
-                self.mm_accumulated.pop(modality_key, None)
+                # DELTA mode: drain modality keys (e.g. audio) so the next
+                # step only sees freshly accumulated data for those keys.
+                if self.output_kind == RequestOutputKind.DELTA:
+                    for modality_key in DRAINABLE_MODALITIES:
+                        self.mm_accumulated.tensors.pop(modality_key, None)

+                return output
+        except (RuntimeError, TypeError, AttributeError):
+            logger.exception("Error creating MultimodalCompletionOutput")
        return base_output

+    def _new_request_output(
+        self,
+        external_req_id: str,
+        outputs: list,
+        finished: bool,
+        kv_transfer_params: dict[str, Any] | None = None,
+    ) -> RequestOutput | PoolingRequestOutput:
+        """Create request output, handling no-detokenizer generation stages.
+
+        Upstream asserts ``self.logprobs_processor is not None`` which fails
+        for generation stages (talker, code2wav) that have no tokenizer.
+        When the logprobs processor is absent we build the RequestOutput
+        directly with ``prompt_logprobs=None``.
+        """
+        if self.logprobs_processor is not None:
+            return super()._new_request_output(
+                external_req_id,
+                outputs,
+                finished,
+                kv_transfer_params,
+            )
+
+        # No-detokenizer path: build RequestOutput directly.
+        prompt_token_ids = self.prompt_token_ids
+        if prompt_token_ids is None and self.prompt_embeds is not None:
+            prompt_token_ids = [0] * len(self.prompt_embeds)
+        if prompt_token_ids is None:
+            prompt_token_ids = []
+
+        return RequestOutput(
+            request_id=external_req_id,
+            lora_request=self.lora_request,
+            prompt=self.prompt,
+            prompt_token_ids=prompt_token_ids,
+            prompt_logprobs=None,
+            outputs=outputs,
+            finished=finished,
+            kv_transfer_params=kv_transfer_params,
+            num_cached_tokens=self.num_cached_tokens,
+            metrics=self.stats,
+        )
+

 class MultimodalOutputProcessor(VLLMOutputProcessor):
-    """Handles multimodal output processing by capturing pooling_output
-    from EngineCoreOutput and accumulating it as multimodal tensors,
-    before delegating to the base vLLM OutputProcessor for text handling.
+    """Handles multimodal output processing.

-    The actual data flow is:
-    1. For each EngineCoreOutput with pooling_output and a detokenizer:
-       - Capture pooling_output into OmniRequestState.add_multimodal_tensor()
-       - Clear eco.pooling_output to force text path in base processor
+    Captures multimodal outputs from OmniEngineCoreOutput and accumulates
+    them as MultimodalPayload in OmniRequestState, before delegating to
+    the base vLLM OutputProcessor for text handling.
+
+    The data flow is:
+    1. For each EngineCoreOutput with multimodal_output:
+       - Capture into OmniRequestState.add_multimodal_tensor()
    2. Base vLLM OutputProcessor handles text detokenization
-    3. On finish, _consolidate_multimodal_tensors() concatenates accumulated tensors
-    4. _new_completion_output() attaches mm_accumulated to CompletionOutput
+    3. On finish, _consolidate_multimodal_tensors() concatenates accumulated
+       tensors using strategy-based dispatch
+    4. _new_completion_output() returns MultimodalCompletionOutput
    """

    def __init__(
@@ -312,6 +463,7 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
        stream_interval: int = 1,
        tracing_enabled: bool = False,
        engine_core_output_type: str | None = None,
+        output_modality: OutputModality = OutputModality.TEXT,
    ):
        """Initialize the multimodal output processor.

@@ -319,9 +471,13 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
            tokenizer: Tokenizer for detokenizing text outputs
            log_stats: Whether to log statistics
            stream_interval: Stream interval for output generation
-            engine_core_output_type: Optional output type specification
-                (e.g., "image", "audio", "latent"). Used to tag multimodal
-                outputs with the correct modality key.
+            engine_core_output_type: Optional output type string (e.g.,
+                "image", "audio", "latent"). Converted to OutputModality
+                internally. Kept for backward compatibility with
+                stage_init_utils.
+            output_modality: Type-safe output modality flag. Used to tag
+                multimodal outputs with the correct modality key when
+                per-output type info is unavailable.
        """
        super().__init__(
            tokenizer=tokenizer,
@@ -329,6 +485,11 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
            stream_interval=stream_interval,
            tracing_enabled=tracing_enabled,
        )
+        # Convert string-based engine_core_output_type to OutputModality
+        if engine_core_output_type is not None:
+            self.output_modality = OutputModality.from_string(engine_core_output_type)
+        else:
+            self.output_modality = output_modality
        self.engine_core_output_type = engine_core_output_type
        self._native_text_metrics_by_request: dict[str, dict[str, Any]] = {}

@@ -426,21 +587,78 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
        engine_core_timestamp: float | None = None,
        iteration_stats: IterationStats | None = None,
    ) -> OutputProcessorOutput:
+        default_mm_type = _modality_to_type_string(self.output_modality)
+
+        # Separate outputs that upstream can handle (has detokenizer or
+        # pooling) from multimodal-only outputs (no detokenizer, no pooling)
+        # that would trigger upstream's `assert detokenizer is not None`.
+        upstream_outputs: list[EngineCoreOutput] = []
+        mm_only_outputs: list[EngineCoreOutput] = []
+
+        for eco in engine_core_outputs:
+            req_state = self.request_states.get(eco.request_id)
+            if req_state is None:
+                continue
+
+            # Accumulate multimodal tensors regardless of path.
+            if isinstance(req_state, OmniRequestState):
+                mm_output = getattr(eco, "multimodal_output", None)
+                if mm_output is not None:
+                    mm_type = getattr(eco, "output_type", None) or default_mm_type
+                    req_state.add_multimodal_tensor(mm_output, mm_type)
+
+            # Route: if no detokenizer and no pooling output, handle locally
+            # to avoid upstream's assert on detokenizer.
+            if req_state.detokenizer is None and eco.pooling_output is None:
+                mm_only_outputs.append(eco)
+            else:
+                upstream_outputs.append(eco)
+
+        # Handle multimodal-only outputs (generation stages) locally.
+        self._process_mm_only_outputs(mm_only_outputs)
+
+        # Delegate text/pooling outputs to upstream.
+        return super().process_outputs(
+            upstream_outputs,
+            engine_core_timestamp=engine_core_timestamp,
+            iteration_stats=iteration_stats,
+        )
+
+    def _process_mm_only_outputs(
+        self,
+        engine_core_outputs: list[EngineCoreOutput],
+    ) -> None:
+        """Handle outputs from generation stages that have no detokenizer.
+
+        These cannot go through upstream process_outputs because it asserts
+        detokenizer is not None when pooling_output is None.
+        """
        for eco in engine_core_outputs:
            req_state = self.request_states.get(eco.request_id)
            if req_state is None or not isinstance(req_state, OmniRequestState):
                continue
-            if eco.pooling_output is not None and req_state.detokenizer is not None:
-                mm_type = (getattr(eco, "output_type", self.engine_core_output_type) or "").lower()
-                req_state.add_multimodal_tensor(eco.pooling_output, mm_type)
-                # Force text path in base processor for multimodal outputs.
-                eco.pooling_output = None

-        return super().process_outputs(
-            engine_core_outputs,
-            engine_core_timestamp=engine_core_timestamp,
-            iteration_stats=iteration_stats,
-        )
+            new_token_ids = eco.new_token_ids
+            finish_reason = eco.finish_reason
+            stop_reason = eco.stop_reason
+            kv_transfer_params = eco.kv_transfer_params
+            routed_experts = eco.routed_experts
+            req_state.num_cached_tokens = eco.num_cached_tokens
+            req_state.is_prefilling = False
+
+            if request_output := req_state.make_request_output(
+                new_token_ids,
+                None,  # pooling_output
+                finish_reason,
+                stop_reason,
+                kv_transfer_params,
+                routed_experts,
+            ):
+                if req_state.queue is not None:
+                    req_state.queue.put(request_output)
+
+            if finish_reason is not None:
+                self._finish_request(req_state)

    def _update_stats_from_output(
        self,
@@ -453,6 +671,14 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
        native_stats = req_state.native_text_stats if isinstance(req_state, OmniRequestState) else None
        previous_last_token_ts = native_stats.last_token_ts if native_stats is not None else 0.0

+        # NOTE: We pass ``None`` for  *iteration_stats* to the parent so that
+        # the upstream's ``_update_stats_from_output`` logs stats via its own
+        # ``RequestStateStats`` (req_state.stats) without touching the omni
+        # ``IterationStats`` object.  Below we re-apply the same update using
+        # *native_stats* (the omni-specific ``RequestStateStats`` attached to
+        # ``OmniRequestState``), which correctly feeds the omni metrics path.
+        # If the upstream adds new logic in ``_update_stats_from_output`` (e.g.
+        # additional stat fields), this override must be kept in sync.
        super()._update_stats_from_output(
            req_state,
            engine_core_output,
@@ -18,6 +18,8 @@ from typing import TYPE_CHECKING, Any
 from urllib.parse import urlparse

 import psutil
+import vllm.v1.engine as _vllm_engine_module
+import vllm.v1.engine.core_client as _vllm_core_client_module
 from vllm.logger import init_logger
 from vllm.v1.engine import EngineCoreRequest
 from vllm.v1.engine.core_client import AsyncMPClient, DPLBAsyncMPClient
@@ -30,6 +32,7 @@ from vllm_omni.distributed.omni_connectors.utils.initialization import (
    KV_TRANSFER_PORT_OFFSET,
 )
 from vllm_omni.distributed.omni_connectors.utils.kv_utils import kv_zmq_port
+from vllm_omni.engine import OmniEngineCoreOutput, OmniEngineCoreOutputs
 from vllm_omni.engine.stage_client import StageClientBase
 from vllm_omni.engine.stage_init_utils import StageMetadata

@@ -174,6 +177,21 @@ class StageEngineCoreClientBase(StageClientBase):
            self.stage_id,
            self.replica_id,
        )
+
+        # Patch the output decoder type so the client decodes
+        # OmniEngineCoreOutputs (which carries multimodal_output per
+        # EngineCoreOutput) instead of the base EngineCoreOutputs.
+        # Must happen BEFORE super().__init__() which creates the decoder.
+        # TODO: Add a defensive assertion after super().__init__() to verify
+        # the decoder uses OmniEngineCoreOutputs, catching import-order regressions.
+        _vllm_engine_module.EngineCoreOutput = OmniEngineCoreOutput
+        _vllm_engine_module.EngineCoreOutputs = OmniEngineCoreOutputs
+        _vllm_core_client_module.EngineCoreOutputs = OmniEngineCoreOutputs
+        logger.debug(
+            "[StageEngineCoreClient] Patched EngineCoreOutputs -> %s",
+            _vllm_core_client_module.EngineCoreOutputs,
+        )
+
        try:
            super().__init__(
                vllm_config,
@@ -14,6 +14,7 @@ from multiprocessing.process import BaseProcess
 from typing import TYPE_CHECKING, Any

 import msgspec
+import vllm.v1.engine.core as _vllm_engine_core_module
 import zmq
 from vllm.logger import init_logger
 from vllm.transformers_utils.config import (
@@ -36,6 +37,7 @@ from vllm.v1.engine.utils import (
 from vllm.v1.utils import shutdown

 from vllm_omni.distributed.omni_coordinator import OmniCoordClientForStage
+from vllm_omni.engine import OmniEngineCoreRequest
 from vllm_omni.engine.stage_init_utils import set_death_signal

 if TYPE_CHECKING:
@@ -122,6 +124,18 @@ class StageEngineCoreProc(EngineCoreProc):
            os.environ.setdefault("FLASHINFER_DISABLE_VERSION_CHECK", "1")
            os.environ["VLLM_OMNI_REPLICA_ID"] = str(max(int(omni_replica_id), 0))

+            # Patch the decoder type so process_input_sockets (started
+            # during __init__) decodes OmniEngineCoreRequest (which
+            # carries additional_information) instead of the base
+            # EngineCoreRequest.  Must happen BEFORE __init__ because
+            # the IO thread creates MsgpackDecoder(EngineCoreRequest)
+            # during __init__.
+            _vllm_engine_core_module.EngineCoreRequest = OmniEngineCoreRequest
+            logger.debug(
+                "[StageEngineCoreProc] Patched EngineCoreRequest -> OmniEngineCoreRequest: %s",
+                _vllm_engine_core_module.EngineCoreRequest,
+            )
+
            engine_core = StageEngineCoreProc(
                *args,
                engine_index=dp_rank,
@@ -4,7 +4,7 @@ import argparse
 import typing

 from vllm.entrypoints.cli.types import CLISubcommand
-from vllm.entrypoints.utils import VLLM_SUBCMD_PARSER_EPILOG
+from vllm.entrypoints.serve.utils.api_utils import VLLM_SUBCMD_PARSER_EPILOG

 from vllm_omni.entrypoints.cli.benchmark.base import OmniBenchmarkSubcommandBase

@@ -23,7 +23,7 @@ def main():
        if "VLLM_LOGGING_COLOR" not in os.environ:
            os.environ["VLLM_LOGGING_COLOR"] = "1"

-        from vllm.entrypoints.utils import VLLM_SUBCMD_PARSER_EPILOG, cli_env_setup
+        from vllm.entrypoints.serve.utils.api_utils import VLLM_SUBCMD_PARSER_EPILOG, cli_env_setup

        import vllm_omni.entrypoints.cli.benchmark.main
        import vllm_omni.entrypoints.cli.serve
@@ -17,7 +17,7 @@ from typing import Any
 import uvloop
 from vllm.entrypoints.cli.types import CLISubcommand
 from vllm.entrypoints.openai.cli_args import make_arg_parser, validate_parsed_serve_args
-from vllm.entrypoints.utils import VLLM_SUBCMD_PARSER_EPILOG
+from vllm.entrypoints.serve.utils.api_utils import VLLM_SUBCMD_PARSER_EPILOG
 from vllm.logger import init_logger

 from vllm_omni.entrypoints.cli.logo import log_logo
@@ -34,7 +34,6 @@ from vllm.engine.protocol import EngineClient
 from vllm.entrypoints.anthropic.serving import AnthropicServingMessages
 from vllm.entrypoints.chat_utils import load_chat_template
 from vllm.entrypoints.launcher import serve_http, terminate_if_errored
-from vllm.entrypoints.logger import RequestLogger
 from vllm.entrypoints.mcp.tool_server import DemoToolServer, MCPToolServer, ToolServer
 from vllm.entrypoints.openai.api_server import build_app as build_openai_app
 from vllm.entrypoints.openai.api_server import setup_server as setup_openai_server
@@ -56,10 +55,7 @@ from vllm.entrypoints.openai.engine.protocol import (
 )
 from vllm.entrypoints.openai.models.protocol import BaseModelPath
 from vllm.entrypoints.openai.models.serving import OpenAIServingModels
-from vllm.entrypoints.openai.orca_metrics import metrics_header
 from vllm.entrypoints.openai.responses.serving import OpenAIServingResponses
-from vllm.entrypoints.openai.server_utils import get_uvicorn_log_config
-from vllm.entrypoints.openai.utils import validate_json_request
 from vllm.entrypoints.pooling.classify.serving import ServingClassification
 from vllm.entrypoints.pooling.embed.serving import ServingEmbedding as OpenAIServingEmbedding
 from vllm.entrypoints.pooling.pooling.serving import ServingPooling
@@ -71,6 +67,16 @@ from vllm.entrypoints.serve.disagg.serving import ServingTokens
 from vllm.entrypoints.serve.instrumentator.basic import base
 from vllm.entrypoints.serve.render.serving import OpenAIServingRender
 from vllm.entrypoints.serve.tokenize.serving import OpenAIServingTokenization
+from vllm.entrypoints.serve.utils.api_utils import (
+    load_aware_call,
+    process_lora_modules,
+    validate_json_request,
+    with_cancellation,
+)
+from vllm.entrypoints.serve.utils.error_response import create_error_response
+from vllm.entrypoints.serve.utils.orca_metrics import metrics_header
+from vllm.entrypoints.serve.utils.request_logger import RequestLogger
+from vllm.entrypoints.serve.utils.server_utils import get_uvicorn_log_config
 from vllm.entrypoints.speech_to_text.realtime.serving import OpenAIServingRealtime
 from vllm.entrypoints.speech_to_text.transcription.serving import (
    OpenAIServingTranscription,
@@ -78,12 +84,6 @@ from vllm.entrypoints.speech_to_text.transcription.serving import (
 from vllm.entrypoints.speech_to_text.translation.serving import (
    OpenAIServingTranslation,
 )
-from vllm.entrypoints.utils import (
-    create_error_response,
-    load_aware_call,
-    process_lora_modules,
-    with_cancellation,
-)
 from vllm.logger import init_logger
 from vllm.tasks import POOLING_TASKS
 from vllm.tool_parsers import ToolParserManager
@@ -3,7 +3,7 @@ from __future__ import annotations
 import asyncio
 import base64
 import json
-from collections.abc import AsyncGenerator
+from collections.abc import AsyncGenerator, Mapping
 from typing import cast
 from uuid import uuid4

@@ -82,7 +82,10 @@ class RealtimeConnection(VllmRealtimeConnection):

    def _extract_audio_chunks(self, output) -> tuple[list[np.ndarray], int]:
        mm = getattr(output, "multimodal_output", None)
-        if not isinstance(mm, dict):
+        if mm is None:
+            return [], 24000
+        # Support both MultimodalPayload and plain dict
+        if not isinstance(mm, Mapping):
            return [], 24000

        sr = mm.get("sr") or mm.get("sample_rate") or mm.get("audio_sample_rate") or 24000
@@ -69,8 +69,8 @@ from vllm.entrypoints.openai.parser.harmony_utils import (
    parse_chat_output,
 )
 from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
-from vllm.entrypoints.openai.utils import maybe_filter_parallel_tool_calls
-from vllm.entrypoints.utils import should_include_usage
+from vllm.entrypoints.serve.utils.api_utils import should_include_usage
+from vllm.entrypoints.serve.utils.tool_calls_utils import maybe_filter_parallel_tool_calls
 from vllm.inputs import PromptType
 from vllm.logger import init_logger
 from vllm.outputs import RequestOutput
@@ -843,6 +843,27 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
            if value is not None and hasattr(obj, key):
                setattr(obj, key, value)

+    def _should_check_for_unstreamed_tool_arg_tokens(
+        self,
+        delta_message: Any,
+        output: Any,
+    ) -> bool:
+        """Check whether the streaming generator should flush unstreamed
+        tool-arg tokens at finish-time.
+
+        This method was moved from OpenAIServingChat to the tool-parser layer
+        by upstream commit 9affc17a05.  Omni's independently-maintained
+        ``chat_completion_stream_generator`` still calls it, so we keep a
+        local copy with the pre-9affc17a05 semantics.
+        """
+        return (
+            output.finish_reason is not None
+            and self.enable_auto_tools
+            and self.tool_parser is not None
+            and delta_message is not None
+            and delta_message.tool_calls
+        )
+
    def _build_sampling_params_list_from_request(
        self,
        request: ChatCompletionRequest,
@@ -2794,6 +2815,16 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                    if isinstance(ar_text, str) and ar_text.strip():
                        cot_output = ar_text

+        req_out = getattr(result, "request_output", None)
+        if req_out:
+            prompt_obj = getattr(req_out, "prompt", None)
+            if isinstance(prompt_obj, dict):
+                extra = prompt_obj.get("extra", {})
+                if isinstance(extra, dict):
+                    ar_text = extra.get("ar_generated_text")
+                    if isinstance(ar_text, str) and ar_text.strip():
+                        cot_output = ar_text
+
        return self._flatten_diffusion_images(images), stage_durations, peak_memory_mb, cot_output

    async def _stream_diffusion_image_chunks(
@@ -44,6 +44,10 @@ from vllm_omni.entrypoints.openai.protocol.audio import (
    SpeechBatchItem,
    SpeechBatchItemResult,
 )
+from vllm_omni.entrypoints.openai.tts_adapters import (
+    SpeechServingContext,
+    resolve_adapter,
+)
 from vllm_omni.entrypoints.utils import coerce_param_message_types
 from vllm_omni.model_executor.models.fish_speech.prompt_utils import (
    build_fish_text_only_prompt_ids,
@@ -80,6 +84,7 @@ _MOSS_TTS_FULL_MODEL_STAGES = {"moss_tts", "moss_tts_codec"}
 _HIGGS_AUDIO_V2_TTS_MODEL_STAGES = {"higgs_audio_v2"}
 _HIGGS_V3_TTS_MODEL_STAGES = {"higgs_audio_v3"}
 _GLM_TTS_MODEL_STAGES = {"glm_tts"}
+_STEP_AUDIO2_TTS_MODEL_STAGES = {"step_audio2_thinker"}
 _TTS_MODEL_STAGES: set[str] = (
    _VOXTRAL_TTS_MODEL_STAGES
    | _QWEN3_TTS_MODEL_STAGES
@@ -94,6 +99,7 @@ _TTS_MODEL_STAGES: set[str] = (
    | _MOSS_TTS_MODEL_STAGES
    | _MOSS_TTS_FULL_MODEL_STAGES
    | _GLM_TTS_MODEL_STAGES
+    | _STEP_AUDIO2_TTS_MODEL_STAGES
 )
 _SAMPLING_MAX_TOKENS_TTS_MODEL_TYPES = {
    "fish_tts",
@@ -472,6 +478,36 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
        self._build_fish_speech_prompt_async = make_async(self._build_fish_speech_prompt, executor=self._tts_executor)
        self._estimate_prompt_len_async = make_async(self._estimate_prompt_len, executor=self._tts_executor)

+        # Resolve the per-model serving adapter (RFC #4327), keyed on the
+        # detected model-type. Every dedicated TTS model has an adapter; the
+        # adapter owns request validation + prompt/param building. Sampling
+        # overrides and the model-type label remain in the orchestrator tail
+        # (keyed on ``_tts_model_type``) during this incremental migration.
+        self._adapter = None
+        if self._tts_stage is not None:
+            adapter_cls = resolve_adapter(self._tts_model_type)
+            if adapter_cls is not None:
+                ctx = SpeechServingContext(server=self, engine_client=self.engine_client)
+                self._adapter = adapter_cls(ctx)
+                logger.info("Resolved TTS serving adapter: %s", adapter_cls.__name__)
+
+    def _get_tts_adapter(self):
+        """Return the per-model serving adapter for the current ``_tts_model_type``.
+
+        Resolved lazily (rebuilt if ``_tts_model_type`` changed since the cached
+        instance was built) so callers that set ``_tts_model_type`` after
+        construction still dispatch to the matching adapter. In production
+        ``_tts_model_type`` is fixed at init, so the cached instance is reused.
+        """
+        adapter_cls = resolve_adapter(self._tts_model_type)
+        if adapter_cls is None:
+            self._adapter = None
+            return None
+        if self._adapter is None or type(self._adapter) is not adapter_cls:
+            ctx = SpeechServingContext(server=self, engine_client=self.engine_client)
+            self._adapter = adapter_cls(ctx)
+        return self._adapter
+
    async def warmup(self) -> None:
        """Run a synthetic speech request to trigger all first-request warmup.

@@ -655,6 +691,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
            return "higgs_audio_v3"
        if model_stage in _GLM_TTS_MODEL_STAGES:
            return "glm_tts"
+        if model_stage in _STEP_AUDIO2_TTS_MODEL_STAGES:
+            return "step_audio2"
        return None

    def _get_custom_voice_dir(self) -> str | None:
@@ -1356,6 +1394,11 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):

    def _validate_tts_request(self, request: OpenAICreateSpeechRequest) -> str | None:
        """Validate TTS request parameters. Returns error message or None."""
+        if self._tts_model_type == "step_audio2":
+            # StepAudio2 only requires non-empty input text
+            if not request.input or not request.input.strip():
+                return "Input text cannot be empty"
+            return None
        if self._tts_model_type == "voxtral_tts":
            return self._validate_voxtral_tts_request(request)
        if self._tts_model_type == "fish_tts":
@@ -2858,6 +2901,33 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                "multi_modal_data": {"audio": [(audio.audio_array, audio.sampling_rate)]},
            }

+    # ---- Step-Audio2 helpers ----
+
+    def _build_step_audio2_prompt(
+        self,
+        request: OpenAICreateSpeechRequest,
+    ) -> dict[str, Any]:
+        """Build prompt for Step-Audio2 TTS.
+
+        Constructs the chat prompt with ``<tts_start>`` as the last token
+        of the assistant turn (without ``<|im_end|>``), so the thinker
+        continues generating audio tokens.
+
+        Prompt format::
+            <|im_start|>system\\n{system_prompt}<|im_end|>\\n
+            <|im_start|>user\\n{input_text}<|im_end|>\\n
+            <|im_start|>assistant\\n<tts_start>
+        """
+        system_prompt = getattr(request, "instructions", None) or "You are a voice assistant. Read the text aloud."
+        text = request.input
+
+        raw_prompt = (
+            f"<|im_start|>system\n{system_prompt}<|im_end|>\n"
+            f"<|im_start|>user\n{text}<|im_end|>\n"
+            f"<|im_start|>assistant\n<tts_start>"
+        )
+        return {"prompt": raw_prompt}
+
    # ---- Fish Speech helpers ----

    def _build_fish_speech_prompt(
@@ -3176,6 +3246,44 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):

    # ---- Common speech generation helpers ----

+    async def _build_qwen3_tts_request(
+        self,
+        request: OpenAICreateSpeechRequest,
+    ) -> tuple[dict[str, Any], dict[str, Any], str | None]:
+        """Build prompt + tts_params for Qwen3-TTS.
+
+        Called from ``Qwen3TTSAdapter.build``. Returns
+        ``(prompt, tts_params, warmup_artifact_key)`` where the warmup key is the
+        Qwen3-TTS ref-audio artifact tracked after ``generate()``.
+        """
+        qwen3_ref_audio_warmup_artifact_key: str | None = None
+        tts_params = self._build_tts_params(request)
+        # Resolve ref_audio (explicit or auto-set for uploaded voices)
+        # to [[wav_list, sr]] so the model doesn't re-decode base64.
+        ref_audio_source = request.ref_audio
+        if ref_audio_source is None and isinstance(tts_params.get("ref_audio"), list):
+            # Uploaded voice: ref_audio was auto-set as [base64_data_url]
+            ref_audio_source = tts_params["ref_audio"][0]
+        if ref_audio_source is not None and isinstance(ref_audio_source, str):
+            wav_list, sr = await self._resolve_ref_audio(ref_audio_source)
+            artifact_key = self._get_resolved_ref_audio_artifact_key(ref_audio_source)
+            if artifact_key:
+                tts_params[_QWEN3_TTS_REF_AUDIO_CACHE_KEY] = [artifact_key]
+            ref_code_length = self._estimate_ref_code_len([wav_list, sr])
+            if ref_code_length is not None:
+                tts_params["ref_code_length"] = [int(ref_code_length)]
+            if self._qwen3_tts_can_use_ref_audio_artifact_only(tts_params, artifact_key):
+                logger.debug("Using Qwen3-TTS ref_audio artifact-only path: %s", artifact_key)
+            else:
+                tts_params["ref_audio"] = [[wav_list, sr]]
+                qwen3_ref_audio_warmup_artifact_key = artifact_key
+
+        ph_len = await self._estimate_prompt_len_async(tts_params)
+        prompt = tokens_input(prompt_token_ids=[1] * ph_len)
+        prompt["additional_information"] = tts_params
+        prompt["cache_salt"] = _conditioning_cache_salt(request, tts_params)
+        return prompt, tts_params, qwen3_ref_audio_warmup_artifact_key
+
    async def _prepare_speech_generation(
        self,
        request: OpenAICreateSpeechRequest,
@@ -3194,203 +3302,35 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
        sampling_params_list = list(self.engine_client.default_sampling_params_list)
        sampling_params_list = coerce_param_message_types(sampling_params_list, request.stream)

-        # Resolve uploaded voice for non-Qwen3 models.
-        # Qwen3 TTS has its own uploaded voice handling in _build_tts_params().
+        # Build prompt + tts_params via the per-model adapter (RFC #4327). Every
+        # dedicated TTS model resolves to an adapter that owns its validation,
+        # uploaded-speaker handling, and prompt/param building. Sampling
+        # overrides and the model-type label remain in the orchestrator tail
+        # below (keyed on ``_tts_model_type``) during this incremental migration.
+        # Non-TTS deployments (no adapter) fall through to the rejection below.
+        # Capture inline-ref-audio status BEFORE validate(): several adapters
+        # apply uploaded speakers inside validate(), which sets request.ref_audio
+        # in place. The builders need to know whether the caller supplied audio
+        # inline vs. via an uploaded voice.
        has_inline_ref_audio = request.ref_audio is not None
-        if self._tts_model_type in ("fish_tts", "cosyvoice3", "moss_tts_nano", "glm_tts"):
-            err = self._apply_uploaded_speaker(request)
-            if err:
-                raise ValueError(err)
-
-        if self._is_fish_speech:
-            validation_error = self._validate_fish_tts_request(request)
+        if self._tts_model_type == "ming_flash_omni_tts":
+            # ming_flash_omni is intentionally NOT migrated onto the adapter
+            # framework in this PR (it has no registered adapter); keep it on the
+            # legacy inline dispatch so serving still works.
+            validation_error = self._validate_ming_flash_omni_tts_request(request)
            if validation_error:
                raise ValueError(validation_error)
-            ref_audio_data = None
-            if request.ref_audio is not None:
-                wav_list, sr = await self._resolve_ref_audio(request.ref_audio)
-                ref_audio_data = (wav_list, sr)
-            prompt = await self._build_fish_speech_prompt_async(
-                request, ref_audio_data=ref_audio_data, has_inline_ref_audio=has_inline_ref_audio
-            )
+            prompt = self._build_ming_flash_omni_prompt(request)
            tts_params = {}
-            # Fish Speech resolves conditioning entirely from raw request fields
-            # (text / ref_audio / ref_text rebuilt into inputs_embeds); unlike
-            # uploaded-voice models it has no resolved tts_params to fold in, so
-            # the salt is derived from the request alone.
-            prompt["cache_salt"] = _conditioning_cache_salt(request, tts_params)
-        elif self._tts_model_type == "omnivoice":
-            if not request.input or not request.input.strip():
-                raise ValueError("Input text cannot be empty")
-            err = self._apply_uploaded_speaker(request)
-            if err:
-                raise ValueError(err)
-            tts_params = {}
-            prompt: dict[str, Any] = {"input": request.input}
-            if request.ref_audio:
-                wav, sr = await self._resolve_ref_audio(request.ref_audio)
-                prompt["ref_audio"] = (np.asarray(wav, dtype=np.float32), sr)
-            if request.ref_text:
-                prompt["ref_text"] = request.ref_text
-            if request.voice:
-                voice_lower = request.voice.lower()
-                if voice_lower in self.uploaded_speakers and not has_inline_ref_audio:
-                    prompt["voice_name"] = voice_lower
-                    prompt["voice_created_at"] = self._voice_created_at(voice_lower)
-            if request.language:
-                prompt["lang"] = request.language
-            if request.instructions:
-                prompt["instruct"] = request.instructions
-        elif self._tts_model_type == "covo_audio":
-            prompt = self._build_covo_audio_prompt(request)
-            tts_params = {}
-        elif self._tts_model_type == "higgs_audio_v2":
-            # Explicit higgs_audio_v2 branch: do NOT fall through to the
-            # Qwen-style generic placeholder path; build prompt_token_ids
-            # directly via the upstream processor + the project's plain-text
-            # helper.
-            #
-            # Resolve uploaded voices first so a voice=<name> request (after
-            # POST /v1/audio/voices) populates ref_audio + ref_text from the
-            # cached speaker entry, mirroring the cosyvoice3 / voxcpm2 flow.
-            err = self._apply_uploaded_speaker(request)
-            if err:
-                raise ValueError(err)
-            validation_error = self._validate_higgs_audio_v2_request(request)
+            qwen3_ref_audio_warmup_artifact_key = None
+        elif (adapter := self._get_tts_adapter()) is not None:
+            validation_error = adapter.validate(request)
            if validation_error:
                raise ValueError(validation_error)
-            prompt = await self._build_higgs_audio_v2_params(request)
-            if request.voice:
-                voice_lower = request.voice.lower()
-                if voice_lower in self.uploaded_speakers and not has_inline_ref_audio:
-                    additional = prompt.setdefault("additional_information", {})
-                    additional["voice_name"] = voice_lower
-                    additional["voice_created_at"] = self._voice_created_at(voice_lower)
-            tts_params = {}
-        elif self._tts_model_type == "higgs_audio_v3":
-            err = self._apply_uploaded_speaker(request)
-            if err:
-                raise ValueError(err)
-            validation_error = self._validate_higgs_audio_v3_request(request)
-            if validation_error:
-                raise ValueError(validation_error)
-            prompt = await self._build_higgs_audio_v3_params(request)
-            tts_params = {}
-        elif self._tts_model_type == "voxcpm2":
-            # voxcpm2 doesn't use `_apply_uploaded_speaker` because the prompt builder needs the
-            # raw waveform tuple for prefill-length accounting, not a base64 data URL.
-            validation_error = self._validate_voxcpm2_request(request)
-            if validation_error:
-                raise ValueError(validation_error)
-
-            uploaded_ref: tuple[np.ndarray, int] | None = None
-            if request.voice:
-                voice_lower = request.voice.lower()
-                if voice_lower in self.uploaded_speakers and not has_inline_ref_audio:
-                    if self.uploaded_speakers[voice_lower].get("embedding_source") == "direct":
-                        raise ValueError(
-                            f"Uploaded voice '{request.voice}' uses a speaker embedding (Qwen3-only). "
-                            f"Re-upload with an audio file for VoxCPM2."
-                        )
-                    if request.ref_audio is None:
-                        uploaded_ref = self._load_uploaded_audio(voice_lower)
-            prompt = await self._build_voxcpm2_prompt(request, uploaded_ref=uploaded_ref)
-            tts_params = {}
-            if request.voice:
-                voice_lower = request.voice.lower()
-                if voice_lower in self.uploaded_speakers or voice_lower in self.precomputed_speakers:
-                    additional = prompt.setdefault("additional_information", {})
-                    additional["voice_name"] = voice_lower
-                    additional["voice_created_at"] = self._voice_created_at(voice_lower)
-        elif self._is_tts:
-            validation_error = self._validate_tts_request(request)
-            if validation_error:
-                raise ValueError(validation_error)
-
-            if self._tts_model_type == "voxtral_tts":
-                prompt = await self._build_voxtral_prompt_async(request)
-                tts_params = {}
-            elif self._tts_model_type == "cosyvoice3":
-                prompt = await self._build_cosyvoice3_prompt(request, has_inline_ref_audio=has_inline_ref_audio)
-                tts_params = {}
-            elif self._tts_model_type == "glm_tts":
-                prompt = await self._build_glm_tts_prompt(request, has_inline_ref_audio=has_inline_ref_audio)
-                tts_params = {}
-            elif self._tts_model_type == "ming_flash_omni_tts":
-                prompt = self._build_ming_flash_omni_prompt(request)
-                tts_params = {}
-            elif self._tts_model_type == "ming_tts":
-                ref_audio_source = request.ref_audio
-                voice_lower = request.voice.lower() if isinstance(request.voice, str) else None
-                if ref_audio_source is None and voice_lower in self.uploaded_speakers:
-                    ref_audio_source = self._get_uploaded_audio_data(request.voice)
-                    if request.ref_text is None:
-                        request.ref_text = self.uploaded_speakers[voice_lower].get("ref_text")
-                ref_audio_data = None
-                if isinstance(ref_audio_source, list):
-                    ref_audio_data = await self._resolve_ref_audio_many(ref_audio_source)
-                    if request.speaker_embedding is None:
-                        request.speaker_embedding = self._extract_ming_speaker_embeddings_from_ref_audio(ref_audio_data)
-                elif ref_audio_source is not None and isinstance(ref_audio_source, str):
-                    wav_list, sr = await self._resolve_ref_audio(ref_audio_source)
-                    ref_audio_data = (wav_list, sr)
-                    if request.speaker_embedding is None:
-                        request.speaker_embedding = self._extract_ming_speaker_embeddings_from_ref_audio(
-                            [ref_audio_data]
-                        )[0]
-                prompt = self._build_ming_dense_prompt(request, ref_audio_data=ref_audio_data)
-                tts_params = prompt.get("additional_information", {})
-            elif self._tts_model_type in ("moss_tts_nano", "moss_tts"):
-                tts_params = await self._build_moss_tts_params(request)
-                if request.voice:
-                    voice_lower = request.voice.lower()
-                    if voice_lower in self.uploaded_speakers and not has_inline_ref_audio:
-                        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]
-                # Full MOSS-TTS family returns real text_ids from the upstream
-                # processor; use them as prompt_token_ids. Nano keeps the
-                # ``[1]`` placeholder (it owns its own tokenisation internally).
-                if isinstance(tts_params.get("prompt_token_ids"), list):
-                    prompt_token_ids = tts_params.pop("prompt_token_ids")
-                    prompt = tokens_input(prompt_token_ids=prompt_token_ids)
-                else:
-                    prompt = tokens_input(prompt_token_ids=[1])
-                prompt["additional_information"] = tts_params
-                prompt["cache_salt"] = _conditioning_cache_salt(request, tts_params)
-            else:
-                tts_params = self._build_tts_params(request)
-                # Resolve ref_audio (explicit or auto-set for uploaded voices)
-                # to [[wav_list, sr]] so the model doesn't re-decode base64.
-                ref_audio_source = request.ref_audio
-                if ref_audio_source is None and isinstance(tts_params.get("ref_audio"), list):
-                    # Uploaded voice: ref_audio was auto-set as [base64_data_url]
-                    ref_audio_source = tts_params["ref_audio"][0]
-                if ref_audio_source is not None and isinstance(ref_audio_source, str):
-                    wav_list, sr = await self._resolve_ref_audio(ref_audio_source)
-                    artifact_key = self._get_resolved_ref_audio_artifact_key(ref_audio_source)
-                    if self._tts_model_type == "qwen3_tts" and artifact_key:
-                        tts_params[_QWEN3_TTS_REF_AUDIO_CACHE_KEY] = [artifact_key]
-                    ref_code_length = self._estimate_ref_code_len([wav_list, sr])
-                    if self._tts_model_type == "qwen3_tts" and ref_code_length is not None:
-                        tts_params["ref_code_length"] = [int(ref_code_length)]
-                    if self._qwen3_tts_can_use_ref_audio_artifact_only(tts_params, artifact_key):
-                        logger.debug("Using Qwen3-TTS ref_audio artifact-only path: %s", artifact_key)
-                    else:
-                        tts_params["ref_audio"] = [[wav_list, sr]]
-                        if self._tts_model_type == "qwen3_tts":
-                            qwen3_ref_audio_warmup_artifact_key = artifact_key
-
-                ph_len = await self._estimate_prompt_len_async(tts_params)
-                prompt = tokens_input(prompt_token_ids=[1] * ph_len)
-                prompt["additional_information"] = tts_params
-                prompt["cache_salt"] = _conditioning_cache_salt(request, tts_params)
+            prepared = await adapter.build(request, sampling_params_list, has_inline_ref_audio)
+            prompt = prepared.prompt
+            tts_params = prepared.tts_params
+            qwen3_ref_audio_warmup_artifact_key = prepared.warmup_artifact_key
        else:
            # Qwen omni models (Qwen3-Omni, Qwen2.5-Omni) use a "talker"
            # stage whose preprocess requires chat-templated tokens.  The
@@ -3413,7 +3353,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
            tts_params = {}
            prompt = {"prompt": request.input}

-        if self._is_fish_speech:
+        if self._tts_model_type == "step_audio2":
+            model_type = "step_audio2"
+        elif self._is_fish_speech:
            model_type = "fish_speech"
        elif self._tts_model_type == "covo_audio":
            model_type = "covo_audio"
@@ -4004,6 +3946,11 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
        merged_requests = [self._merge_batch_item(batch_request, item) for item in batch_request.items]

        async def _run_item(idx: int, req: OpenAICreateSpeechRequest) -> SpeechBatchItemResult:
+            # Batch validation still goes through _validate_tts_request directly
+            # (not the adapter). The single-request path validates via the
+            # adapter; both ultimately call the same per-model validators, so the
+            # surfaces stay in sync. Routing batch through the adapter (and its
+            # uploaded-speaker handling) is a follow-up (RFC #4327).
            validation_error = self._validate_tts_request(req)
            if validation_error is not None:
                return SpeechBatchItemResult(index=idx, status="error", error=validation_error)
@@ -5,6 +5,7 @@ from __future__ import annotations

 import copy
 import time
+from collections.abc import Mapping
 from dataclasses import dataclass
 from http import HTTPStatus
 from typing import Any, cast
@@ -620,7 +621,7 @@ class OmniOpenAIServingVideo:
    def _resolve_fps(result: Any) -> int | None:
        """Extract fps from multimodal_output if the model reported it."""
        multimodal_output = getattr(result, "multimodal_output", None)
-        if isinstance(multimodal_output, dict):
+        if isinstance(multimodal_output, Mapping):
            fps = multimodal_output.get("fps")
            if fps is not None:
                try:
@@ -633,7 +634,7 @@ class OmniOpenAIServingVideo:
        request_output = getattr(result, "request_output", None)
        if isinstance(request_output, dict):
            mm = request_output.get("multimodal_output") or {}
-            if isinstance(mm, dict):
+            if isinstance(mm, Mapping):
                fps = mm.get("fps")
                if fps is not None:
                    try:
@@ -644,7 +645,7 @@ class OmniOpenAIServingVideo:
                        pass
        elif hasattr(request_output, "multimodal_output"):
            mm = getattr(request_output, "multimodal_output", None)
-            if isinstance(mm, dict):
+            if isinstance(mm, Mapping):
                fps = mm.get("fps")
                if fps is not None:
                    try:
@@ -659,7 +660,7 @@ class OmniOpenAIServingVideo:
    @classmethod
    def _extract_audio_sample_rate_from_result(cls, result: Any) -> int | None:
        multimodal_output = getattr(result, "multimodal_output", None)
-        if isinstance(multimodal_output, dict):
+        if isinstance(multimodal_output, Mapping):
            sample_rate = cls._coerce_audio_sample_rate(
                multimodal_output.get("audio_sample_rate")
                or multimodal_output.get("sample_rate")
@@ -672,7 +673,7 @@ class OmniOpenAIServingVideo:
        request_output = getattr(result, "request_output", None)
        if isinstance(request_output, dict):
            multimodal_output = request_output.get("multimodal_output") or {}
-            if isinstance(multimodal_output, dict):
+            if isinstance(multimodal_output, Mapping):
                return cls._coerce_audio_sample_rate(
                    multimodal_output.get("audio_sample_rate")
                    or multimodal_output.get("sample_rate")
@@ -681,7 +682,7 @@ class OmniOpenAIServingVideo:
                )
        elif hasattr(request_output, "multimodal_output"):
            multimodal_output = getattr(request_output, "multimodal_output", None)
-            if isinstance(multimodal_output, dict):
+            if isinstance(multimodal_output, Mapping):
                return cls._coerce_audio_sample_rate(
                    multimodal_output.get("audio_sample_rate")
                    or multimodal_output.get("sample_rate")
@@ -32,6 +32,7 @@ import json
 import time as _time
 import uuid
 import wave
+from collections.abc import Mapping
 from typing import Any

 import torch
@@ -796,7 +797,7 @@ class OmniStreamingVideoHandler:
        if not isinstance(outputs, list) or not outputs:
            return None
        mm_output = getattr(outputs[0], "multimodal_output", None)
-        if not isinstance(mm_output, dict):
+        if not isinstance(mm_output, Mapping):
            return None
        return mm_output.get("audio")

@@ -0,0 +1,80 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Registry of TTS serving adapters.
+
+Adapters register themselves by their ``name`` (the model-type discriminator
+returned by ``OmniOpenAIServingSpeech._detect_tts_model_type``) via
+``@register_tts_adapter``. Resolution is by that name: detection already encodes
+the (order-sensitive, partly ``model_arch``-based) stage logic, so keying the
+registry on the resolved type is both simpler and correct by construction.
+``stage_keys`` is retained on each adapter as documentation metadata.
+"""
+
+from vllm.logger import init_logger
+
+from vllm_omni.entrypoints.openai.tts_adapters.base import (
+    ARTTSAdapter,
+    DiffusionTTSAdapter,
+    OutputPolicy,
+    PreparedRequest,
+    SpeechServingContext,
+    TTSModelAdapter,
+)
+
+logger = init_logger(__name__)
+
+TTS_ADAPTER_REGISTRY: dict[str, type[TTSModelAdapter]] = {}
+
+
+def register_tts_adapter(cls: type[TTSModelAdapter]) -> type[TTSModelAdapter]:
+    """Class decorator: index ``cls`` under its ``name`` (model-type)."""
+    if cls.name in TTS_ADAPTER_REGISTRY:
+        raise ValueError(
+            f"TTS adapter name {cls.name!r} already registered to "
+            f"{TTS_ADAPTER_REGISTRY[cls.name].__qualname__}; {cls.__qualname__} conflicts."
+        )
+    TTS_ADAPTER_REGISTRY[cls.name] = cls
+    return cls
+
+
+def all_tts_model_types() -> frozenset[str]:
+    """All registered model-type names."""
+    return frozenset(TTS_ADAPTER_REGISTRY)
+
+
+def resolve_adapter(model_type: str | None) -> type[TTSModelAdapter] | None:
+    """Return the adapter for a detected model-type, or ``None``."""
+    if model_type is None:
+        return None
+    return TTS_ADAPTER_REGISTRY.get(model_type)
+
+
+# Import adapter modules for their registration side effects. Keep at the bottom
+# so the registry helpers above are defined first.
+from vllm_omni.entrypoints.openai.tts_adapters import (  # noqa: E402,F401
+    cosyvoice3,
+    covo_audio,
+    fish_speech,
+    glm_tts,
+    higgs_audio_v2,
+    higgs_audio_v3,
+    ming_tts,
+    moss_tts,
+    omnivoice,
+    qwen3_tts,
+    step_audio2,
+    voxcpm2,
+    voxtral,
+)
+
+__all__ = [
+    "ARTTSAdapter",
+    "DiffusionTTSAdapter",
+    "OutputPolicy",
+    "PreparedRequest",
+    "SpeechServingContext",
+    "TTSModelAdapter",
+    "TTS_ADAPTER_REGISTRY",
+    "all_tts_model_types",
+    "register_tts_adapter",
+    "resolve_adapter",
+]
@@ -0,0 +1,199 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Base contract for per-model TTS serving adapters.
+
+This package factors the per-model ``if self._tts_model_type == ...`` dispatch
+in ``serving_speech.py`` into one adapter class per model. Each adapter owns its
+model's request normalization, validation, prompt/param building, sampling
+overrides, and output policy, so adding a model means writing one adapter file
+instead of editing the shared serving module in ~10 scattered places.
+
+See the RFC for the full design (issue #4327). This is the foundation landed in
+the first migration PR; Qwen3-TTS is the first model routed through it while the
+remaining models stay on the legacy path until individually migrated.
+"""
+
+from abc import ABC, abstractmethod
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, ClassVar
+
+if TYPE_CHECKING:
+    from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest
+
+
+_conditioning_cache_salt_fn: "Callable[..., str] | None" = None
+
+
+def conditioning_cache_salt(request: "OpenAICreateSpeechRequest", tts_params: dict) -> str:
+    """Return the conditioning cache salt for ``request`` + ``tts_params``.
+
+    Lazily imports and caches ``serving_speech._conditioning_cache_salt`` on first
+    use: the import is deferred to break the adapters<->serving_speech import
+    cycle, and cached so it resolves once instead of on every ``build()`` call.
+    """
+    global _conditioning_cache_salt_fn
+    if _conditioning_cache_salt_fn is None:
+        from vllm_omni.entrypoints.openai.serving_speech import _conditioning_cache_salt
+
+        _conditioning_cache_salt_fn = _conditioning_cache_salt
+    return _conditioning_cache_salt_fn(request, tts_params)
+
+
+@dataclass
+class OutputPolicy:
+    """How the orchestrator aggregates engine output for a model.
+
+    ``accumulate_nonstreaming`` enables MOSS-style cross-step accumulation in the
+    non-streaming path. Streaming cumulative/delta semantics stay engine-side,
+    keyed by request id, and are unaffected by this flag.
+    """
+
+    accumulate_nonstreaming: bool = False
+
+
+@dataclass
+class PreparedRequest:
+    """Everything the generic orchestrator needs to call ``<engine>.generate()``.
+
+    The fields mirror what ``_prepare_speech_generation`` assembled inline:
+    ``prompt`` is the engine prompt dict, ``tts_params`` the per-model parameter
+    dict, ``model_type`` the discriminator used for logging, and
+    ``output_policy`` controls non-streaming aggregation.
+    """
+
+    prompt: dict[str, Any]
+    tts_params: dict[str, Any] = field(default_factory=dict)
+    model_type: str = "generic"
+    output_policy: OutputPolicy = field(default_factory=OutputPolicy)
+    #: Cross-cutting per-request state the orchestrator still owns (e.g. the
+    #: Qwen3-TTS ref-audio warmup artifact key tracked after ``generate()``).
+    warmup_artifact_key: str | None = None
+
+
+@dataclass
+class SpeechServingContext:
+    """Shared state and helpers an adapter may use.
+
+    During the incremental migration this holds a back-reference to the owning
+    ``OmniOpenAIServingSpeech`` instance (``server``) so adapters can reuse the
+    existing, battle-tested helper implementations (ref-audio resolution,
+    uploaded-speaker handling, prompt-length estimation, speaker storage)
+    without relocating them. Follow-up PRs may narrow this to explicit fields as
+    more models migrate. Exactly one of ``engine_client`` / ``diffusion_engine``
+    is set, matching the deployment's serving backend.
+    """
+
+    server: Any
+    engine_client: Any | None = None
+    diffusion_engine: Any | None = None
+
+
+class TTSModelAdapter(ABC):
+    """Mandatory base class for a TTS model served via ``/v1/audio/speech``.
+
+    One concrete subclass per model, registered by stage key. The serving layer
+    resolves exactly one adapter at startup and dispatches all per-model
+    decisions to it. Adapters reuse shared helpers through ``ctx`` rather than
+    re-implementing them.
+    """
+
+    #: Stable discriminator string (the model-type from detection); registry key.
+    name: ClassVar[str]
+    #: Engine ``model_stage`` key(s) this model uses, for documentation only.
+    stage_keys: ClassVar[frozenset[str]] = frozenset()
+    #: Serving backend: ``"ar"`` (engine_client) or ``"diffusion"``.
+    backend: ClassVar[str] = "ar"
+
+    def __init__(self, ctx: SpeechServingContext) -> None:
+        self.ctx = ctx
+
+    def normalize(self, request: "OpenAICreateSpeechRequest") -> None:
+        """In-place request normalization/mutation (e.g. infer task type,
+        lowercase voice). Default: no-op."""
+
+    def validate(self, request: "OpenAICreateSpeechRequest") -> str | None:
+        """Return an error string, or ``None`` if the request is valid.
+
+        Should be free of new side effects beyond what ``normalize`` did.
+        """
+        return None
+
+    @abstractmethod
+    async def build(
+        self,
+        request: "OpenAICreateSpeechRequest",
+        sampling_params_list: list,
+        has_inline_ref_audio: bool,
+    ) -> PreparedRequest:
+        """Build the engine prompt + tts_params for this request.
+
+        ``sampling_params_list`` is passed read-only for models (e.g. MOSS) that
+        fold the resolved seed into ``additional_information`` at build time.
+
+        ``has_inline_ref_audio`` is captured by the orchestrator *before*
+        ``validate()`` runs, because ``_apply_uploaded_speaker`` (invoked inside
+        several adapters' ``validate``) sets ``request.ref_audio`` in place.
+        Recomputing it here would misclassify uploaded voices as inline and drop
+        the ``voice_name`` / ``voice_created_at`` metadata.
+        """
+
+    def apply_sampling_overrides(
+        self,
+        sampling_params_list: list,
+        request: "OpenAICreateSpeechRequest",
+    ) -> list:
+        """Apply model-specific sampling mutations.
+
+        The orchestrator guarantees the call order:
+        stream-coercion -> extra_params -> THIS -> seed. Default: identity.
+        """
+        return sampling_params_list
+
+
+class ARTTSAdapter(TTSModelAdapter):
+    """Adapter for models served through the AR ``engine_client`` path.
+
+    Covers pure-AR codec models as well as AR-base-LM + diffusion-side-computation
+    hybrids (e.g. VoxCPM2, Ming) whose internal diffusion is invisible to the
+    serving layer.
+    """
+
+    backend: ClassVar[str] = "ar"
+
+
+class DiffusionTTSAdapter(TTSModelAdapter):
+    """Adapter for pure-diffusion pipelines served through the diffusion engine.
+
+    In scope today: OmniVoice. Bridges to the diffusion pipeline parameter
+    contract (#3572) when present; see ``extra_body_params``.
+    """
+
+    backend: ClassVar[str] = "diffusion"
+
+    #: Backing diffusion pipeline class (for EXTRA_BODY_PARAMS lookup).
+    pipeline_cls: ClassVar[type | None] = None
+
+    @classmethod
+    def extra_body_params(cls) -> frozenset[str]:
+        """Fallback-safe access to the pipeline's declared body params.
+
+        Returns the pipeline's ``EXTRA_BODY_PARAMS`` if the #3572 contract is
+        present, else an empty frozenset (the adapter then uses its own inline
+        parameter logic).
+        """
+        params = getattr(cls.pipeline_cls, "EXTRA_BODY_PARAMS", None)
+        return frozenset(params) if params is not None else frozenset()
+
+
+# Re-exported here to avoid import cycles at call sites.
+__all__ = [
+    "ARTTSAdapter",
+    "DiffusionTTSAdapter",
+    "OutputPolicy",
+    "PreparedRequest",
+    "SpeechServingContext",
+    "TTSModelAdapter",
+]
+
+# Type alias kept for readability in adapter signatures.
+NormalizeFn = Callable[["OpenAICreateSpeechRequest"], None]
@@ -0,0 +1,30 @@
+# SPDX-License-Identifier: Apache-2.0
+"""CosyVoice3 TTS serving adapter."""
+
+from typing import TYPE_CHECKING
+
+from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter
+from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest
+
+if TYPE_CHECKING:
+    from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest
+
+
+@register_tts_adapter
+class CosyVoice3Adapter(ARTTSAdapter):
+    stage_keys = frozenset({"cosyvoice3_talker"})
+    name = "cosyvoice3"
+
+    def validate(self, request: "OpenAICreateSpeechRequest") -> str | None:
+        err = self.ctx.server._apply_uploaded_speaker(request)
+        if err:
+            return err
+        return self.ctx.server._validate_cosyvoice3_request(request)
+
+    async def build(
+        self, request: "OpenAICreateSpeechRequest", sampling_params_list: list, has_inline_ref_audio: bool
+    ) -> PreparedRequest:
+        prompt = await self.ctx.server._build_cosyvoice3_prompt(request, has_inline_ref_audio=has_inline_ref_audio)
+        # NOTE: CosyVoice3 dynamic-token sampling stays in the orchestrator tail
+        # (keyed on _tts_model_type) during this incremental migration.
+        return PreparedRequest(prompt=prompt, tts_params={}, model_type="cosyvoice3")
@@ -0,0 +1,22 @@
+# SPDX-License-Identifier: Apache-2.0
+"""CoVo-Audio serving adapter."""
+
+from typing import TYPE_CHECKING
+
+from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter
+from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest
+
+if TYPE_CHECKING:
+    from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest
+
+
+@register_tts_adapter
+class CovoAudioAdapter(ARTTSAdapter):
+    stage_keys = frozenset({"fused_thinker_talker"})
+    name = "covo_audio"
+
+    async def build(
+        self, request: "OpenAICreateSpeechRequest", sampling_params_list: list, has_inline_ref_audio: bool
+    ) -> PreparedRequest:
+        prompt = self.ctx.server._build_covo_audio_prompt(request)
+        return PreparedRequest(prompt=prompt, tts_params={}, model_type="covo_audio")
@@ -0,0 +1,37 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Fish Speech serving adapter (retires the legacy ``_is_fish_speech`` flag)."""
+
+from typing import TYPE_CHECKING
+
+from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter
+from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest, conditioning_cache_salt
+
+if TYPE_CHECKING:
+    from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest
+
+
+@register_tts_adapter
+class FishSpeechAdapter(ARTTSAdapter):
+    stage_keys = frozenset({"fish_speech_slow_ar"})
+    name = "fish_tts"
+
+    def validate(self, request: "OpenAICreateSpeechRequest") -> str | None:
+        err = self.ctx.server._apply_uploaded_speaker(request)
+        if err:
+            return err
+        return self.ctx.server._validate_fish_tts_request(request)
+
+    async def build(
+        self, request: "OpenAICreateSpeechRequest", sampling_params_list: list, has_inline_ref_audio: bool
+    ) -> PreparedRequest:
+        server = self.ctx.server
+        ref_audio_data = None
+        if request.ref_audio is not None:
+            wav_list, sr = await server._resolve_ref_audio(request.ref_audio)
+            ref_audio_data = (wav_list, sr)
+        prompt = await server._build_fish_speech_prompt_async(
+            request, ref_audio_data=ref_audio_data, has_inline_ref_audio=has_inline_ref_audio
+        )
+        tts_params = {}
+        prompt["cache_salt"] = conditioning_cache_salt(request, tts_params)
+        return PreparedRequest(prompt=prompt, tts_params=tts_params, model_type="fish_speech")
@@ -0,0 +1,30 @@
+# SPDX-License-Identifier: Apache-2.0
+"""GLM-TTS serving adapter."""
+
+from typing import TYPE_CHECKING
+
+from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter
+from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest
+
+if TYPE_CHECKING:
+    from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest
+
+
+@register_tts_adapter
+class GlmTTSAdapter(ARTTSAdapter):
+    stage_keys = frozenset({"glm_tts"})
+    name = "glm_tts"
+
+    def validate(self, request: "OpenAICreateSpeechRequest") -> str | None:
+        err = self.ctx.server._apply_uploaded_speaker(request)
+        if err:
+            return err
+        return self.ctx.server._validate_glm_tts_request(request)
+
+    async def build(
+        self, request: "OpenAICreateSpeechRequest", sampling_params_list: list, has_inline_ref_audio: bool
+    ) -> PreparedRequest:
+        prompt = await self.ctx.server._build_glm_tts_prompt(request, has_inline_ref_audio=has_inline_ref_audio)
+        # GLM-TTS dynamic-token sampling stays in the orchestrator tail
+        # (keyed on _tts_model_type) during this incremental migration.
+        return PreparedRequest(prompt=prompt, tts_params={}, model_type="glm_tts")
@@ -0,0 +1,35 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Higgs-Audio v2 serving adapter."""
+
+from typing import TYPE_CHECKING
+
+from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter
+from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest
+
+if TYPE_CHECKING:
+    from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest
+
+
+@register_tts_adapter
+class HiggsAudioV2Adapter(ARTTSAdapter):
+    stage_keys = frozenset({"higgs_audio_v2"})
+    name = "higgs_audio_v2"
+
+    def validate(self, request: "OpenAICreateSpeechRequest") -> str | None:
+        err = self.ctx.server._apply_uploaded_speaker(request)
+        if err:
+            return err
+        return self.ctx.server._validate_higgs_audio_v2_request(request)
+
+    async def build(
+        self, request: "OpenAICreateSpeechRequest", sampling_params_list: list, has_inline_ref_audio: bool
+    ) -> PreparedRequest:
+        server = self.ctx.server
+        prompt = await server._build_higgs_audio_v2_params(request)
+        if request.voice:
+            voice_lower = request.voice.lower()
+            if voice_lower in server.uploaded_speakers and not has_inline_ref_audio:
+                additional = prompt.setdefault("additional_information", {})
+                additional["voice_name"] = voice_lower
+                additional["voice_created_at"] = server._voice_created_at(voice_lower)
+        return PreparedRequest(prompt=prompt, tts_params={}, model_type="higgs_audio_v2")
@@ -0,0 +1,28 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Higgs-Audio v3 serving adapter."""
+
+from typing import TYPE_CHECKING
+
+from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter
+from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest
+
+if TYPE_CHECKING:
+    from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest
+
+
+@register_tts_adapter
+class HiggsAudioV3Adapter(ARTTSAdapter):
+    stage_keys = frozenset({"higgs_audio_v3"})
+    name = "higgs_audio_v3"
+
+    def validate(self, request: "OpenAICreateSpeechRequest") -> str | None:
+        err = self.ctx.server._apply_uploaded_speaker(request)
+        if err:
+            return err
+        return self.ctx.server._validate_higgs_audio_v3_request(request)
+
+    async def build(
+        self, request: "OpenAICreateSpeechRequest", sampling_params_list: list, has_inline_ref_audio: bool
+    ) -> PreparedRequest:
+        prompt = await self.ctx.server._build_higgs_audio_v3_params(request)
+        return PreparedRequest(prompt=prompt, tts_params={}, model_type="higgs_audio_v3")
@@ -0,0 +1,44 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Ming-TTS (dense) serving adapter."""
+
+from typing import TYPE_CHECKING
+
+from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter
+from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest
+
+if TYPE_CHECKING:
+    from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest
+
+
+@register_tts_adapter
+class MingTTSAdapter(ARTTSAdapter):
+    # Detected by model_arch (MingTTSForConditionalGeneration), not stage key.
+    name = "ming_tts"
+
+    def validate(self, request: "OpenAICreateSpeechRequest") -> str | None:
+        return self.ctx.server._validate_ming_tts_request(request)
+
+    async def build(
+        self, request: "OpenAICreateSpeechRequest", sampling_params_list: list, has_inline_ref_audio: bool
+    ) -> PreparedRequest:
+        server = self.ctx.server
+        ref_audio_source = request.ref_audio
+        voice_lower = request.voice.lower() if isinstance(request.voice, str) else None
+        if ref_audio_source is None and voice_lower in server.uploaded_speakers:
+            ref_audio_source = server._get_uploaded_audio_data(request.voice)
+            if request.ref_text is None:
+                request.ref_text = server.uploaded_speakers[voice_lower].get("ref_text")
+        ref_audio_data = None
+        if isinstance(ref_audio_source, list):
+            ref_audio_data = await server._resolve_ref_audio_many(ref_audio_source)
+            if request.speaker_embedding is None:
+                request.speaker_embedding = server._extract_ming_speaker_embeddings_from_ref_audio(ref_audio_data)
+        elif ref_audio_source is not None and isinstance(ref_audio_source, str):
+            wav_list, sr = await server._resolve_ref_audio(ref_audio_source)
+            ref_audio_data = (wav_list, sr)
+            if request.speaker_embedding is None:
+                request.speaker_embedding = server._extract_ming_speaker_embeddings_from_ref_audio([ref_audio_data])[0]
+        prompt = server._build_ming_dense_prompt(request, ref_audio_data=ref_audio_data)
+        tts_params = prompt.get("additional_information", {})
+        # Ming stop-token / max_tokens sampling stays in the orchestrator tail.
+        return PreparedRequest(prompt=prompt, tts_params=tts_params, model_type="ming_tts")
@@ -0,0 +1,58 @@
+# SPDX-License-Identifier: Apache-2.0
+"""MOSS-TTS serving adapters (Nano + full family).
+
+Both variants share the same build/validate flow (``_build_moss_tts_params``
+handles each); they are registered under distinct model-type names.
+"""
+
+from typing import TYPE_CHECKING
+
+from vllm.inputs import tokens_input
+
+from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter
+from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest, conditioning_cache_salt
+
+if TYPE_CHECKING:
+    from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest
+
+
+class _MossTTSAdapterBase(ARTTSAdapter):
+    def validate(self, request: "OpenAICreateSpeechRequest") -> str | None:
+        err = self.ctx.server._apply_uploaded_speaker(request)
+        if err:
+            return err
+        return self.ctx.server._validate_moss_tts_request(request)
+
+    async def build(
+        self, request: "OpenAICreateSpeechRequest", sampling_params_list: list, has_inline_ref_audio: bool
+    ) -> PreparedRequest:
+        server = self.ctx.server
+        tts_params = await server._build_moss_tts_params(request)
+        if request.voice:
+            voice_lower = request.voice.lower()
+            if voice_lower in server.uploaded_speakers and not has_inline_ref_audio:
+                tts_params["voice_name"] = [voice_lower]
+                tts_params["voice_created_at"] = [server._voice_created_at(voice_lower)]
+        # MOSS reads the resolved seed at build time (it samples internally).
+        if sampling_params_list and getattr(sampling_params_list[0], "seed", None) is not None:
+            tts_params["seed"] = [sampling_params_list[0].seed]
+        if isinstance(tts_params.get("prompt_token_ids"), list):
+            prompt_token_ids = tts_params.pop("prompt_token_ids")
+            prompt = tokens_input(prompt_token_ids=prompt_token_ids)
+        else:
+            prompt = tokens_input(prompt_token_ids=[1])
+        prompt["additional_information"] = tts_params
+        prompt["cache_salt"] = conditioning_cache_salt(request, tts_params)
+        return PreparedRequest(prompt=prompt, tts_params=tts_params, model_type=self.name)
+
+
+@register_tts_adapter
+class MossTTSNanoAdapter(_MossTTSAdapterBase):
+    stage_keys = frozenset({"moss_tts_nano"})
+    name = "moss_tts_nano"
+
+
+@register_tts_adapter
+class MossTTSAdapter(_MossTTSAdapterBase):
+    stage_keys = frozenset({"moss_tts", "moss_tts_codec"})
+    name = "moss_tts"
Author	SHA1	Message	Date
Zhou Taichang	7b837944a1	[Rebase] Rebase to vllm 0.23.0 (#4286 ) Signed-off-by: tzhouam <tzhouam@connect.ust.hk>	2026-06-13 16:43:02 +00:00
Chendi.Xue	a9ee92ddec	[XPU] update Dreamzero to support any HW (#4399 ) Signed-off-by: Chendi Xue <chendi.xue@intel.com>	2026-06-13 15:32:39 +00:00
Yuanheng Zhao	31886d6ece	[Doc][Recipe] Update CUDA verifications for inclusionAI/Ming-omni-tts-0.5B (#4324 ) Signed-off-by: Yuanheng <jonathan.zhaoyh@gmail.com>	2026-06-13 18:12:37 +08:00
Hongsheng Liu	fa94f41807	[Bugfix] Surface all-rank diffusion RPC failures (#4403 ) Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>	2026-06-13 15:42:16 +08:00
Yueqian Lin	d75a6433b4	[Refactor] TTS serving adapter framework + migrate TTS models (excl. ming_flash_omni) (#4330 ) Signed-off-by: Yueqian Lin <linyueqian@outlook.com>	2026-06-13 04:22:54 +00:00
Chendi.Xue	6506ed3d4d	[XPU][CI] Fix docker build slowness (#4402 ) Signed-off-by: Chendi Xue <chendi.xue@intel.com>	2026-06-13 10:59:15 +08:00
akshatvishu	d4f2a49a68	[Bugfix] Ming-omni-tts: generalize dict checks to Mapping (#4397 ) Signed-off-by: akshatvishu <akshatnayak197@gmail.com>	2026-06-13 09:37:57 +08:00
wuli666 linyueqian wuli666	9bc8705c98	[New Model] Step-Audio2 (#464 ) Signed-off-by: linyueqian <linyueqian@outlook.com> Signed-off-by: wuli666 <djjpro975@gmail.com> Co-authored-by: linyueqian <linyueqian@outlook.com> Co-authored-by: wuli666 <djjpro975@gmail.com>	2026-06-13 01:00:33 +00:00
Henry J jjw Claude Fable 5 Yueqian Lin	1733831f78	[Bugfix] Fix cross-request codes.ref leak in Qwen3-TTS make_omni_output (#4373 ) Signed-off-by: jjw <m@jjw.in> Co-authored-by: jjw <m@jjw.in> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>	2026-06-13 00:41:49 +00:00
Joni JuvonenandHongsheng Liu	20b7a67b0d	[Perf][Wan2.2] Skip attention mask for zero-padded SP sequences to avoid varlen path (#3763 ) Signed-off-by: jjuvonen-amd <joni.juvonen@amd.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-06-13 07:50:39 +08:00
wangyu	9577ccba50	[CI] skip MOSS-TTS-Nano E2E tests pending issue#4361 (#4391 )	2026-06-13 00:07:57 +08:00
wangyuandHongsheng Liu	54a9d241c0	[CI] Diff-gate L2/L3 E2E jobs and migrate single-GPU tests to gpu_1_queue (#4365 ) Signed-off-by: wangyu <410167048@qq.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-06-12 19:05:43 +08:00
zyz111222andCanlin Guo	6de0ca36bb	[Hardware][Ascend] Adapt Qwen3 TTS for 310P (#4283 ) Signed-off-by: zouyizhou <zouyizhou@huawei.com> Co-authored-by: Canlin Guo <canlinguosdu@gmail.com>	2026-06-12 09:39:21 +00:00
ChenWenjing	293396737d	[ci]skip voxcpm2 pcm hnr test (#4375 ) Signed-off-by: CHEN <116010019@link.cuhk.edu.cn>	2026-06-12 17:14:33 +08:00
tanhaoan333	aa59dc50c7	[bugfix]VoxCPM2 audio encoder adapt other than CUDA (#4374 ) Signed-off-by: tanhaoan333 <tanhaoan@huawei.com>	2026-06-12 08:54:29 +00:00
Samit	f44f3ef9f7	[platform] fix: set UnspecifiedOmniPlatform device_type to cpu (#4357 ) Signed-off-by: samithuang <285365963@qq.com>	2026-06-12 08:40:00 +00:00
Megha Agarwal Copilot Hongsheng Liu meghaagr13 Boao Shi Taichang Zhou Cursor	8f869c8ea5	[Refactor] Output Processor Phase 2: separate multimodal output channel (#1601 ) (#2744 ) Signed-off-by: Megha Agarwal <agarwalmegha1308@gmail.com> Signed-off-by: meghaagr13 <divyanshsinghvi@gmail.com> Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com> Co-authored-by: meghaagr13 <divyanshsinghvi@gmail.com> Co-authored-by: Boao Shi <aoibosh@connect.hku.hk> Co-authored-by: Taichang Zhou <tzhouam@connect.ust.hk> Co-authored-by: Cursor <cursoragent@cursor.com>	2026-06-12 15:19:14 +08:00