Format diffusion RPC guard

Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>
Address diffusion RPC review feedback
2026-06-13 14:43:32 +08:00 · 2026-06-13 14:39:53 +08:00 · 2026-06-13 12:32:44 +08:00 · 2026-06-13 12:30:44 +08:00 · 2026-06-13 04:22:54 +00:00 · 2026-06-13 10:59:15 +08:00
124 changed files with 8786 additions and 715 deletions
@@ -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
@@ -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***:

@@ -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 = [
@@ -1,3 +1,3 @@
 -r common.txt
 onnxruntime-cann>=1.23.2
-torchaudio==2.9.0
+torchaudio==2.10.0
@@ -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
@@ -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)
@@ -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
@@ -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)
@@ -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)
@@ -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
@@ -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.

@@ -41,10 +128,12 @@ class OmniRequestState(RequestState):
        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 +142,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 +150,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 +277,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 +300,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 +330,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 +461,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 +469,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 +483,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 +585,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,
@@ -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,
@@ -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
@@ -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"
@@ -0,0 +1,49 @@
+# SPDX-License-Identifier: Apache-2.0
+"""OmniVoice serving adapter (AR engine_client path).
+
+OmniVoice can also be served through the pure-diffusion engine via
+``for_diffusion`` (``_create_diffusion_speech``); unifying that path is a
+follow-up. This adapter covers the AR-stage deployment.
+"""
+
+from typing import TYPE_CHECKING
+
+import numpy as np
+
+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 OmniVoiceAdapter(ARTTSAdapter):
+    stage_keys = frozenset({"omnivoice_generator"})
+    name = "omnivoice"
+
+    def validate(self, request: "OpenAICreateSpeechRequest") -> str | None:
+        if not request.input or not request.input.strip():
+            return "Input text cannot be empty"
+        return self.ctx.server._apply_uploaded_speaker(request)
+
+    async def build(
+        self, request: "OpenAICreateSpeechRequest", sampling_params_list: list, has_inline_ref_audio: bool
+    ) -> PreparedRequest:
+        server = self.ctx.server
+        prompt: dict = {"input": request.input}
+        if request.ref_audio:
+            wav, sr = await server._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 server.uploaded_speakers and not has_inline_ref_audio:
+                prompt["voice_name"] = voice_lower
+                prompt["voice_created_at"] = server._voice_created_at(voice_lower)
+        if request.language:
+            prompt["lang"] = request.language
+        if request.instructions:
+            prompt["instruct"] = request.instructions
+        return PreparedRequest(prompt=prompt, tts_params={}, model_type="omnivoice")
@@ -0,0 +1,56 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Qwen3-TTS serving adapter.
+
+First model migrated onto the adapter framework (RFC #4327). It is the dispatch
+surface for Qwen3-TTS: the serving orchestrator routes prompt/param building
+through ``build()`` instead of the inline ``else`` branch in
+``_prepare_speech_generation``.
+
+To stay behaviour-identical and avoid divergence from the batch path (which
+shares ``_validate_qwen_tts_request`` and ``_build_tts_params``), the adapter
+reuses the single-source helper implementations on the serving instance through
+``ctx.server`` rather than copying them. Subsequent PRs relocate those bodies as
+more models migrate and the batch path also routes through adapters.
+"""
+
+from typing import TYPE_CHECKING
+
+from vllm.logger import init_logger
+
+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
+
+logger = init_logger(__name__)
+
+
+@register_tts_adapter
+class Qwen3TTSAdapter(ARTTSAdapter):
+    """Adapter for Qwen3-TTS (AR ``engine_client`` backend)."""
+
+    stage_keys = frozenset({"qwen3_tts"})
+    name = "qwen3_tts"
+
+    def normalize(self, request: "OpenAICreateSpeechRequest") -> None:
+        """Qwen3-TTS normalization (Base-task inference, voice lowercasing) is
+        performed inside ``validate`` today; kept fused for a strict behaviour
+        match. See ``_validate_qwen_tts_request``."""
+
+    def validate(self, request: "OpenAICreateSpeechRequest") -> str | None:
+        # Go through the shared dispatcher (which routes to
+        # _validate_qwen_tts_request for this model-type) rather than the leaf
+        # validator directly, matching the pre-adapter qwen3 path.
+        return self.ctx.server._validate_tts_request(request)
+
+    async def build(
+        self, request: "OpenAICreateSpeechRequest", sampling_params_list: list, has_inline_ref_audio: bool
+    ) -> PreparedRequest:
+        prompt, tts_params, warmup_key = await self.ctx.server._build_qwen3_tts_request(request)
+        return PreparedRequest(
+            prompt=prompt,
+            tts_params=tts_params,
+            model_type=tts_params.get("task_type", ["unknown"])[0],
+            warmup_artifact_key=warmup_key,
+        )
@@ -0,0 +1,33 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Step-Audio2 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 StepAudio2Adapter(ARTTSAdapter):
+    """Adapter for Step-Audio2 (AR ``engine_client`` backend).
+
+    Step-Audio2 runs a single thinker stage that emits audio tokens after a
+    ``<tts_start>`` marker; prompt building is a synchronous chat-template
+    construction with no uploaded-voice / ref-audio handling.
+    """
+
+    stage_keys = frozenset({"step_audio2_thinker"})
+    name = "step_audio2"
+
+    def validate(self, request: "OpenAICreateSpeechRequest") -> str | None:
+        # Dispatcher routes to the step_audio2 case (non-empty input check).
+        return self.ctx.server._validate_tts_request(request)
+
+    async def build(
+        self, request: "OpenAICreateSpeechRequest", sampling_params_list: list, has_inline_ref_audio: bool
+    ) -> PreparedRequest:
+        prompt = self.ctx.server._build_step_audio2_prompt(request)
+        return PreparedRequest(prompt=prompt, tts_params={}, model_type="step_audio2")
@@ -0,0 +1,49 @@
+# SPDX-License-Identifier: Apache-2.0
+"""VoxCPM2 serving adapter (AR base-LM + diffusion side-computation)."""
+
+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 VoxCPM2Adapter(ARTTSAdapter):
+    """VoxCPM2 shares ``latent_generator`` with VoxCPM; selected when no ``vae``
+    stage is present (and/or via ``model_arch``)."""
+
+    stage_keys = frozenset({"latent_generator"})
+    name = "voxcpm2"
+
+    def validate(self, request: "OpenAICreateSpeechRequest") -> str | None:
+        return self.ctx.server._validate_voxcpm2_request(request)
+
+    async def build(
+        self, request: "OpenAICreateSpeechRequest", sampling_params_list: list, has_inline_ref_audio: bool
+    ) -> PreparedRequest:
+        server = self.ctx.server
+        # VoxCPM2 needs the raw waveform tuple for prefill-length accounting, so
+        # it loads uploaded audio directly rather than via _apply_uploaded_speaker.
+        uploaded_ref = None
+        if request.voice:
+            voice_lower = request.voice.lower()
+            if voice_lower in server.uploaded_speakers and not has_inline_ref_audio:
+                if server.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 = server._load_uploaded_audio(voice_lower)
+        prompt = await server._build_voxcpm2_prompt(request, uploaded_ref=uploaded_ref)
+        tts_params = {}
+        if request.voice:
+            voice_lower = request.voice.lower()
+            if voice_lower in server.uploaded_speakers or voice_lower in server.precomputed_speakers:
+                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=tts_params, model_type="voxcpm2")
@@ -0,0 +1,25 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Voxtral 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 VoxtralTTSAdapter(ARTTSAdapter):
+    stage_keys = frozenset({"audio_generation"})
+    name = "voxtral_tts"
+
+    def validate(self, request: "OpenAICreateSpeechRequest") -> str | None:
+        return self.ctx.server._validate_voxtral_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_voxtral_prompt_async(request)
+        return PreparedRequest(prompt=prompt, tts_params={}, model_type="voxtral_tts")
@@ -21,6 +21,7 @@ to avoid OOM from vLLM's post-init KV-cache profiling. Mirrors qwen3_tts

 from __future__ import annotations

+import contextlib
 import tempfile
 import threading
 from collections.abc import Iterable
@@ -38,6 +39,40 @@ from vllm_omni.model_executor.models.utils import reinit_rotary_inv_freq
 logger = init_logger(__name__)


+@contextlib.contextmanager
+def _hf_load_without_tp_warmup():
+    """Work around a transformers 5.8.x crash loading some trust_remote_code LMs.
+
+    ``caching_allocator_warmup`` does ``len(model._tp_plan)`` when
+    ``torch.distributed`` is initialized, but the MOSS-TTS-Nano remote model
+    leaves ``_tp_plan`` as ``None`` (it survives ``PreTrainedModel.__init__``),
+    so the warmup raises ``TypeError: object of type 'NoneType' has no len()``.
+    vLLM initializes ``torch.distributed`` even at TP=1, so the crash fires.
+
+    MOSS-TTS-Nano runs single-GPU (no tensor parallelism), so forcing the
+    warmup down its non-distributed branch is byte-for-byte equivalent and
+    sidesteps the ``None`` plan. Scoped to the wrapped ``from_pretrained`` call.
+
+    Targets the private ``transformers.modeling_utils._is_torch_distributed_initialized``.
+    Guarded by ``hasattr`` so that if a future transformers renames or removes it
+    (and presumably the crash with it), we no-op instead of raising AttributeError.
+    """
+    import transformers.modeling_utils as _tmu
+
+    if not hasattr(_tmu, "_is_torch_distributed_initialized"):
+        # Internal hook gone (transformers refactor); assume the upstream crash
+        # is gone too and run the load unpatched.
+        yield
+        return
+
+    original = _tmu._is_torch_distributed_initialized
+    _tmu._is_torch_distributed_initialized = lambda: False
+    try:
+        yield
+    finally:
+        _tmu._is_torch_distributed_initialized = original
+
+
 def _patch_torchaudio_load() -> None:
    """Patch torchaudio.load to use soundfile if torchcodec is unavailable."""
    try:
@@ -158,11 +193,12 @@ class MossTTSNanoForGeneration(nn.Module):
        logger.info("Loading MOSS-TTS-Nano LM from %s (dtype=%s)", self.model_path, tts_dtype)
        from transformers import AutoModelForCausalLM

-        lm = AutoModelForCausalLM.from_pretrained(
-            self.model_path,
-            trust_remote_code=True,
-            torch_dtype=tts_dtype,
-        )
+        with _hf_load_without_tp_warmup():
+            lm = AutoModelForCausalLM.from_pretrained(
+                self.model_path,
+                trust_remote_code=True,
+                torch_dtype=tts_dtype,
+            )
        if device.type == "cuda":
            try:
                import flash_attn  # noqa: F401
@@ -197,11 +233,12 @@ class MossTTSNanoForGeneration(nn.Module):
        logger.info("Loading MOSS-Audio-Tokenizer-Nano from %s", codec_path)
        from transformers import AutoModel

-        audio_tokenizer = AutoModel.from_pretrained(
-            codec_path,
-            trust_remote_code=True,
-            torch_dtype=torch.float32,
-        )
+        with _hf_load_without_tp_warmup():
+            audio_tokenizer = AutoModel.from_pretrained(
+                codec_path,
+                trust_remote_code=True,
+                torch_dtype=torch.float32,
+            )
        audio_tokenizer.to(device=device)
        audio_tokenizer.eval()
        self._audio_tokenizer: nn.Module = audio_tokenizer
@@ -329,6 +329,7 @@ class Qwen3TTSPromptEmbedsBuilder:
        self._tts_pad_embed_buffer = tts_pad_embed
        self._encode_ref_audio_batch_fn = encode_ref_audio_batch
        self._speaker_cache = speaker_cache
+        self._embedding_dtype = torch.bfloat16

        self._text_tokenizer: Any | None = None

@@ -598,7 +599,7 @@ class Qwen3TTSPromptEmbedsBuilder:
        if isinstance(ref_code, torch.Tensor):
            entry["ref_code"] = ref_code.detach().to("cpu", dtype=torch.long).contiguous()
        if isinstance(ref_spk_embedding, torch.Tensor):
-            entry["ref_spk_embedding"] = ref_spk_embedding.detach().to("cpu", dtype=torch.bfloat16).reshape(-1)
+            entry["ref_spk_embedding"] = ref_spk_embedding.detach().to("cpu", dtype=self._embedding_dtype).reshape(-1)
        if not entry:
            return
        self._ref_audio_artifact_cache[cache_key] = entry
@@ -613,10 +614,11 @@ class Qwen3TTSPromptEmbedsBuilder:
        # CUDA. Ensure the speaker encoder is on the same device/dtype as the
        # main model before running it.
        dev = self._device()
+        dtype = self._embedding_dtype
        try:
            spk_param = next(self._speaker_encoder.parameters())
-            if spk_param.device != dev or spk_param.dtype != torch.bfloat16:
-                self._speaker_encoder.to(device=dev, dtype=torch.bfloat16)
+            if spk_param.device != dev or spk_param.dtype != dtype:
+                self._speaker_encoder.to(device=dev, dtype=dtype)
        except StopIteration:
            pass

@@ -641,8 +643,8 @@ class Qwen3TTSPromptEmbedsBuilder:
            fmin=0,
            fmax=12000,
        ).transpose(1, 2)
-        spk = self._speaker_encoder(mels.to(dtype=torch.bfloat16))[0]
-        return spk.to(dtype=torch.bfloat16)
+        spk = self._speaker_encoder(mels.to(dtype=dtype))[0]
+        return spk.to(dtype=dtype)

    def encode_ref_audio_batch(
        self,
@@ -1187,17 +1189,18 @@ class Qwen3TTSPromptEmbedsBuilder:
            # NOTE: Do NOT use _as_singleton here — the embedding may be a plain
            # float list (from API via msgspec IPC) that _as_singleton would
            # destructively unwrap to a single scalar.
+            speaker_dtype = self._embedding_dtype
            spk = None
            if voice_clone_prompt is not None:
                spk = voice_clone_prompt.get("ref_spk_embedding")
            if isinstance(spk, torch.Tensor):
-                speaker_embed = spk.to(device=input_ids.device, dtype=torch.bfloat16).view(1, 1, -1)
+                speaker_embed = spk.to(device=input_ids.device, dtype=speaker_dtype).view(1, 1, -1)
            elif isinstance(spk, (list, np.ndarray)):
-                speaker_embed = torch.tensor(spk, dtype=torch.bfloat16, device=input_ids.device).view(1, 1, -1)
+                speaker_embed = torch.tensor(spk, dtype=speaker_dtype, device=input_ids.device).view(1, 1, -1)
            elif cached_artifacts is not None and isinstance(cached_artifacts.get("ref_spk_embedding"), torch.Tensor):
                speaker_embed = (
                    cached_artifacts["ref_spk_embedding"]
-                    .to(device=input_ids.device, dtype=torch.bfloat16)
+                    .to(device=input_ids.device, dtype=speaker_dtype)
                    .view(1, 1, -1)
                )
            elif artifact_only:
@@ -400,6 +400,7 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
            torch.zeros(1, int(self.talker_config.hidden_size), dtype=model_dtype),
            persistent=False,
        )
+        self._embedding_dtype = torch.bfloat16

        tokenizer_config = Qwen3TTSTokenizerV2Config.from_pretrained(
            self.model_path,
@@ -537,10 +538,12 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
            logger.warning_once("runtime_additional_information is deprecated, use model_intermediate_buffer")
        audio_codes_list: list[torch.Tensor] = []
        ref_code_len_list: list[torch.Tensor] = []
-        ref_code_tensor: torch.Tensor | None = None
+        ref_code_list: list[torch.Tensor] = []
+        has_ref_code = False
        codec_streaming_list: list[torch.Tensor] = []
        for info in info_dicts:
            if not isinstance(info, dict):
+                ref_code_list.append(torch.empty(0, dtype=torch.long))
                continue
            codes = info.get("codes", {})
            meta = info.get("meta", {})
@@ -554,7 +557,10 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
                    )
            ref_code = codes.get("ref")
            if isinstance(ref_code, torch.Tensor) and ref_code.numel() > 0:
-                ref_code_tensor = ref_code
+                ref_code_list.append(ref_code)
+                has_ref_code = True
+            else:
+                ref_code_list.append(torch.empty(0, dtype=torch.long))
            ref_len = meta.get("ref_code_len")
            if ref_len is None:
                continue
@@ -584,8 +590,11 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
        mm: OmniPayload = {"codes": {"audio": audio_codes}}
        if ref_code_len_list:
            mm.setdefault("meta", {})["ref_code_len"] = torch.cat(ref_code_len_list, dim=0)[:span_len]
-        if ref_code_tensor is not None:
-            mm.setdefault("codes", {})["ref"] = [ref_code_tensor]
+        if has_ref_code:
+            # Batch-aligned, one entry per request: ``to_payload_element``
+            # indexes ``element[idx]`` per request, and a shorter list would
+            # silently broadcast one request's reference codes to the others.
+            mm.setdefault("codes", {})["ref"] = ref_code_list
        if codec_streaming_list:
            mm.setdefault("meta", {})["codec_streaming"] = torch.cat(codec_streaming_list, dim=0)[:span_len]
        return OmniOutput(text_hidden_states=hidden, multimodal_outputs=mm)
@@ -642,7 +651,8 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
        # :meth:`_init_runtime_buffers`. Materialize once on the right
        # device/dtype and reuse for both the prefill placeholder padding
        # and the decode text-step fallback below.
-        tts_pad_embed = self._tts_pad_embed.to(device=input_ids.device, dtype=torch.bfloat16).reshape(1, -1)
+        dtype = self._embedding_dtype
+        tts_pad_embed = self._tts_pad_embed.to(device=input_ids.device, dtype=dtype).reshape(1, -1)

        if is_prefill:
            # Prefill (prompt embeddings)
@@ -689,7 +699,7 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
                pad_n = int(span_len - int(take.shape[0]))
                pad_rows = tts_pad_embed.reshape(1, -1).to("cpu").expand(pad_n, -1)
                take = torch.cat([take, pad_rows], dim=0)
-            prompt_embeds = take.to(device=input_ids.device, dtype=torch.bfloat16)
+            prompt_embeds = take.to(device=input_ids.device, dtype=dtype)
            info_update["meta"]["talker_prefill_offset"] = int(offset + span_len)

            # When inputs_embeds is set, token ids are ignored by the model but must stay in-vocab for vLLM bookkeeping.
@@ -719,7 +729,7 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
                    tail[text_offset : text_offset + 1]
                    .to(
                        device=input_ids.device,
-                        dtype=torch.bfloat16,
+                        dtype=dtype,
                    )
                    .reshape(1, -1)
                )
@@ -744,14 +754,14 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):

        last_hidden = hs.get("last")
        if isinstance(last_hidden, torch.Tensor):
-            past_hidden = last_hidden.to(device=input_ids.device, dtype=torch.bfloat16).reshape(1, -1)
+            past_hidden = last_hidden.to(device=input_ids.device, dtype=dtype).reshape(1, -1)
        else:
            # Defensive: EOS step row is zeroed by the invalid-layer-0 mask and filtered downstream.
            past_hidden = torch.zeros_like(text_step)

        # Use OmniGPUModelRunner talker_mtp fast-path for residual codebooks and per-step inputs_embeds update.
        last_id_hidden = self.embed_input_ids(input_ids.reshape(1, 1).to(torch.long)).to(
-            device=input_ids.device, dtype=torch.bfloat16
+            device=input_ids.device, dtype=dtype
        )
        inputs_embeds_out = last_id_hidden.reshape(1, -1)

@@ -784,7 +794,7 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
            )

        device = input_ids_flat.device
-        dtype = torch.bfloat16
+        dtype = self._embedding_dtype
        # Request-independent constant (see :meth:`_init_runtime_buffers`) —
        # compute once for the batch instead of fetching it per request from
        # ``info_dict["embed"]["tts_pad"]``.
@@ -1043,11 +1053,12 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
        bsz = int(input_ids.shape[0])
        q = int(self.talker_config.num_code_groups)
        dev = input_embeds.device
+        dtype = self._embedding_dtype

        input_ids = input_ids.reshape(bsz, 1).to(dtype=torch.long, device=dev)
-        last_id_hidden = input_embeds.reshape(bsz, 1, -1).to(dtype=torch.bfloat16, device=dev)
-        past_hidden = last_talker_hidden.reshape(bsz, 1, -1).to(dtype=torch.bfloat16, device=dev)
-        text_step = text_step.reshape(bsz, 1, -1).to(dtype=torch.bfloat16, device=dev)
+        last_id_hidden = input_embeds.reshape(bsz, 1, -1).to(dtype=dtype, device=dev)
+        past_hidden = last_talker_hidden.reshape(bsz, 1, -1).to(dtype=dtype, device=dev)
+        text_step = text_step.reshape(bsz, 1, -1).to(dtype=dtype, device=dev)

        # Residual predictor runs fixed-length (Q-1) steps via the vLLM-native code_predictor.
        max_steps = q - 1
@@ -53,6 +53,32 @@ _OMNI_MODELS = {
        "qwen3_omni_code2wav",
        "Qwen3OmniMoeCode2Wav",
    ),
+    # Step-Audio2 models
+    "StepAudio2ForCausalLM": (
+        "step_audio2",
+        "step_audio2",
+        "StepAudio2ForConditionalGeneration",
+    ),
+    "StepAudio2ForConditionalGeneration": (
+        "step_audio2",
+        "step_audio2",
+        "StepAudio2ForConditionalGeneration",
+    ),
+    "StepAudio2ThinkerForConditionalGeneration": (
+        "step_audio2",
+        "step_audio2_thinker",
+        "StepAudio2ThinkerForConditionalGeneration",
+    ),
+    "StepAudio2Token2WavModel": (
+        "step_audio2",
+        "step_audio2_token2wav",
+        "StepAudio2Token2WavForConditionalGeneration",
+    ),
+    "StepAudio2Token2WavForConditionalGeneration": (
+        "step_audio2",
+        "step_audio2_token2wav",
+        "StepAudio2Token2WavForConditionalGeneration",
+    ),
    "CosyVoice3Model": (
        "cosyvoice3",
        "cosyvoice3",
@@ -0,0 +1,3 @@
+from .step_audio2 import StepAudio2ForConditionalGeneration
+
+__all__ = ["StepAudio2ForConditionalGeneration"]
@@ -0,0 +1,195 @@
+import torch
+import torch.nn as nn
+from vllm.config import VllmConfig
+from vllm.logger import init_logger
+from vllm.model_executor.models.interfaces import SupportsMultiModal, SupportsPP
+from vllm.model_executor.models.utils import init_vllm_registered_model, maybe_prefix
+from vllm.multimodal import MULTIMODAL_REGISTRY
+from vllm.sequence import IntermediateTensors
+
+from .step_audio2_thinker import (
+    StepAudio2DummyInputsBuilder,
+    StepAudio2MultiModalProcessor,
+    StepAudio2ProcessingInfo,
+)
+
+logger = init_logger(__name__)
+
+
+@MULTIMODAL_REGISTRY.register_processor(
+    StepAudio2MultiModalProcessor,
+    info=StepAudio2ProcessingInfo,
+    dummy_inputs=StepAudio2DummyInputsBuilder,
+)
+class StepAudio2ForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):
+    """
+    Step-Audio2 Main Controller
+
+    Manages two-stage inference pipeline:
+    - Stage 1 (Thinker): Audio understanding and token generation
+    - Stage 2 (Token2Wav): Audio token to waveform synthesis
+
+    Usage:
+        # Stage 1: Thinker
+        model = StepAudio2ForConditionalGeneration(
+            vllm_config=config,
+            model_stage="thinker"
+        )
+
+        # Stage 2: Token2Wav
+        model = StepAudio2ForConditionalGeneration(
+            vllm_config=config,
+            model_stage="token2wav"
+        )
+    """
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        self.have_multimodal_outputs = True
+
+        config = vllm_config.model_config.hf_config
+        multimodal_config = vllm_config.model_config.multimodal_config
+
+        self.config = config
+        self.multimodal_config = multimodal_config
+        self.vllm_config = vllm_config
+
+        raw_model_stage = vllm_config.model_config.model_stage
+        # Normalise: "step_audio2_thinker" → "thinker"
+        self.model_stage = "thinker" if raw_model_stage in ("thinker", "step_audio2_thinker") else raw_model_stage
+
+        if self.model_stage == "thinker":
+            self.thinker = init_vllm_registered_model(
+                vllm_config=vllm_config,
+                prefix=maybe_prefix(prefix, "thinker"),
+                hf_config=config,
+                architectures=["StepAudio2ThinkerForConditionalGeneration"],
+            )
+            self.model = self.thinker
+            self.token2wav = None
+
+            logger.info("Initialized Step-Audio2 Thinker (Stage 0)")
+
+        elif self.model_stage == "token2wav":
+            self.thinker = None
+            self.token2wav = init_vllm_registered_model(
+                vllm_config=vllm_config,
+                prefix=maybe_prefix(prefix, "token2wav"),
+                hf_config=config,
+                architectures=["StepAudio2Token2WavForConditionalGeneration"],
+            )
+            self.model = self.token2wav
+
+            logger.info("Initialized Step-Audio2 Token2Wav (Stage 1)")
+
+        else:
+            raise ValueError(
+                f"Invalid model_stage: {self.model_stage}. Must be 'thinker'/'step_audio2_thinker' or 'token2wav'"
+            )
+
+        self.make_empty_intermediate_tensors = (
+            self.thinker.make_empty_intermediate_tensors if self.model_stage == "thinker" else lambda: None
+        )
+
+    @classmethod
+    def get_placeholder_str(cls, modality: str, i: int) -> str | None:
+        """Get placeholder string for a modality
+
+        Returns:
+            For audio: "<audio_patch>" (matches processor's audio_token)
+        """
+        if modality == "audio":
+            return "<audio_patch>"
+        return None
+
+    def get_language_model(self) -> nn.Module:
+        """Get the underlying language model."""
+        if self.model_stage == "thinker":
+            return self.thinker.get_language_model()
+        # Token2Wav is not a real language model. Return the wrapper itself so
+        # vLLM interface helpers resolve embed_input_ids() on the adapter layer.
+        return self.token2wav
+
+    def get_multimodal_embeddings(self, **kwargs):
+        """Get multimodal embeddings - only used in Thinker stage."""
+        if self.model_stage == "thinker":
+            return self.thinker.get_multimodal_embeddings(**kwargs)
+        return None
+
+    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
+        """Compatibility helper used by older call sites."""
+        if self.model_stage == "thinker":
+            return self.model.get_input_embeddings(input_ids)
+        return self.model.embed_input_ids(input_ids)
+
+    def embed_input_ids(
+        self,
+        input_ids: torch.Tensor,
+        multimodal_embeddings=None,
+        is_multimodal=None,
+    ) -> torch.Tensor:
+        """Explicit vLLM embedding hook for both stages."""
+        if self.model_stage == "token2wav":
+            return self.model.embed_input_ids(input_ids)
+        return self.model.get_input_embeddings(input_ids, multimodal_embeddings)
+
+    def embed_multimodal(self, **kwargs):
+        """Delegate multimodal embedding to thinker stage only."""
+        if self.model_stage == "thinker":
+            return self.model.embed_multimodal(**kwargs)
+        return None
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        **kwargs,
+    ):
+        """
+        Forward pass through the model
+
+        For Thinker:
+            Returns hidden states/logits
+        For Token2Wav:
+            Returns waveform
+        """
+        return self.model.forward(
+            input_ids=input_ids,
+            positions=positions,
+            intermediate_tensors=intermediate_tensors,
+            inputs_embeds=inputs_embeds,
+            **kwargs,
+        )
+
+    def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
+        """Compute logits from hidden states"""
+        return self.model.compute_logits(hidden_states)
+
+    def load_weights(self, weights):
+        """Load weights"""
+        return self.model.load_weights(weights)
+
+    def move_submodules_to_devices(
+        self,
+        *,
+        thinker_device: str | torch.device | None = None,
+        token2wav_device: str | torch.device | None = None,
+    ) -> None:
+        """
+        Optionally move thinker/token2wav to different devices
+
+        Example:
+            model.move_submodules_to_devices(
+                thinker_device='cuda:0',
+                token2wav_device='cuda:1',
+            )
+        """
+        if thinker_device is not None and self.thinker is not None:
+            self.thinker.to(torch.device(thinker_device))
+            logger.info(f"Moved Thinker to {thinker_device}")
+
+        if token2wav_device is not None and self.token2wav is not None:
+            self.token2wav.to(torch.device(token2wav_device))
+            logger.info(f"Moved Token2Wav to {token2wav_device}")
@@ -0,0 +1,181 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Step-Audio2 configuration constants - Single Source of Truth."""
+
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class StepAudio2TokenConfig:
+    """
+    Step-Audio2 token ID ranges and vocabulary configuration.
+
+    Token ID Layout (based on vocab_size=158720):
+    - Text tokens: 0 - 151688 (standard Qwen tokenizer)
+    - Special tokens: 151689 - 151695 (reserved)
+    - Audio tokens: 151696 - 158257 (vocab size 6562)
+
+    These values are architecture constants for Step-Audio2 and do not
+    change between model variants (mini/7B).
+    """
+
+    # Text token range
+    text_max: int = 151688
+    """Maximum text token ID (inclusive)"""
+
+    # Audio token range (absolute IDs in full vocabulary)
+    audio_start: int = 151696
+    """First audio token ID in absolute vocabulary"""
+
+    audio_vocab_size: int = 6562
+    """Total number of audio tokens"""
+
+    audio_eos: int = 6561
+    """Audio EOS token ID (relative to audio_start, used for padding)"""
+
+    # Special tokens
+    audio_patch_token_id: int = 151690
+    """<audio_patch> placeholder token ID"""
+
+    @property
+    def audio_end(self) -> int:
+        """Last audio token ID in absolute vocabulary"""
+        return self.audio_start + self.audio_vocab_size - 1
+
+
+@dataclass(frozen=True)
+class StepAudio2EncoderConfig:
+    """
+    Step-Audio2 audio encoder architecture configuration.
+
+    These values define the encoder structure and should match your model weights.
+    """
+
+    # Mel spectrogram parameters
+    n_mels: int = 128
+    """Number of mel frequency bins"""
+
+    # Encoder architecture
+    n_audio_ctx: int = 1500
+    """Audio context length (max sequence length)"""
+
+    n_audio_state: int = 512
+    """Audio encoder hidden state dimension"""
+
+    n_audio_head: int = 8
+    """Number of attention heads in encoder"""
+
+    n_audio_layer: int = 6
+    """Number of encoder layers"""
+
+    # Adapter parameters
+    kernel_size: int = 3
+    """Convolution kernel size for adapter"""
+
+    adapter_stride: int = 2
+    """Stride for adapter convolution (downsampling factor)"""
+
+
+@dataclass(frozen=True)
+class StepAudio2Token2WavConfig:
+    """Token2Wav synthesis configuration."""
+
+    n_timesteps: int = 10
+    """Number of diffusion timesteps for flow model inference"""
+
+    sample_rate: int = 24000
+    """Output audio sample rate"""
+
+
+@dataclass(frozen=True)
+class StepAudio2StreamConfig:
+    """Streaming (chunked) inference configuration for Token2Wav."""
+
+    chunk_size: int = 25
+    """Number of audio tokens consumed per chunk."""
+
+    pre_lookahead_len: int = 3
+    """Conformer encoder lookahead tokens beyond the chunk boundary."""
+
+    up_rate: int = 2
+    """Token-to-mel upsampling rate inside the flow model."""
+
+    mel_cache_len: int = 8
+    """HiFT overlap mel frames (160 ms at 50 Hz mel rate)."""
+
+    n_timesteps: int = 10
+    """Flow model ODE solver steps."""
+
+    estimator_cache_keep: int = 100
+    """Attention-cache pruning window for the flow estimator."""
+
+
+@dataclass(frozen=True)
+class StepAudio2ModelConfig:
+    """
+    Step-Audio2 complete model configuration - Single Source of Truth.
+
+    All architecture constants are defined here. The stage YAML and model config.json
+    can override these values if needed, but this serves as the default.
+    """
+
+    # LLM configuration (Qwen2)
+    hidden_size: int = 4096
+    """LLM hidden dimension (Qwen2-7B: 4096, Qwen2-1.5B: 1536, Qwen2-0.5B: 896)"""
+
+    # Token configuration
+    text_max: int = 151688
+    audio_start: int = 151696
+    audio_vocab_size: int = 6562
+    audio_eos: int = 6561
+    audio_patch_token_id: int = 151690
+
+    # Encoder configuration
+    n_mels: int = 128
+    n_audio_ctx: int = 1500
+    n_audio_state: int = 512
+    n_audio_head: int = 8
+    n_audio_layer: int = 6
+    kernel_size: int = 3
+    adapter_stride: int = 2
+
+
+# Default configuration instances
+DEFAULT_TOKEN_CONFIG = StepAudio2TokenConfig()
+DEFAULT_ENCODER_CONFIG = StepAudio2EncoderConfig()
+DEFAULT_TOKEN2WAV_CONFIG = StepAudio2Token2WavConfig()
+DEFAULT_STREAM_CONFIG = StepAudio2StreamConfig()
+DEFAULT_MODEL_CONFIG = StepAudio2ModelConfig()
+
+STREAM_SOURCE_CACHE_LEN = DEFAULT_STREAM_CONFIG.mel_cache_len * 480  # 3840 samples
+
+# Export constants for backward compatibility
+STEP_AUDIO2_TEXT_MAX = DEFAULT_TOKEN_CONFIG.text_max
+STEP_AUDIO2_AUDIO_START = DEFAULT_TOKEN_CONFIG.audio_start
+STEP_AUDIO2_AUDIO_VOCAB_SIZE = DEFAULT_TOKEN_CONFIG.audio_vocab_size
+STEP_AUDIO2_AUDIO_EOS = DEFAULT_TOKEN_CONFIG.audio_eos
+STEP_AUDIO2_AUDIO_PATCH_TOKEN_ID = DEFAULT_TOKEN_CONFIG.audio_patch_token_id
+STEP_AUDIO2_AUDIO_END = DEFAULT_TOKEN_CONFIG.audio_end
+
+# Default prompt wav path for Token2Wav (can be overridden by env var)
+STEP_AUDIO2_DEFAULT_PROMPT_WAV = "default_female.wav"
+
+
+__all__ = [
+    "StepAudio2TokenConfig",
+    "StepAudio2EncoderConfig",
+    "StepAudio2Token2WavConfig",
+    "StepAudio2StreamConfig",
+    "StepAudio2ModelConfig",
+    "DEFAULT_TOKEN_CONFIG",
+    "DEFAULT_ENCODER_CONFIG",
+    "DEFAULT_TOKEN2WAV_CONFIG",
+    "DEFAULT_STREAM_CONFIG",
+    "DEFAULT_MODEL_CONFIG",
+    "STREAM_SOURCE_CACHE_LEN",
+    "STEP_AUDIO2_TEXT_MAX",
+    "STEP_AUDIO2_AUDIO_START",
+    "STEP_AUDIO2_AUDIO_VOCAB_SIZE",
+    "STEP_AUDIO2_AUDIO_EOS",
+    "STEP_AUDIO2_AUDIO_PATCH_TOKEN_ID",
+    "STEP_AUDIO2_AUDIO_END",
+]
@@ -0,0 +1,898 @@
+"""
+Step-Audio2 Thinker - Stage 1 LLM for Audio Understanding
+
+This file contains:
+- Audio encoder components (AudioEncoder, Adaptor)
+- Audio preprocessing utilities (mel-spectrogram extraction)
+- Multi-modal processing (StepAudio2MultiModalProcessor)
+- Main thinker model (StepAudio2ThinkerForConditionalGeneration)
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterable, Mapping, Sequence
+from typing import TYPE_CHECKING, Any, TypedDict
+
+import librosa  # noqa: TID251
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.utils.rnn import pad_sequence
+from transformers import AutoTokenizer
+from transformers.processing_utils import ProcessorMixin
+from vllm.config import VllmConfig
+from vllm.inputs import MultiModalDataDict
+from vllm.model_executor.models.interfaces import SupportsMultiModal, SupportsPP
+from vllm.model_executor.models.utils import (
+    _merge_multimodal_embeddings,
+    flatten_bn,
+    init_vllm_registered_model,
+    maybe_prefix,
+)
+from vllm.multimodal import MULTIMODAL_REGISTRY
+from vllm.multimodal.inputs import (
+    MultiModalFieldConfig,
+    MultiModalKwargsItems,
+    NestedTensors,
+)
+from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser
+from vllm.multimodal.processing import (
+    BaseDummyInputsBuilder,
+    BaseMultiModalProcessor,
+    BaseProcessingInfo,
+    PromptReplacement,
+    PromptUpdate,
+)
+from vllm.sequence import IntermediateTensors
+
+from vllm_omni.model_executor.models.output_templates import OmniOutput
+from vllm_omni.model_executor.models.step_audio2.step_audio2_constants import (
+    DEFAULT_MODEL_CONFIG,
+    DEFAULT_TOKEN_CONFIG,
+    STEP_AUDIO2_AUDIO_PATCH_TOKEN_ID,
+)
+
+if TYPE_CHECKING:
+    pass
+
+
+# ============================================================================
+# Audio Preprocessing Utilities
+# ============================================================================
+
+
+def _mel_filters(n_mels: int) -> torch.Tensor:
+    """Generate mel filter banks"""
+    assert n_mels in {80, 128}, f"Unsupported n_mels: {n_mels}"
+    return torch.from_numpy(librosa.filters.mel(sr=16000, n_fft=400, n_mels=n_mels))
+
+
+def _normalize_audio(audio: Any) -> torch.Tensor:
+    """Normalize audio to torch.Tensor"""
+    if isinstance(audio, tuple):
+        audio = audio[0]
+    if isinstance(audio, np.ndarray):
+        audio = torch.from_numpy(audio)
+    elif not torch.is_tensor(audio):
+        audio = torch.tensor(audio, dtype=torch.float32)
+    return audio.float()
+
+
+def log_mel_spectrogram(audio: Any, n_mels: int = 128, padding: int = 479) -> torch.Tensor:
+    """
+    Convert audio to log mel-spectrogram
+
+    Args:
+        audio: Audio waveform (any format)
+        n_mels: Number of mel frequency bins (default: 128)
+        padding: Right padding for STFT (default: 479)
+
+    Returns:
+        Log mel-spectrogram tensor of shape (n_mels, T)
+    """
+    audio = _normalize_audio(audio)
+    if padding > 0:
+        audio = F.pad(audio, (0, padding))
+    window = torch.hann_window(400, device=audio.device)
+    stft = torch.stft(audio, 400, 160, window=window, return_complex=True)
+    magnitudes = stft[..., :-1].abs() ** 2
+    filters = _mel_filters(n_mels).to(audio.device)
+    mel_spec = filters @ magnitudes
+
+    log_spec = torch.clamp(mel_spec, min=1e-10).log10()
+    log_spec = torch.maximum(log_spec, log_spec.max() - 8.0)
+    log_spec = (log_spec + 4.0) / 4.0
+    return log_spec
+
+
+def padding_mels(data: list[torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Pad mel spectrograms to same length
+
+    Args:
+        data: List of mel spectrograms, each of shape (n_mels, T)
+
+    Returns:
+        padded_feats: Padded tensor of shape (batch, n_mels, max_T)
+        feats_lengths: Lengths of each sequence
+    """
+    feats_lengths = torch.tensor([s.size(1) - 2 for s in data], dtype=torch.int32)
+    feats = [s.t() for s in data]
+    padded_feats = pad_sequence(feats, batch_first=True, padding_value=0)
+    return padded_feats.transpose(1, 2), feats_lengths
+
+
+# ============================================================================
+# Encoder Utilities
+# ============================================================================
+
+
+def make_non_pad_mask(lengths: torch.Tensor, max_len: int = 0) -> torch.Tensor:
+    """Create non-padding mask from lengths"""
+    batch_size = lengths.size(0)
+    max_len = max_len if max_len > 0 else lengths.max().item()
+    seq_range = torch.arange(0, max_len, dtype=torch.int64, device=lengths.device)
+    seq_range_expand = seq_range.unsqueeze(0).expand(batch_size, max_len)
+    seq_length_expand = lengths.unsqueeze(-1)
+    mask = seq_range_expand >= seq_length_expand
+    return ~mask
+
+
+def mask_to_bias(mask: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
+    """Convert mask to attention bias"""
+    assert mask.dtype == torch.bool
+    assert dtype in [torch.float32, torch.bfloat16, torch.float16]
+    mask = mask.to(dtype)
+    mask = (1.0 - mask) * -1.0e10
+    return mask
+
+
+# ============================================================================
+# Audio Encoder Components
+# ============================================================================
+
+
+class MultiHeadAttention(nn.Module):
+    """Multi-head attention for audio encoder"""
+
+    def __init__(self, n_state: int, n_head: int):
+        super().__init__()
+        self.n_head = n_head
+        self.query = nn.Linear(n_state, n_state)
+        self.key = nn.Linear(n_state, n_state, bias=False)
+        self.value = nn.Linear(n_state, n_state)
+        self.out = nn.Linear(n_state, n_state)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        mask: torch.Tensor | None = None,
+    ):
+        q = self.query(x)
+        k = self.key(x)
+        v = self.value(x)
+
+        wv, qk = self.qkv_attention(q, k, v, mask)
+        return self.out(wv), qk
+
+    def qkv_attention(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, mask: torch.Tensor | None = None):
+        B, T, D = q.shape
+        head_dim = D // self.n_head
+
+        # Reshape: (B, T, D) -> (B, n_head, T, head_dim)
+        q = q.view(B, T, self.n_head, head_dim).transpose(1, 2)
+        k = k.view(B, T, self.n_head, head_dim).transpose(1, 2)
+        v = v.view(B, T, self.n_head, head_dim).transpose(1, 2)
+
+        # Use F.scaled_dot_product_attention
+        # mask is additive bias (0 or -1e10), pass directly as attn_mask
+        out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
+
+        # Reshape back: (B, n_head, T, head_dim) -> (B, T, D)
+        out = out.transpose(1, 2).reshape(B, T, D)
+        return out, None
+
+
+class ResidualAttentionBlock(nn.Module):
+    """Residual attention block for audio encoder"""
+
+    def __init__(self, n_state: int, n_head: int):
+        super().__init__()
+
+        self.attn = MultiHeadAttention(n_state, n_head)
+        self.attn_ln = nn.LayerNorm(n_state)
+
+        n_mlp = n_state * 4
+        self.mlp = nn.Sequential(nn.Linear(n_state, n_mlp), nn.GELU(), nn.Linear(n_mlp, n_state))
+        self.mlp_ln = nn.LayerNorm(n_state)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        mask: torch.Tensor | None = None,
+    ):
+        x = x + self.attn(self.attn_ln(x), mask=mask)[0]
+        x = x + self.mlp(self.mlp_ln(x))
+        return x
+
+
+class AudioEncoder(nn.Module):
+    """
+    Step-Audio2 Audio Encoder
+
+    Lightweight audio encoder (6 layers, 512 hidden)
+    Optimized for 25s audio chunks at 25Hz
+    """
+
+    def __init__(self, n_mels: int, n_ctx: int, n_state: int, n_head: int, n_layer: int):
+        super().__init__()
+        self.conv1 = nn.Conv1d(n_mels, n_state, kernel_size=3, padding=1)
+        self.conv2 = nn.Conv1d(n_state, n_state, kernel_size=3, stride=2, padding=1)
+        self.positional_embedding = nn.Embedding(n_ctx, n_state)
+
+        self.blocks: Iterable[ResidualAttentionBlock] = nn.ModuleList(
+            [ResidualAttentionBlock(n_state, n_head) for _ in range(n_layer)]
+        )
+        self.avg_pooler = nn.AvgPool1d(2, stride=2)
+        self.after_norm = nn.LayerNorm(n_state)
+
+    def forward(self, x: torch.Tensor, x_len: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Args:
+            x: mel spectrogram, shape (batch_size, n_mels, T)
+            x_len: length of each audio, shape (batch_size,)
+
+        Returns:
+            x: encoded features, shape (batch_size, T//4, n_state)
+            x_len: updated lengths, shape (batch_size,)
+        """
+        T = x.size(-1)
+        x = F.gelu(self.conv1(x))
+        x = F.gelu(self.conv2(x))
+        x = x.permute(0, 2, 1)  # (B, T // 2, n_state)
+
+        # Create attention mask
+        mask = make_non_pad_mask(x_len, T).unsqueeze(1)  # (B, 1, T)
+        mask = mask_to_bias(mask[:, :, (T + 1) % 2 :: 2], x.dtype)  # (B, 1, T // 2)
+
+        # Add positional embedding
+        x = (x + self.positional_embedding.weight[: x.shape[1], :]).to(x.dtype).contiguous()
+
+        # Apply attention blocks
+        for block in self.blocks:
+            x = block(x, mask.unsqueeze(1))
+
+        # Pool and normalize
+        x = x.permute(0, 2, 1)
+        x = self.avg_pooler(x)
+        x = x.permute(0, 2, 1)
+        x_len = (x_len + 1) // 2 // 2
+        x = self.after_norm(x)
+
+        return x, x_len
+
+
+class Adaptor(nn.Module):
+    """
+    Adaptor to project audio features to LLM dimension
+
+    Maps from n_state (512) to n_hidden (LLM dimension, e.g., 4096)
+    with optional downsampling via convolution
+    """
+
+    def __init__(
+        self, n_state: int = 512, n_hidden: int = 4096, kernel_size: int = 3, stride: int = 2, adapter_state: int = 2048
+    ):
+        super().__init__()
+        self.stride = stride
+
+        if self.stride != -1:
+            self.conv = nn.Conv1d(n_state, n_state, kernel_size, stride, padding=1)
+
+        self.linear1 = nn.Linear(n_state, adapter_state)
+        self.relu = nn.ReLU()
+        self.linear2 = nn.Linear(adapter_state, n_hidden)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: audio features, shape (batch_size, T, n_state)
+
+        Returns:
+            x: projected features, shape (batch_size, T//stride, n_hidden)
+        """
+        if self.stride != -1:
+            x = x.permute(0, 2, 1)  # (B, n_state, T)
+            x = F.gelu(self.conv(x))
+            x = x.permute(0, 2, 1)  # (B, T//stride, n_state)
+
+        x = self.linear1(x)
+        x = self.relu(x)
+        x = self.linear2(x)
+
+        return x
+
+
+# ============================================================================
+# Feature Length Calculation Utility
+# ============================================================================
+
+
+def calculate_audio_feature_length(mel_length: int) -> int:
+    """
+    Calculate output feature length after encoder + adapter processing.
+
+    Processing chain:
+    1. Mel spectrogram input: length = mel_length
+    2. Encoder (AudioEncoder):
+       - conv1: kernel=3, padding=1, stride=1 → no change
+       - conv2: kernel=3, padding=1, stride=2 → T // 2
+       - avg_pool: kernel=2, stride=2 → T // 4
+       - Output length: (mel_length + 1) // 2 // 2 ≈ mel_length // 4
+    3. Adapter (Adaptor):
+       - conv: kernel=3, padding=1, stride=2 → T // 2
+       - Output length: (encoder_output - 1) // 2 + 1
+
+    Args:
+        mel_length: Length of mel spectrogram (time steps)
+
+    Returns:
+        Final feature length after encoder + adapter
+
+    Example:
+        >>> calculate_audio_feature_length(1000)
+        125
+    """
+    encoder_output_len = (mel_length + 1) // 4
+
+    adapter_output_len = (encoder_output_len - 1) // 2 + 1
+
+    return max(1, adapter_output_len)
+
+
+class StepAudio2Processor(ProcessorMixin):
+    """Processor for Step-Audio2 that handles text tokenization and audio mel-spectrograms."""
+
+    attributes = ["tokenizer"]
+    attribute_class = {"tokenizer": "PreTrainedTokenizerBase"}
+
+    def __init__(self, tokenizer):
+        self.tokenizer = tokenizer
+        self.audio_token = "<audio_patch>"
+
+    def __call__(self, text=None, audio=None, **kwargs):
+        """Process text and audio inputs."""
+        if text is None:
+            text = ""
+
+        allowed = {
+            "padding",
+            "truncation",
+            "max_length",
+            "return_tensors",
+            "add_special_tokens",
+            "pad_to_multiple_of",
+            "stride",
+            "return_attention_mask",
+            "return_token_type_ids",
+            "return_overflowing_tokens",
+            "return_offsets_mapping",
+            "return_special_tokens_mask",
+            "verbose",
+            "is_split_into_words",
+        }
+        tok_kwargs = {k: v for k, v in kwargs.items() if k in allowed}
+        if "return_tensors" not in tok_kwargs:
+            tok_kwargs["return_tensors"] = "pt"
+
+        encoded = self.tokenizer(text, **tok_kwargs)
+
+        if audio is None:
+            encoded["audio_mels"] = torch.empty((0, 128, 0))
+            encoded["audio_lens"] = torch.tensor([], dtype=torch.int32)
+            return encoded
+
+        audio_list = list(audio) if isinstance(audio, (list, tuple)) else [audio]
+
+        mels = [log_mel_spectrogram(a) for a in audio_list]
+        audio_mels, audio_lens = padding_mels(mels)
+
+        encoded["audio_mels"] = audio_mels
+        encoded["audio_lens"] = audio_lens
+        return encoded
+
+
+class StepAudio2AudioInputs(TypedDict):
+    """Audio inputs for Step-Audio2"""
+
+    audio_mels: torch.Tensor
+    """Shape: (batch, num_mel_bins, time_steps)"""
+
+    audio_lens: list[int]
+    """Shape: (batch,) - length of each audio in time steps"""
+
+
+class StepAudio2ProcessingInfo(BaseProcessingInfo):
+    """Processing info for Step-Audio2"""
+
+    def get_tokenizer(self, **kwargs):
+        return AutoTokenizer.from_pretrained(
+            self.ctx.model_config.model,
+            trust_remote_code=True,
+            fix_mistral_regex=True,
+            **kwargs,
+        )
+
+    def get_hf_processor(self, **kwargs: object):
+        return StepAudio2Processor(self.get_tokenizer())
+
+    def build_data_parser(self) -> MultiModalDataParser:
+        return MultiModalDataParser(target_sr=16000)
+
+    # Backward/branch compatibility:
+    # some code paths still call get_data_parser() on ProcessingInfo.
+    def get_data_parser(self) -> MultiModalDataParser:
+        return MultiModalDataParser(target_sr=16000)
+
+    def get_supported_mm_limits(self) -> Mapping[str, int | None]:
+        return {"audio": None}
+
+    def get_mm_max_tokens_per_item(
+        self,
+        seq_len: int,
+        mm_counts: Mapping[str, int],
+    ) -> Mapping[str, int]:
+        max_audio_tokens = 250
+        return {"audio": max_audio_tokens}
+
+
+class StepAudio2DummyInputsBuilder(BaseDummyInputsBuilder[StepAudio2ProcessingInfo]):
+    """Dummy inputs builder for Step-Audio2"""
+
+    def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
+        num_audios = mm_counts.get("audio", 0)
+        return "<audio_patch>" * num_audios
+
+    def get_dummy_mm_data(
+        self,
+        seq_len: int,
+        mm_counts: Mapping[str, int],
+        mm_options: Mapping[str, object] | None = None,
+    ) -> MultiModalDataDict:
+        audio_len = 16000 * 25
+        num_audios = mm_counts.get("audio", 0)
+        return {"audio": self._get_dummy_audios(length=audio_len, num_audios=num_audios)}
+
+
+class StepAudio2MultiModalProcessor(BaseMultiModalProcessor[StepAudio2ProcessingInfo]):
+    """Multi-modal processor for Step-Audio2"""
+
+    def _get_mm_fields_config(
+        self,
+        hf_inputs,
+        hf_processor_mm_kwargs: Mapping[str, object],
+    ) -> Mapping[str, MultiModalFieldConfig]:
+        """Map processor outputs to vLLM multimodal fields.
+
+        audio_mels is a padded 3D tensor per audio item with matching audio_lens.
+        Use flat_from_sizes to keep one mm item per audio, so downstream mm_kwargs
+        alignment does not collapse to a single item.
+        """
+        audio_lens = hf_inputs.get("audio_lens", torch.empty(0))
+        if isinstance(audio_lens, torch.Tensor):
+            lens_tensor = audio_lens.flatten()
+        elif audio_lens is None:
+            lens_tensor = torch.empty(0, dtype=torch.int32)
+        else:
+            lens_tensor = torch.tensor(
+                list(audio_lens) if hasattr(audio_lens, "__iter__") else [int(audio_lens)], dtype=torch.int32
+            )
+
+        return dict(
+            audio_mels=MultiModalFieldConfig.flat_from_sizes("audio", lens_tensor),
+            audio_lens=MultiModalFieldConfig.batched("audio"),
+        )
+
+    def _get_prompt_updates(
+        self,
+        mm_items: MultiModalDataItems,
+        hf_processor_mm_kwargs: Mapping[str, object],
+        out_mm_kwargs: MultiModalKwargsItems,
+    ) -> Sequence[PromptUpdate]:
+        processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
+
+        audio_token = getattr(processor, "audio_token", "<audio_patch>")
+
+        audio_items = mm_items.get("audio", [])
+        num_audio_items = len(audio_items) if audio_items else 0
+
+        out_mm_data = out_mm_kwargs.get_data()
+        audio_lens = out_mm_data.get("audio_lens")
+
+        feature_lens = []
+        if audio_lens is not None:
+            if isinstance(audio_lens, torch.Tensor):
+                audio_lens_list = audio_lens.tolist()
+            else:
+                audio_lens_list = list(audio_lens) if hasattr(audio_lens, "__iter__") else [audio_lens]
+
+            for length in audio_lens_list:
+                if length > 0:
+                    feature_len = calculate_audio_feature_length(int(length))
+                    feature_lens.append(feature_len)
+
+        # CRITICAL: Align feature_lens with mm_items['audio'] count
+        if num_audio_items > 0:
+            if len(feature_lens) < num_audio_items:
+                default_feature_len = 250
+                pad_count = num_audio_items - len(feature_lens)
+                feature_lens.extend([default_feature_len] * pad_count)
+            elif len(feature_lens) > num_audio_items:
+                feature_lens = feature_lens[:num_audio_items]
+        elif not feature_lens:
+            feature_lens = [250]
+
+        def get_replacement_audio(item_idx: int):
+            """Generate replacement tokens for audio placeholder.
+
+            Following Qwen2.5-Omni pattern: returns list[int] of token IDs.
+            """
+            if item_idx >= len(feature_lens):
+                num_features = 1
+            else:
+                num_features = feature_lens[item_idx]
+
+            return [STEP_AUDIO2_AUDIO_PATCH_TOKEN_ID] * num_features
+
+        return [
+            PromptReplacement(
+                modality="audio",
+                target=audio_token,
+                replacement=get_replacement_audio,
+            )
+        ]
+
+    def _call_hf_processor(
+        self,
+        prompt: str,
+        mm_data: Mapping[str, object],
+        mm_kwargs: Mapping[str, object] | None = None,
+        tok_kwargs: Mapping[str, object] | None = None,
+        **_: object,
+    ):
+        """Call HF processor and post-process outputs"""
+        mm_data = dict(mm_data)
+        audios = mm_data.pop("audios", [])
+
+        mm_kwargs = mm_kwargs or {}
+        tok_kwargs = tok_kwargs or {}
+
+        # CRITICAL: Ensure audios is ALWAYS a list to prevent string iteration
+        if audios:
+            if isinstance(audios, str):
+                mm_data["audio"] = [audios]
+            elif isinstance(audios, (list, tuple)):
+                mm_data["audio"] = audios
+            else:
+                mm_data["audio"] = [audios]
+
+        hf_inputs = super()._call_hf_processor(
+            prompt=prompt,
+            mm_data=mm_data,
+            mm_kwargs=mm_kwargs,
+            tok_kwargs=tok_kwargs,
+        )
+
+        if "audio_mels" not in hf_inputs:
+            hf_inputs["audio_mels"] = torch.empty((0, 128, 0))
+        if "audio_lens" not in hf_inputs:
+            hf_inputs["audio_lens"] = torch.tensor([], dtype=torch.int32)
+
+        return hf_inputs
+
+
+# ============================================================================
+# Main Thinker Model
+# ============================================================================
+
+
+@MULTIMODAL_REGISTRY.register_processor(
+    StepAudio2MultiModalProcessor,
+    info=StepAudio2ProcessingInfo,
+    dummy_inputs=StepAudio2DummyInputsBuilder,
+)
+class StepAudio2ThinkerForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):
+    """
+    Step-Audio2 Thinker - Stage 1 LLM
+
+    Architecture:
+        AudioEncoder (6 layers, 512 hidden) →
+        Adaptor (512 → LLM dim) →
+        Qwen2 LLM (vocab=64012)
+    """
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        self.have_multimodal_outputs = True
+
+        config = vllm_config.model_config.hf_config
+        multimodal_config = vllm_config.model_config.multimodal_config
+
+        self.config = config
+        self.multimodal_config = multimodal_config
+
+        audio_enc_cfg = getattr(config, "audio_encoder_config", {})
+
+        def _get_config_value(cfg, attr_name: str, default_value):
+            """Get value from config object or dict"""
+            if isinstance(cfg, dict):
+                return cfg.get(attr_name, default_value)
+            else:
+                return getattr(cfg, attr_name, default_value)
+
+        n_mels = _get_config_value(audio_enc_cfg, "n_mels", DEFAULT_MODEL_CONFIG.n_mels)
+        n_audio_ctx = _get_config_value(audio_enc_cfg, "n_audio_ctx", DEFAULT_MODEL_CONFIG.n_audio_ctx)
+        n_audio_state = _get_config_value(audio_enc_cfg, "n_audio_state", DEFAULT_MODEL_CONFIG.n_audio_state)
+        n_audio_head = _get_config_value(audio_enc_cfg, "n_audio_head", DEFAULT_MODEL_CONFIG.n_audio_head)
+        n_audio_layer = _get_config_value(audio_enc_cfg, "n_audio_layer", DEFAULT_MODEL_CONFIG.n_audio_layer)
+        kernel_size = _get_config_value(audio_enc_cfg, "kernel_size", DEFAULT_MODEL_CONFIG.kernel_size)
+        adapter_stride = _get_config_value(audio_enc_cfg, "adapter_stride", DEFAULT_MODEL_CONFIG.adapter_stride)
+
+        n_hidden = None
+
+        if audio_enc_cfg:
+            n_hidden = _get_config_value(audio_enc_cfg, "llm_dim", None)
+
+        if n_hidden is None:
+            text_config = getattr(config, "text_config", None)
+            if text_config is not None:
+                n_hidden = _get_config_value(text_config, "hidden_size", None)
+
+        if n_hidden is None:
+            n_hidden = _get_config_value(config, "hidden_size", None)
+
+        if n_hidden is None:
+            n_hidden = DEFAULT_MODEL_CONFIG.hidden_size
+
+        self.encoder = AudioEncoder(
+            n_mels=n_mels,
+            n_ctx=n_audio_ctx,
+            n_state=n_audio_state,
+            n_head=n_audio_head,
+            n_layer=n_audio_layer,
+        )
+
+        self.adapter = Adaptor(
+            n_state=n_audio_state,
+            n_hidden=n_hidden,
+            kernel_size=kernel_size,
+            stride=adapter_stride,
+        )
+
+        from transformers import PretrainedConfig
+
+        text_config = getattr(config, "text_config", None)
+        lm_config = text_config if isinstance(text_config, PretrainedConfig) else config
+        architectures = getattr(lm_config, "architectures", None) or getattr(config, "architectures", None)
+        if not architectures:
+            model_type = getattr(lm_config, "model_type", None) or getattr(config, "model_type", None)
+            model_type_to_arch = {
+                "qwen2": "Qwen2ForCausalLM",
+                "qwen2_5": "Qwen2_5ForCausalLM",
+            }
+            arch = model_type_to_arch.get(model_type)
+            architectures = [arch] if arch else ["Qwen2ForCausalLM"]
+        lm_vllm_config = vllm_config.with_hf_config(lm_config, architectures=architectures)
+        self.language_model = init_vllm_registered_model(
+            vllm_config=lm_vllm_config, hf_config=lm_config, prefix=maybe_prefix(prefix, "language_model")
+        )
+
+        self.make_empty_intermediate_tensors = self.language_model.make_empty_intermediate_tensors
+
+    @classmethod
+    def get_placeholder_str(cls, modality: str, i: int) -> str | None:
+        if modality.startswith("audio"):
+            return "<audio_patch>"
+        raise ValueError("Only audio modality is supported")
+
+    @property
+    def device(self):
+        return next(self.parameters()).device
+
+    @property
+    def dtype(self):
+        return next(self.parameters()).dtype
+
+    def _parse_and_validate_audio_input(self, **kwargs: object) -> StepAudio2AudioInputs | None:
+        """Parse audio inputs from kwargs"""
+        audio_mels = kwargs.get("audio_mels", None)
+        audio_lens = kwargs.get("audio_lens", None)
+
+        if audio_mels is None:
+            return None
+
+        if isinstance(audio_mels, torch.Tensor):
+            if audio_mels.dim() >= 4:
+                audio_mels = audio_mels.flatten(0, 1)
+            elif audio_mels.dim() == 3:
+                # Already (B, n_mels, T)
+                pass
+            elif audio_mels.dim() == 2:
+                audio_mels = audio_mels.unsqueeze(0)
+            else:
+                return None
+        else:
+            audio_mels = flatten_bn(audio_mels, concat=True)
+
+        if isinstance(audio_lens, torch.Tensor):
+            if audio_lens.dim() >= 2:
+                audio_lens = audio_lens.flatten(0, 1)
+            elif audio_lens.dim() == 0:
+                audio_lens = audio_lens.unsqueeze(0)
+            # dim == 1 is already fine
+        else:
+            audio_lens = flatten_bn(audio_lens, concat=True)
+
+        expected_n_mels = 128
+        if audio_mels.ndim == 3 and audio_mels.size(1) != expected_n_mels:
+            return None
+
+        return StepAudio2AudioInputs(
+            audio_mels=audio_mels.to(self.dtype).to(self.device),
+            audio_lens=audio_lens,
+        )
+
+    def _process_audio_input(self, audio_input: StepAudio2AudioInputs) -> tuple[torch.Tensor, ...]:
+        """Process audio mels through encoder and adapter"""
+        audio_mels = audio_input["audio_mels"]
+        audio_lens = audio_input["audio_lens"]
+        if not isinstance(audio_lens, torch.Tensor):
+            audio_lens = torch.tensor(audio_lens, device=self.device)
+        else:
+            audio_lens = audio_lens.to(self.device)
+
+        audio_features, audio_lens = self.encoder(audio_mels, audio_lens)
+
+        audio_features = self.adapter(audio_features)
+
+        audio_feature_lens = (audio_lens - 1) // 2 + 1
+
+        audio_feature_list = [audio_features[i, : audio_feature_lens[i]] for i in range(audio_features.size(0))]
+
+        return audio_feature_list
+
+    def get_multimodal_embeddings(self, **kwargs) -> NestedTensors | None:
+        """Get multimodal embeddings"""
+        audio_input = self._parse_and_validate_audio_input(**kwargs)
+        if audio_input is None:
+            return None
+        else:
+            audio_embeddings = self._process_audio_input(audio_input)
+            return audio_embeddings
+
+    def embed_multimodal(self, **kwargs: object) -> NestedTensors:
+        """vLLM multimodal encoder entrypoint."""
+        audio_embeddings = self.get_multimodal_embeddings(**kwargs)
+        if audio_embeddings is None:
+            return []
+        return audio_embeddings
+
+    def get_input_embeddings(
+        self,
+        input_ids: torch.Tensor,
+        multimodal_embeddings: NestedTensors | None = None,
+    ) -> torch.Tensor:
+        """Get input embeddings with multimodal fusion"""
+        # Under newer vLLM versions the registered LM wrapper exposes
+        # get_input_embeddings(), while the inner `.model` may be a bare
+        # Qwen2Model without that compatibility method.
+        inputs_embeds = self.language_model.embed_input_ids(input_ids)
+        if multimodal_embeddings is not None:
+            inputs_embeds = _merge_multimodal_embeddings(
+                input_ids, inputs_embeds, multimodal_embeddings, STEP_AUDIO2_AUDIO_PATCH_TOKEN_ID
+            )
+        return inputs_embeds
+
+    def get_language_model(self) -> nn.Module:
+        return self.language_model
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        **kwargs: object,
+    ):
+        """Forward pass - returns OmniOutput for multi-stage pipeline"""
+        if intermediate_tensors is not None:
+            inputs_embeds = None
+        elif inputs_embeds is None:
+            audio_embeddings = self.get_multimodal_embeddings(**kwargs)
+            inputs_embeds = self.get_input_embeddings(input_ids, audio_embeddings)
+            input_ids = None
+
+        hidden_states = self.language_model(
+            input_ids,
+            positions,
+            intermediate_tensors,
+            inputs_embeds=inputs_embeds,
+        )
+
+        return OmniOutput(
+            text_hidden_states=hidden_states,
+            multimodal_outputs={},
+        )
+
+    def compute_logits(
+        self,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor | None:
+        return self.language_model.compute_logits(hidden_states)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        """Load weights with name mapping from HuggingFace to vLLM structure"""
+        from vllm.model_executor.models.utils import AutoWeightsLoader
+
+        # Map HuggingFace weight names to our model structure
+        # HF: model.*, lm_head.* → Our: language_model.model.*, language_model.lm_head.*
+        # HF: encoder.*, adapter.* → Our: encoder.*, adapter.* (no change)
+        def weight_mapper(name: str) -> str:
+            if name.startswith("lm_head."):
+                return name.replace("lm_head.", "language_model.lm_head.")
+            elif name.startswith("model."):
+                return name.replace("model.", "language_model.model.")
+            else:
+                # encoder.* and adapter.* remain unchanged
+                return name
+
+        # Apply mapping to weights
+        mapped_weights = [(weight_mapper(name), tensor) for name, tensor in weights]
+
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(mapped_weights)
+
+    @staticmethod
+    def separate_tokens(
+        token_ids: list[int],
+        text_max: int = DEFAULT_TOKEN_CONFIG.text_max,
+        audio_start: int = DEFAULT_TOKEN_CONFIG.audio_start,
+    ) -> tuple[list[int], list[int]]:
+        """Separate generated tokens into text and audio tokens"""
+        text_tokens = [tid for tid in token_ids if tid < text_max]
+        audio_tokens = [tid - audio_start for tid in token_ids if tid >= audio_start]
+        return text_tokens, audio_tokens
+
+    @staticmethod
+    def has_audio_output(token_ids: list[int], audio_start: int = DEFAULT_TOKEN_CONFIG.audio_start) -> bool:
+        """Check if generated tokens contain audio tokens"""
+        return any(tid >= audio_start for tid in token_ids)
+
+
+class StepAudio2OutputProcessor:
+    """Helper class to process Step-Audio2 outputs"""
+
+    @staticmethod
+    def process_output(output_ids: torch.Tensor, tokenizer, remove_audio_padding: bool = True) -> dict:
+        """Process model output and separate text/audio tokens"""
+        if isinstance(output_ids, torch.Tensor):
+            output_ids = output_ids.squeeze().tolist()
+
+        # Separate tokens
+        text_tokens, audio_tokens = StepAudio2ThinkerForConditionalGeneration.separate_tokens(output_ids)
+
+        # Remove audio padding if requested
+        if remove_audio_padding and audio_tokens:
+            audio_tokens = [t for t in audio_tokens if t < DEFAULT_TOKEN_CONFIG.audio_eos]
+
+        # Decode text
+        text = tokenizer.decode(text_tokens, skip_special_tokens=False)
+
+        return {
+            "text": text,
+            "text_tokens": text_tokens,
+            "audio_tokens": audio_tokens,
+            "has_audio": len(audio_tokens) > 0,
+            "all_tokens": output_ids,
+        }
@@ -0,0 +1,759 @@
+import io
+import os
+from collections.abc import Iterable
+from dataclasses import dataclass, field
+from typing import Any
+
+import librosa  # noqa: TID251
+import numpy as np
+import onnxruntime
+import s3tokenizer
+import soundfile as sf
+import torch
+import torch.nn as nn
+import torchaudio
+import torchaudio.compliance.kaldi as kaldi
+from flashcosyvoice.modules.hifigan import HiFTGenerator
+from flashcosyvoice.utils.audio import mel_spectrogram
+from hyperpyyaml import load_hyperpyyaml
+from vllm.config import VllmConfig
+from vllm.logger import init_logger
+from vllm.model_executor.models.interfaces import SupportsPP
+from vllm.model_executor.models.utils import WeightsMapper
+from vllm.sequence import IntermediateTensors
+from vllm.v1.outputs import SamplerOutput
+from vllm.v1.sample.metadata import SamplingMetadata
+from vllm.v1.sample.sampler import Sampler
+
+from vllm_omni.model_executor.models.output_templates import OmniOutput
+from vllm_omni.model_executor.models.step_audio2.step_audio2_constants import (
+    DEFAULT_STREAM_CONFIG,
+    DEFAULT_TOKEN2WAV_CONFIG,
+    STEP_AUDIO2_DEFAULT_PROMPT_WAV,
+    STREAM_SOURCE_CACHE_LEN,
+)
+
+logger = init_logger(__name__)
+
+
+def fade_in_out(
+    fade_in_mel: torch.Tensor,
+    fade_out_mel: torch.Tensor,
+    window: torch.Tensor,
+) -> torch.Tensor:
+    """Cross-fade two overlapping waveform segments using a Hamming window.
+
+    The window is split in half: the first half ramps *up* (fade-in) and
+    the second half ramps *down* (fade-out).  The overlap region of
+    ``fade_in_mel`` is blended with the tail of ``fade_out_mel``.
+    """
+    mel_overlap_len = int(window.shape[0] / 2)
+    fade_in_mel = fade_in_mel.clone()
+    fade_in_mel[..., :mel_overlap_len] = (
+        fade_in_mel[..., :mel_overlap_len] * window[:mel_overlap_len]
+        + fade_out_mel[..., -mel_overlap_len:] * window[mel_overlap_len:]
+    )
+    return fade_in_mel
+
+
+@dataclass
+class _StreamState:
+    """Per-request streaming state for concurrent Token2Wav sessions."""
+
+    setup_done: bool = False
+    stream_cache: dict | None = None
+    hift_cache_dict: dict = field(default_factory=dict)
+    finished: bool = False
+
+
+class StepAudio2Token2WavCore(nn.Module):
+    """Core Token2Wav model - handles token to waveform conversion"""
+
+    def __init__(
+        self,
+        model_path: str,
+        float16: bool = False,
+        device: str = "cuda",
+        n_timesteps: int = DEFAULT_TOKEN2WAV_CONFIG.n_timesteps,
+    ):
+        super().__init__()
+        self.model_path = model_path
+        self.float16 = float16
+        self.device = torch.device(device)
+        self.n_timesteps = n_timesteps
+
+        self._models_loaded = False
+        self._audio_tokenizer = None
+        self._spk_model = None
+        self._flow = None
+        self._hift = None
+
+        self.cache = {}
+
+        # Streaming state (constants from centralised config)
+        self.mel_cache_len = DEFAULT_STREAM_CONFIG.mel_cache_len
+        self.source_cache_len = STREAM_SOURCE_CACHE_LEN
+        self.speech_window: torch.Tensor | None = None  # created lazily on device
+        self.stream_cache: dict | None = None  # flow model causal cache
+        self.hift_cache_dict: dict = {}  # HiFT vocoder overlap cache
+
+    def _ensure_models_loaded(self):
+        """Lazy load models on first use"""
+        if self._models_loaded:
+            return
+
+        from vllm.logger import init_logger
+
+        logger = init_logger(__name__)
+        logger.info(f"Loading Token2Wav models from: {self.model_path}")
+
+        self._audio_tokenizer = (
+            s3tokenizer.load_model(f"{self.model_path}/speech_tokenizer_v2_25hz.onnx").to(self.device).eval()
+        )
+
+        option = onnxruntime.SessionOptions()
+        option.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
+        option.intra_op_num_threads = 1
+        self._spk_model = onnxruntime.InferenceSession(
+            f"{self.model_path}/campplus.onnx", sess_options=option, providers=["CPUExecutionProvider"]
+        )
+
+        with open(f"{self.model_path}/flow.yaml") as f:
+            configs = load_hyperpyyaml(f)
+        self._flow = configs["flow"]
+        if self.float16:
+            self._flow.half()
+        self._flow.load_state_dict(
+            torch.load(f"{self.model_path}/flow.pt", map_location="cpu", weights_only=True), strict=True
+        )
+        self._flow.to(self.device).eval()
+
+        self._hift = HiFTGenerator()
+        hift_state_dict = {
+            k.replace("generator.", ""): v
+            for k, v in torch.load(f"{self.model_path}/hift.pt", map_location="cpu", weights_only=True).items()
+        }
+        self._hift.load_state_dict(hift_state_dict, strict=True)
+        self._hift.to(self.device).eval()
+
+        self._models_loaded = True
+        logger.info("Token2Wav models loaded successfully")
+
+    @property
+    def audio_tokenizer(self):
+        """Lazy-loaded audio tokenizer"""
+        self._ensure_models_loaded()
+        return self._audio_tokenizer
+
+    @property
+    def spk_model(self):
+        """Lazy-loaded speaker model"""
+        self._ensure_models_loaded()
+        return self._spk_model
+
+    @property
+    def flow(self):
+        """Lazy-loaded flow model"""
+        self._ensure_models_loaded()
+        return self._flow
+
+    @property
+    def hift(self):
+        """Lazy-loaded HiFT vocoder"""
+        self._ensure_models_loaded()
+        return self._hift
+
+    def _prepare_prompt(self, prompt_wav: str):
+        """Prepare prompt audio for conditioning"""
+        # Prefer soundfile/librosa path to avoid torchaudio->torchcodec runtime coupling.
+        try:
+            audio_np_raw, sample_rate = sf.read(prompt_wav, dtype="float32", always_2d=False)
+        except Exception:
+            audio_np_raw, sample_rate = librosa.load(prompt_wav, sr=None, mono=True)
+        if isinstance(audio_np_raw, np.ndarray) and audio_np_raw.ndim > 1:
+            audio_np_raw = np.mean(audio_np_raw, axis=-1)
+        sample_rate = int(sample_rate)
+
+        # 16 kHz branch: s3tokenizer + speaker embedding
+        audio_16k = audio_np_raw.astype(np.float32)
+        if sample_rate != 16000:
+            audio_16k = librosa.resample(y=audio_16k, orig_sr=sample_rate, target_sr=16000)
+        audio = torch.from_numpy(audio_16k)
+        mels = s3tokenizer.log_mel_spectrogram(audio)
+        mels, mels_lens = s3tokenizer.padding([mels])
+        prompt_speech_tokens, prompt_speech_tokens_lens = self.audio_tokenizer.quantize(
+            mels.to(self.device), mels_lens.to(self.device)
+        )
+
+        spk_feat = kaldi.fbank(audio.unsqueeze(0), num_mel_bins=80, dither=0, sample_frequency=16000)
+        spk_feat = spk_feat - spk_feat.mean(dim=0, keepdim=True)
+        spk_emb = torch.tensor(
+            self.spk_model.run(None, {self.spk_model.get_inputs()[0].name: spk_feat.unsqueeze(dim=0).cpu().numpy()})[0],
+            device=self.device,
+        )
+
+        # 24 kHz branch: mel spectrogram for flow model conditioning
+        # Must resample from the ORIGINAL audio, not the 16 kHz version.
+        audio_24k = audio_np_raw.astype(np.float32)
+        if sample_rate != 24000:
+            audio_24k = librosa.resample(y=audio_24k, orig_sr=sample_rate, target_sr=24000)
+        audio_24k_t = torch.from_numpy(audio_24k).unsqueeze(0)  # [1, T]
+        prompt_mel = mel_spectrogram(audio_24k_t).transpose(1, 2).squeeze(0)  # [T, num_mels]
+        prompt_mels = prompt_mel.unsqueeze(0).to(self.device)
+        prompt_mels_lens = torch.tensor([prompt_mels.shape[1]], dtype=torch.int32, device=self.device)
+        prompt_mels = torch.nn.functional.pad(
+            prompt_mels,
+            (0, 0, 0, prompt_speech_tokens.shape[1] * self.flow.up_rate - prompt_mels.shape[1]),
+            mode="replicate",
+        )
+
+        return prompt_speech_tokens, prompt_speech_tokens_lens, spk_emb, prompt_mels, prompt_mels_lens
+
+    def forward(
+        self, generated_speech_tokens: list, prompt_wav: str, return_bytes: bool = True
+    ) -> torch.Tensor | bytes:
+        """
+        Convert audio tokens to waveform
+
+        Args:
+            generated_speech_tokens: List of audio token IDs (0-6560)
+            prompt_wav: Path to prompt wav file for speaker conditioning
+            return_bytes: If True, return WAV bytes; else return tensor
+
+        Returns:
+            Audio waveform as bytes or tensor
+        """
+        if prompt_wav not in self.cache:
+            self.cache[prompt_wav] = self._prepare_prompt(prompt_wav)
+        prompt_speech_tokens, prompt_speech_tokens_lens, spk_emb, prompt_mels, prompt_mels_lens = self.cache[prompt_wav]
+
+        generated_speech_tokens = torch.tensor([generated_speech_tokens], dtype=torch.int32, device=self.device)
+        generated_speech_tokens_lens = torch.tensor(
+            [generated_speech_tokens.shape[1]], dtype=torch.int32, device=self.device
+        )
+
+        with torch.amp.autocast(str(self.device.type), dtype=torch.float16 if self.float16 else torch.float32):
+            mel = self.flow.inference(
+                generated_speech_tokens,
+                generated_speech_tokens_lens,
+                prompt_speech_tokens,
+                prompt_speech_tokens_lens,
+                prompt_mels,
+                prompt_mels_lens,
+                spk_emb,
+                self.n_timesteps,
+            )
+
+        wav, _ = self.hift(speech_feat=mel)
+
+        if return_bytes:
+            output = io.BytesIO()
+            torchaudio.save(output, wav.cpu(), sample_rate=DEFAULT_TOKEN2WAV_CONFIG.sample_rate, format="wav")
+            return output.getvalue()
+        else:
+            return wav
+
+    # ------------------------------------------------------------------
+    # Streaming (chunked) inference
+    # ------------------------------------------------------------------
+
+    def setup_stream(self, prompt_wav: str) -> None:
+        """Initialize flow + HiFT caches for chunked streaming inference.
+
+        Must be called once before the first ``stream_chunk()`` call for a
+        given utterance.
+        """
+        if prompt_wav not in self.cache:
+            self.cache[prompt_wav] = self._prepare_prompt(prompt_wav)
+        (
+            prompt_speech_tokens,
+            _,
+            spk_emb,
+            prompt_mels,
+            _,
+        ) = self.cache[prompt_wav]
+
+        # Lazily create the Hamming window on the correct device
+        if self.speech_window is None or self.speech_window.device != self.device:
+            self.speech_window = torch.from_numpy(np.hamming(2 * self.source_cache_len)).to(
+                device=self.device, dtype=torch.float32
+            )
+
+        # Flow model: setup causal cache with prompt tokens + lookahead
+        pre_lookahead = DEFAULT_STREAM_CONFIG.pre_lookahead_len
+        self.stream_cache = self.flow.setup_cache(
+            torch.cat([prompt_speech_tokens, prompt_speech_tokens[:, :pre_lookahead]], dim=1),
+            prompt_mels,
+            spk_emb,
+            n_timesteps=self.n_timesteps,
+        )
+
+        # HiFT vocoder: empty overlap buffers
+        self.hift_cache_dict = {
+            "mel": torch.zeros(1, prompt_mels.shape[2], 0, device=self.device),
+            "source": torch.zeros(1, 1, 0, device=self.device),
+            "speech": torch.zeros(1, 0, device=self.device),
+        }
+
+    def stream_chunk(
+        self,
+        audio_tokens: list[int],
+        prompt_wav: str,
+        last_chunk: bool = False,
+    ) -> torch.Tensor:
+        """Process one chunk of audio tokens and return a waveform segment.
+
+        Args:
+            audio_tokens: Token IDs for this chunk (includes lookahead for
+                non-last chunks).
+            prompt_wav: Path to the prompt wav (for cached speaker embeddings).
+            last_chunk: Whether this is the final chunk of the utterance.
+
+        Returns:
+            1-D float tensor of audio samples (24 kHz).
+        """
+        if self.stream_cache is None:
+            raise ValueError("stream_cache not initialised – call setup_stream() first")
+
+        if prompt_wav not in self.cache:
+            self.cache[prompt_wav] = self._prepare_prompt(prompt_wav)
+        _, _, spk_emb, prompt_mels, _ = self.cache[prompt_wav]
+
+        token_tensor = torch.tensor([audio_tokens], dtype=torch.int32, device=self.device)
+
+        with torch.amp.autocast(
+            str(self.device.type),
+            dtype=torch.float16 if self.float16 else torch.float32,
+        ):
+            chunk_mel, self.stream_cache = self.flow.inference_chunk(
+                token=token_tensor,
+                spk=spk_emb,
+                cache=self.stream_cache,
+                last_chunk=last_chunk,
+                n_timesteps=self.n_timesteps,
+            )
+
+        # Trim estimator attention cache to avoid unbounded growth
+        keep = DEFAULT_STREAM_CONFIG.estimator_cache_keep
+        est_att = self.stream_cache.get("estimator_att_cache")
+        if est_att is not None and est_att.shape[4] > (prompt_mels.shape[1] + keep):
+            self.stream_cache["estimator_att_cache"] = torch.cat(
+                [est_att[:, :, :, :, : prompt_mels.shape[1]], est_att[:, :, :, :, -keep:]],
+                dim=4,
+            )
+
+        # ---- HiFT vocoder with overlap-add ----
+        hift_cache_mel = self.hift_cache_dict["mel"]
+        hift_cache_source = self.hift_cache_dict["source"]
+        hift_cache_speech = self.hift_cache_dict["speech"]
+
+        mel = torch.cat([hift_cache_mel, chunk_mel], dim=2)
+        speech, source = self.hift(mel, hift_cache_source)
+
+        # Cross-fade overlap region for smooth concatenation
+        if hift_cache_speech.shape[-1] > 0:
+            speech = fade_in_out(speech, hift_cache_speech, self.speech_window)
+
+        # Update vocoder cache for the next chunk
+        self.hift_cache_dict = {
+            "mel": mel[..., -self.mel_cache_len :].clone().detach(),
+            "source": source[:, :, -self.source_cache_len :].clone().detach(),
+            "speech": speech[:, -self.source_cache_len :].clone().detach(),
+        }
+
+        # For non-last chunks, trim the tail that will be re-synthesised
+        if not last_chunk:
+            speech = speech[:, : -self.source_cache_len]
+
+        return speech.squeeze(0)  # [samples]
+
+    def reset_stream(self) -> None:
+        """Clear all streaming state after an utterance is complete."""
+        self.stream_cache = None
+        self.hift_cache_dict = {}
+
+    # ------------------------------------------------------------------
+    # Per-request streaming (does NOT mutate self – uses external state)
+    # ------------------------------------------------------------------
+
+    def setup_stream_for(self, prompt_wav: str, state: _StreamState) -> None:
+        """Initialise flow + HiFT caches into *state* (no self mutation)."""
+        if prompt_wav not in self.cache:
+            self.cache[prompt_wav] = self._prepare_prompt(prompt_wav)
+        prompt_speech_tokens, _, spk_emb, prompt_mels, _ = self.cache[prompt_wav]
+
+        if self.speech_window is None or self.speech_window.device != self.device:
+            self.speech_window = torch.from_numpy(np.hamming(2 * self.source_cache_len)).to(
+                device=self.device, dtype=torch.float32
+            )
+
+        pre_lookahead = DEFAULT_STREAM_CONFIG.pre_lookahead_len
+        state.stream_cache = self.flow.setup_cache(
+            torch.cat(
+                [prompt_speech_tokens, prompt_speech_tokens[:, :pre_lookahead]],
+                dim=1,
+            ),
+            prompt_mels,
+            spk_emb,
+            n_timesteps=self.n_timesteps,
+        )
+        state.hift_cache_dict = {
+            "mel": torch.zeros(1, prompt_mels.shape[2], 0, device=self.device),
+            "source": torch.zeros(1, 1, 0, device=self.device),
+            "speech": torch.zeros(1, 0, device=self.device),
+        }
+        state.setup_done = True
+
+    def stream_chunk_for(
+        self,
+        audio_tokens: list[int],
+        prompt_wav: str,
+        last_chunk: bool,
+        state: _StreamState,
+    ) -> torch.Tensor:
+        """Process one chunk using *state* (no self mutation except speech_window)."""
+        if state.stream_cache is None:
+            raise ValueError("stream_cache not initialised – call setup_stream_for() first")
+
+        if prompt_wav not in self.cache:
+            self.cache[prompt_wav] = self._prepare_prompt(prompt_wav)
+        _, _, spk_emb, prompt_mels, _ = self.cache[prompt_wav]
+
+        token_tensor = torch.tensor([audio_tokens], dtype=torch.int32, device=self.device)
+
+        with torch.amp.autocast(
+            str(self.device.type),
+            dtype=torch.float16 if self.float16 else torch.float32,
+        ):
+            chunk_mel, state.stream_cache = self.flow.inference_chunk(
+                token=token_tensor,
+                spk=spk_emb,
+                cache=state.stream_cache,
+                last_chunk=last_chunk,
+                n_timesteps=self.n_timesteps,
+            )
+
+        # Trim estimator attention cache to avoid unbounded growth
+        keep = DEFAULT_STREAM_CONFIG.estimator_cache_keep
+        est_att = state.stream_cache.get("estimator_att_cache")
+        if est_att is not None and est_att.shape[4] > (prompt_mels.shape[1] + keep):
+            state.stream_cache["estimator_att_cache"] = torch.cat(
+                [
+                    est_att[:, :, :, :, : prompt_mels.shape[1]],
+                    est_att[:, :, :, :, -keep:],
+                ],
+                dim=4,
+            )
+
+        # ---- HiFT vocoder with overlap-add ----
+        hift_cache_mel = state.hift_cache_dict["mel"]
+        hift_cache_source = state.hift_cache_dict["source"]
+        hift_cache_speech = state.hift_cache_dict["speech"]
+
+        mel = torch.cat([hift_cache_mel, chunk_mel], dim=2)
+        speech, source = self.hift(mel, hift_cache_source)
+
+        if hift_cache_speech.shape[-1] > 0:
+            speech = fade_in_out(speech, hift_cache_speech, self.speech_window)
+
+        state.hift_cache_dict = {
+            "mel": mel[..., -self.mel_cache_len :].clone().detach(),
+            "source": source[:, :, -self.source_cache_len :].clone().detach(),
+            "speech": speech[:, -self.source_cache_len :].clone().detach(),
+        }
+
+        if not last_chunk:
+            speech = speech[:, : -self.source_cache_len]
+
+        return speech.squeeze(0)
+
+    @staticmethod
+    def reset_stream_for(state: _StreamState) -> None:
+        """Clear streaming caches for the given state."""
+        state.stream_cache = None
+        state.hift_cache_dict = {}
+        state.setup_done = False
+        state.finished = True
+
+
+class StepAudio2Token2WavForConditionalGeneration(nn.Module, SupportsPP):
+    """vLLM-compatible wrapper for Step-Audio2 Token2Wav"""
+
+    hf_to_vllm_mapper = WeightsMapper(
+        orig_to_new_prefix={
+            "model.": "token2wav.",
+        }
+    )
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        self.vllm_config = vllm_config
+        self.config = vllm_config.model_config.hf_config
+
+        model_path = getattr(self.config, "token2wav_path", None)
+        model_name_or_path = vllm_config.model_config.model
+        if model_path is None:
+            # Resolve HF repo names to local cache path
+            if not os.path.isdir(model_name_or_path):
+                from huggingface_hub import snapshot_download
+
+                model_name_or_path = snapshot_download(model_name_or_path)
+            model_path = f"{model_name_or_path}/token2wav"
+
+        float16 = getattr(self.config, "token2wav_float16", False)
+
+        device = "cuda"
+        if hasattr(vllm_config, "device_config") and vllm_config.device_config:
+            device = str(vllm_config.device_config.device)
+
+        self.token2wav = StepAudio2Token2WavCore(model_path=model_path, float16=float16, device=device)
+
+        self.have_multimodal_outputs = True
+        # Required for the runner to pass async_chunk info via
+        # runtime_additional_information on every forward step.
+        self.enable_update_additional_information = True
+
+        # Per-request streaming states (ordered list matching batch order).
+        self._stream_states: list[_StreamState] = []
+
+        # Resolve default prompt wav path from model directory.
+        # Priority: env var > {model_dir}/assets/default_female.wav > bare filename (cwd)
+        self._default_prompt_wav = self._resolve_prompt_wav(model_name_or_path)
+
+        self.make_empty_intermediate_tensors = lambda: None
+
+    @staticmethod
+    def _resolve_prompt_wav(model_root: str) -> str:
+        """Find the default prompt wav, searching model assets directory."""
+        env_val = os.environ.get("STEP_AUDIO2_DEFAULT_PROMPT_WAV")
+        if env_val and os.path.exists(env_val):
+            return env_val
+        # Search common locations: model dir, then bundled assets in this package
+        pkg_dir = os.path.dirname(__file__)
+        candidates = [
+            os.path.join(model_root, "assets", STEP_AUDIO2_DEFAULT_PROMPT_WAV),
+            os.path.join(model_root, STEP_AUDIO2_DEFAULT_PROMPT_WAV),
+            os.path.join(pkg_dir, "assets", STEP_AUDIO2_DEFAULT_PROMPT_WAV),
+            os.path.join(pkg_dir, STEP_AUDIO2_DEFAULT_PROMPT_WAV),
+            STEP_AUDIO2_DEFAULT_PROMPT_WAV,
+        ]
+        for path in candidates:
+            if os.path.exists(path):
+                return path
+        # Return the env var (even if not found) or the bare filename as fallback
+        return env_val or STEP_AUDIO2_DEFAULT_PROMPT_WAV
+
+    def get_language_model(self) -> torch.nn.Module:
+        # Return the wrapper so vLLM interface helpers resolve
+        # embed_input_ids() on this adapter instead of the bare core model.
+        return self
+
+    @property
+    def sampler(self):
+        return Sampler()
+
+    def embed_input_ids(
+        self,
+        input_ids: torch.Tensor,
+        multimodal_embeddings=None,
+        is_multimodal=None,
+    ) -> torch.Tensor:
+        hidden_size = self.vllm_config.model_config.get_hidden_size()
+        return torch.zeros_like(input_ids).reshape(-1, 1).repeat(1, hidden_size)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        additional: dict | None = None,
+        runtime_additional_information: list[dict[str, Any]] | None = None,
+        **kwargs: Any,
+    ) -> torch.Tensor | bytes:
+        """
+        Forward pass for Token2Wav
+
+        Args:
+            input_ids: Audio token IDs (from stage input processor)
+            positions: Position IDs (unused)
+            intermediate_tensors: Intermediate tensors (unused)
+            inputs_embeds: Input embeddings (unused)
+            additional: Additional information dict (deprecated)
+            runtime_additional_information: Per-request dicts from the
+                async_chunk processor.  Each dict contains ``audio_tokens``
+                (list[int]) and ``last_chunk`` (bool).
+            **kwargs: Other kwargs
+
+        Returns:
+            OmniOutput with waveform tensor
+        """
+        # ----- async_chunk streaming path -----
+        # Only enter when at least one entry carries ``left_context_size``
+        # (the field the framework forwards from the async_chunk payload).
+        # The runner always passes runtime_additional_information (possibly
+        # [{}]), so a bare truthiness check would incorrectly enter this
+        # branch in synchronous mode.
+        if runtime_additional_information and any(
+            "left_context_size" in info for info in runtime_additional_information
+        ):
+            return self._forward_async_chunk(input_ids, runtime_additional_information, **kwargs)
+
+        # ----- profiling / warmup guard -----
+        # Profile/warmup passes use all-zero token IDs. Real audio tokens
+        # are non-zero, so check sum == 0 rather than a length threshold
+        # (short utterances may have fewer than 50 tokens).
+        if isinstance(input_ids, torch.Tensor):
+            is_profile_mode = input_ids.numel() == 0 or input_ids.sum().item() == 0
+        else:
+            is_profile_mode = not input_ids if not hasattr(input_ids, "__len__") else len(input_ids) == 0
+
+        if is_profile_mode:
+            device = torch.device("cpu")
+            if isinstance(inputs_embeds, torch.Tensor) and inputs_embeds.numel() > 0:
+                device = inputs_embeds.device
+            elif isinstance(input_ids, torch.Tensor) and input_ids.numel() > 0:
+                device = input_ids.device
+
+            dummy_audio = torch.zeros(1000, device=device, dtype=torch.float32)
+            return OmniOutput(
+                text_hidden_states=None,
+                multimodal_outputs={"model_outputs": [dummy_audio]},
+            )
+
+        # ----- synchronous (full-sequence) path -----
+        if isinstance(input_ids, torch.Tensor):
+            if input_ids.dim() == 1:
+                input_ids_list = input_ids.cpu().tolist()
+            else:
+                input_ids_list = input_ids[0].cpu().tolist()
+        else:
+            input_ids_list = input_ids if isinstance(input_ids, list) else [input_ids]
+
+        prompt_wav = self._default_prompt_wav
+
+        if not os.path.exists(prompt_wav):
+            raise FileNotFoundError(f"Token2Wav: prompt_wav file not found: {prompt_wav}")
+
+        waveform_tensor = self.token2wav(input_ids_list, prompt_wav=prompt_wav, return_bytes=False)
+
+        if waveform_tensor.dim() == 2 and waveform_tensor.shape[0] == 1:
+            waveform_tensor = waveform_tensor.squeeze(0)
+
+        waveform_tensor_cpu = waveform_tensor.detach().cpu().contiguous()
+
+        return OmniOutput(
+            text_hidden_states=None,
+            multimodal_outputs={"model_outputs": [waveform_tensor_cpu]},
+        )
+
+    # ------------------------------------------------------------------
+    # Async-chunk streaming forward
+    # ------------------------------------------------------------------
+
+    def _forward_async_chunk(
+        self,
+        input_ids: torch.Tensor,
+        runtime_additional_information: list[dict[str, Any]],
+        **kwargs: Any,
+    ) -> OmniOutput:
+        """Handle one async_chunk forward step (batch=1).
+
+        The end-of-stream flag is encoded as ``left_context_size == 1``
+        (the only field the framework transparently forwards from the
+        async_chunk payload to ``request.additional_information``).
+        """
+        # --- batch=1 guard ---
+        batch_size = sum(1 for info in runtime_additional_information if "left_context_size" in info)
+        if batch_size != 1:
+            raise RuntimeError(
+                f"Token2Wav async_chunk only supports batch=1, got {batch_size}. "
+                "Batch>1 requires framework support for batch_req_ids."
+            )
+
+        info = next(info for info in runtime_additional_information if "left_context_size" in info)
+        last_chunk = info.get("left_context_size", 0) == 1
+
+        # --- Manage single stream state ---
+        # If the previous request was preempted (setup_done but not finished),
+        # its stale caches must be discarded before we start a new session.
+        if self._stream_states and not self._stream_states[0].finished:
+            state = self._stream_states[0]
+            if state.setup_done and not last_chunk:
+                # Heuristic: a non-last chunk arriving into an already-active
+                # state that we did NOT set up this call means a new request
+                # reused the slot.  For a genuinely continuing request the
+                # state is fine; we cannot distinguish without request IDs,
+                # so we only forcibly reset on an explicit new-session signal
+                # (setup_done=True but stream_cache already consumed all
+                # previous data -- detected below when stream_chunk fails).
+                pass
+        if not self._stream_states or self._stream_states[0].finished:
+            self._stream_states = [_StreamState()]
+        state = self._stream_states[0]
+
+        # --- Extract audio tokens ---
+        audio_tokens = input_ids.flatten().cpu().tolist()
+
+        # Empty chunk (e.g. EOF with no audio tokens)
+        if not audio_tokens:
+            if state.setup_done:
+                self.token2wav.reset_stream_for(state)
+            else:
+                state.finished = True
+            return OmniOutput(
+                text_hidden_states=None,
+                multimodal_outputs={"model_outputs": [torch.zeros(1, dtype=torch.float32)]},
+            )
+
+        prompt_wav = self._default_prompt_wav
+        if not os.path.exists(prompt_wav):
+            raise FileNotFoundError(f"Token2Wav: prompt_wav file not found: {prompt_wav}")
+
+        try:
+            if not state.setup_done:
+                logger.info("Token2Wav: initialising streaming caches")
+                self.token2wav.setup_stream_for(prompt_wav, state)
+
+            waveform = self.token2wav.stream_chunk_for(
+                audio_tokens,
+                prompt_wav=prompt_wav,
+                last_chunk=last_chunk,
+                state=state,
+            )
+        except Exception:
+            logger.exception("Token2Wav: stream error, resetting state")
+            self.token2wav.reset_stream_for(state)
+            raise
+
+        if last_chunk:
+            self.token2wav.reset_stream_for(state)
+
+        return OmniOutput(
+            text_hidden_states=None,
+            multimodal_outputs={"model_outputs": [waveform.detach().cpu().contiguous()]},
+        )
+
+    def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
+        return hidden_states
+
+    def sample(
+        self,
+        logits: torch.Tensor,
+        sampling_metadata: SamplingMetadata,
+    ) -> SamplerOutput | None:
+        return None
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        loaded_params = set()
+        for name, _ in self.token2wav.named_parameters():
+            loaded_params.add(f"token2wav.{name}")
+        for name, _ in self.token2wav.named_buffers():
+            loaded_params.add(f"token2wav.{name}")
+        return loaded_params
+
+    def load_weights_without_buffers(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        loaded_params = set()
+        for name, _ in self.token2wav.named_parameters():
+            loaded_params.add(f"token2wav.{name}")
+        return loaded_params
@@ -283,7 +283,8 @@ def _encode_raw_audio(
        pad = (padding_size, 0) if padding_mode == "left" else (0, padding_size)
        audio = torch.nn.functional.pad(audio, pad)

-    feat = tts.audio_vae.encode(audio.to(tts.device), encode_sr).cpu()
+    vae_device = next(tts.audio_vae.parameters()).device
+    feat = tts.audio_vae.encode(audio.to(vae_device), encode_sr).cpu()
    return feat.view(tts.audio_vae.latent_dim, -1, tts.patch_size).permute(1, 2, 0)


@@ -1,4 +1,4 @@
-from collections.abc import Iterable
+from collections.abc import Iterable, Mapping
 from dataclasses import replace
 from functools import cached_property
 from pathlib import Path
@@ -220,7 +220,7 @@ class VoxtralTTSForConditionalGeneration(

    def tts_postprocess(self, hidden_states: torch.Tensor, multimodal_outputs: object, **info_dict: object | None):
        update_dict = {}
-        if isinstance(multimodal_outputs, dict) and "audio" in multimodal_outputs:
+        if isinstance(multimodal_outputs, Mapping) and "audio" in multimodal_outputs:
            assert self.post_process_idx < len(multimodal_outputs["audio"]), (
                f"Expect {self.post_process_idx=} < {len(multimodal_outputs['audio'])=}"
            )
@@ -0,0 +1,122 @@
+# Stage config for Step-Audio2 with vLLM-Omni
+#
+# Two-stage pipeline:
+#   Stage 0 (Thinker): Audio understanding → Text + Audio tokens
+#   Stage 1 (Token2Wav): Audio tokens → Waveform (24kHz)
+#
+# Hardware requirement: 1-2 GPUs (can run on single GPU with devices: "0" for both)
+
+# Model configuration reference (for documentation)
+# Actual values are set in hf_overrides below
+#
+# Token configuration:
+#   text_max: 151688              # Maximum text token ID (inclusive)
+#   audio_start: 151696           # First audio token ID in absolute vocabulary
+#   audio_vocab_size: 6562        # Total number of audio tokens
+#   audio_eos: 6561               # Audio EOS token ID (relative to audio_start)
+#   audio_patch_token_id: 151690  # <audio_patch> placeholder token ID
+#
+# Encoder configuration:
+#   n_mels: 128                   # Number of mel frequency bins
+#   n_audio_ctx: 1500             # Audio context length
+#   n_audio_state: 512            # Audio encoder hidden state dimension
+#   n_audio_head: 8               # Number of attention heads
+#   n_audio_layer: 6              # Number of encoder layers
+#   kernel_size: 3                # Convolution kernel size for adapter
+#   adapter_stride: 2             # Stride for adapter convolution
+
+stage_args:
+  # Stage 0: Thinker (Audio Understanding + Token Generation)
+  - stage_id: 0
+    runtime:
+      process: true              # Run in separate process
+      devices: "0,1,2,3"          # Use 3 GPUs for tensor parallelism
+      max_batch_size: 1         # Single sample at a time
+    engine_args:
+      model_stage: step_audio2_thinker
+      model_arch: StepAudio2ThinkerForConditionalGeneration
+      tensor_parallel_size: 4   # Split model across 3 GPUs
+        # NOTE: All default values are defined in step_audio2_constants.py (DEFAULT_MODEL_CONFIG)
+        # Only override here if your model differs from the defaults
+        #
+        # Default values (for reference, from constants.py):
+        #   hidden_size: 4096              # Qwen2-7B (change to 1536 for 1.5B, 896 for 0.5B)
+        #   text_max: 151688
+        #   audio_start: 151696
+        #   audio_vocab_size: 6562
+        #   audio_eos: 6561
+        #   audio_patch_token_id: 151690
+        #   n_mels: 128
+        #   n_audio_ctx: 1500
+        #   n_audio_state: 512
+        #   n_audio_head: 8
+        #   n_audio_layer: 6
+        #   kernel_size: 3
+        #   adapter_stride: 2
+        #
+        # Uncomment and modify only if needed:
+        # hidden_size: 4096
+      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      gpu_memory_utilization: 0.8
+      enforce_eager: true        # Required for now
+      trust_remote_code: true
+      enable_prefix_caching: false
+      max_num_batched_tokens: 8192
+      engine_output_type: text   # Outputs text + audio tokens
+    is_comprehension: true       # This stage does comprehension
+    final_output: true          # Can be final output (for ASR mode)
+    final_output_type: text
+    default_sampling_params:
+      temperature: 0.7
+      top_p: 0.9
+      top_k: -1
+      max_tokens: 1024
+      seed: 42
+      detokenize: true
+      repetition_penalty: 1.05
+
+  # Stage 1: Token2Wav (Audio Synthesis)
+  - stage_id: 1
+    runtime:
+      process: true
+      devices: "3"              # Use GPU 3 (separate from Stage 0's 0-2)
+      max_batch_size: 1
+    engine_args:
+      model_stage: token2wav
+      model_arch: StepAudio2Token2WavModel
+      hf_overrides:
+        architectures: ["StepAudio2Token2WavForConditionalGeneration"]
+      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
+      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
+
+      gpu_memory_utilization: 0.3  # Token2Wav needs less memory
+      enforce_eager: true
+      trust_remote_code: true
+      async_scheduling: false
+      enable_prefix_caching: false
+      max_num_batched_tokens: 4096
+      engine_output_type: audio    # Outputs waveform
+    engine_input_source: [0]       # Takes input from Stage 0
+    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.step_audio2.thinker2token2wav
+    final_output: true            # Final output stage
+    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 configuration
+runtime:
+  enabled: true
+  defaults:
+    window_size: -1             # Wait for full input before triggering
+    max_inflight: 1             # Process one at a time
+
+  edges:
+    - from: 0                   # Thinker → Token2Wav
+      to: 1
+      window_size: -1           # Wait for complete output from Thinker
@@ -0,0 +1,35 @@
+# Stage config for Step-Audio2 ASR (speech-to-text only)
+#
+# Single-stage pipeline: only Thinker, no Token2Wav.
+# Use this for audio_to_text tasks where no audio output is needed.
+
+stage_args:
+  - stage_id: 0
+    stage_type: llm
+    runtime:
+      process: true
+      devices: "0,1"
+      max_batch_size: 1
+    engine_args:
+      model_stage: step_audio2_thinker
+      model_arch: StepAudio2ThinkerForConditionalGeneration
+      worker_type: ar
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      gpu_memory_utilization: 0.8
+      enforce_eager: true
+      trust_remote_code: true
+      enable_prefix_caching: false
+      max_num_batched_tokens: 2048
+      tensor_parallel_size: 2
+      engine_output_type: text
+    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: 1024
+      seed: 42
+      detokenize: true
+      repetition_penalty: 1.05
@@ -0,0 +1,78 @@
+# Stage config for Step-Audio2 with async_chunk support
+#
+# Two-stage pipeline with streaming between stages:
+#   Stage 0 (Thinker): Audio understanding → Text + Audio tokens (streamed in chunks)
+#   Stage 1 (Token2Wav): Audio tokens → Waveform (24kHz)
+#
+# Key difference from step_audio_2.yaml:
+#   - async_chunk: true enables streaming from Thinker to Token2Wav
+#   - Thinker uses custom_process_next_stage_input_func to send audio token chunks
+#   - Token2Wav receives chunks incrementally instead of waiting for full output
+#
+# Hardware requirement: 1-2 GPUs (can run on single GPU with devices: "0" for both)
+
+async_chunk: true
+stage_args:
+  # Stage 0: Thinker (Audio Understanding + Token Generation)
+  - stage_id: 0
+    stage_type: llm
+    runtime:
+      process: true
+      devices: "0,1"
+      max_batch_size: 1
+    engine_args:
+      model_stage: step_audio2_thinker
+      model_arch: StepAudio2ThinkerForConditionalGeneration
+      worker_type: ar
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      gpu_memory_utilization: 0.8
+      enforce_eager: true
+      trust_remote_code: true
+      enable_prefix_caching: false
+      max_num_batched_tokens: 1024
+      tensor_parallel_size: 2
+      engine_output_type: latent
+      custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.step_audio2.thinker2token2wav_async_chunk
+    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: 1024
+      seed: 42
+      detokenize: true
+      repetition_penalty: 1.05
+
+  # Stage 1: Token2Wav (Audio Synthesis)
+  - stage_id: 1
+    stage_type: llm
+    runtime:
+      process: true
+      devices: "3"
+      max_batch_size: 1
+    engine_args:
+      model_stage: token2wav
+      model_arch: StepAudio2Token2WavModel
+      worker_type: generation
+      hf_overrides:
+        architectures: ["StepAudio2Token2WavForConditionalGeneration"]
+      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
+      gpu_memory_utilization: 0.3
+      enforce_eager: true
+      trust_remote_code: true
+      async_scheduling: false
+      enable_prefix_caching: false
+      max_num_batched_tokens: 2048
+      engine_output_type: audio
+    engine_input_source: [0]
+    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
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from collections import defaultdict
+from collections.abc import Mapping
 from contextlib import nullcontext
 from typing import Any

@@ -162,7 +163,7 @@ def text2flow(

 def talker2code2wav_async_chunk(
    transfer_manager: Any,
-    pooling_output: dict[str, Any] | None,
+    multimodal_output: dict[str, Any] | None,
    request: Any,
    is_finished: bool = False,
 ) -> OmniPayloadStruct | None:
@@ -198,12 +199,14 @@ def talker2code2wav_async_chunk(
                    value = _to_cpu_tensor(info_embed.get(key))
                    if value is not None:
                        prompt_payload[key] = value
-                if isinstance(pooling_output, dict):
-                    po_embed = pooling_output.get("embed", {}) if isinstance(pooling_output.get("embed"), dict) else {}
+                if isinstance(multimodal_output, Mapping):
+                    mm_embed = multimodal_output.get("embed", {})
+                    if not isinstance(mm_embed, Mapping):
+                        mm_embed = multimodal_output
                    for key in cond_keys:
                        if key in prompt_payload:
                            continue
-                        value = _to_cpu_tensor(po_embed.get(key))
+                        value = _to_cpu_tensor(mm_embed.get(key))
                        if value is not None:
                            prompt_payload[key] = value
                # Drop any right-padding carried from batched talker emission so
@@ -1,5 +1,6 @@
 """Stage input processor for Fish Speech S2 Pro: Slow AR → DAC Decoder."""

+from collections.abc import Mapping
 from typing import Any

 import torch
@@ -91,14 +92,14 @@ def _select_backlog_chunk_size(
    return backlog_chunk_size


-def _extract_last_frame(pooling_output: dict[str, Any]) -> torch.Tensor | None:
-    """Extract the last frame of audio codes from the pooling output."""
-    audio_codes = pooling_output.get("audio_codes")
+def _extract_last_frame(multimodal_output: dict[str, Any]) -> torch.Tensor | None:
+    """Extract the last frame of audio codes from the multimodal output."""
+    audio_codes = multimodal_output.get("audio_codes")
    if not isinstance(audio_codes, torch.Tensor) or audio_codes.numel() == 0:
        return None
    if audio_codes.ndim == 2:
        frame = audio_codes[-1]
-        valid = pooling_output.get("audio_code_valid")
+        valid = multimodal_output.get("audio_code_valid")
        if isinstance(valid, torch.Tensor) and valid.numel() > 0:
            is_valid = bool(valid.reshape(-1)[-1].item())
        elif valid is not None:
@@ -146,7 +147,7 @@ def slow_ar_to_dac_decoder(

 def slow_ar_to_dac_decoder_async_chunk(
    transfer_manager: Any,
-    pooling_output: dict[str, Any] | None,
+    multimodal_output: dict[str, Any] | None,
    request: Any,
    is_finished: bool = False,
 ) -> OmniPayloadStruct | None:
@@ -160,8 +161,8 @@ def slow_ar_to_dac_decoder_async_chunk(
    finished = bool(is_finished or request.is_finished())
    cfg = _get_connector_extra(transfer_manager)

-    if isinstance(pooling_output, dict):
-        frame = _extract_last_frame(pooling_output)
+    if isinstance(multimodal_output, Mapping):
+        frame = _extract_last_frame(multimodal_output)
        if frame is not None:
            transfer_manager.code_prompt_token_ids[request_id].append(frame.detach())
    elif not finished:
@@ -4,6 +4,7 @@

 import math
 import time
+from collections.abc import Mapping
 from typing import Any

 import torch
@@ -23,7 +24,7 @@ def _has_source_image(mm_data: Any) -> bool:
    - `img2img`: legacy single-image key
    - `images`: list or single image
    """
-    if not isinstance(mm_data, dict):
+    if not isinstance(mm_data, Mapping):
        return False
    if mm_data.get("image") is not None:
        return True
@@ -35,7 +36,7 @@ def _has_source_image(mm_data: Any) -> bool:

 def _first_source_image(mm_data: Any) -> Any:
    """Get first source image from normalized multimodal keys."""
-    if not isinstance(mm_data, dict):
+    if not isinstance(mm_data, Mapping):
        return None

    image = mm_data.get("image")
@@ -285,7 +286,7 @@ def ar2diffusion(

        if hasattr(ar_output, "multimodal_output") and ar_output.multimodal_output:
            mm_output = ar_output.multimodal_output
-            if isinstance(mm_output, dict) and mm_output.get("ids", {}).get("prior_image") is not None:
+            if isinstance(mm_output, Mapping) and mm_output.get("ids", {}).get("prior_image") is not None:
                is_i2i = True
        _dt_mode = (time.perf_counter() - _t_mode) * 1000

@@ -322,7 +323,7 @@ def ar2diffusion(
        # Check ar_output (RequestOutput) for multimodal_output - this is the correct location
        if hasattr(ar_output, "multimodal_output") and ar_output.multimodal_output:
            mm_output = ar_output.multimodal_output
-            if isinstance(mm_output, dict):
+            if isinstance(mm_output, Mapping):
                raw_prior_image_ids = mm_output.get("ids", {}).get("prior_image")
                if raw_prior_image_ids is not None:
                    # Handle different formats:
@@ -348,7 +349,7 @@ def ar2diffusion(
            if hasattr(output, "multimodal_output") and output.multimodal_output:
                mm_output = output.multimodal_output
                logger.debug(f"[ar2diffusion] Request {i}: found multimodal_output on CompletionOutput (fallback)")
-                if isinstance(mm_output, dict):
+                if isinstance(mm_output, Mapping):
                    raw_prior_image_ids = mm_output.get("ids", {}).get("prior_image")
                    if raw_prior_image_ids is not None:
                        if isinstance(raw_prior_image_ids, torch.Tensor):
@@ -188,7 +188,7 @@ def _to_cpu_tensor(value: Any) -> torch.Tensor | None:

 def ar_to_dit_async_chunk(
    transfer_manager: Any,
-    pooling_output: dict[str, Any] | None,
+    multimodal_output: dict[str, Any] | None,
    request: Any,
    is_finished: bool = False,
 ) -> OmniPayloadStruct | None:
@@ -216,6 +216,7 @@ def ar_to_dit_async_chunk(
    if request_id is None:
        raise ValueError("GLM-TTS async chunk request is missing request id")
    finished = bool(is_finished or request.is_finished())
+    pooling_output = multimodal_output

    # Read connector chunk config (supports progressive list or single int)
    connector = getattr(transfer_manager, "connector", None)
@@ -176,7 +176,7 @@ def _extract_last_frame(pooling_output: OmniPayload) -> torch.Tensor | None:

 def talker2code2wav_async_chunk(
    transfer_manager: Any,
-    pooling_output: OmniPayload | None,
+    multimodal_output: OmniPayload | None,
    request: Any,
    is_finished: bool = False,
 ) -> OmniPayloadStruct | None:
@@ -188,6 +188,7 @@ def talker2code2wav_async_chunk(
    """
    request_id = request.external_req_id
    finished = bool(is_finished or request.is_finished())
+    pooling_output = multimodal_output

    if isinstance(pooling_output, dict):
        frame = _extract_last_frame(pooling_output)
Author	SHA1	Message	Date
hsliuustc0106	2b3a1c8200	Format diffusion RPC guard Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>	2026-06-13 14:43:32 +08:00
hsliuustc0106	08b2f694f3	Address diffusion RPC review feedback Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>	2026-06-13 14:39:53 +08:00
Hongsheng Liu	170f98c847	Merge branch 'main' into codex/fix-diffusion-allrank-rpc	2026-06-13 12:32:44 +08:00
hsliuustc0106	3805852a71	[Bugfix] Surface all-rank diffusion RPC failures Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>	2026-06-13 12:30:44 +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