[Rebase] Rebase to vllm 0.23.0 (#4286 )

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
[XPU] update Dreamzero to support any HW (#4399 )
2026-06-13 16:43:02 +00:00 · 2026-06-13 15:32:39 +00:00 · 2026-06-13 18:12:37 +08:00 · 2026-06-13 15:42:16 +08:00 · 2026-06-13 04:22:54 +00:00 · 2026-06-13 10:59:15 +08:00
234 changed files with 16904 additions and 1907 deletions
@@ -10,7 +10,7 @@ steps:
          DOCKER_BUILDKIT: "1"
          # Buildkite will automatically replace this with the actual commit hash
          VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}"
-          VLLM_VERSION: "v0.22.0"
+          VLLM_VERSION: "v0.23.0"
        priority: 100
        timeout_in_minutes: 60
        soft_fail: false
@@ -31,7 +31,7 @@ steps:
    key: upload-ready-pipeline
    if: __UPLOAD_READY_IF__
    commands:
-      - buildkite-agent pipeline upload .buildkite/test-ready.yml
+      - python3 .buildkite/scripts/upload_pipeline.py --upload .buildkite/test-ready.yml
    agents:
      queue: "cpu_queue_premerge"

@@ -14,7 +14,7 @@ steps:
    depends_on: image-build
    key: upload-ready-pipeline
    commands:
-      - buildkite-agent pipeline upload .buildkite/test-ready.yml
+      - python3 .buildkite/scripts/upload_pipeline.py --upload .buildkite/test-ready.yml
    agents:
      queue: "cpu_queue_premerge"

@@ -44,9 +44,6 @@ steps:
    depends_on: upload-merge-pipeline
    steps:
      - label: "Diffusion · LoRA Test"
-        source_file_dependencies:
-          - tests/diffusion/lora/
-          - vllm_omni/diffusion/lora/
        timeout_in_minutes: 30
        commands:
          - pytest -sv tests/diffusion/lora/ -m 'advanced_model and cuda' --run-level "advanced_model"
@@ -64,9 +61,6 @@ steps:
                - "/fsx/hf_cache:/fsx/hf_cache"

  - label: "Distributed Test"
-    source_file_dependencies:
-      - tests/distributed/
-      - vllm_omni/distributed/
    timeout_in_minutes: 60
    commands:
      - |
@@ -111,9 +105,6 @@ steps:

  - group: ":card_index_dividers: Entrypoints Test"
    depends_on: upload-merge-pipeline
-    source_file_dependencies:
-      - tests/entrypoints/
-      - vllm_omni/entrypoints/
    steps:
      - label: "Entrypoint Test with H100"
        timeout_in_minutes: 60
@@ -158,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
@@ -173,6 +164,28 @@ 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
@@ -180,7 +193,6 @@ steps:
          - vllm_omni/model_executor/models/qwen3_tts/
          - vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
          - vllm_omni/deploy/qwen3_tts.yaml
-          - vllm_omni/model_executor/stage_configs/qwen3_tts_uniproc.yaml
        commands:
          - |
            timeout 20m bash -c "
@@ -189,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
@@ -216,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
@@ -229,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:
          - |
@@ -568,6 +568,55 @@ steps:
                      path: /mnt/hf-cache
                      type: DirectoryOrCreate

+  - label: ":full_moon: TTS · Higgs-Audio-V3 · Perf Test"
+    key: nightly-tts-performance-higgs-audio-v3
+    depends_on: upload-nightly-pipeline
+    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" || build.pull_request.labels includes "tts-test"
+    timeout_in_minutes: 180
+    commands:
+      - export BENCHMARK_DIR=tests/dfx/perf/results
+      - export VLLM_USE_DEEP_GEMM="0"
+      - export VLLM_MOE_USE_DEEP_GEMM="0"
+      - |
+        set +e
+        pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_higgs_audio_v3.json
+        EXIT=$$?
+        buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
+        exit $$EXIT
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                resources:
+                  limits:
+                    nvidia.com/gpu: 1
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: hf-cache
+                    mountPath: /root/.cache/huggingface
+                env:
+                  - name: HF_HOME
+                    value: /root/.cache/huggingface
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
+            nodeSelector:
+              node.kubernetes.io/instance-type: gpu-h100-sxm
+            volumes:
+              - name: devshm
+                emptyDir:
+                  medium: Memory
+              - name: hf-cache
+                hostPath:
+                  path: /mnt/hf-cache
+                  type: DirectoryOrCreate
+
  # Diffusion X2I suite: x2i / x2a / x2t and related non-video paths; x2v is only in "Diffusion X2V Model Test" below.
  - group: ":card_index_dividers: Diffusion X2I(&A&T) Model Test"
    key: nightly-diffusion-x2iat-group
@@ -1463,6 +1512,7 @@ steps:
      - nightly-omni-performance-vllm-text
      - nightly-omni-performance-multi-replicas
      - nightly-tts-performance
+      - nightly-tts-performance-higgs-audio-v3
      - nightly-diffusion-x2iat-performance-qwen-image
      - nightly-diffusion-x2iat-performance-qwen-image-edit
      - nightly-diffusion-x2iat-performance-qwen-image-edit-2511
@@ -1476,6 +1526,7 @@ steps:
      - export DEFAULT_INPUT_DIR=tests/dfx/perf/results
      - export DEFAULT_OUTPUT_DIR=tests/dfx/perf/results
      - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-tts-performance
+      - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-tts-performance-higgs-audio-v3
      - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-no-async-chunk
      - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-async-chunk
      - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-vllm-text
@@ -215,10 +215,78 @@ steps:
                  path: /mnt/hf-cache
                  type: DirectoryOrCreate

+  - group: ":card_index_dividers: Entrypoints Test"
+    depends_on: upload-ready-pipeline
+    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
+                    emptyDir:
+                      medium: Memory
+                  - name: hf-cache
+                    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 "
@@ -239,6 +307,11 @@ steps:
                - "/fsx/hf_cache:/fsx/hf_cache"

      - label: "Omni · Qwen3-Omni Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_qwen3_omni.py
+          - vllm_omni/model_executor/models/qwen3_omni/
+          - vllm_omni/model_executor/stage_input_processors/qwen3_omni.py
+          - vllm_omni/deploy/qwen3_omni_moe.yaml
        commands:
          - |
            timeout 20m bash -c "
@@ -273,50 +346,18 @@ steps:
                      path: /mnt/hf-cache
                      type: DirectoryOrCreate

-      - label: "Omni · Qwen3-Omni Realtime WebSocket"
-        commands:
-          - |
-            timeout 30m bash -c "
-              export VLLM_WORKER_MULTIPROC_METHOD=spawn
-              pytest -s -v tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py::TestQwen3OmniRealtimeWebSocket::test_streaming_audio_input_pcm_output_async_chunk_smoke -m 'core_model' --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
-                    emptyDir:
-                      medium: Memory
-                  - name: hf-cache
-                    hostPath:
-                      path: /mnt/hf-cache
-                      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"
@@ -333,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 "
@@ -340,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
@@ -354,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 "
@@ -363,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
@@ -377,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 "
@@ -386,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
@@ -400,6 +456,9 @@ steps:
                - "/fsx/hf_cache:/fsx/hf_cache"

      - label: "Diffusion · Qwen Image Edit Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_qwen_image_edit.py
+          - vllm_omni/diffusion/models/qwen_image/
        commands:
          - |
            timeout 20m bash -c "
@@ -435,6 +494,12 @@ steps:
                      type: DirectoryOrCreate

      - label: "Diffusion · Bagel Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_bagel.py
+          - vllm_omni/model_executor/models/bagel/
+          - vllm_omni/diffusion/models/bagel/
+          - vllm_omni/model_executor/stage_input_processors/bagel.py
+          - vllm_omni/deploy/bagel.yaml
        commands:
          - |
            timeout 40m bash -c "
@@ -476,6 +541,9 @@ steps:
                      type: DirectoryOrCreate

      - label: "Diffusion · Wan22 Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_wan22_t2v.py
+          - vllm_omni/diffusion/models/wan2_2/
        commands:
          - |
            timeout 40m bash -c "
@@ -518,6 +586,9 @@ steps:


      - label: "Diffusion · Qwen Image Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_qwen_image.py
+          - vllm_omni/diffusion/models/qwen_image/
        commands:
          - |
            timeout 40m bash -c "
@@ -559,6 +630,9 @@ steps:
                      type: DirectoryOrCreate

      - label: "Diffusion · Qwen Image Layered Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_qwen_image_layered.py
+          - vllm_omni/diffusion/models/qwen_image/
        commands:
          - |
            timeout 40m bash -c "
@@ -0,0 +1,103 @@
+---
+name: quantization
+description: Work on vLLM-Omni quantization for diffusion, autoregressive, omni, or multi-stage models. Use when choosing or adding methods such as fp8, int8, gguf, mxfp8, mxfp4, mxfp4_dualscale, ModelOpt, AutoRound, INC, msModelSlim, awq, or gptq; debugging quantized loading; or validating memory, speed, and output quality.
+---
+
+# vLLM-Omni Quantization
+
+Use this skill for quantization work in `vllm-omni`. Start from the local
+code and docs, then pick the smallest validated path for the model, method,
+hardware, and quality target.
+
+## First Checks
+
+Before changing code or recommending a command, identify:
+
+- Model and task: AR text, omni/TTS, text-to-image, text-to-video, or multi-stage diffusion.
+- Hardware backend: CUDA generation, Ascend NPU, Intel XPU, or another platform.
+- Quantization mode: online from BF16/FP16 weights, pre-quantized checkpoint, runtime KV/FA quantization, or per-component routing.
+- Scope: transformer/DiT, thinker language model, text encoder, VAE, audio/vision encoder, talker, or a specific stage.
+- Validation target: memory reduction, latency/throughput, output quality, or all three.
+
+Do not assume a method is supported just because the CLI accepts the string.
+Check `docs/user_guide/quantization/` and the closest implementation first.
+
+## Quick Decision
+
+| Task | Start With |
+|------|------------|
+| Choose a method or command | `references/methods.md` and `references/modality-compat.md` |
+| Use `build_quant_config()` or per-component routing | `references/methods.md` |
+| Work on diffusion quantization | `references/diffusion.md` |
+| Add quantization to a new model | `references/adding-models.md` |
+| Convert or load ModelOpt FP8 checkpoints | `references/modelopt-fp8.md` |
+| Debug ModelOpt, GGUF, AutoRound, MXFP, or serialized Int8 loading | `references/diffusion.md` and method docs under `docs/user_guide/quantization/` |
+
+## Current Runtime Shape
+
+The unified entrypoint is:
+
+```python
+from vllm_omni.quantization import build_quant_config
+```
+
+It supports method strings, flat method dictionaries, per-component
+dictionaries, existing `QuantizationConfig` objects, and `None`. The factory
+delegates generic methods to upstream `vllm` and keeps vLLM-Omni overrides for
+diffusion or omni-specific routing:
+
+- `gguf`
+- `int8`
+- `mxfp8`
+- `mxfp4`
+- `mxfp4_dualscale`
+- `inc`, `auto-round`, `auto_round`
+- ModelOpt auto-detected configs: `modelopt`, `modelopt_fp4`, `modelopt_mixed`
+
+Use `ComponentQuantizationConfig` when only one stage or component should be
+quantized. Pre-quantized ModelOpt-style checkpoints should not spill into
+vision/audio encoders that have no corresponding scale tensors.
+
+## Ownership Boundary
+
+- Upstream `vllm` owns generic quantization configs, kernels, loader semantics,
+  hardware capability rules, and generic AR quantization methods.
+- `vllm-omni` owns unified config routing, diffusion-specific wrappers,
+  component scoping, GGUF or ModelOpt checkpoint adapters, model-specific
+  prefix mapping, docs, examples, and validation.
+
+If a new method needs missing generic kernels or loader behavior, fix upstream
+`vllm` first. In `vllm-omni`, add thin integration and model wiring.
+
+## Working Loop
+
+1. Read the local method doc under `docs/user_guide/quantization/`.
+2. Inspect the active factory and config class in `vllm_omni/quantization/`.
+3. Inspect the target model pipeline and transformer for stable prefixes and
+   whether every relevant vLLM linear layer receives `quant_config`.
+4. For pre-quantized checkpoints, inspect `transformer/config.json` and the
+   checkpoint tensor names before running full generation.
+5. Validate with the same prompt, seed, size, scheduler, steps, dtype, eager or
+   non-eager mode, and parallelism as the BF16 baseline.
+6. Report quality, latency, and memory separately. Do not call a method
+   supported until quality is checked.
+
+## Common Mistakes
+
+| Symptom | Likely Cause | Fix |
+|---------|--------------|-----|
+| `--quantization` has no visible effect | Wrong scope or unsupported model path | Check component routing and method docs |
+| Some layers stay BF16 unexpectedly | `quant_config` was not threaded into all vLLM linear layers | Audit transformer constructors and prefixes |
+| Quality collapses but loading succeeds | Too many sensitive layers were quantized | Add model-specific `ignored_layers` and compare to BF16 |
+| ModelOpt checkpoint loads but output is corrupted | Prefixes, packed-module mapping, scale routing, or BF16 fallback is wrong | Use `references/modelopt-fp8.md` |
+| GGUF shape or tensor mismatch | Missing architecture-specific adapter | Add explicit adapter mapping; avoid generic fallback |
+| Online and offline MXFP4 disagree | Wrong `mxfp4` vs `mxfp4_dualscale` mode | Check `docs/user_guide/quantization/mxfp4.md` |
+| Quantized path is slower than BF16 | Kernel path, compile mode, dtype, or shape mismatch | Compare eager-to-eager and non-eager-to-non-eager |
+
+## References
+
+- General methods and config routing: [references/methods.md](references/methods.md)
+- Modality and model scope: [references/modality-compat.md](references/modality-compat.md)
+- Diffusion implementation patterns: [references/diffusion.md](references/diffusion.md)
+- Adding model support: [references/adding-models.md](references/adding-models.md)
+- ModelOpt FP8 conversion and loading: [references/modelopt-fp8.md](references/modelopt-fp8.md)
@@ -0,0 +1,70 @@
+# Adding Quantization Support to a Model
+
+Use this checklist when extending quantization to a new architecture or stage.
+
+## Classify Ownership
+
+| Case | Owner |
+|------|-------|
+| New generic AR method, kernel, or loader semantics | upstream `vllm` first |
+| Existing upstream method needs diffusion wiring | `vllm-omni` model and config integration |
+| GGUF or ModelOpt tensor mapping for a diffusion architecture | `vllm-omni` adapter |
+| NPU MXFP or msModelSlim diffusion path | `vllm-omni` platform/method wrapper and docs |
+| Multi-stage omni checkpoint quantizes only the thinker LM | `vllm-omni` component scoping |
+
+Avoid building a private quantization stack in model code.
+
+## Model Wiring Checklist
+
+1. Find the closest supported model implementation.
+2. Confirm the model uses vLLM linear layers that accept `quant_config`.
+3. Add `quant_config` to the transformer constructor.
+4. Pass stable `prefix` values through every linear layer.
+5. Handle fused modules and packed mappings explicitly.
+6. Ensure `load_weights()` and any `WeightsMapper` remap `ignored_layers` or
+   `modules_to_not_convert`.
+7. Keep tokenizer, scheduler, VAE, text encoder, audio/vision encoders, and
+   talker BF16 unless method-specific docs validate them.
+8. Add or update a method doc under `docs/user_guide/quantization/`.
+
+## Checkpoint Adapter Checklist
+
+For pre-quantized checkpoints:
+
+- Inspect `transformer/config.json` first.
+- Check `quant_method`, `method`, `producer.name`, `quant_algo`, serialized
+  flags, `ignored_layers`, and any method-specific fields.
+- Verify tensor name mapping before generation.
+- Verify packed modules such as `to_qkv`, `add_kv_proj`, and `w13`.
+- Confirm scale tensor names and shapes.
+- Allow dequantization back to BF16 for intentionally skipped target layers.
+
+## Sensitive Layer Discovery
+
+Start conservative. Quantize the backbone first and keep output-side or routing
+layers BF16 until quality is proven.
+
+Frequently sensitive areas:
+
+- attention output projections such as `to_out`
+- FFN down projections such as `w2`
+- final projections such as `proj_out`, `final_layer`, and output heads
+- modulation, timestep, position, image, or text embedders
+- refiner blocks
+- MoE routers, experts selected by routing, and layernorm-adjacent modules
+- VAE, text/vision towers, audio encoders, and codec/talker stages
+
+Do not reuse one model family's ignore list blindly.
+
+## Tests and Evidence
+
+Minimum evidence for a PR:
+
+- loader/config test or smoke command
+- fixed-seed BF16 vs quantized output comparison
+- quality metrics and saved outputs for precision-changing work
+- memory and latency comparison under identical execution mode
+- failure behavior for unsupported models or missing checkpoint fields
+
+When the PR adds a method to a support table, include the exact model,
+hardware, command, quality threshold, and artifact path used for validation.
@@ -0,0 +1,113 @@
+# Diffusion Quantization Reference
+
+Use this file for DiT quantization, config routing, checkpoint adapters, and
+quality or loader debugging.
+
+## Supported Method Families
+
+The current local diffusion quantization families are:
+
+- CUDA/GPU online or checkpoint: `fp8`, `int8`, `gguf`, ModelOpt
+- Ascend NPU: `int8`, `mxfp8`, `mxfp4`, `mxfp4_dualscale`, msModelSlim paths
+- Intel XPU: AutoRound/INC and online MXFP8 paths where documented
+
+Always cross-check the method doc before claiming model support. The CLI
+accepting a method string is not enough.
+
+## Implementation Starting Points
+
+| Area | Files |
+|------|-------|
+| Unified factory | `vllm_omni/quantization/factory.py` |
+| Per-component routing | `vllm_omni/quantization/component_config.py` |
+| GGUF diffusion method | `vllm_omni/quantization/gguf_config.py` |
+| Int8 diffusion method | `vllm_omni/quantization/int8_config.py` |
+| MXFP8 diffusion method | `vllm_omni/quantization/mxfp8_config.py` |
+| MXFP4 and dual-scale method | `vllm_omni/quantization/mxfp4_config.py` |
+| AutoRound/INC omni extension | `vllm_omni/quantization/inc_config.py` |
+| ModelOpt checkpoint adapters | `vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py` |
+| FP8 scaled-mm runtime patches | `vllm_omni/patch.py` |
+| Similarity tool | `vllm_omni/quantization/tools/compare_diffusion_trajectory_similarity.py` |
+
+## Threading `quant_config`
+
+For a diffusion transformer:
+
+1. Accept `quant_config: QuantizationConfig | None` in the transformer.
+2. Pass `quant_config` and stable `prefix` values to every vLLM linear layer.
+3. Preserve packed-module mappings for fused QKV, FFN, or model-specific fused
+   projections.
+4. Call `apply_vllm_mapper()` on configs that need HF-to-vLLM name remapping.
+5. For component configs, keep prefixes specific enough to avoid quantizing VAE
+   or encoders by accident.
+
+Common skipped areas are `proj_out`, output projections, modulation layers,
+embedders, refiner blocks, MoE routers, and layernorm-adjacent paths.
+
+## Online vs Offline
+
+Online quantization:
+
+- starts from BF16/FP16 checkpoint weights
+- computes or stores quantized weights at load time
+- typically uses `--quantization <method>` or `quantization="<method>"`
+
+Offline or pre-quantized checkpoints:
+
+- include quantized weights and scales on disk
+- usually carry `quantization_config` in `transformer/config.json`
+- should auto-detect when the config is complete
+- may require model-specific adapter logic for tensor names and packed modules
+
+Do not mix online flags with incompatible offline checkpoint configs. If the
+active config and disk config differ, `resolve_quant_config_from_disk()` should
+raise or rebuild rather than silently loading corrupted weights.
+
+## GGUF Rules
+
+- Require `quantization_config.gguf_model`.
+- Use the normal base model for tokenizer, scheduler, text encoder, and VAE.
+- Add explicit architecture adapters for tensor-name mapping.
+- Use `named_parameters()` and `named_buffers()`, not `state_dict()`, to
+  discover loadable names.
+- Guard fused QKV/KV rewrites so tensors are not rewritten twice.
+- Flatten N-D activations around the GGUF matmul and restore shape.
+
+Unsupported architectures should fail clearly instead of falling back to a fake
+generic mapper.
+
+## MXFP Rules
+
+- `mxfp8` is W8A8 microscaling, currently centered on NPU and documented XPU
+  behavior.
+- `mxfp4` is online single-scale W4A4 on NPU.
+- `mxfp4_dualscale` is online/offline dual-scale W4A4 with BF16 fallback; the
+  offline path is the production path when calibrated `mul_scale` exists.
+- Offline dual-scale checkpoints must inject per-transformer `ignored_layers`
+  into config because cascade transformers can have different BF16 fallback
+  sets.
+
+Do not load an offline dual-scale checkpoint with the online single-scale path.
+
+## Common Failures
+
+| Symptom | Likely Cause | Fix |
+|---------|--------------|-----|
+| `Unknown quantization method` | Factory override or upstream registration is missing | Add the right override or upstream method |
+| Missing scale tensors | Wrong scope, wrong adapter, or BF16 layer not allowed to dequantize | Check config scope and adapter scale routing |
+| Output rank collapses from 3-D to 2-D | FP8 scaled-mm kernel ignored output shape | Check the patch in `vllm_omni/patch.py` |
+| Quality degrades silently | Sensitive layers quantized | Add `ignored_layers`; compare to BF16 |
+| Only one cascade transformer works | Disk config differs per transformer but active config is reused | Rebuild from each transformer's config |
+| Quantized path slower than BF16 | Kernel, dtype, compile, or benchmark mode mismatch | Compare eager-to-eager and non-eager-to-non-eager |
+
+## Verification
+
+For a new or changed diffusion quantization path, run:
+
+1. Config and loader smoke test.
+2. One fixed-seed generation against BF16.
+3. Numeric comparison with the similarity tool or equivalent metrics.
+4. Serving test if the path is user-facing online.
+5. Memory and latency check under the same execution mode as the baseline.
+
+Keep output artifacts when precision changes are part of the PR.
@@ -0,0 +1,82 @@
+# General Quantization Methods
+
+Use this file for method selection and unified `vllm-omni` config syntax. For
+diffusion-specific wiring, see `diffusion.md`. For ModelOpt checkpoint export
+and loading details, see `modelopt-fp8.md`.
+
+## Unified Entry Point
+
+`vllm_omni.quantization.build_quant_config()` accepts:
+
+- `None` or `"none"` to disable quantization
+- a method string such as `"fp8"`, `"int8"`, `"mxfp8"`, `"awq"`, or `"auto-round"`
+- a flat dict with `"method"` or `"quant_method"`
+- a per-component dict such as `{"transformer": {"method": "fp8"}, "vae": None}`
+- an already-built `QuantizationConfig`
+
+Per-component dicts are routed by `ComponentQuantizationConfig` using
+longest-prefix matching. Prefix stability matters: if a model remaps HF names
+to vLLM names, apply the mapper to `ignored_layers` or component prefixes.
+
+## Local Overrides
+
+The local factory delegates most methods to upstream `vllm`, but has vLLM-Omni
+overrides for methods that need diffusion, NPU, XPU, or component behavior:
+
+| Method | Local Config | Main Use |
+|--------|--------------|----------|
+| `gguf` | `DiffusionGGUFConfig` | Pre-quantized diffusion transformer weights |
+| `int8` | `DiffusionInt8Config` | Online or serialized W8A8 diffusion transformers |
+| `mxfp8` | `DiffusionMXFP8Config` | W8A8 MXFP8 diffusion on NPU and online XPU path |
+| `mxfp4` | `DiffusionMXFP4Config` | Online W4A4 MXFP4 diffusion on NPU |
+| `mxfp4_dualscale` | `DiffusionMXFP4DualScaleMixedConfig` | Online/offline W4A4 dual-scale plus BF16 fallback on NPU |
+| `inc`, `auto-round`, `auto_round` | `OmniINCConfig` | AutoRound/INC checkpoint loading, including MXFP8 checkpoint metadata |
+| ModelOpt config dicts | upstream ModelOpt configs through auto-detect | FP8, NVFP4, or mixed ModelOpt checkpoints |
+
+Check `SUPPORTED_QUANTIZATION_METHODS` in `vllm_omni/quantization/factory.py`
+before adding a new method name.
+
+## AR and Generic Quantization
+
+For AR-backed models, upstream `vllm` methods usually apply:
+
+- `awq` and `gptq`: 4-bit or 8-bit weight quantization with calibration,
+  commonly used for text and AR stages.
+- `fp8`: FP8 weight path or KV-cache FP8 depending on flags and hardware.
+- `inc` or `auto-round`: pre-quantized checkpoints with Intel Neural
+  Compressor / AutoRound metadata.
+
+Use these for the AR language model stage only when the checkpoint and model
+guide support it. Omni and TTS models often keep audio, vision, talker, and
+codec stages in BF16.
+
+## Diffusion and Offline Methods
+
+For diffusion transformers, prefer the method-specific docs:
+
+- `docs/user_guide/quantization/fp8.md`
+- `docs/user_guide/quantization/int8.md`
+- `docs/user_guide/quantization/gguf.md`
+- `docs/user_guide/quantization/modelopt.md`
+- `docs/user_guide/quantization/mxfp8.md`
+- `docs/user_guide/quantization/mxfp4.md`
+- `docs/user_guide/quantization/autoround.md`
+- `docs/user_guide/quantization/msmodelslim.md`
+
+Pre-quantized checkpoints should normally auto-detect from config. Runtime
+flags such as `--force-cutlass-fp8`, `--linear-backend cutlass`, and
+`--moe-backend cutlass` select kernels; they do not quantize a BF16 checkpoint.
+
+## Validation
+
+Use fixed-seed A/B comparisons and record:
+
+- model path and quantized checkpoint path
+- method config and `ignored_layers`
+- prompt, size, frame count, steps, scheduler, dtype, and seed
+- eager/non-eager mode, compile flags, and parallelism
+- output metrics such as PSNR, MAE, cosine similarity, SSIM, LPIPS, plus visual artifacts
+- latency and peak memory, reported separately from quality
+
+The helper `vllm_omni.quantization.tools.compare_diffusion_trajectory_similarity`
+is the default lightweight image/video comparison tool for diffusion paths.
@@ -0,0 +1,47 @@
+# Quantization Compatibility by Modality
+
+Use this as a routing guide. The source of truth is the local method docs under
+`docs/user_guide/quantization/` and the current model implementation.
+
+## Summary
+
+- AR and omni language-model quantization mostly follows upstream `vllm`, with
+  vLLM-Omni scoping for multi-stage models.
+- Diffusion quantization is method-specific and model-specific. It usually
+  targets the diffusion transformer only.
+- Pre-quantized checkpoints should quantize only the component that has
+  quantized weights and scale tensors.
+
+## Model Type Matrix
+
+| Model Type | Typical Quantization Scope | Notes |
+|------------|----------------------------|-------|
+| AR text models | Language model | Upstream `vllm` methods usually apply |
+| Omni/TTS models | Thinker or AR language-model stage | Audio, vision, talker, and codec stages usually stay BF16 |
+| DiT image models | Diffusion transformer | Check FP8, Int8, GGUF, ModelOpt, AutoRound docs |
+| DiT video models | Diffusion transformer | HunyuanVideo FP8 and Wan2.2 MXFP paths are documented; validate other video paths |
+| Multi-stage diffusion | Selected stage only | Use per-component routing and per-stage validation |
+| Audio diffusion | Usually unsupported | Prefer dtype, offload, cache, or parallelism unless docs say otherwise |
+
+## Method Hints
+
+| Goal | Candidate |
+|------|-----------|
+| AR memory reduction | AWQ, GPTQ, AutoRound, ModelOpt where checkpoint supports it |
+| CUDA diffusion W8A8 | FP8 or Int8 |
+| Pre-quantized CUDA diffusion checkpoint | ModelOpt or GGUF |
+| Ascend NPU diffusion | Int8, MXFP8, MXFP4, MXFP4 DualScale, msModelSlim |
+| Intel XPU checkpoint | AutoRound/INC paths |
+| Conservative image quality | selective FP8/Int8 with `ignored_layers` |
+| Unsupported diffusion model | cache acceleration, CPU/layerwise offload, dtype, TP/SP/CFG parallelism |
+
+## Validation Rule
+
+Compatibility means more than successful loading. A method is compatible only
+after it has:
+
+- correct component scope
+- complete checkpoint/config metadata if offline
+- fixed-seed quality comparison against BF16
+- acceptable latency and memory behavior
+- docs or examples that state model, method, hardware, and caveats
@@ -0,0 +1,122 @@
+# ModelOpt FP8 Conversion and Loading
+
+Use this when exporting or loading ModelOpt FP8 checkpoints for diffusion
+transformers. This reference captures the practical conversion lessons from the
+2026-04-21 ModelOpt FP8 paste and aligns them with the current
+`vllm-omni` ModelOpt adapter.
+
+## Core Rule
+
+Treat ModelOpt FP8 as selective checkpoint quantization, not one-click
+whole-model FP8. Quantize the main transformer or DiT backbone, keep sensitive
+or peripheral modules BF16, serialize a complete ModelOpt checkpoint, and let
+vLLM-Omni auto-detect it from config.
+
+Do not pass `--quantization fp8` to force a ModelOpt checkpoint path. The
+checkpoint's `quantization_config` should select `modelopt`.
+
+## Conversion Strategy
+
+Prefer:
+
+- backbone attention Q/K/V and FFN expansion projections
+- conservative per-model ignore lists
+- MHA quantizers disabled until basic linear quantization is stable
+- full module-name prefixes in ignore/routing logic
+- BF16 fallback for output-side and routing-sensitive layers
+
+Avoid quantizing by default:
+
+- VAE
+- text, vision, or audio encoders
+- embedders and modulation layers
+- final heads and output projections
+- refiner modules
+- MoE routers and layernorm-sensitive blocks
+
+## Model-Specific Starting Points
+
+| Model Family | Quantize First | Keep BF16 First |
+|--------------|----------------|-----------------|
+| Z-Image | attention q/k/v, FFN w1/w3 | `to_out`, `w2`, `noise_refiner`, `context_refiner`, embedders, modulation, final/output |
+| Qwen-Image | transformer backbone linear layers | `proj_out`, `img_in`, `txt_in`, norms, modulation, position/time/embed layers |
+| FLUX2-dev / FLUX2-klein | backbone linear layers | `proj_out`, context/x embedders, `norm_out`, stream modulation, time embeddings |
+| HunyuanImage-3.0 | `model.model` backbone | VAE, vision tower, token/time/image embedders, final image head, LM head, routers, layernorm-sensitive blocks |
+
+These are starting points, not support claims. Validate with BF16 outputs.
+
+## Required Checkpoint Metadata
+
+`transformer/config.json` must include a complete `quantization_config`.
+At minimum it should identify ModelOpt and the algorithm:
+
+```json
+{
+  "quantization_config": {
+    "producer": {"name": "modelopt"},
+    "quant_method": "modelopt",
+    "quant_algo": "FP8"
+  }
+}
+```
+
+Current vLLM-Omni auto-detects:
+
+- `FP8` and `FP8_PER_CHANNEL_PER_TOKEN` -> `modelopt`
+- `NVFP4` -> `modelopt_fp4`
+- `MIXED_PRECISION` -> `modelopt_mixed`
+
+Packed modules and scales must map correctly:
+
+- `to_qkv` from `to_q`, `to_k`, `to_v`
+- `add_kv_proj` from `add_q_proj`, `add_k_proj`, `add_v_proj`
+- `w13` from `w1`, `w3`
+- `.input_scale`, `.weight_scale`, `.weight_scale_2`, `.weight_scale_inv`
+
+For selective FP8, target layers that remain BF16 must be allowed to
+dequantize from FP8 weights using the matching `weight_scale`, or the adapter
+must skip the quantized scale tensors for that BF16 target.
+
+## Runtime Signals
+
+Healthy loading should show ModelOpt auto-detection and a ModelOpt linear
+kernel path. Useful signals include:
+
+- auto-detected quantization from config
+- selected ModelOpt FP8 linear method
+- CUTLASS or FlashInfer FP8 scaled-mm backend when enabled and supported
+- non-eager path compiling the transformer when expected
+
+If online serving fails while offline inference works, inspect warmup and
+serving initialization before changing checkpoint conversion.
+
+## Validation Order
+
+1. Check config completeness.
+2. Run an export checker such as `examples/quantization/check_modelopt_fp8_export.py`.
+3. Confirm actual FP8 tensors and expected disk-size reduction.
+4. Run one offline generation with fixed prompt, seed, size, and steps.
+5. Run serving with `vllm serve --omni` and verify `/v1/images/generations`.
+6. Compare against BF16 using visual review plus metrics such as LPIPS, PSNR,
+   MAE, cosine similarity, and SSIM.
+7. Compare performance only under matched modes:
+   - BF16 eager vs FP8 eager
+   - BF16 non-eager vs FP8 non-eager
+
+FP8 is not automatically faster. Check whether execution reaches FP8 kernels,
+whether compile mode is comparable, and whether the workload shape fits the
+kernel path.
+
+## Triage
+
+| Symptom | First Check |
+|---------|-------------|
+| Startup failure | `config.json`, `model_index.json`, `quantization_config`, missing scales |
+| Loads but output is corrupted | Over-quantized sensitive layers or wrong full-name prefixes |
+| Offline works but online fails | Dummy warmup, serving init, compile path |
+| FP8 is not faster | Kernel selection, dtype, compile mode, workload shape |
+| BF16 fallback fails | Missing `weight_scale` for dequantization or wrong packed-module mapping |
+
+The success criterion is stable, correctly scoped checkpoint loading with
+quality preserved enough for the target use case, not maximizing the number of
+FP8 layers.
@@ -21,6 +21,8 @@ include:
 - `diffusion-perf-opt`: guides diffusion model performance optimization,
  including profiling traces, parallel strategies, stage timing analysis, and
  benchmark-driven tuning
+- `quantization`: guides quantization method selection, model integration,
+  checkpoint loading, and quality/performance validation for vLLM-Omni
 - `add-omni-model`: covers addition of new omni-modality model support
 - `add-tts-model`: covers integration of new TTS models and related serving
  workflows
@@ -1,24 +1,15 @@
 <!-- markdownlint-disable -->
-**Before submitting:** run the precheck-pr skill with code agent for a self-check against project conventions — catches dead code, missing benchmarks, and title format issues.
-
-PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTTOM) HAVE BEEN CONSIDERED.
+PLEASE FILL IN THE PR DESCRIPTION HERE.

 ## Purpose

 ## Test Plan

+**vLLM Version:**
+
+**vLLM-Omni Commit:**
+
 ## Test Result

---
-<details>
-<summary> Essential Elements of an Effective PR Description Checklist </summary>
-
- [ ] The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
- [ ] The test plan. Please provide the test scripts & test commands. Please state the reasons if your codes don't require additional test scripts. For test file guidelines, please check the [test style doc](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/tests_style/)
- [ ] The test results. Please paste the results comparison before and after, or the e2e results.
- [ ] (Optional) The necessary documentation update, such as updating `supported_models.md` and `examples` for a new model. **Please run `mkdocs serve` to sync the documentation editions to `./docs`.**
- [ ] (Optional) Release notes update. If your change is user-facing, please update the release notes draft.
-</details>
-
-**BEFORE SUBMITTING,** please read <https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md>
+**BEFORE SUBMITTING:** read [CONTRIBUTING.md](https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md) and run the [precheck-pr skill](https://github.com/vllm-project/vllm-omni/blob/main/.claude/skills/precheck-pr/SKILL.md) with the code agent for a self-check against project conventions.
 (anything written below this line will be removed by GitHub Actions)
@@ -1,3 +1,5 @@
 # Contributing to vLLM-Omni

 You may find information about contributing to vLLM-Omni on [Contributing](https://vllm-omni.readthedocs.io/en/latest/contributing/)
+
+Before submitting a PR, run the [precheck-pr skill](https://github.com/vllm-project/vllm-omni/blob/main/.claude/skills/precheck-pr/SKILL.md) with the code agent for a self-check against project conventions.
@@ -33,6 +33,15 @@ models:
      voice_clone:
        task_type: Base

+  bosonai/higgs-audio-v3-tts-4b:
+    supported_tasks: [default_voice, voice_clone]
+    backend: openai-audio-speech
+    endpoint: /v1/audio/speech
+    trust_remote_code: true
+    task_extra_body:
+      default_voice: {}
+      voice_clone: {}
+
  openbmb/VoxCPM2:
    supported_tasks: [voice_clone, default_voice]
    trust_remote_code: true
@@ -2,7 +2,7 @@
 # The pinned VLLM_PRECOMPILED_WHEEL_COMMIT and VLLM_BASE_TAG below must
 # be updated on every rebase to match the new upstream vllm commit/tag.
 ARG VLLM_BASE_IMAGE=vllm/vllm-openai
-ARG VLLM_BASE_TAG=v0.22.0
+ARG VLLM_BASE_TAG=v0.23.0
 FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG}
 ARG APP_DIR=/workspace/vllm-omni
 WORKDIR ${APP_DIR}
@@ -16,6 +16,7 @@ RUN apt-get update && \

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

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

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

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

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

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

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

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

 FROM vllm-omni AS vllm-omni-openai
@@ -87,6 +87,8 @@ nav:
      - AutoRound: user_guide/quantization/autoround.md
      - msModelSlim: user_guide/quantization/msmodelslim.md
    - ComfyUI: features/comfyui.md
+  - Fault Tolerance:
+    - Failure Modes: user_guide/fault_injection_reliability_matrix.md
 - Developer Guide:
  - General:
    - contributing/README.md
@@ -118,17 +118,13 @@ Only specific types of PRs will be reviewed. The PR title is prefixed appropriat

 ### Pre-Check Before Submitting

-Before submitting a PR, run the precheck-pr skill with code agent for a self-review against project conventions:
-
-```
-/precheck-pr
-```
+Before submitting a PR, run the [precheck-pr skill](https://github.com/vllm-project/vllm-omni/blob/main/.claude/skills/precheck-pr/SKILL.md) with the code agent for a self-review against project conventions:

 The skill offers two modes:
 - **Quick (~3 min):** catches showstoppers — PR title format, missing benchmark claims, rebase status
 - **Full (~10 min):** thorough maintainer-grade review — dead code scan, copy-paste detection, import hygiene

-The precheck covers five PR types: Bug Fix, Performance, New Model, Diffusion Model, and General. Each type has a tailored checklist that validates evidence quality (repro steps, A/B benchmarks, registry entries, etc.). See the skill in `.claude/skills/precheck-pr/` for the full checklist.
+The precheck covers five PR types: Bug Fix, Performance, New Model, Diffusion Model, and General. Each type has a tailored checklist that validates evidence quality (repro steps, A/B benchmarks, registry entries, etc.). See the [precheck-pr skill](https://github.com/vllm-project/vllm-omni/blob/main/.claude/skills/precheck-pr/SKILL.md) for the full checklist.

 ### Local Test
 Please run the L1 and L2 test cases locally first and attach the results before contacting us to add the "ready" label. Please refer to the [test instructions](./ci/test_guide.md) for running the test cases.
@@ -268,6 +268,120 @@ 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.

+## 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`.
+
+#### What `source_file_dependencies` is
+
+- A **uploader-only** YAML key on a Buildkite step or group. **Buildkite does not understand it**; `upload_pipeline.py` always **removes** it from the YAML that is uploaded.
+- A list of path **prefixes** (directories or individual files). If **any** changed file in the diff equals a prefix or starts with `prefix/`, the step is kept; otherwise the step (or entire group) is **omitted** from the uploaded pipeline.
+
+#### When filtering runs
+
+| Build context | Changed files used for matching |
+| --- | --- |
+| Pull request | `git diff --name-only origin/<base>...<BUILDKITE_COMMIT>` |
+| `main` branch push | `git diff --name-only <commit>^..<commit>` |
+| Other (e.g. local dry-run, non-PR branch) | Diff cannot be resolved → **no filtering**; all steps are uploaded and `source_file_dependencies` is still stripped |
+
+Docs-only PRs are handled earlier in bootstrap (`pipeline.yml`) via skip-ci logic; `source_file_dependencies` applies only to the **uploaded** L2/L3 test pipelines.
+
+#### Where it is configured
+
+| Level | Pipeline file | Upload entry |
+| --- | --- | --- |
+| L2 | `.buildkite/test-ready.yml` | `upload_pipeline.py --upload .buildkite/test-ready.yml` (from `pipeline.yml`) |
+| L3 | `.buildkite/test-merge.yml` | `upload_pipeline.py --upload .buildkite/test-merge.yml` (from `pipeline.yml`) |
+
+Steps **without** `source_file_dependencies` are always uploaded (subject to the usual label conditions: `ready` for L2, `merge-test` for L3).
+
+#### Current skip policy (L2 / L3)
+
+To balance CI cost and coverage:
+
+- **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 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 examples
+
+H100 E2E (kubernetes / `mithril-h100-pool`):
+
+```yaml
+      - label: "Diffusion · Qwen Image Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_qwen_image.py
+          - vllm_omni/diffusion/models/qwen_image/
+        commands:
+          - pytest -s -v tests/e2e/online_serving/test_qwen_image.py -m 'core_model' ...
+        agents:
+          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
+
+```bash
+# Render filtered YAML to stdout (no upload)
+python3 .buildkite/scripts/upload_pipeline.py .buildkite/test-ready.yml
+
+# Confirm uploader-only keys are stripped
+python3 .buildkite/scripts/upload_pipeline.py .buildkite/test-merge.yml | grep source_file_dependencies
+# (no output expected)
+```
+
+On a PR build, Buildkite logs from `upload_pipeline.py` include lines such as `skip '…' (no changes under …)` for omitted steps.
+
+#### Related
+
+- 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.
+
+| Variable | Accepted values | Description |
+| -------- | --------------- | ----------- |
+| `VLLM_OMNI_KEEP_REQUEST_MEDIA` | `1`, `true`, `yes` (case-insensitive) | When enabled, temporary WAV files created by `tests.helpers.media.convert_audio_bytes_to_text` are **not** deleted when the pytest process exits. By default, each call writes a unique file under the system temp directory via `tempfile.mkstemp` and registers `atexit` cleanup. Use this when debugging audio output validation (Whisper transcription, keyword checks, text–audio similarity). The saved path is logged as `audio data is saved: <path>`. |
+
+Example (Linux / macOS):
+
+```bash
+export VLLM_OMNI_KEEP_REQUEST_MEDIA=1
+pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -k test_mix_to_text_audio
+```
+
+Example (Windows PowerShell):
+
+```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
+```
+
 ## Chapter 1: L1 & L2 Level Testing - Unit Testing and Basic End-to-End Verification

 ### 1.1 Testing Purpose
@@ -302,6 +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, **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***:

@@ -404,7 +519,8 @@ 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.
+-   ***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—**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***:

@@ -520,6 +636,8 @@ L3 level testing executes after code is merged into the main branch. Its core pu

    **Implicit Validation**: The `send_omni_request` and `send_diffusion_request` methods internally includes validation logic dynamically selected based on the --run-level parameter: core_model performs basic validation, while advanced_model and full_model perform deep validation.

+    **Audio output debugging**: Deep validation may transcribe returned audio via `convert_audio_bytes_to_text` (Whisper). If an audio keyword or text–audio similarity assertion fails, set `VLLM_OMNI_KEEP_REQUEST_MEDIA=1` before running pytest to keep the intermediate WAV files for inspection (see [Test helper environment variables](#test-helper-environment-variables)).
+
 -   ***Run Command (L3 merge)***: `pytest -s -v /tests/e2e/online_serving/test_{model_name}.py -m advanced_model --run-level=advanced_model`

 -   ***Run Command (L4 nightly expansion)***: `pytest -s -v /tests/e2e/online_serving/test_{model_name}_expansion.py -m full_model --run-level=full_model`
@@ -69,6 +69,8 @@ If the failure is already tracked, leave a comment on your PR noting that the fa

 If the failure appears to be new, investigate whether your changes caused it.

+For resilience behavior under injected process signals/OOM pressure (validation scenarios, not root-cause equivalence), see [Failure Mode Injection Scenarios and Expected Behavior Matrix](../../user_guide/fault_injection_reliability_matrix.md).
+
 ### Reading Buildkite Logs

 1. Click **Details** next to the Buildkite check on your PR.
@@ -34,7 +34,7 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr

 ### Pre-built wheels

-Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, 0.21.0, and 0.22.0. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
+Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, 0.21.0, 0.22.0, and 0.23.0. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).

 === "NVIDIA CUDA"

@@ -20,7 +20,7 @@ Therefore, it is recommended to install vLLM and vLLM-Omni with a **fresh new**

 vLLM-Omni is built based on vLLM. Please install it with command below.
 ```bash
-uv pip install vllm==0.22.0 --torch-backend=auto
+uv pip install vllm==0.23.0 --torch-backend=auto
 ```

 #### Installation of vLLM-Omni
@@ -39,13 +39,13 @@ uv pip install 'vllm-omni[demo]'
 # --8<-- [start:build-wheel-from-source]

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

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

-The 0.22.0 release of vLLM ships CUDA 12.9-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
+The 0.23.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.

 #### Installation of vLLM-Omni
 Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -66,12 +66,12 @@ If you want to check, modify or debug with source code of vLLM, install the libr
 ```bash
 git clone https://github.com/vllm-project/vllm.git
 cd vllm
-git checkout v0.22.0
+git checkout v0.23.0
 ```
 Set up environment variables to get pre-built wheels. If there are internet problems, just download the whl file manually. And set `VLLM_PRECOMPILED_WHEEL_LOCATION` as your local absolute path of whl file.
 ```bash
-#For CUDA 12.9 (the default for v0.22.0; the wheel filename has no `+cu129` suffix)
-export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.22.0/vllm-0.22.0-cp38-abi3-manylinux_2_28_x86_64.whl
+#For CUDA 13.0 (the default for v0.23.0; the wheel filename has no `+cu130` suffix)
+export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.23.0/vllm-0.23.0-cp38-abi3-manylinux_2_28_x86_64.whl
 ```
 Install vllm with command below (If you have no existing PyTorch).
 ```bash
@@ -102,7 +102,7 @@ docker run --runtime nvidia --gpus 2 \
    --env "HF_TOKEN=$HF_TOKEN" \
    -p 8091:8091 \
    --ipc=host \
-    vllm/vllm-omni:v0.22.0 \
+    vllm/vllm-omni:v0.23.0 \
    vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 ```

@@ -13,7 +13,7 @@ vLLM-Omni current recommends the steps in under setup through Docker Images.

 vLLM-Omni is built based on vLLM. Please install it with command below.
 ```bash
-uv pip install vllm==0.22.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.22.0/rocm722
+uv pip install vllm==0.23.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.23.0/rocm722
 ```

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

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

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

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

 #### Installation of vLLM-Omni
 Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -58,7 +58,7 @@ If you want to check, modify or debug with source code of vLLM, install the libr
 ```bash
 git clone https://github.com/vllm-project/vllm.git
 cd vllm
-git checkout v0.22.0
+git checkout v0.23.0
 python3 -m pip install -r requirements/rocm.txt
 python3 setup.py develop
 ```
@@ -130,7 +130,7 @@ docker run --rm \
  -v ~/.cache/huggingface:/root/.cache/huggingface \
  --env "HF_TOKEN=$HF_TOKEN" \
  -p 8091:8091 \
-  vllm/vllm-omni-rocm:v0.22.0 \
+  vllm/vllm-omni-rocm:v0.23.0 \
  --model Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 ```

@@ -149,7 +149,7 @@ docker run --rm -it \
  -v ~/.cache/huggingface:/root/.cache/huggingface \
  --env "HF_TOKEN=$HF_TOKEN" \
  --entrypoint bash \
-  vllm/vllm-omni-rocm:v0.22.0
+  vllm/vllm-omni-rocm:v0.23.0
 ```

 # --8<-- [end:pre-built-images]
@@ -19,10 +19,10 @@ uv venv --python 3.12 --seed
 source .venv/bin/activate

 # On CUDA
-uv pip install vllm==0.22.0 --torch-backend=auto
+uv pip install vllm==0.23.0 --torch-backend=auto

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

 git clone https://github.com/vllm-project/vllm-omni.git
 cd vllm-omni
@@ -35,7 +35,7 @@ For additional installation methods — please see the [installation guide](inst
 !!! note
    It is important to install the same major & minor version of vLLM and vLLM Omni, otherwise things may not work as expected. If the versions are misaligned, you will see a warning when you import vLLM Omni.

-    If you are seeing strange behavior with the `vllm` command not handling the `--omni` flag correctly, you most likely have a version mismatch with vLLM < `0.22.0` and vLLM Omni `0.22.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue.
+    If you are seeing strange behavior with the `vllm` command not handling the `--omni` flag correctly, you most likely have a version mismatch with vLLM < `0.23.0` and vLLM Omni `0.23.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue.

 ## Offline Inference

@@ -33,7 +33,7 @@ th {
 | `ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, V2V, T2V with sound, action policy | `nvidia/Cosmos3-Nano` | ✅︎ | | | |
+| `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, V2V, T2V with sound, action policy | `nvidia/Cosmos3-Nano` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `WanSpeechToVideoPipeline` | Wan2.2-S2V | `Wan-AI/Wan2.2-S2V-14B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `LTX2Pipeline` | LTX-2-T2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | |
@@ -42,7 +42,7 @@ th {
 | `LTX2ImageToVideoTwoStagesPipeline` | LTX-2-I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |
 | `LTX23Pipeline` | LTX-2.3-T2V | `dg845/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
 | `LTX23ImageToVideoPipeline` | LTX-2.3-I2V | `dg845/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
-| `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | | | |
+| `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | |
 | `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | |
 | `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | ✅︎ | ✅︎ | | ✅︎ |
@@ -114,6 +114,7 @@ Content-Type: application/json
 | `instructions` | string | "" | Voice style/emotion instructions |
 | `max_new_tokens` | integer | 2048 | Maximum tokens to generate |
 | `initial_codec_chunk_frames` | integer | null | Per-request initial chunk size override for TTFA tuning. When null, IC is computed dynamically based on server load. |
+| `non_streaming_mode` | bool | null | Qwen3-TTS prompt construction mode override. Does not affect HTTP response streaming or async-chunk pipelining. When null, Qwen3-TTS uses model defaults: Base=false, CustomVoice/VoiceDesign=true. |
 | `stream` | bool | false | Stream raw PCM chunks as they are decoded (requires `response_format="pcm"`) |

 **Supported languages:** Only applicable to Qwen3-TTS. Derived from the model configuration (`talker_config.codec_language_id` in the checkpoint's `config.json`), plus `Auto`, which is always accepted. Official Qwen3-TTS checkpoints support: Auto, Chinese, English, Japanese, Korean, German, French, Russian, Portuguese, Spanish, Italian.
@@ -332,7 +333,8 @@ curl -X POST http://localhost:8091/v1/audio/speech \
        "input": "Hello, this is a cloned voice",
        "task_type": "Base",
        "ref_audio": "https://example.com/reference.wav",
-        "ref_text": "Original transcript of the reference audio"
+        "ref_text": "Original transcript of the reference audio",
+        "non_streaming_mode": true
    }' --output cloned.wav
 ```

@@ -454,6 +456,7 @@ Content-Type: application/json
 | `ref_audio` | string | null | Default reference audio (Base task) |
 | `ref_text` | string | null | Default reference transcript (Base task) |
 | `max_new_tokens` | integer | null | Default max tokens |
+| `non_streaming_mode` | bool | null | Default Qwen3-TTS prompt construction mode override. Does not affect HTTP response streaming or async-chunk pipelining. When null, Qwen3-TTS uses model defaults: Base=false, CustomVoice/VoiceDesign=true. |

 Each item in the `items` array requires only `input` (the text). All other fields are optional and override the batch-level defaults when set:

@@ -469,6 +472,7 @@ Each item in the `items` array requires only `input` (the text). All other field
 | `ref_audio` | string | Override reference audio |
 | `ref_text` | string | Override reference transcript |
 | `max_new_tokens` | integer | Override max tokens |
+| `non_streaming_mode` | bool | Override Qwen3-TTS prompt construction mode. Does not affect HTTP response streaming or async-chunk pipelining. When null, inherits the batch-level value (then the model default). |

 ### Response Format

@@ -70,6 +70,7 @@ curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o output.mp4
 | `input_reference` | file | null | Uploaded reference image or video for image-to-video/video-to-video requests |
 | `image_reference` | string | null | JSON-encoded reference image payload; do not combine with `input_reference` or `video_reference` |
 | `video_reference` | string | null | JSON-encoded reference video payload; do not combine with `input_reference` or `image_reference` |
+| `audio_reference` | string | null | JSON-encoded audio reference for speech-to-video: `{"audio_url": "..."}` — supports HTTP(s) URLs or base64 data URLs |
 | `width` | integer | model default | Output video width |
 | `height` | integer | model default | Output video height |
 | `num_frames` | integer | model default | Number of generated frames |
@@ -161,6 +162,25 @@ are not implemented yet. Models may expose additional V2V controls through
 `condition_frame_indexes_vision` and `condition_video_keep` to select which
 decoded reference frames are used as clean conditioning.

+### Speech-to-Video
+
+For models that support audio-driven generation (e.g., Wan2.2-S2V), pass both
+an image reference and an audio reference. The `audio_reference` field accepts a
+JSON string with `audio_url` pointing to an HTTP(s) URL or base64 data URL.
+
+```bash
+curl -s http://localhost:8091/v1/videos \
+  -F "prompt=A person singing" \
+  -F 'image_reference={"image_url": "https://example.com/face.png"}' \
+  -F 'audio_reference={"audio_url": "https://example.com/speech.mp3"}' \
+  -F "width=832" \
+  -F "height=480" \
+  -F "num_inference_steps=40" \
+  -F "guidance_scale=4.5" \
+  -F "fps=16"
+```
+
+
 ### Synchronous Generation

 ```bash
@@ -189,5 +209,7 @@ walkthroughs, see:

 - [Text-to-Video](../user_guide/examples/online_serving/text_to_video.md)
 - [Image-to-Video](../user_guide/examples/online_serving/image_to_video.md)
+- [Speech-to-Video](../user_guide/examples/online_serving/speech_to_video.md)
+  for Wan2.2-S2V audio-driven lip-sync generation
 - [Cosmos3 recipes](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-Nano.md)
  for model-specific video-to-video examples and conditioning controls
@@ -0,0 +1,128 @@
+# Speech-To-Video
+
+Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/speech_to_video>.
+
+This example demonstrates how to deploy the Wan2.2 speech-to-video (S2V) model for online video generation using vLLM-Omni.
+
+## Supported Models
+
+| Model | Model ID |
+|-------|----------|
+| Wan2.2 S2V (14B) | `Wan-AI/Wan2.2-S2V-14B` |
+
+## Start Server
+
+### Basic Start
+
+```bash
+VLLM_WORKER_MULTIPROC_METHOD=spawn \
+vllm serve Wan-AI/Wan2.2-S2V-14B --omni \
+  --model-class-name WanS2VPipeline \
+  --tensor-parallel-size 2 \
+  --flow-shift 3.0 \
+  --vae-use-slicing --vae-use-tiling \
+  --cache-backend cache_dit \
+  --port 8091
+```
+
+### Start with Script
+
+```bash
+bash run_server.sh
+```
+
+The script allows overriding:
+- `MODEL` (default: `Wan-AI/Wan2.2-S2V-14B`)
+- `PORT` (default: `8091`)
+- `FLOW_SHIFT` (default: `3.0`)
+- `TP` (default: `2`)
+- `CACHE_BACKEND` (default: `cache_dit`)
+
+### Start with 4 GPUs
+
+```bash
+TP=4 bash run_server.sh
+```
+
+## API Endpoints
+
+The server provides OpenAI-compatible video generation endpoints:
+
+- `POST /v1/videos`: create a video generation job (async)
+- `POST /v1/videos/sync`: generate a video and return raw bytes (sync, for benchmarks)
+- `GET /v1/videos/{video_id}`: retrieve the current job status and metadata
+- `GET /v1/videos`: list stored video jobs
+- `GET /v1/videos/{video_id}/content`: download the generated video file
+- `DELETE /v1/videos/{video_id}`: delete the job and any stored output
+
+## Sync API (Recommended for Testing)
+
+`POST /v1/videos/sync` blocks until generation completes and returns the raw
+video bytes (`video/mp4`) directly in the response body.
+
+```bash
+bash run_curl_speech_to_video.sh
+```
+
+## Async Job API
+
+`POST /v1/videos` creates an asynchronous job and returns immediately.
+
+```bash
+no_proxy=127.0.0.1 \
+create_response=$(curl -sS -X POST http://127.0.0.1:8091/v1/videos \
+  -H "Accept: application/json" \
+  -F "prompt=A person singing" \
+  -F 'image_reference={"image_url": "https://raw.githubusercontent.com/Wan-Video/Wan2.2/main/examples/Five%20Hundred%20Miles.png"}' \
+  -F 'audio_reference={"audio_url": "https://raw.githubusercontent.com/Wan-Video/Wan2.2/main/examples/Five%20Hundred%20Miles.MP3"}' \
+  -F "width=832" -F "height=480" \
+  -F "num_inference_steps=40" \
+  -F "guidance_scale=4.5" \
+  -F "fps=16")
+
+video_id=$(echo "$create_response" | jq -r '.id')
+
+# Poll until complete
+while true; do
+  status=$(curl -s "http://127.0.0.1:8091/v1/videos/${video_id}" | jq -r '.status')
+  if [ "$status" = "completed" ]; then
+    break
+  fi
+  if [ "$status" = "failed" ]; then
+    echo "Video generation failed"
+    exit 1
+  fi
+  sleep 2
+done
+
+# Download
+curl -L "http://127.0.0.1:8091/v1/videos/${video_id}/content" -o s2v_output.mp4
+```
+
+## Request Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `prompt` | str | - | Text description of the desired video |
+| `image_reference` | JSON str | - | Image reference: `{"image_url": "..."}` — supports HTTP(s) URLs or base64 data URLs |
+| `audio_reference` | JSON str | - | Audio reference: `{"audio_url": "..."}` — supports HTTP(s) URLs or base64 data URLs |
+| `width` | int | None | Video width in pixels |
+| `height` | int | None | Video height in pixels |
+| `fps` | int | None | Output video frame rate |
+| `num_inference_steps` | int | None | Number of denoising steps (40 recommended) |
+| `guidance_scale` | float | None | CFG guidance scale |
+| `seed` | int | None | Random seed for reproducibility |
+
+## Notes
+
+- S2V requires both a reference image (`image_reference`) and an audio reference
+  (`audio_reference`). The generated video will show a person matching the reference
+  image with lip movements synchronized to the audio.
+- `audio_reference` accepts a JSON string: `{"audio_url": "..."}` where the URL can be:
+  - An HTTP/HTTPS URL (e.g., `https://example.com/audio.mp3`)
+  - A base64 data URL (e.g., `data:audio/mp3;base64,...`)
+- `--model-class-name WanS2VPipeline` is required on the server to select the
+  S2V pipeline (distinct from the T2V/I2V pipelines).
+- `--cache-backend cache_dit` enables DiT caching for ~2x speedup on cached steps.
+- Audio is muxed into the output MP4 automatically.
+- Always pass `fps=16` to ensure correct video/audio alignment.
@@ -0,0 +1,117 @@
+# vLLM-Omni Failure Mode Injection Scenarios and Expected Behavior Matrix
+
+## Background and Goals
+
+The online serving pipeline of vLLM-Omni can be abstracted as: `Serve` (service entry process) -> `Engine` (model orchestration unit) -> `Worker` (execution process) -> `GPU` (compute resource).
+
+- `Serve` provides external APIs, accepts requests, and manages lifecycle.
+- `Engine` handles request scheduling and inference orchestration; the number and organization of engines can vary by model type.
+- `Worker` performs concrete model computation, usually bound to underlying GPU resources.
+- `GPU` is the resource layer that ultimately carries inference compute and memory usage.
+
+<div markdown style="text-align: center">
+
+```mermaid
+%%{init: {'flowchart': {'useMaxWidth': false}}}%%
+flowchart TD
+    R(request)--> A(serve)
+    A --> B(engine)
+    A --> C(engine)
+    A --> D(engine)
+    B --> E(worker)
+    C --> F(worker)
+    D --> I(worker)
+    E --> G(GPU)
+    F --> H(GPU)
+    I --> L(GPU)
+```
+
+</div>
+
+The diagram above shows an abstract relationship. Omni models and Diffusion models share the same hierarchy; differences mainly come from engine decomposition, worker count, and scheduling strategy.
+
+In real customer environments, the system will inevitably encounter accidental operations or environmental disturbances, such as unintended process termination, container/node jitter, or short-term resource contention. From the system perspective, many of these events map to two fault categories: "process receives abnormal signals" or "GPU memory is squeezed." For example, pressing `Ctrl+C` essentially sends a `SIGINT` signal to the `serve` process; running `kill -15 <pid>` maps to graceful termination via `SIGTERM`; running `kill -9 <pid>` maps to forced termination via `SIGKILL`; when `docker kill` is executed or a Pod is deleted, the container runtime usually sends `SIGKILL`.
+
+This document currently covers two core failure modes, and the system is expected to show predictable and observable behavior under each mode:
+
+- **Failure Mode 1: Process receives abnormal signals (`SIGINT` / `SIGTERM` / `SIGKILL`)**  
+  Focus on three aspects: whether processes exit as expected and complete cleanup, whether requests fail fast or connections are interrupted, and whether GPU resources are eventually fully released with no residue.
+- **Failure Mode 2: OOM (occupy all free GPU memory via an extra process)**  
+  Focus on two aspects: whether service health degrades to the expected state (for example, `503`), and whether requests fail within an acceptable time without hanging.
+
+More failure modes (for example, network jitter and network interruption) will be added in future iterations to improve end-to-end reliability validation coverage.
+
+> Note: independent fault injection on the `engine` component is not covered in the current version and will be added gradually in future versions.
+
+## Fault Injection Scenario Matrix
+
+| Scenario | Fault Type | System Behavior | Current Status |
+|------|----------------------------|----------|----------|
+| No load | Send `SIGKILL` to Worker | Worker process is killed immediately; main service detects child-process loss and turns unavailable; API enters stable 5xx |  |
+| No load | Send `SIGTERM` to Worker | Worker exits after receiving termination signal; main service is marked unavailable; API enters stable 5xx |  |
+| No load | Send `SIGKILL` to serve main process | serve main process exits instantly; request connections are interrupted; related child processes are cleaned up with no residue; GPU memory is released quickly | [#3725](https://github.com/vllm-project/vllm-omni/issues/3725) <br>[#43060](https://github.com/vllm-project/vllm/issues/43060) |
+| No load | Send `SIGTERM` to serve main process | serve enters graceful shutdown and stops serving; then exits and completes cleanup; GPU memory is released |  |
+| No load | Send `SIGINT` to serve main process (equivalent to `Ctrl+C`) | Triggers serve shutdown path; service stops responding and becomes unavailable; related child processes exit and resources are released |  |
+| No load | Send `SIGKILL` to all related processes | All related processes terminate immediately; service becomes unavailable at once; no residual processes remain; GPU memory is released quickly |  |
+| No load | Send `SIGTERM` to all related processes | All processes enter exit flow and complete shutdown; service becomes unavailable; resource release is completed |  |
+| Under load | Send `SIGKILL` to Worker | In-flight requests are hard interrupted (5xx/connection drop); main service becomes unavailable | [#3683](https://github.com/vllm-project/vllm-omni/issues/3683) |
+| Under load | Send `SIGTERM` to Worker | In-flight requests are canceled or fail fast; main service becomes unavailable |  |
+| Under load | Send `SIGKILL` to serve main process | serve is hard-killed and current connections are interrupted; in-flight requests fail; after cleanup there are no residual processes and GPU memory is released | [#3683](https://github.com/vllm-project/vllm-omni/issues/3683) |
+| Under load | Send `SIGTERM` to serve main process | serve stops accepting new requests and executes shutdown flow; in-flight requests fail; no residue remains and GPU memory is released | [#3683](https://github.com/vllm-project/vllm-omni/issues/3683) |
+| Under load | Send `SIGINT` to serve main process (equivalent to `Ctrl+C`) | `Ctrl+C`-style serve shutdown; in-flight requests fail (5xx/connection interruption); service unavailable; after exit there is no residue and GPU memory is released | [#3683](https://github.com/vllm-project/vllm-omni/issues/3683) |
+| Under load | Send `SIGKILL` to all related processes | All processes terminate instantly; all in-flight requests fail; service becomes unavailable immediately; GPU memory is released quickly |  |
+| Under load | Send `SIGTERM` to all related processes | All processes exit gracefully; in-flight requests fail; service unavailable; GPU memory is released | [#3683](https://github.com/vllm-project/vllm-omni/issues/3683) |
+| OOM | Occupy all free GPU memory via an extra process | After OOM injection process starts, GPU memory is continuously saturated; service enters unavailable/degraded state and health check drops to 503; different request types (chat/speech, etc.) fail fast within a fixed time and return 500 (no hanging) | [#4285](https://github.com/vllm-project/vllm-omni/issues/4285) |
+
+## Source of Conclusions
+
+The behaviors and conclusions above are summarized from current fault injection validation results on `Qwen3-Omni` and `Wan2.2`.
+
+### Example: `SIGTERM` fault injection log (Qwen3-Omni)
+
+The following full log is from a real run where a `SIGTERM` signal was sent to the serve root process.
+
+```text
+Launching OmniServer with: /workspace/.venv/bin/python3 -m vllm_omni.entrypoints.cli.main serve Qwen/Qwen3-Omni-30B-A3B-Instruct --host 127.0.0.1 --port 60675 --omni --async-chunk --stage-init-timeout 600 --init-timeout 900 --log-stats --stage-configs-path vllm-omni/vllm_omni/deploy/qwen3_omni_moe.yaml
+Server ready on 127.0.0.1:60675 (OmniServer startup took 206.023s)
+OmniServer started successfully
+[reliability][process-kill] current_server_proc pid=556855 name=python3 cmdline=/workspace/.venv/bin/python3 -m vllm_omni.entrypoints.cli.main serve Qwen/Qwen3-Omni-30B-A3B-Instruct --host 127.0.0.1 --port 60675 --omni --async-chunk --stage-init-timeout 600 --init-timeout 900 --log-stats --stage-configs-path vllm-omni/vllm_omni/deploy/qwen3_omni_moe.yaml
+[reliability][process-kill] current_server_proc pid=557118 name=python3 cmdline=/workspace/.venv/bin/python3 -c from multiprocessing.resource_tracker import main;main(57)
+[reliability][process-kill] current_server_proc pid=557119 name=VLLM::StageEngineCoreProc_noid_replica0_DP0 cmdline=VLLM::StageEngineCoreProc_noid_replica0_DP0
+[reliability][process-kill] current_server_proc pid=557122 name=VLLM::StageEngineCoreProc_noid_replica0_DP0 cmdline=VLLM::StageEngineCoreProc_noid_replica0_DP0
+[reliability][process-kill] current_server_proc pid=558201 name=VLLM::StageEngineCoreProc_noid_replica0_DP0 cmdline=VLLM::StageEngineCoreProc_noid_replica0_DP0
+
+[reliability][process-kill] root-kill pid=556855 name=python3 signal=SIGTERM cmdline=/workspace/.venv/bin/python3 -m vllm_omni.entrypoints.cli.main serve Qwen/Qwen3-Omni-30B-A3B-Instruct --host 127.0.0.1 --port 60675 --omni --async-chunk --stage-init-timeout 600 --init-timeout 900 --log-stats --stage-configs-path vllm-omni/vllm_omni/deploy/qwen3_omni_moe.yaml
+[0;36m(APIServer pid=556855)[0;0m INFO 05-19 12:24:51 [omni_base.py:463] [AsyncOmni] Shutting down
+[0;36m(APIServer pid=556855)[0;0m INFO 05-19 12:24:51 [async_omni_engine.py:2397] [AsyncOmniEngine] Shutting down Orchestrator
+[0;36m(APIServer pid=556855)[0;0m INFO 05-19 12:24:51 [orchestrator.py:351] [Orchestrator] Received shutdown signal
+[0;36m(APIServer pid=556855)[0;0m INFO 05-19 12:24:51 [orchestrator.py:1359] [Orchestrator] Shutting down all 3 client(s)
+
+[0;36m(StageEngineCoreProc pid=557119)[0;0m [32mINFO[0m [90m05-19 12:24:51[0m [90m[core.py:1242][0m Shutdown initiated (timeout=0)
+[0;36m(StageEngineCoreProc pid=557119)[0;0m [32mINFO[0m [90m05-19 12:24:51[0m [90m[core.py:1265][0m Shutdown complete
+[0;36m(APIServer pid=556855)[0;0m INFO:     Shutting down
+[0;36m(APIServer pid=556855)[0;0m INFO:     Waiting for application shutdown.
+[0;36m(APIServer pid=556855)[0;0m INFO:     Application shutdown complete.
+[0;36m(APIServer pid=556855)[0;0m INFO:     Finished server process [556855]
+[0;36m(APIServer pid=556855)[0;0m INFO 05-19 12:24:51 [omni_base.py:463] [AsyncOmni] Shutting down
+[rank0]:[W519 12:24:53.427894833 ProcessGroupNCCL.cpp:1575] Warning: WARNING: destroy_process_group() was not called before program exit, which can leak resources. For more info, please see https://pytorch.org/docs/stable/distributed.html#shutdown (function operator())
+[0;36m(APIServer pid=556855)[0;0m INFO 05-19 12:24:56 [stage_pool.py:731] [StagePool] Stage 0 replica 0 shut down
+[0;36m(StageEngineCoreProc pid=557122)[0;0m [32mINFO[0m [90m05-19 12:24:56[0m [90m[core.py:1242][0m Shutdown initiated (timeout=0)
+[0;36m(StageEngineCoreProc pid=557122)[0;0m [32mINFO[0m [90m05-19 12:24:56[0m [90m[core.py:1265][0m Shutdown complete
+[rank0]:[W519 12:24:57.569425065 ProcessGroupNCCL.cpp:1575] Warning: WARNING: destroy_process_group() was not called before program exit, which can leak resources. For more info, please see https://pytorch.org/docs/stable/distributed.html#shutdown (function operator())
+[0;36m(APIServer pid=556855)[0;0m INFO 05-19 12:25:00 [stage_pool.py:731] [StagePool] Stage 1 replica 0 shut down
+[0;36m(StageEngineCoreProc pid=558201)[0;0m [32mINFO[0m [90m05-19 12:25:00[0m [90m[core.py:1242][0m Shutdown initiated (timeout=0)
+[0;36m(StageEngineCoreProc pid=558201)[0;0m [32mINFO[0m [90m05-19 12:25:00[0m [90m[core.py:1265][0m Shutdown complete
+[0;36m(APIServer pid=556855)[0;0m WARNING 05-19 12:25:01 [async_omni_engine.py:2406] [AsyncOmniEngine] Orchestrator thread did not exit in time
+[rank0]:[W519 12:25:01.867242016 ProcessGroupNCCL.cpp:1575] Warning: WARNING: destroy_process_group() was not called before program exit, which can leak resources. For more info, please see https://pytorch.org/docs/stable/distributed.html#shutdown (function operator())
+```
+
+### How this maps to the matrix
+
+- This log corresponds to the `Send SIGTERM to serve main process` scenario in the matrix.
+- The `root-kill ... signal=SIGTERM` line confirms the fault injection action is applied as expected.
+- The sequence from `Shutting down Orchestrator` to stage-level `Shutdown complete` and `Stage ... shut down` matches the expected graceful shutdown path and child-process cleanup behavior described in the matrix.
+
+## Related Documentation
+
+- For practical CI debugging procedures and triage guidance, see [CI Failures](../contributing/ci/failures.md).
@@ -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,115 @@
+# Speech-To-Video
+
+This example demonstrates how to deploy the Wan2.2 speech-to-video (S2V) model for online video generation using vLLM-Omni.
+
+## Supported Models
+
+| Model | Model ID |
+|-------|----------|
+| Wan2.2 S2V (14B) | `Wan-AI/Wan2.2-S2V-14B` |
+
+## Start Server
+
+### Basic Start
+
+```bash
+VLLM_WORKER_MULTIPROC_METHOD=spawn \
+vllm serve Wan-AI/Wan2.2-S2V-14B --omni \
+  --model-class-name WanS2VPipeline \
+  --tensor-parallel-size 2 \
+  --flow-shift 3.0 \
+  --vae-use-slicing --vae-use-tiling \
+  --cache-backend cache_dit \
+  --port 8091
+```
+
+### Start with Script
+
+```bash
+bash run_server.sh
+```
+
+The script allows overriding:
+- `MODEL` (default: `Wan-AI/Wan2.2-S2V-14B`)
+- `PORT` (default: `8091`)
+- `FLOW_SHIFT` (default: `3.0`)
+- `TP` (default: `2`)
+- `CACHE_BACKEND` (default: `cache_dit`)
+
+### Start with 4 GPUs
+
+```bash
+TP=4 bash run_server.sh
+```
+
+## Sync API (Recommended for Testing)
+
+`POST /v1/videos/sync` blocks until generation completes and returns the raw
+video bytes (`video/mp4`) directly in the response body.
+
+```bash
+bash run_curl_speech_to_video.sh
+```
+
+## Async Job API
+
+`POST /v1/videos` creates an asynchronous job and returns immediately.
+
+```bash
+no_proxy=127.0.0.1 \
+create_response=$(curl -sS -X POST http://127.0.0.1:8091/v1/videos \
+  -H "Accept: application/json" \
+  -F "prompt=A person singing" \
+  -F 'image_reference={"image_url": "https://raw.githubusercontent.com/Wan-Video/Wan2.2/main/examples/Five%20Hundred%20Miles.png"}' \
+  -F 'audio_reference={"audio_url": "https://raw.githubusercontent.com/Wan-Video/Wan2.2/main/examples/Five%20Hundred%20Miles.MP3"}' \
+  -F "width=832" -F "height=480" \
+  -F "num_inference_steps=40" \
+  -F "guidance_scale=4.5" \
+  -F "fps=16")
+
+video_id=$(echo "$create_response" | jq -r '.id')
+
+# Poll until complete
+while true; do
+  status=$(curl -s "http://127.0.0.1:8091/v1/videos/${video_id}" | jq -r '.status')
+  if [ "$status" = "completed" ]; then
+    break
+  fi
+  if [ "$status" = "failed" ]; then
+    echo "Video generation failed"
+    exit 1
+  fi
+  sleep 2
+done
+
+# Download
+curl -L "http://127.0.0.1:8091/v1/videos/${video_id}/content" -o s2v_output.mp4
+```
+
+## Request Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `prompt` | str | - | Text description of the desired video |
+| `image_reference` | JSON str | - | Image reference: `{"image_url": "..."}` — supports HTTP(s) URLs or base64 data URLs |
+| `audio_reference` | JSON str | - | Audio reference: `{"audio_url": "..."}` — supports HTTP(s) URLs or base64 data URLs |
+| `width` | int | None | Video width in pixels |
+| `height` | int | None | Video height in pixels |
+| `fps` | int | None | Output video frame rate |
+| `num_inference_steps` | int | None | Number of denoising steps (40 recommended) |
+| `guidance_scale` | float | None | CFG guidance scale |
+| `seed` | int | None | Random seed for reproducibility |
+
+## Notes
+
+- S2V requires both a reference image (`image_reference`) and an audio reference
+  (`audio_reference`). The generated video will show a person matching the reference
+  image with lip movements synchronized to the audio.
+- `audio_reference` accepts a JSON string: `{"audio_url": "..."}` where the URL can be:
+  - An HTTP/HTTPS URL (e.g., `https://example.com/audio.mp3`)
+  - A base64 data URL (e.g., `data:audio/mp3;base64,...`)
+- `--model-class-name WanS2VPipeline` is required on the server to select the
+  S2V pipeline (distinct from the T2V/I2V pipelines).
+- `--cache-backend cache_dit` enables DiT caching for ~2x speedup on cached steps.
+- Audio is muxed into the output MP4 automatically.
+- Always pass `fps=16` to ensure correct video/audio alignment.
@@ -0,0 +1,44 @@
+#!/bin/bash
+# Wan2.2 S2V (speech-to-video) curl example using the sync video API.
+
+set -euo pipefail
+
+BASE_URL="${BASE_URL:-http://localhost:8091}"
+OUTPUT_PATH="${OUTPUT_PATH:-s2v_480p_serve.mp4}"
+IMAGE_URL="${IMAGE_URL:-https://raw.githubusercontent.com/Wan-Video/Wan2.2/main/examples/Five%20Hundred%20Miles.png}"
+AUDIO_URL="${AUDIO_URL:-https://raw.githubusercontent.com/Wan-Video/Wan2.2/main/examples/Five%20Hundred%20Miles.MP3}"
+PROMPT="${PROMPT:-A person singing}"
+WIDTH="${WIDTH:-832}"
+HEIGHT="${HEIGHT:-480}"
+NUM_INFERENCE_STEPS="${NUM_INFERENCE_STEPS:-40}"
+GUIDANCE_SCALE="${GUIDANCE_SCALE:-4.5}"
+FPS="${FPS:-16}"
+
+echo "Sending S2V request..."
+echo "  Image URL: $IMAGE_URL"
+echo "  Audio URL: $AUDIO_URL"
+echo "  Prompt: $PROMPT"
+echo "  Resolution: ${WIDTH}x${HEIGHT}"
+echo "  Steps: $NUM_INFERENCE_STEPS"
+echo "  FPS: $FPS"
+
+IMAGE_REF_JSON="{\"image_url\": \"${IMAGE_URL}\"}"
+AUDIO_REF_JSON="{\"audio_url\": \"${AUDIO_URL}\"}"
+
+no_proxy=127.0.0.1 \
+curl -X POST "${BASE_URL}/v1/videos/sync" \
+  -F "prompt=${PROMPT}" \
+  -F "image_reference=${IMAGE_REF_JSON}" \
+  -F "audio_reference=${AUDIO_REF_JSON}" \
+  -F "width=${WIDTH}" -F "height=${HEIGHT}" \
+  -F "num_inference_steps=${NUM_INFERENCE_STEPS}" \
+  -F "guidance_scale=${GUIDANCE_SCALE}" \
+  -F "fps=${FPS}" \
+  --output "${OUTPUT_PATH}"
+
+if [ -f "$OUTPUT_PATH" ] && [ -s "$OUTPUT_PATH" ]; then
+    echo "Saved video to ${OUTPUT_PATH} ($(du -h "$OUTPUT_PATH" | cut -f1))"
+else
+    echo "ERROR: Output file is empty or missing"
+    exit 1
+fi
@@ -0,0 +1,29 @@
+#!/bin/bash
+# Wan2.2 S2V (speech-to-video) online serving startup script
+
+MODEL="${MODEL:-Wan-AI/Wan2.2-S2V-14B}"
+PORT="${PORT:-8091}"
+FLOW_SHIFT="${FLOW_SHIFT:-3.0}"
+TP="${TP:-2}"
+CACHE_BACKEND="${CACHE_BACKEND:-cache_dit}"
+
+echo "Starting Wan2.2 S2V server..."
+echo "Model: $MODEL"
+echo "Port: $PORT"
+echo "Flow shift: $FLOW_SHIFT"
+echo "Tensor parallel size: $TP"
+echo "Cache backend: $CACHE_BACKEND"
+
+CACHE_BACKEND_FLAG=""
+if [ "$CACHE_BACKEND" != "none" ]; then
+    CACHE_BACKEND_FLAG="--cache-backend $CACHE_BACKEND"
+fi
+
+VLLM_WORKER_MULTIPROC_METHOD=spawn \
+vllm serve "$MODEL" --omni \
+    --model-class-name WanS2VPipeline \
+    --tensor-parallel-size "$TP" \
+    --flow-shift "$FLOW_SHIFT" \
+    --vae-use-slicing --vae-use-tiling \
+    --port "$PORT" \
+    $CACHE_BACKEND_FLAG
@@ -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!"
@@ -408,7 +408,10 @@ def create_app(api_base: str):
                await resp.aclose()
                await client.aclose()

-        return StreamingResponse(relay(), media_type="application/octet-stream")
+        # Use audio/pcm to bypass browsers (e.g. Chromium) MIME type sniffing, which
+        # buffers application/octet-stream responses before delivering chunks to
+        # fetch().body.getReader(), inflating TTFP to total_ms.
+        return StreamingResponse(relay(), media_type="audio/pcm")

    # ── Gradio UI ────────────────────────────────────────────────
    voices = fetch_voices(api_base)
@@ -66,7 +66,8 @@ dev = [
    "voxcpm>=2.0.2", # for VoxCPM2 path-locating in tests
    "torchdiffeq>=0.2.5", # for covo_audio Test
    "FlagEmbedding", # for hunyuan_image3 tests
-    "kernels<0.15.0" # for accuracy tests comparison with diffusers output (0.15.0 incompatible with diffusers 0.38.0)
+    "kernels<0.15.0", # for accuracy tests comparison with diffusers output (0.15.0 incompatible with diffusers 0.38.0)
+    "s3tokenizer>=0.3.0" # for tts test
 ]

 demo = [
@@ -280,6 +281,7 @@ extend-ignore-identifiers-re = [
    ".*[Oo]no_[Aa]nna.*",
    ".*cann.*",
    ".*[Nn][Dd]3.*",
+    ".*updator.*",
    "byt5.*",
 ]
 extend-ignore-re = [
@@ -292,5 +294,9 @@ fullset = "fullset"
 Vai = "Vai"
 tockens = "tockens"
 CANN = "CANN"
+# Step-Audio reasoning parser streams `</think>` tokens split across chunks
+# (e.g. "</" + "thi" + "nk>"). The fragment "thi" appears literally in its
+# parser logic and tests and is not a typo of "the"/"this".
+thi = "thi"
 ND = "ND"
 byt5 = "byt5"
@@ -52,6 +52,7 @@ vllm-omni serve bosonai/higgs-audio-v3-tts-4b \

 The default deploy config `vllm_omni/deploy/higgs_multimodal_qwen3.yaml` is
 loaded automatically by model registry (HF `model_type=higgs_multimodal_qwen3`).
+It matches the high-throughput profile.

 **Offline batch inference:**

@@ -103,8 +104,28 @@ python examples/online_serving/text_to_speech/higgs_audio_v3/batch_speech_client
 - Voice cloning: `ref_audio` accepts WAV/FLAC/MP3; `ref_text` is optional but improves fidelity.
 - Deploy config: `vllm_omni/deploy/higgs_multimodal_qwen3.yaml` (auto-discovered from `model_type`).
  - `max_num_seqs=16` for both stages.
-  - `enforce_eager=true` for both stages (CUDA graph available for stage 0 but no throughput gain at batch>1).
+  - Stage 0 and Stage 1 default to the same device (`0`) for single-GPU serving.
+  - Stage 0 intentionally keeps `enforce_eager=true`. This preserves the Higgs-specific local MLP CUDA graph path, which is the current high-throughput default.
+  - Stage 1 remains `enforce_eager=true` for the codec decoder.
+- Deploy profiles:
+  - High throughput: `vllm_omni/deploy/higgs_multimodal_qwen3_high_throughput.yaml`.
+    Use this for medium/high concurrency. The auto-discovered
+    `higgs_multimodal_qwen3.yaml` is kept as a compatibility/default alias for
+    this profile.
+  - Low latency: `vllm_omni/deploy/higgs_multimodal_qwen3_low_latency.yaml`.
+    Use this for low-concurrency serving (for example c1-c4). It sets Stage 0
+    `enforce_eager=false` and enables vLLM `FULL_DECODE_ONLY` CUDA graph through
+    YAML `compilation_config`; no environment variable is required.
+  - Profile details: `vllm_omni/deploy/README_higgs_audio_v3.md`.
+- Performance note:
+  - Do not switch the auto-discovered default Stage 0 profile to vLLM
+    `FULL_DECODE_ONLY` CUDA graph without an end-to-end throughput and
+    audio-quality revalidation. On the H20 SeedTTS c16/full-dataset benchmark,
+    the eager Stage 0 path with Higgs local MLP CUDA graph reproduced
+    ~35 audio_s/s (`1088/1088` OK per run, three runs around 134-136s). A
+    separate H20 low-concurrency smoke after the FULL_DECODE audio-feedback
+    capture fix showed lower c1/c4 latency, so the FULL_DECODE profile is kept
+    as an explicit low-latency option rather than the throughput default.
 - Known limitations:
-  - CUDA graph for stage 0 helps batch=1 (-13% RTF) but not batch>1 due to model-owned sampler running outside graph.
  - Stage 1 (code2wav) must use `enforce_eager=true` (`@torch.inference_mode` incompatible with graph capture).
-  - Async streaming (chunk-based) not yet implemented; sync-only pipeline.
+  - Stage 0 full-decode CUDA graph is experimental; sampler, delay-state updates, staging, and request postprocess remain outside the graph.
@@ -54,7 +54,6 @@ recipes/
 | [`StabilityAI/Stable-Diffusion-3.5.md`](./StabilityAI/Stable-Diffusion-3.5.md) | Text-to-image serving (SD 3.5-medium and SD 3.5-large) | 1x RTX A6000 48GB |
 | [`zai-org/GLM-TTS.md`](./zai-org/GLM-TTS.md) | Online serving for Chinese/English zero-shot voice-cloned TTS | 1x A40 48GB |
 | [`GLM/GLM-Image.md`](./GLM/GLM-Image.md) | Online serving for image generation | 1x A800 80GB / 2x A800 80GB |
-|

 Within a single recipe file, include different hardware support sections such
 as `GPU`, `ROCm`, and `NPU`, and add concrete tested configurations like
@@ -171,6 +171,34 @@ python examples/offline_inference/speech_to_video/speech_to_video.py \
 ```


+### Online Serving
+
+#### Server
+
+```bash
+vllm serve Wan-AI/Wan2.2-S2V-14B --omni \
+  --model-class-name WanS2VPipeline \
+  --tensor-parallel-size 2 \
+  --flow-shift 3.0 \
+  --vae-use-slicing --vae-use-tiling \
+  --port 8091
+```
+
+#### Client
+
+```bash
+no_proxy=127.0.0.1 \
+curl -X POST http://127.0.0.1:8091/v1/videos/sync \
+  -F "prompt=A person singing" \
+  -F 'image_reference={"image_url": "https://raw.githubusercontent.com/Wan-Video/Wan2.2/main/examples/Five%20Hundred%20Miles.png"}' \
+  -F 'audio_reference={"audio_url": "https://raw.githubusercontent.com/Wan-Video/Wan2.2/main/examples/Five%20Hundred%20Miles.MP3"}' \
+  -F "width=832" -F "height=480" \
+  -F "num_inference_steps=40" \
+  -F "guidance_scale=4.5" \
+  -F "fps=16" \
+  --output s2v_480p_serve.mp4
+```
+
 #### Verification

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

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

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

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

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

                async def close(self):
                    return None

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

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

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

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

-    assert result["completed"] == sent_requests == num_prompts
-    assert requested_urls
-    assert all(url == expected_url for url in requested_urls), f"Unexpected target URLs: {requested_urls}"
+    assert result["completed"] == sent_bench_requests == num_prompts, (
+        f"completed={result['completed']}, bench_requests={sent_bench_requests}, all_requested_urls={requested_urls}"
+    )
+    assert bench_requests
+    assert all(url == expected_url for url in bench_requests), f"Unexpected target URLs: {bench_requests}"
@@ -0,0 +1,74 @@
+[
+  {
+    "test_name": "test_higgs_audio_v3",
+    "server_params": {
+      "model": "bosonai/higgs-audio-v3-tts-4b",
+      "stage_overrides": {
+        "0": {
+          "gpu_memory_utilization": 0.3
+        },
+        "1": {
+          "gpu_memory_utilization": 0.2
+        }
+      },
+      "extra_cli_args": [
+        "--deploy-config",
+        "vllm_omni/deploy/higgs_multimodal_qwen3.yaml",
+        "--trust-remote-code"
+      ]
+    },
+    "benchmark_params": [
+      {
+        "task": "default_voice",
+        "eval_phase": "latency",
+        "dataset_name": "seed-tts-text",
+        "backend": "openai-audio-speech",
+        "endpoint": "/v1/audio/speech",
+        "dataset_path": "benchmarks/build_dataset/seed_tts_smoke",
+        "num_prompts": [
+          20
+        ],
+        "max_concurrency": [
+          1
+        ],
+        "seed_tts_locale": "en",
+        "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
+        "baseline": {
+          "median_audio_ttfp_ms": [
+            3000
+          ],
+          "median_audio_rtf": [
+            0.5
+          ]
+        }
+      },
+      {
+        "task": "default_voice",
+        "eval_phase": "throughput",
+        "dataset_name": "seed-tts-text",
+        "backend": "openai-audio-speech",
+        "endpoint": "/v1/audio/speech",
+        "dataset_path": "benchmarks/build_dataset/seed_tts_smoke",
+        "num_prompts": [
+          500
+        ],
+        "max_concurrency": [
+          16
+        ],
+        "seed_tts_locale": "en",
+        "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
+        "baseline": {
+          "median_audio_ttfp_ms": [
+            5000
+          ],
+          "median_audio_rtf": [
+            1.2
+          ],
+          "audio_throughput": [
+            30
+          ]
+        }
+      }
+    ]
+  }
+]
@@ -198,19 +198,78 @@ def test_video_chat_stream_invalid_requests(
 # WS /v1/realtime
 # ─────────────────────────────────────────────────────────────────────────────

+# Resolved in ``test_realtime_invalid_requests`` (skip ``session.created`` handshake).
+_REALTIME_WS_MODEL_MISMATCH = object()
+_REALTIME_WS_INVALID_AUDIO_APPEND = object()
+

@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.parametrize(
+    "send_frames_spec, ws_error_code, err_message",
+    [
+        pytest.param("{not-json", "invalid_json", "Invalid JSON", id="invalid_json"),
+        pytest.param(
+            json.dumps({"type": "session.update"}),
+            "invalid_event",
+            ("model", "Missing required field"),
+            id="session_update_missing_model",
+        ),
+        pytest.param(
+            _REALTIME_WS_MODEL_MISMATCH,
+            "model_not_found",
+            ("model", "does not exist"),
+            id="session_update_model_mismatch",
+        ),
+        pytest.param(
+            json.dumps({"type": "input_audio_buffer.commit", "final": False}),
+            "model_not_validated",
+            ("session.update", "validate the model"),
+            id="commit_without_model_validation",
+        ),
+        pytest.param(
+            json.dumps({"type": "unknown.event"}),
+            "unknown_event",
+            "Unknown event type",
+            id="unknown_event_type",
+        ),
+        pytest.param(
+            _REALTIME_WS_INVALID_AUDIO_APPEND,
+            "invalid_audio",
+            "Invalid audio data",
+            id="invalid_audio_append",
+        ),
+    ],
+)
@pytest.mark.parametrize("omni_server", _QWEN3_OMNI_SERVER, indirect=True)
-def test_realtime_rejected_when_async_chunk_enabled(
+def test_realtime_invalid_requests(
    omni_server: OmniServer,
    openai_client: OpenAIClientHandler,
+    send_frames_spec: Any,
+    ws_error_code: str,
+    err_message: str | tuple[str, ...],
 ) -> None:
+    if send_frames_spec is _REALTIME_WS_MODEL_MISMATCH:
+        send_frames = json.dumps(
+            {
+                "type": "session.update",
+                "model": "this-model-is-not-served-on-this-instance",
+            }
+        )
+    elif send_frames_spec is _REALTIME_WS_INVALID_AUDIO_APPEND:
+        send_frames = [
+            json.dumps({"type": "session.update", "model": omni_server.model}),
+            json.dumps({"type": "input_audio_buffer.append", "audio": "not-valid-base64!!!"}),
+        ]
+    else:
+        send_frames = send_frames_spec
    openai_client.send_realtime_ws_request(
        {
+            "send_frames": send_frames,
            "timeout": 120,
            "ws_max_size": None,
+            "ws_skip_types": ["session.created"],
            "ws_json_type": "error",
-            "ws_error_code": "unsupported",
-            "err_message": "async_chunk",
+            "ws_error_code": ws_error_code,
+            "err_message": err_message,
        }
    )
@@ -578,12 +578,27 @@ def test_rmsnorm_fused_native_parity_2d_input():


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

    hidden_size = 64
    eps = 1e-6

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

-        torch.testing.assert_close(out_fused, out_native, atol=1e-3, rtol=1e-3)
+        torch.testing.assert_close(out_fused, out_native, atol=atol_map[dtype], rtol=rtol_map[dtype])
@@ -131,6 +131,38 @@ def test_forward_returns_video_prediction(monkeypatch: pytest.MonkeyPatch) -> No
    assert tuple(output.shape) == (1, 2, 1, 2, 2)


+def test_forward_gathers_gen_tokens_before_unpatchify(monkeypatch: pytest.MonkeyPatch) -> None:
+    from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
+
+    class RecordingGather(nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.calls = 0
+
+        def forward(self, hidden_gen: torch.Tensor) -> torch.Tensor:
+            self.calls += 1
+            return hidden_gen
+
+    monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (1, 0, None))
+    model = transformer_cosmos3.Cosmos3VFMTransformer(
+        SimpleNamespace(tf_model_config=_tiny_cosmos3_config(), dtype=torch.float32)
+    )
+    gather = RecordingGather()
+    model.gen_sp_gather = gather
+
+    output = model(
+        hidden_states=torch.zeros(1, 2, 1, 2, 2),
+        timestep=torch.tensor([1.0]),
+        text_ids=torch.tensor([[1, 2]], dtype=torch.long),
+        text_mask=torch.ones(1, 2, dtype=torch.long),
+        video_shape=(1, 2, 2),
+        fps=24.0,
+    )
+
+    assert gather.calls == 1
+    assert tuple(output.shape) == (1, 2, 1, 2, 2)
+
+
 def test_sound_and_action_modules_follow_config() -> None:
    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer

@@ -0,0 +1,728 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for multi-resolution support in HunyuanImage3.
+
+Covers three areas:
+
+1. ``ResolutionGroup._calc_by_step`` — training-bucket generation,
+   verifying that each resolution has correct width×height and that
+   indices match expected ratio ordering.
+
+2. ``HUNYUAN_IMAGE3_EXTRA_RESOLUTIONS`` — extra resolution matching
+   via ``Resolution.append`` and new-entry insertion in
+   ``ResolutionGroup.__init__``.
+
+3. ``get_cached_resolution_group`` and ``ar2diffusion`` — caching
+   behaviour and correct ``ResolutionGroup`` construction for the
+   AR→DiT stage transition.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Any
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_image3_transformer import (
+    HUNYUAN_IMAGE3_EXTRA_RESOLUTIONS,
+    Resolution,
+    ResolutionGroup,
+    get_cached_resolution_group,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+# ---------------------------------------------------------------------------
+# helpers
+# ---------------------------------------------------------------------------
+
+
+def _make_reso(desc: str | tuple[int, int]) -> Resolution:
+    """Build a ``Resolution`` from a ``"HxW"`` string or ``(h, w)`` tuple."""
+    if isinstance(desc, str):
+        h, w = desc.split("x")
+        return Resolution(int(h), int(w))
+    return Resolution(desc[0], desc[1])
+
+
+def _reso_set(rg: ResolutionGroup) -> set[tuple[int, int]]:
+    """Return all ``(h, w)`` tuples reachable through a ``ResolutionGroup``.
+
+    For entries that carry ``extra_res`` this includes every alternative
+    resolution in that set.
+    """
+    out: set[tuple[int, int]] = set()
+    for r in rg.data:
+        out.add((r.h, r.w))
+        out |= r.extra_res
+    return out
+
+
+# ---------------------------------------------------------------------------
+# 1.  _calc_by_step — training-bucket resolution parsing
+# ---------------------------------------------------------------------------
+
+_BASE = 1024
+_STEP = _BASE // 16  # 64
+_ALIGN = 16
+_MIN_SIDE = _BASE // 2  # 512
+_MAX_SIDE = _BASE * 2  # 2048
+
+
+def _build_plain_group(
+    base_size: int = _BASE,
+    step: int | None = None,
+    align: int = _ALIGN,
+) -> ResolutionGroup:
+    """Return a ``ResolutionGroup`` with *no* extra resolutions."""
+    return ResolutionGroup(base_size=base_size, step=step, align=align, extra_resolutions=None)
+
+
+class TestCalcByStep:
+    """Verify ``_calc_by_step`` generates the expected training buckets."""
+
+    # -- count & range -------------------------------------------------------
+
+    def test_count_for_base_1024(self):
+        """33 buckets: 1 centre + 2×16 from the stepping loops."""
+        rg = _build_plain_group()
+        # 1 centre + 16 portrait + 16 landscape
+        assert len(rg.data) == 33, f"expected 33 training buckets, got {len(rg.data)}"
+
+    def test_all_sizes_within_bounds(self):
+        rg = _build_plain_group()
+        for r in rg.data:
+            assert _MIN_SIDE <= r.h <= _MAX_SIDE, f"height {r.h} outside [{_MIN_SIDE}, {_MAX_SIDE}]"
+            assert _MIN_SIDE <= r.w <= _MAX_SIDE, f"width {r.w} outside [{_MIN_SIDE}, {_MAX_SIDE}]"
+
+    def test_all_sizes_aligned(self):
+        rg = _build_plain_group()
+        for r in rg.data:
+            assert r.h % _ALIGN == 0, f"height {r.h} not {_ALIGN}-aligned"
+            assert r.w % _ALIGN == 0, f"width {r.w} not {_ALIGN}-aligned"
+
+    # -- ratio ordering ------------------------------------------------------
+
+    def test_sorted_by_ratio_ascending(self):
+        rg = _build_plain_group()
+        ratios = [r.ratio for r in rg.data]
+        for i in range(1, len(ratios)):
+            assert ratios[i - 1] <= ratios[i], f"ratio[{i - 1}]={ratios[i - 1]:.4f} > ratio[{i}]={ratios[i]:.4f}"
+
+    def test_ratio_index_matches_numpy_array(self):
+        rg = _build_plain_group()
+        for i, r in enumerate(rg.data):
+            assert math.isclose(rg.ratio[i], r.ratio, rel_tol=1e-9), (
+                f"rg.ratio[{i}]={rg.ratio[i]:.6f} != data[{i}].ratio={r.ratio:.6f}"
+            )
+
+    # -- key index positions -------------------------------------------------
+
+    @pytest.mark.parametrize(
+        "idx,expected_hw",
+        [
+            # First (most portrait) → 1:4 ratio bucket
+            (0, (512, 2048)),
+            # Centre square bucket
+            (16, (1024, 1024)),
+            # Most landscape bucket → 4:1 ratio
+            (32, (2048, 512)),
+        ],
+    )
+    def test_landmark_index(self, idx, expected_hw):
+        rg = _build_plain_group()
+        r = rg.data[idx]
+        assert (r.h, r.w) == expected_hw, f"idx {idx}: expected {expected_hw}, got ({r.h}, {r.w})"
+
+    def test_center_is_square(self):
+        rg = _build_plain_group()
+        r = rg.data[16]
+        assert r.h == r.w == _BASE
+
+    def test_first_and_last_are_extremes(self):
+        rg = _build_plain_group()
+        # first  → 512×2048  (ratio 0.25)
+        assert rg.data[0].ratio < 0.3
+        # last   → 2048×512  (ratio 4.0)
+        assert rg.data[-1].ratio > 3.9
+
+    # -- width × height matching ---------------------------------------------
+
+    def test_every_resolution_has_same_area_product(self):
+        """Symmetry: stepping loop produces the same product set in both
+        directions, so (h × w) tuples should appear symmetrically."""
+        rg = _build_plain_group()
+        products = {(r.h, r.w) for r in rg.data}
+        # For every (h, w) there should exist (w, h) or a nearby equivalent
+        for h, w in products:
+            if h == w:
+                continue  # square — trivially symmetric
+            assert (w, h) in products or any(abs(rr.h - w) <= _STEP and abs(rr.w - h) <= _STEP for rr in rg.data), (
+                f"no symmetric counterpart for {h}×{w}"
+            )
+
+    def test_product_of_wh_divisible_by_256(self):
+        """Each bucket's area should be a multiple of 256 (16×16),
+        which is the VAE down-sample grid."""
+        rg = _build_plain_group()
+        for r in rg.data:
+            assert (r.h * r.w) % 256 == 0, f"{r.h}×{r.w} = {r.h * r.w} not divisible by 256"
+
+    # -- custom base_size / step / align -------------------------------------
+
+    @pytest.mark.parametrize("base_size", [512, 768, 1280])
+    def test_count_scales_with_base_size(self, base_size):
+        step = base_size // 16
+        rg = ResolutionGroup(base_size=base_size, step=step, align=16, extra_resolutions=None)
+        expected = 1 + 2 * (base_size // step)  # centre + 2×phases
+        assert len(rg.data) == expected, f"base_size={base_size}: expected {expected}, got {len(rg.data)}"
+
+    def test_custom_step(self):
+        """A smaller step produces more buckets."""
+        rg_small = ResolutionGroup(base_size=1024, step=32, align=16, extra_resolutions=None)
+        rg_default = _build_plain_group()
+        assert len(rg_small.data) > len(rg_default.data), "smaller step should produce more buckets"
+
+
+# ---------------------------------------------------------------------------
+# 2.  HUNYUAN_IMAGE3_EXTRA_RESOLUTIONS matching
+# ---------------------------------------------------------------------------
+
+_EXTRA_SQUARES = {"512x512", "640x640", "768x768", "896x896"}
+_EXTRA_NON_SQUARES = {"1024x768", "1280x720", "768x1024", "720x1280"}
+
+
+def _build_extra_group() -> ResolutionGroup:
+    return ResolutionGroup(
+        base_size=_BASE,
+        extra_resolutions=[_make_reso(s) for s in HUNYUAN_IMAGE3_EXTRA_RESOLUTIONS],
+    )
+
+
+class TestExtraResolutionMatching:
+    """Verify extra resolutions are correctly integrated into
+    the ``ResolutionGroup``."""
+
+    # -- square extras (ratio = 1.0) -----------------------------------------
+
+    def test_square_extras_appended_to_base_square(self):
+        """512², 640², 768², 896² all have ratio 1.0 and should be
+        appended to the (1024, 1024) Resolution via ``append()``."""
+        rg = _build_extra_group()
+        # Find the 1024×1024 entry
+        square_reso = next(r for r in rg.data if r.h == 1024 and r.w == 1024)
+        expected_extra = {(512, 512), (640, 640), (768, 768), (896, 896)}
+        assert square_reso.extra_res == expected_extra, f"square extra_res: {square_reso.extra_res}"
+
+    def test_base_1024_unchanged(self):
+        """The 1024×1024 entry itself still reports h=1024, w=1024."""
+        rg = _build_extra_group()
+        r = next(r for r in rg.data if r.h == 1024 and r.w == 1024)
+        assert r.h == 1024 and r.w == 1024
+
+    def test_squares_not_standalone_entries(self):
+        """Square extras must NOT appear as top-level ``Resolution`` entries;
+        they live inside ``extra_res`` only."""
+        rg = _build_extra_group()
+        standalone = {(r.h, r.w) for r in rg.data}
+        for square in [(512, 512), (640, 640), (768, 768), (896, 896)]:
+            assert square not in standalone, f"{square} should be in extra_res, not a standalone entry"
+
+    # -- non-square extras (ratio ≠ 1.0) ------------------------------------
+
+    @pytest.mark.parametrize("reso_str", sorted(_EXTRA_NON_SQUARES))
+    def test_non_square_extra_is_standalone_entry(self, reso_str):
+        """1024×768, 1280×720, 768×1024, 720×1280 are appended as
+        new ``Resolution`` entries because no training bucket shares
+        their exact ratio."""
+        r = _make_reso(reso_str)
+        rg = _build_extra_group()
+        matching = [d for d in rg.data if d.h == r.h and d.w == r.w]
+        assert len(matching) == 1, f"{reso_str} should be in rg.data exactly once"
+        assert matching[0].extra_res == set(), f"{reso_str} should have empty extra_res, got {matching[0].extra_res}"
+
+    def test_non_square_extras_empty_extra_res(self):
+        """Non-square extras have no own extra_res."""
+        rg = _build_extra_group()
+        for r in rg.data:
+            if (r.h, r.w) in {(1024, 768), (1280, 720), (768, 1024), (720, 1280)}:
+                assert r.extra_res == set()
+
+    # -- total count ---------------------------------------------------------
+
+    def test_total_count(self):
+        """33 training + 4 non-square extras = 37."""
+        rg = _build_extra_group()
+        assert len(rg.data) == 37, f"expected 37, got {len(rg.data)}"
+
+    def test_extra_count(self):
+        """Exactly 8 extra resolutions (4 square + 4 non-square)."""
+        rg = _build_extra_group()
+        reso_set = _reso_set(rg)
+        extra_in_set = reso_set - {(r.h, r.w) for r in _build_plain_group().data}
+        assert len(extra_in_set) == 8, f"expected 8 extra, got {len(extra_in_set)}: {extra_in_set}"
+
+    # -- index stability -----------------------------------------------------
+
+    def test_indices_0_32_unchanged(self):
+        """Indices 0–32 remain the training buckets (sorted by ratio)."""
+        plain = _build_plain_group()
+        rg = _build_extra_group()
+        for i in range(33):
+            assert (rg.data[i].h, rg.data[i].w) == (plain.data[i].h, plain.data[i].w), (
+                f"idx {i}: extra group has ({rg.data[i].h},{rg.data[i].w}), "
+                f"expected ({plain.data[i].h},{plain.data[i].w})"
+            )
+
+    @pytest.mark.parametrize(
+        "reso_str,expected_idx",
+        [
+            ("1024x768", 33),
+            ("1280x720", 34),
+            ("768x1024", 35),
+            ("720x1280", 36),
+        ],
+    )
+    def test_non_square_extra_indices(self, reso_str, expected_idx):
+        """Non-square extras are appended in order at indices 33-36."""
+        r = _make_reso(reso_str)
+        rg = _build_extra_group()
+        idx = next(i for i, d in enumerate(rg.data) if d.h == r.h and d.w == r.w)
+        assert idx == expected_idx, f"{reso_str}: expected idx {expected_idx}, got {idx}"
+
+    # -- ratio lookup via get_base_size_and_ratio_index ----------------------
+
+    @pytest.mark.parametrize(
+        "h,w,expected_idx",
+        [
+            (1024, 768, 33),
+            (1280, 720, 34),
+            (768, 1024, 35),
+            (720, 1280, 36),
+            (1024, 1024, 16),
+            (512, 512, 16),
+            (640, 640, 16),
+            (768, 768, 16),
+            (896, 896, 16),
+        ],
+    )
+    def test_ratio_index_for_resolutions(self, h, w, expected_idx):
+        """``get_base_size_and_ratio_index`` must return the correct
+        ratio index for both training-bucket and extra resolutions."""
+        rg = _build_extra_group()
+        _, idx = rg.get_base_size_and_ratio_index(w, h)
+        assert idx == expected_idx, f"({w},{h}): expected idx {expected_idx}, got {idx}"
+
+
+# ---------------------------------------------------------------------------
+# Resolution.match — area-based matching within a ratio group
+# ---------------------------------------------------------------------------
+
+
+class TestResolutionMatch:
+    """Verify ``Resolution.match`` picks the best-fit alternative from
+    ``extra_res`` based on area proximity."""
+
+    def test_no_extra_returns_self(self):
+        r = _make_reso("1024x768")
+        assert r.match(800, 600) == (768, 1024)
+
+    @pytest.mark.parametrize(
+        "target_w,target_h,expected",
+        [
+            # Exact area match in extra_res
+            (512, 512, (512, 512)),
+            (640, 640, (640, 640)),
+            (768, 768, (768, 768)),
+            (896, 896, (896, 896)),
+            # Closer to 1024² than to any extra square
+            (1024, 1024, (1024, 1024)),
+            (1000, 1000, (1024, 1024)),
+            # Small square picks 512²
+            (256, 256, (512, 512)),
+            (400, 400, (512, 512)),
+            # 700² → closest to 640² (diff 129600) vs 768² (diff 99856)
+            # 768*768 - 700*700 = 589824 - 490000 = 99824
+            # 640*640 - 700*700 = 409600 - 490000 = -80400 → abs=80400
+            # So 640² is closer in area
+            (700, 700, (640, 640)),
+            # 800² → closest to 768² (diff 49984) vs 896² (diff 162944)
+            (800, 800, (768, 768)),
+        ],
+    )
+    def test_square_extra_match(self, target_w, target_h, expected):
+        """Build a Resolution like the 1024² entry with square extras,
+        then verify area-based matching."""
+        r = Resolution(1024, 1024)
+        for s in [(512, 512), (640, 640), (768, 768), (896, 896)]:
+            r.extra_res.add(s)
+        result = r.match(target_w, target_h)
+        assert result == expected, f"match({target_w},{target_h}) = {result}, expected {expected}"
+
+    def test_match_returns_tuple(self):
+        r = Resolution(1024, 1024)
+        r.extra_res.add((512, 512))
+        result = r.match(512, 512)
+        assert isinstance(result, tuple)
+        assert len(result) == 2
+        assert all(isinstance(v, int) for v in result)
+
+    def test_match_does_not_mutate_resolution(self):
+        r = Resolution(1024, 1024)
+        r.extra_res.add((512, 512))
+        extra_before = set(r.extra_res)
+        r.match(300, 300)
+        assert r.extra_res == extra_before
+
+
+# ---------------------------------------------------------------------------
+# ResolutionGroup.get_target_size — end-to-end resolution selection
+# ---------------------------------------------------------------------------
+
+
+class TestGetTargetSize:
+    """Verify ``get_target_size`` uses ratio-then-area logic correctly."""
+
+    @pytest.fixture(scope="class")
+    def rg(self) -> ResolutionGroup:
+        return _build_extra_group()
+
+    @pytest.mark.parametrize(
+        "w,h,expected_min_area,expected_max_area",
+        [
+            # Square targets should land on square variants
+            (1024, 1024, 1024 * 1024, 1024 * 1024),
+            (512, 512, 512 * 512, 512 * 512),
+            (640, 640, 640 * 640, 640 * 640),
+            (768, 768, 768 * 768, 768 * 768),
+            (896, 896, 896 * 896, 896 * 896),
+        ],
+    )
+    def test_square_target_returns_square(self, rg, w, h, expected_min_area, expected_max_area):
+        rw, rh = rg.get_target_size(w, h)
+        area = rw * rh
+        assert expected_min_area <= area <= expected_max_area, f"({w},{h}) → ({rw},{rh}) area={area}"
+
+    def test_training_bucket_self_maps(self, rg):
+        """A training bucket should map to itself (or an extra_res variant
+        with the same ratio)."""
+        for r in rg.data:
+            rw, rh = rg.get_target_size(r.w, r.h)
+            # Either exact match or the ratio must be preserved
+            assert abs(rh / rw - r.ratio) < 0.01, f"({r.w},{r.h}) → ({rw},{rh}) ratio drifted from {r.ratio:.4f}"
+
+    def test_wild_ratio_still_finds_bucket(self, rg):
+        """Even an unusual aspect ratio should map to some valid bucket."""
+        for w, h in [(800, 600), (1200, 900), (1600, 400), (400, 1600)]:
+            rw, rh = rg.get_target_size(w, h)
+            assert rw > 0 and rh > 0
+            assert rw % 16 == 0 and rh % 16 == 0
+
+
+# ---------------------------------------------------------------------------
+# 3.  get_cached_resolution_group & ar2diffusion integration
+# ---------------------------------------------------------------------------
+
+
+class TestGetCachedResolutionGroup:
+    """Verify ``get_cached_resolution_group`` returns a correctly
+    configured ``ResolutionGroup`` and respects the LRU cache."""
+
+    def test_returns_resolution_group(self):
+        rg = get_cached_resolution_group(base_size=1024)
+        assert isinstance(rg, ResolutionGroup)
+
+    def test_base_size_propagated(self):
+        rg = get_cached_resolution_group(base_size=1024)
+        assert rg.base_size == 1024
+
+    def test_has_extra_resolutions(self):
+        rg = get_cached_resolution_group(base_size=1024)
+        reso_set = _reso_set(rg)
+        for s in HUNYUAN_IMAGE3_EXTRA_RESOLUTIONS:
+            r = _make_reso(s)
+            assert (r.h, r.w) in reso_set, f"{s} missing from cached group"
+
+    def test_cache_same_base_size_same_object(self):
+        """LRU cache: same base_size → same object identity."""
+        # Clear any prior cache entries
+        get_cached_resolution_group.cache_clear()
+        rg1 = get_cached_resolution_group(base_size=1024)
+        rg2 = get_cached_resolution_group(base_size=1024)
+        assert rg1 is rg2, "same base_size must return cached object"
+
+    def test_cache_different_base_size_different_object(self):
+        get_cached_resolution_group.cache_clear()
+        rg1 = get_cached_resolution_group(base_size=1024)
+        rg2 = get_cached_resolution_group(base_size=768)
+        assert rg1 is not rg2, "different base_size must return different objects"
+
+    def test_cache_size_limit(self):
+        """Cache is bounded at 4 entries (``maxsize=4``)."""
+        get_cached_resolution_group.cache_clear()
+        for bs in [512, 768, 1024, 1280, 1536]:
+            get_cached_resolution_group(base_size=bs)
+        info = get_cached_resolution_group.cache_info()
+        assert info.currsize <= 4, f"cache grew to {info.currsize}, max is 4"
+        # All 5 unique base_size values caused cache misses
+        assert info.misses >= 5, f"expected >=5 misses for 5 unique base_size values, got {info.misses}"
+
+    def test_resolution_indexing_works(self):
+        """Indices 0..len(rg)-1 must be accessible."""
+        rg = get_cached_resolution_group(base_size=1024)
+        for i in range(len(rg)):
+            r = rg[i]
+            assert isinstance(r, Resolution)
+            assert r.h > 0 and r.w > 0
+
+    @pytest.mark.parametrize("base_size", [512, 1024])
+    def test_ratio_index_roundtrip_via_group_length(self, base_size):
+        """Every ratio index in [0, len(rg)) must resolve to a valid
+        Resolution with height/width within bounds."""
+        rg = get_cached_resolution_group(base_size=base_size)
+        for i in range(len(rg)):
+            r = rg[i]
+            assert isinstance(r.height, int) and r.height > 0
+            assert isinstance(r.width, int) and r.width > 0
+
+
+# ---------------------------------------------------------------------------
+# ar2diffusion — integration tests
+# ---------------------------------------------------------------------------
+
+
+# Token IDs matching the real prompt_utils values
+_RATIO_0 = 128044
+_RATIO_32 = 128076
+_RATIO_33 = 130103
+_RATIO_36 = 130106
+
+
+def _fake_ar_output(
+    *,
+    cumulative_token_ids: list[int] | None = None,
+    text: str = "",
+) -> MagicMock:
+    """Minimal stub for a vLLM ``RequestOutput`` with one completion."""
+    choice = MagicMock()
+    choice.cumulative_token_ids = cumulative_token_ids or []
+    choice.text = text
+    output = MagicMock()
+    output.outputs = [choice]
+    return output
+
+
+class TestAr2DiffusionIntegration:
+    """Verify that ``ar2diffusion`` calls ``get_cached_resolution_group``
+    and obtains a correctly configured ``ResolutionGroup`` whose indices map
+    to the expected training-bucket or extra-resolution sizes."""
+
+    def _call_ar2diffusion(
+        self,
+        token_ids: list[int],
+        prompt_kwargs: dict[str, Any] | None = None,
+    ) -> tuple[Any, Any]:
+        """Invoke ``ar2diffusion`` with a single fake AR output.
+
+        Returns ``(diffusion_inputs, ar_predicted_sizes)`` so callers can
+        assert on the resolved height/width.
+        """
+        from vllm_omni.model_executor.stage_input_processors.hunyuan_image3 import (
+            ar2diffusion,
+        )
+
+        pk = dict(prompt_kwargs or {})
+        pk.setdefault("prompt", "a beautiful sunset")
+        pk.setdefault("height", 1024)
+        pk.setdefault("width", 1024)
+        pk.setdefault("image_base_size", 1024)
+
+        # Build a prompt-like object that behaves like the real OmniTokensPrompt
+        class _FakePrompt:
+            def __init__(self, d):
+                self._d = d
+
+            def __getitem__(self, key):
+                return self._d[key]
+
+            def get(self, key, default=None):
+                return self._d.get(key, default)
+
+            def _asdict(self):
+                return dict(self._d)
+
+        prompt = _FakePrompt(pk)
+        ar_output = _fake_ar_output(cumulative_token_ids=token_ids)
+        result = ar2diffusion(
+            source_outputs=[ar_output],
+            prompt=[prompt],
+            requires_multimodal_data=False,
+        )
+        return result
+
+    # -- ratio token ↔ resolution resolution ---------------------------------
+
+    @pytest.mark.parametrize(
+        "ratio_idx,expected_hw",
+        [
+            # Training buckets
+            (0, (512, 2048)),
+            (16, (1024, 1024)),
+            (32, (2048, 512)),
+            # Non-square extras
+            (33, (1024, 768)),
+            (34, (1280, 720)),
+            (35, (768, 1024)),
+            (36, (720, 1280)),
+        ],
+    )
+    def test_ratio_index_resolves_to_correct_size(self, ratio_idx, expected_hw):
+        """When the AR emits ``<img_ratio_N>``, the DiT input must carry
+        the training-bucket width/height for that index."""
+        rg = get_cached_resolution_group(base_size=1024)
+        reso = rg[ratio_idx]
+        assert (reso.height, reso.width) == expected_hw, (
+            f"ratio_idx {ratio_idx}: expected {expected_hw}, got ({reso.height}, {reso.width})"
+        )
+
+    @pytest.mark.parametrize("ratio_idx", [0, 16, 32, 33, 34, 35, 36])
+    def test_ar2diffusion_extracts_height_width(self, ratio_idx):
+        """End-to-end: simulate AR output containing ``<img_ratio_{idx}>``
+        and assert the DiT inputs carry the expected size."""
+        rg = get_cached_resolution_group(base_size=1024)
+        expected_h = rg[ratio_idx].height
+        expected_w = rg[ratio_idx].width
+
+        token_id = _RATIO_0 + ratio_idx if ratio_idx <= 32 else _RATIO_33 + (ratio_idx - 33)
+        result = self._call_ar2diffusion(
+            token_ids=[150000, 150001, token_id],
+            prompt_kwargs={"height": 1024, "width": 1024, "image_base_size": 1024},
+        )
+        assert len(result) == 1
+        diffusion_input = result[0]
+        assert diffusion_input["height"] == expected_h, (
+            f"ratio_idx={ratio_idx}: expected height {expected_h}, got {diffusion_input['height']}"
+        )
+        assert diffusion_input["width"] == expected_w, (
+            f"ratio_idx={ratio_idx}: expected width {expected_w}, got {diffusion_input['width']}"
+        )
+
+    def test_no_ratio_token_falls_back_to_prompt_size(self):
+        """If no ratio token is present, height/width from the prompt
+        are kept unchanged."""
+        result = self._call_ar2diffusion(
+            token_ids=[150000, 150001],  # no ratio tokens
+            prompt_kwargs={"height": 768, "width": 1024, "image_base_size": 1024},
+        )
+        assert len(result) == 1
+        diffusion_input = result[0]
+        assert diffusion_input["height"] == 768
+        assert diffusion_input["width"] == 1024
+
+    def test_out_of_range_ratio_index_falls_back(self):
+        """If ratio_idx ≥ len(reso_group), ar2diffusion must keep the
+        prompt size (fallback path)."""
+        with patch(
+            "vllm_omni.model_executor.stage_input_processors.hunyuan_image3._extract_ratio_index",
+            return_value=999,  # out of range
+        ):
+            result = self._call_ar2diffusion(
+                token_ids=[150000],
+                prompt_kwargs={"height": 512, "width": 512, "image_base_size": 1024},
+            )
+        assert len(result) == 1
+        assert result[0]["height"] == 512
+        assert result[0]["width"] == 512
+
+    # -- cache behaviour -----------------------------------------------------
+
+    def test_ar2diffusion_uses_cached_group(self):
+        """Each invocation of ``ar2diffusion`` calls
+        ``get_cached_resolution_group``, which returns a shared cached object."""
+        get_cached_resolution_group.cache_clear()
+        rg_before = get_cached_resolution_group(base_size=1024)
+
+        token_id = _RATIO_0 + 16  # <img_ratio_16> = 1024×1024
+        self._call_ar2diffusion(
+            token_ids=[token_id],
+            prompt_kwargs={"height": 1024, "width": 1024, "image_base_size": 1024},
+        )
+
+        rg_after = get_cached_resolution_group(base_size=1024)
+        assert rg_before is rg_after, "get_cached_resolution_group should return the same cached object"
+
+    def test_correct_reso_group_used_for_base_size(self):
+        """Verify ``ar2diffusion`` constructs or retrieves the
+        ResolutionGroup for the correct ``image_base_size``."""
+        get_cached_resolution_group.cache_clear()
+
+        # First request with a non-default base_size
+        token_id = _RATIO_0  # <img_ratio_0>
+        self._call_ar2diffusion(
+            token_ids=[token_id],
+            prompt_kwargs={"height": 512, "width": 512, "image_base_size": 512},
+        )
+
+        # The cached group for base_size=512 should exist
+        rg512 = get_cached_resolution_group(base_size=512)
+        assert rg512.base_size == 512
+
+        # And base_size=1024 should also be separately accessible
+        rg1024 = get_cached_resolution_group(base_size=1024)
+        assert rg1024.base_size == 1024
+        assert rg512 is not rg1024
+
+
+# ---------------------------------------------------------------------------
+# Edge cases
+# ---------------------------------------------------------------------------
+
+
+class TestEdgeCases:
+    def test_resolution_string_parsing(self):
+        """Resolution accepts both "HxW" strings and (h, w) tuples."""
+        r1 = Resolution("512x2048")
+        assert r1.h == 512 and r1.w == 2048
+
+        r2 = Resolution(768, 1024)
+        assert r2.h == 768 and r2.w == 1024
+
+        r3 = Resolution(1024)  # single int → square
+        assert r3.h == 1024 and r3.w == 1024
+
+    def test_resolution_extra_res_initially_empty(self):
+        r = Resolution(1024, 1024)
+        assert r.extra_res == set()
+        assert r.match(512, 512) == (1024, 1024)
+
+    def test_resolution_append_deduplicates(self):
+        """Appending the same resolution twice is a no-op on the set."""
+        r = Resolution(1024, 1024)
+        r.append(Resolution(512, 512))
+        r.append(Resolution(512, 512))
+        assert len(r.extra_res) == 1
+
+    def test_all_final_resolutions_are_16_aligned(self):
+        """Every reachable resolution (including extras) must be
+        divisible by the VAE down-sample factor of 16."""
+        rg = _build_extra_group()
+        for h, w in _reso_set(rg):
+            assert h % 16 == 0, f"height {h} not 16-aligned in {h}×{w}"
+            assert w % 16 == 0, f"width {w} not 16-aligned in {h}×{w}"
+
+    def test_ratio_array_length_matches_data(self):
+        rg = _build_extra_group()
+        assert len(rg.ratio) == len(rg.data)
+
+    def test_resolution_repr(self):
+        r = Resolution(1024, 1024)
+        rep = repr(r)
+        assert "1024x1024" in rep
+
+        r2 = Resolution(1024, 1024)
+        r2.extra_res.add((512, 512))
+        rep2 = repr(r2)
+        assert "512x512" in rep2
@@ -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 ─────


@@ -145,7 +145,7 @@ def test_shm_connector_flow(mocker: MockerFixture):
    Using real SharedMemoryConnector (inline mode for simplicity).
    """
    # 1. Setup Connector
-    config = {"shm_threshold_bytes": 1024}  # Large threshold to use inline
+    config = {"shm_threshold_bytes": 1024, "inline_small_payloads": True}
    connector = SharedMemoryConnector(config)
    connectors_map = {("0", "1"): connector}

@@ -75,18 +75,21 @@ def test_create_unknown_connector():

@pytest.fixture
 def shm_connector():
-    config = {"shm_threshold_bytes": 100}  # Small threshold for testing
+    config = {"shm_threshold_bytes": 100, "inline_small_payloads": True}
    return SharedMemoryConnector(config)


 def test_put_get_inline(shm_connector):
-    """Test transfer for small data (currently always uses SHM)."""
+    """Test inline transfer for small data."""
    data = {"small": "data"}

    success, size, metadata = shm_connector.put("stage_0", "stage_1", "req_1", data)
    assert success is True
-    assert "shm" in metadata
+    assert "inline_bytes" in metadata
+    assert "shm" not in metadata
    assert "size" in metadata
+    assert shm_connector._metrics["inline_writes"] == 1
+    assert shm_connector._metrics["shm_writes"] == 0

    # Retrieve
    retrieved_data, ret_size = shm_connector.get("stage_0", "stage_1", "req_1", metadata)
@@ -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
@@ -13,7 +13,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]

@pytest.fixture()
 def connector():
-    c = SharedMemoryConnector({"shm_threshold_bytes": 64})
+    c = SharedMemoryConnector({"shm_threshold_bytes": 0})
    yield c
    c.close()

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

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

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


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


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


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

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

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

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

 import os
+from collections.abc import Mapping

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

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

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

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

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

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

@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
+@pytest.mark.skip(reason="#4335")
 def test_response_format_001(omni_server, openai_client) -> None:
    """
    Test voice-clone non-stream PCM output via OpenAI API.
@@ -0,0 +1,59 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for _ensure_tensor_values boundary enforcement."""
+
+import pytest
+import torch
+
+from vllm_omni.worker.gpu_ar_model_runner import _ensure_tensor_values
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_tensors_pass_through():
+    """Tensor values pass through unchanged."""
+    t = torch.randn(10)
+    result = _ensure_tensor_values({"audio": t})
+    assert torch.equal(result["audio"], t)
+
+
+def test_scalar_int_wrapped():
+    """Integer scalars are wrapped in torch.tensor."""
+    result = _ensure_tensor_values({"sr": 24000})
+    assert isinstance(result["sr"], torch.Tensor)
+    assert result["sr"].item() == 24000
+
+
+def test_scalar_float_wrapped():
+    """Float scalars are wrapped in torch.tensor."""
+    result = _ensure_tensor_values({"gain": 0.5})
+    assert isinstance(result["gain"], torch.Tensor)
+    assert result["gain"].item() == pytest.approx(0.5)
+
+
+def test_list_converted():
+    """Lists of numbers are converted to tensors."""
+    result = _ensure_tensor_values({"ids": [1, 2, 3]})
+    assert isinstance(result["ids"], torch.Tensor)
+    assert list(result["ids"]) == [1, 2, 3]
+
+
+def test_non_tensorizable_dropped():
+    """Values that cannot be tensorized are dropped with warning."""
+    result = _ensure_tensor_values({"ok": torch.randn(5), "bad": {"nested": "dict"}})
+    assert "ok" in result
+    assert "bad" not in result
+
+
+def test_mixed_payload():
+    """Mixed tensor + scalar + list payload is properly sanitized."""
+    payload = {
+        "audio": torch.randn(100),
+        "sr": 24000,
+        "codes": [1, 2, 3],
+        "hidden": torch.randn(10, 64),
+    }
+    result = _ensure_tensor_values(payload)
+    assert len(result) == 4
+    for v in result.values():
+        assert isinstance(v, torch.Tensor)
@@ -161,3 +161,113 @@ async def test_diffusion_error_output_routed_as_finished(orchestrator_factory) -
    finally:
        orchestrator_fixture.request_sync_q.put_nowait(ShutdownRequestMessage())
        orchestrator_fixture.thread.join(timeout=5)
+
+
+@pytest.mark.asyncio
+async def test_diffusion_client_error_output_propagates_status_code(orchestrator_factory) -> None:
+    """A client-error OmniRequestOutput (e.g. a guardrail 4xx) from a diffusion
+    stage must surface as an ErrorMessage that carries status_code/error_type,
+    so the frontend can map it to the correct HTTP response instead of a 500.
+    """
+    stage0 = FakeStageClient(stage_type="diffusion", final_output=True, final_output_type="image")
+    orchestrator_fixture = orchestrator_factory([stage0])
+    params = OmniDiffusionSamplingParams()
+
+    try:
+        await _enqueue_add_request(
+            orchestrator_fixture,
+            request_id="req-blocked",
+            prompt={"prompt": "blocked"},
+            original_prompt={"prompt": "blocked"},
+            sampling_params_list=[params],
+            final_stage_id=0,
+        )
+
+        await _wait_for(lambda: len(stage0.add_request_calls) == 1)
+
+        # Push a client-error output from the diffusion stage.
+        stage0.push_diffusion_output(
+            OmniRequestOutput.from_error(
+                "req-blocked",
+                "Input was blocked by Cosmos3 guardrails.",
+                status_code=400,
+                error_type="BadRequestError",
+            )
+        )
+
+        msg = await _get_any_output_message(orchestrator_fixture)
+
+        assert isinstance(msg, ErrorMessage)
+        assert msg.request_id == "req-blocked"
+        assert msg.stage_id == 0
+        assert msg.fatal is False
+        assert msg.status_code == 400
+        assert msg.error_type == "BadRequestError"
+        assert msg.error == "Input was blocked by Cosmos3 guardrails."
+
+        # Request state should be cleaned up.
+        await _wait_for(lambda: "req-blocked" not in orchestrator_fixture.orchestrator.request_states)
+    finally:
+        orchestrator_fixture.request_sync_q.put_nowait(ShutdownRequestMessage())
+        orchestrator_fixture.thread.join(timeout=5)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("status_code", "error_type"),
+    [
+        pytest.param(429, "RateLimitError", id="429-too-many-requests"),
+        pytest.param(413, "PayloadTooLargeError", id="413-payload-too-large"),
+        pytest.param(422, "UnprocessableEntityError", id="422-unprocessable-entity"),
+        pytest.param(403, "PermissionDeniedError", id="403-forbidden"),
+    ],
+)
+async def test_diffusion_client_error_output_propagates_non_400_status(
+    orchestrator_factory,
+    status_code: int,
+    error_type: str,
+) -> None:
+    """A diffusion-stage client error with a non-400 4xx status must be serialized
+    into an ErrorMessage that preserves that exact status_code/error_type.
+    """
+    stage0 = FakeStageClient(stage_type="diffusion", final_output=True, final_output_type="image")
+    orchestrator_fixture = orchestrator_factory([stage0])
+    params = OmniDiffusionSamplingParams()
+
+    try:
+        await _enqueue_add_request(
+            orchestrator_fixture,
+            request_id="req-blocked",
+            prompt={"prompt": "blocked"},
+            original_prompt={"prompt": "blocked"},
+            sampling_params_list=[params],
+            final_stage_id=0,
+        )
+
+        await _wait_for(lambda: len(stage0.add_request_calls) == 1)
+
+        # Push a non-400 client-error output from the diffusion stage.
+        stage0.push_diffusion_output(
+            OmniRequestOutput.from_error(
+                "req-blocked",
+                "client error from diffusion stage",
+                status_code=status_code,
+                error_type=error_type,
+            )
+        )
+
+        msg = await _get_any_output_message(orchestrator_fixture)
+
+        assert isinstance(msg, ErrorMessage)
+        assert msg.request_id == "req-blocked"
+        assert msg.stage_id == 0
+        assert msg.fatal is False
+        assert msg.status_code == status_code
+        assert msg.error_type == error_type
+        assert msg.error == "client error from diffusion stage"
+
+        # Request state should be cleaned up.
+        await _wait_for(lambda: "req-blocked" not in orchestrator_fixture.orchestrator.request_states)
+    finally:
+        orchestrator_fixture.request_sync_q.put_nowait(ShutdownRequestMessage())
+        orchestrator_fixture.thread.join(timeout=5)
@@ -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,187 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for decode_audio_url helper."""
+
+import asyncio
+import base64
+import os
+
+import httpx
+import pytest
+
+from vllm_omni.entrypoints.openai.errors import InvalidInputReferenceError
+from vllm_omni.entrypoints.openai.video_api_utils import decode_audio_url
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _run(coro):
+    return asyncio.run(coro)
+
+
+def _cleanup(path):
+    if path and os.path.exists(path):
+        os.unlink(path)
+
+
+class TestDecodeAudioUrlDataUrl:
+    """Tests for base64 data URL decoding."""
+
+    def test_valid_mp3_data_url(self):
+        audio_bytes = b"\xff\xfb\x90\x00" * 100
+        b64 = base64.b64encode(audio_bytes).decode()
+        url = f"data:audio/mp3;base64,{b64}"
+
+        path = _run(decode_audio_url(url))
+        try:
+            assert os.path.isfile(path)
+            assert path.endswith(".mp3")
+            with open(path, "rb") as f:
+                assert f.read() == audio_bytes
+        finally:
+            _cleanup(path)
+
+    def test_valid_wav_data_url(self):
+        audio_bytes = b"RIFF" + b"\x00" * 100
+        b64 = base64.b64encode(audio_bytes).decode()
+        url = f"data:audio/wav;base64,{b64}"
+
+        path = _run(decode_audio_url(url))
+        try:
+            assert path.endswith(".wav")
+            with open(path, "rb") as f:
+                assert f.read() == audio_bytes
+        finally:
+            _cleanup(path)
+
+    def test_valid_mpeg_data_url_uses_mp3_suffix(self):
+        audio_bytes = b"\xff\xfb\x90\x00" * 50
+        b64 = base64.b64encode(audio_bytes).decode()
+        url = f"data:audio/mpeg;base64,{b64}"
+
+        path = _run(decode_audio_url(url))
+        try:
+            assert path.endswith(".mp3")
+        finally:
+            _cleanup(path)
+
+    def test_valid_ogg_data_url(self):
+        audio_bytes = b"OggS" + b"\x00" * 100
+        b64 = base64.b64encode(audio_bytes).decode()
+        url = f"data:audio/ogg;base64,{b64}"
+
+        path = _run(decode_audio_url(url))
+        try:
+            assert path.endswith(".ogg")
+        finally:
+            _cleanup(path)
+
+    def test_invalid_base64_raises(self):
+        url = "data:audio/mp3;base64,!!!not-valid-base64!!!"
+
+        with pytest.raises(InvalidInputReferenceError, match="not valid base64"):
+            _run(decode_audio_url(url))
+
+    def test_empty_audio_data_raises(self):
+        url = "data:audio/mp3;base64,"
+
+        with pytest.raises(InvalidInputReferenceError, match="audio data is empty"):
+            _run(decode_audio_url(url))
+
+
+class TestDecodeAudioUrlSuffixSanitization:
+    """Tests for MIME extension sanitization (path traversal prevention)."""
+
+    def test_path_traversal_in_mime_is_neutralized(self):
+        audio_bytes = b"\x00" * 100
+        b64 = base64.b64encode(audio_bytes).decode()
+        url = f"data:audio/../..;base64,{b64}"
+
+        path = _run(decode_audio_url(url))
+        try:
+            # ".." contains non-alnum dots, so falls back to .wav
+            assert path.endswith(".wav")
+            assert os.path.dirname(path) == "/tmp"
+        finally:
+            _cleanup(path)
+
+    def test_non_alnum_extension_falls_back_to_wav(self):
+        audio_bytes = b"\x00" * 100
+        b64 = base64.b64encode(audio_bytes).decode()
+        url = f"data:audio/foo-bar;base64,{b64}"
+
+        path = _run(decode_audio_url(url))
+        try:
+            assert path.endswith(".wav")
+        finally:
+            _cleanup(path)
+
+    def test_too_long_extension_falls_back_to_wav(self):
+        audio_bytes = b"\x00" * 100
+        b64 = base64.b64encode(audio_bytes).decode()
+        url = f"data:audio/verylongextension;base64,{b64}"
+
+        path = _run(decode_audio_url(url))
+        try:
+            assert path.endswith(".wav")
+        finally:
+            _cleanup(path)
+
+    def test_valid_short_alnum_extension_is_used(self):
+        audio_bytes = b"\x00" * 100
+        b64 = base64.b64encode(audio_bytes).decode()
+        url = f"data:audio/flac;base64,{b64}"
+
+        path = _run(decode_audio_url(url))
+        try:
+            assert path.endswith(".flac")
+        finally:
+            _cleanup(path)
+
+
+class TestDecodeAudioUrlHttp:
+    """Tests for HTTP URL decoding."""
+
+    def test_valid_http_url(self, monkeypatch):
+        audio_bytes = b"\xff\xfb\x90\x00" * 50
+        fake_resp = httpx.Response(200, content=audio_bytes)
+
+        async def _mock_get(self, url):
+            return fake_resp
+
+        monkeypatch.setattr(httpx.AsyncClient, "get", _mock_get)
+        monkeypatch.setattr(httpx.Response, "raise_for_status", lambda self: None)
+
+        path = _run(decode_audio_url("https://example.com/audio.mp3"))
+        try:
+            assert os.path.isfile(path)
+            assert path.endswith(".wav")
+            with open(path, "rb") as f:
+                assert f.read() == audio_bytes
+        finally:
+            _cleanup(path)
+
+    def test_http_error_raises(self, monkeypatch):
+        async def _mock_get(self, url):
+            raise httpx.HTTPStatusError("Not Found", request=httpx.Request("GET", url), response=httpx.Response(404))
+
+        monkeypatch.setattr(httpx.AsyncClient, "get", _mock_get)
+
+        with pytest.raises(InvalidInputReferenceError, match="failed to download audio"):
+            _run(decode_audio_url("https://example.com/missing.mp3"))
+
+
+class TestDecodeAudioUrlInvalid:
+    """Tests for invalid URL schemes."""
+
+    def test_local_file_path_raises(self):
+        with pytest.raises(InvalidInputReferenceError, match="must be an http.*URL or data URL"):
+            _run(decode_audio_url("/path/to/local/file.mp3"))
+
+    def test_ftp_url_raises(self):
+        with pytest.raises(InvalidInputReferenceError, match="must be an http.*URL or data URL"):
+            _run(decode_audio_url("ftp://example.com/file.mp3"))
+
+    def test_empty_string_raises(self):
+        with pytest.raises(InvalidInputReferenceError, match="must be an http.*URL or data URL"):
+            _run(decode_audio_url(""))
@@ -29,6 +29,7 @@ from vllm_omni.entrypoints.openai.image_api_utils import (
    parse_size,
 )
 from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
+from vllm_omni.errors import GuardrailViolationError
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams

 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -587,6 +588,24 @@ def test_generate_single_image(test_client):
    assert test_client.app.state.engine_client.captured_prompt["modalities"] == ["image"]


+def test_generate_images_guardrail_error_returns_400(test_client, mock_async_diffusion):
+    async def blocked_generate(**kwargs):
+        raise GuardrailViolationError("Input was blocked by Cosmos3 guardrails.")
+        yield MockGenerationResult([])  # pragma: no cover
+
+    mock_async_diffusion.generate = blocked_generate
+    response = test_client.post(
+        "/v1/images/generations",
+        json={
+            "prompt": "blocked prompt",
+            "n": 1,
+            "size": "1024x1024",
+        },
+    )
+    assert response.status_code == 400
+    assert response.json()["detail"] == "Input was blocked by Cosmos3 guardrails."
+
+
 def test_generate_images_async_omni_sampling_params(async_omni_test_client):
    """Test AsyncOmni path uses per-stage sampling params."""
    response = async_omni_test_client.post(
@@ -891,6 +910,31 @@ def test_image_edits_streaming_errors_on_empty_final_image(streaming_image_edit_
    assert payloads[2] == "[DONE]"


+def test_image_edits_streaming_guardrail_error_uses_400(streaming_image_edit_client):
+    async def generate_guardrail_error(prompt, request_id, sampling_params=None, sampling_params_list=None, **kwargs):
+        raise GuardrailViolationError("Input was blocked by Cosmos3 guardrails.")
+        yield MockStageResult(stage_id=1, final_output_type="image", images=[])  # pragma: no cover
+
+    streaming_image_edit_client.app.state.engine_client.generate = generate_guardrail_error
+    img_bytes = make_test_image_bytes((16, 16))
+    response = streaming_image_edit_client.post(
+        "/v1/images/edits",
+        files=[("image", img_bytes)],
+        data={
+            "prompt": "blocked prompt",
+            "stream": "true",
+        },
+    )
+
+    assert response.status_code == 200
+    payloads = _parse_sse_payloads(response.text)
+    assert payloads[0]["object"] == "error"
+    assert payloads[0]["error"]["message"] == "Input was blocked by Cosmos3 guardrails."
+    assert payloads[0]["error"]["type"] == "BadRequestError"
+    assert payloads[0]["error"]["code"] == 400
+    assert payloads[1] == "[DONE]"
+
+
 def test_image_edits_streaming_rejects_single_stage(test_client):
    img_bytes = make_test_image_bytes((16, 16))
    response = test_client.post(
@@ -31,6 +31,7 @@ from vllm_omni.entrypoints.openai.protocol.audio import (
    CreateAudio,
    OpenAICreateSpeechRequest,
    SpeechBatchItem,
+    StreamingSpeechSessionConfig,
 )
 from vllm_omni.entrypoints.openai.serving_speech import (
    _TTS_LANGUAGES,
@@ -1235,6 +1236,49 @@ class TestTTSMethods:
        assert params["language"] == ["English"]
        assert params["task_type"] == ["CustomVoice"]

+    def test_build_tts_params_base_non_streaming_mode_true(self, speech_server):
+        """Base task should pass through an explicit non_streaming_mode override."""
+        req = OpenAICreateSpeechRequest(
+            input="Hello",
+            task_type="Base",
+            ref_audio="data:audio/wav;base64,abc",
+            ref_text="reference",
+            non_streaming_mode=True,
+        )
+
+        params = speech_server._build_tts_params(req)
+
+        assert params["task_type"] == ["Base"]
+        assert params["non_streaming_mode"] == [True]
+
+    def test_build_tts_params_base_omits_non_streaming_mode_by_default(self, speech_server):
+        """Base task should keep using the model default when no override is sent."""
+        req = OpenAICreateSpeechRequest(
+            input="Hello",
+            task_type="Base",
+            ref_audio="data:audio/wav;base64,abc",
+            ref_text="reference",
+        )
+
+        params = speech_server._build_tts_params(req)
+
+        assert params["task_type"] == ["Base"]
+        assert "non_streaming_mode" not in params
+
+    def test_build_tts_params_explicit_non_streaming_mode_overrides_voicedesign_default(self, speech_server):
+        """Explicit false should not be replaced by the VoiceDesign fallback."""
+        req = OpenAICreateSpeechRequest(
+            input="Hello",
+            task_type="VoiceDesign",
+            instructions="warm and calm",
+            non_streaming_mode=False,
+        )
+
+        params = speech_server._build_tts_params(req)
+
+        assert params["task_type"] == ["VoiceDesign"]
+        assert params["non_streaming_mode"] == [False]
+
    def test_load_supported_speakers(self, mocker: MockerFixture):
        """Test _load_supported_speakers."""
        mock_engine_client = mocker.MagicMock()
@@ -2023,6 +2067,34 @@ class TestMergeBatchItem:
        assert merged.task_type == "CustomVoice"
        assert merged.max_new_tokens == 512

+    def test_non_streaming_mode_batch_default_used(self):
+        """Batch-level non_streaming_mode should be used when item doesn't specify one."""
+        batch = BatchSpeechRequest(
+            items=[SpeechBatchItem(input="hi")],
+            non_streaming_mode=True,
+        )
+
+        merged = OmniOpenAIServingSpeech._merge_batch_item(batch, batch.items[0])
+
+        assert merged.non_streaming_mode is True
+
+    def test_non_streaming_mode_item_override_wins(self):
+        """Per-item false should override a true batch-level default."""
+        batch = BatchSpeechRequest(
+            items=[SpeechBatchItem(input="hi", non_streaming_mode=False)],
+            non_streaming_mode=True,
+        )
+
+        merged = OmniOpenAIServingSpeech._merge_batch_item(batch, batch.items[0])
+
+        assert merged.non_streaming_mode is False
+
+
+def test_streaming_speech_session_config_accepts_non_streaming_mode():
+    config = StreamingSpeechSessionConfig(non_streaming_mode=True)
+
+    assert config.non_streaming_mode is True
+

 class TestAsyncOmniSupportedTasks:
    """Test that AsyncOmni reports correct supported tasks based on output modalities."""
@@ -2811,6 +2883,52 @@ class TestTTSAsyncOffloading:
        qwen3_tts_server._build_tts_params.assert_called_once()
        qwen3_tts_server._estimate_prompt_len_async.assert_awaited_once()

+    def test_prepare_speech_generation_qwen3_voicedesign_non_streaming_mode_false(
+        self, qwen3_tts_server, mocker: MockerFixture
+    ):
+        """VoiceDesign explicit false should reach the model prompt additional_information."""
+        qwen3_tts_server._validate_tts_request = mocker.MagicMock(return_value=None)
+        qwen3_tts_server._estimate_prompt_len_async = mocker.AsyncMock(return_value=512)
+
+        request = OpenAICreateSpeechRequest(
+            input="hello",
+            task_type="VoiceDesign",
+            instructions="warm and calm",
+            non_streaming_mode=False,
+        )
+        _request_id, _generator, tts_params = asyncio.run(qwen3_tts_server._prepare_speech_generation(request))
+
+        assert tts_params["task_type"] == ["VoiceDesign"]
+        assert tts_params["non_streaming_mode"] == [False]
+        prompt = qwen3_tts_server.engine_client.generate.call_args.kwargs["prompt"]
+        assert prompt["additional_information"] is tts_params
+        assert prompt["additional_information"]["non_streaming_mode"] == [False]
+
+    def test_prepare_speech_generation_qwen3_base_non_streaming_mode_true(
+        self, qwen3_tts_server, mocker: MockerFixture
+    ):
+        """Base explicit true should reach the model prompt additional_information."""
+        qwen3_tts_server._validate_tts_request = mocker.MagicMock(return_value=None)
+        qwen3_tts_server._resolve_ref_audio = mocker.AsyncMock(return_value=([0.0] * 48000, 24000))
+        qwen3_tts_server._get_resolved_ref_audio_artifact_key = mocker.MagicMock(return_value=None)
+        qwen3_tts_server._estimate_prompt_len_async = mocker.AsyncMock(return_value=512)
+
+        request = OpenAICreateSpeechRequest(
+            input="hello",
+            task_type="Base",
+            ref_audio="data:audio/wav;base64,abc",
+            ref_text="reference transcript",
+            non_streaming_mode=True,
+        )
+        _request_id, _generator, tts_params = asyncio.run(qwen3_tts_server._prepare_speech_generation(request))
+
+        assert tts_params["task_type"] == ["Base"]
+        assert tts_params["ref_text"] == ["reference transcript"]
+        assert tts_params["non_streaming_mode"] == [True]
+        prompt = qwen3_tts_server.engine_client.generate.call_args.kwargs["prompt"]
+        assert prompt["additional_information"] is tts_params
+        assert prompt["additional_information"]["non_streaming_mode"] == [True]
+
    def test_qwen3_repeated_ref_audio_hot_path_sends_cache_key_without_waveform(self, qwen3_tts_server):
        """After a ref artifact is marked ready, repeated requests avoid ref_audio payload IPC."""
        wav_list = [0.0] * 48000
@@ -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"])
@@ -31,6 +31,7 @@ from vllm_omni.entrypoints.openai.protocol.videos import (
 from vllm_omni.entrypoints.openai.serving_video import OmniOpenAIServingVideo
 from vllm_omni.entrypoints.openai.storage import LocalStorageManager
 from vllm_omni.entrypoints.openai.stores import AsyncDictStore, TaskRegistry
+from vllm_omni.errors import GuardrailViolationError
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams

 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -86,8 +87,10 @@ class BlockingVideoHandler:
        if self.stage_configs is None:
            self.stage_configs = stage_configs

-    async def generate_video_bytes(self, request, reference_id, *, reference_image=None, reference_video=None):
-        del request, reference_id, reference_image, reference_video
+    async def generate_video_bytes(
+        self, request, reference_id, *, reference_image=None, reference_video=None, reference_audio=None
+    ):
+        del request, reference_id, reference_image, reference_video, reference_audio
        self.started.set()
        try:
            await asyncio.Future()
@@ -457,6 +460,41 @@ def test_cosmos3_reference_video_limit_caps_condition_frames_to_output_frames():
    assert api_server._reference_video_decode_spec(request, _cosmos3_stage_configs()).max_frames == 5


+def test_s2v_video_generation_with_audio_reference_form(test_client, mocker: MockerFixture):
+    """Speech-to-video: image + audio_reference (base64 data URL) passes audio path to multi_modal_data."""
+    audio_bytes = b"\xff\xfb\x90\x00" * 50
+    audio_b64 = base64.b64encode(audio_bytes).decode()
+    audio_ref = json.dumps({"audio_url": f"data:audio/mp3;base64,{audio_b64}"})
+
+    mocker.patch(
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
+        return_value=b"fake-video",
+    )
+    response = test_client.post(
+        "/v1/videos",
+        data={
+            "prompt": "A person singing",
+            "audio_reference": audio_ref,
+            "width": "832",
+            "height": "480",
+        },
+        files={"input_reference": ("face.png", _make_test_image_bytes((64, 64)), "image/png")},
+    )
+
+    assert response.status_code == 200
+    video_id = response.json()["id"]
+    _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
+
+    engine = test_client.app.state.openai_serving_video._engine_client
+    prompt = engine.captured_prompt
+    assert "multi_modal_data" in prompt
+    assert "image" in prompt["multi_modal_data"]
+    assert "audio" in prompt["multi_modal_data"]
+    audio_path = prompt["multi_modal_data"]["audio"]
+    assert isinstance(audio_path, str)
+    assert audio_path.endswith(".mp3")
+
+
 def test_seconds_defaults_fps_and_frames(test_client, mocker: MockerFixture):
    fps_values = []

@@ -982,6 +1020,25 @@ def test_invalid_lora_returns_400(test_client):
    assert "lora object" in failed["error"]["message"].lower()


+def test_async_guardrail_error_returns_400_on_retrieve(test_client, mocker: MockerFixture):
+    mocker.patch.object(
+        OmniOpenAIServingVideo,
+        "generate_video_bytes",
+        side_effect=GuardrailViolationError("Input was blocked by Cosmos3 guardrails."),
+    )
+    response = test_client.post("/v1/videos", data={"prompt": "blocked prompt"})
+    assert response.status_code == 200
+
+    video_id = response.json()["id"]
+    failed = _wait_for_status(test_client, video_id, VideoGenerationStatus.FAILED.value)
+    assert failed["error"]["code"] == 400
+    assert failed["error"]["message"] == "Input was blocked by Cosmos3 guardrails."
+
+    retrieve = test_client.get(f"/v1/videos/{video_id}")
+    assert retrieve.status_code == 400
+    assert retrieve.json()["error"]["code"] == 400
+
+
 def test_unsupported_image_reference_file_id_returns_400(test_client):
    response = test_client.post(
        "/v1/videos",
@@ -1484,6 +1541,20 @@ def test_sync_generation_error_returns_500(test_client, mocker: MockerFixture):
    assert "GPU exploded" in response.json()["detail"]


+def test_sync_guardrail_error_returns_400(test_client, mocker: MockerFixture):
+    mocker.patch.object(
+        OmniOpenAIServingVideo,
+        "generate_video_bytes",
+        side_effect=GuardrailViolationError("Input was blocked by Cosmos3 guardrails."),
+    )
+    response = test_client.post(
+        "/v1/videos/sync",
+        data={"prompt": "blocked prompt"},
+    )
+    assert response.status_code == 400
+    assert response.json()["detail"] == "Input was blocked by Cosmos3 guardrails."
+
+
 def test_sync_does_not_create_store_entry(test_client, mocker: MockerFixture):
    """The sync endpoint should NOT leave any record in VIDEO_STORE — it is
    stateless by design."""
@@ -18,6 +18,12 @@ from vllm_omni.entrypoints.async_omni import AsyncOmni
 from vllm_omni.entrypoints.client_request_state import ClientRequestState
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.entrypoints.omni_base import OmniEngineDeadError
+from vllm_omni.errors import (
+    OmniClientError,
+    client_error_from_metadata,
+    client_error_metadata,
+    is_client_error_status,
+)
 from vllm_omni.outputs import OmniRequestOutput

 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -604,6 +610,71 @@ async def test_async_omni_propagates_fatal_error_context(monkeypatch: pytest.Mon
    assert getattr(exc_info.value, "error_stage_id") == 2


+def _enqueue_client_error_message(engine: FakeAsyncOmniEngine, msg: dict[str, Any]) -> None:
+    engine.output_q.put_nowait(
+        ErrorMessage(
+            request_id=msg["request_id"],
+            stage_id=2,
+            error="Input was blocked by Cosmos3 guardrails.",
+            status_code=400,
+            error_type="BadRequestError",
+        )
+    )
+
+
+@pytest.mark.asyncio
+async def test_async_omni_propagates_client_error_status(monkeypatch: pytest.MonkeyPatch):
+    """A non-fatal client error from the orchestrator must be routed to the
+    requesting generate() call (not raised in the shared dispatcher) and
+    surface as an OmniClientError carrying status_code/error_type."""
+    engine = FakeAsyncOmniEngine(stage_metadata=THREE_STAGE_META, on_add_request=_enqueue_client_error_message)
+    _patch_engine(monkeypatch, engine)
+
+    app = AsyncOmni("dummy-model")
+    try:
+        with pytest.raises(OmniClientError) as exc_info:
+            async for _ in app.generate(prompt="blocked", request_id="req-1"):
+                pass
+    finally:
+        app.shutdown()
+
+    assert exc_info.value.status_code == 400
+    assert exc_info.value.error_type == "BadRequestError"
+    assert str(exc_info.value) == "Input was blocked by Cosmos3 guardrails."
+
+
+def _enqueue_server_error_message(engine: FakeAsyncOmniEngine, msg: dict[str, Any]) -> None:
+    engine.output_q.put_nowait(
+        ErrorMessage(
+            request_id=msg["request_id"],
+            stage_id=2,
+            error="GPU exploded",
+        )
+    )
+
+
+@pytest.mark.asyncio
+async def test_async_omni_propagates_server_error_as_runtime(monkeypatch: pytest.MonkeyPatch):
+    """A non-fatal error WITHOUT a 4xx status_code is a server fault: it must
+    surface as a RuntimeError (-> HTTP 500), not an OmniClientError (-> 400).
+    Guards against the fix over-broadly mapping every error to a client error."""
+    engine = FakeAsyncOmniEngine(stage_metadata=THREE_STAGE_META, on_add_request=_enqueue_server_error_message)
+    _patch_engine(monkeypatch, engine)
+
+    app = AsyncOmni("dummy-model")
+    try:
+        # OmniClientError subclasses ValueError, not RuntimeError, so matching
+        # RuntimeError here also proves it was NOT raised as a client error.
+        with pytest.raises(RuntimeError) as exc_info:
+            async for _ in app.generate(prompt="boom", request_id="req-1"):
+                pass
+    finally:
+        app.shutdown()
+
+    assert not isinstance(exc_info.value, OmniClientError)
+    assert str(exc_info.value) == "GPU exploded"
+
+
 def test_omni_generate_py_generator_yields_final_outputs_for_each_request(monkeypatch: pytest.MonkeyPatch):
    sampling_params = [SamplingParams(max_tokens=8) for _ in range(3)]
    engine = FakeAsyncOmniEngine(
@@ -761,6 +832,140 @@ def test_non_fatal_error_raises_runtime():
        base._handle_output_message(msg)


+def test_non_fatal_client_error_raises_omni_client_error():
+    """A non-fatal ErrorMessage carrying a 4xx status_code (e.g. a guardrail
+    block) must surface as an OmniClientError with the metadata intact, not a
+    bare RuntimeError. This covers the offline/sync Omni consumer path."""
+    base = _make_base()
+    msg = ErrorMessage(
+        error="Input was blocked by Cosmos3 guardrails.",
+        status_code=400,
+        error_type="BadRequestError",
+    )
+
+    with pytest.raises(OmniClientError) as exc_info:
+        base._handle_output_message(msg)
+
+    assert exc_info.value.status_code == 400
+    assert exc_info.value.error_type == "BadRequestError"
+    assert str(exc_info.value) == "Input was blocked by Cosmos3 guardrails."
+
+
+_NON_400_CLIENT_ERRORS = [
+    pytest.param(429, "RateLimitError", id="429-too-many-requests"),
+    pytest.param(413, "PayloadTooLargeError", id="413-payload-too-large"),
+    pytest.param(422, "UnprocessableEntityError", id="422-unprocessable-entity"),
+    pytest.param(403, "PermissionDeniedError", id="403-forbidden"),
+]
+
+
+@pytest.mark.parametrize(
+    ("status_code", "error_type"), [pytest.param(400, "BadRequestError", id="400")] + _NON_400_CLIENT_ERRORS
+)
+def test_client_error_metadata_round_trip_preserves_4xx(status_code: int, error_type: str):
+    original = OmniClientError("blocked", status_code=status_code, error_type=error_type)
+
+    # Outbound: what the broad `except Exception` handlers serialize.
+    carried_status, carried_type = client_error_metadata(original)
+    assert carried_status == status_code
+    assert carried_type == error_type
+    assert is_client_error_status(carried_status)
+
+    # Inbound: reconstruction at the consuming end.
+    rebuilt = client_error_from_metadata("blocked", status_code=carried_status, error_type=carried_type)
+    assert isinstance(rebuilt, OmniClientError)
+    assert rebuilt.status_code == status_code
+    assert rebuilt.error_type == error_type
+    assert str(rebuilt) == "blocked"
+
+
+@pytest.mark.parametrize(("status_code", "error_type"), _NON_400_CLIENT_ERRORS)
+def test_non_fatal_client_error_preserves_non_400_status(status_code: int, error_type: str):
+    base = _make_base()
+    msg = ErrorMessage(
+        error="client side failure",
+        status_code=status_code,
+        error_type=error_type,
+    )
+
+    with pytest.raises(OmniClientError) as exc_info:
+        base._handle_output_message(msg)
+
+    assert exc_info.value.status_code == status_code
+    assert exc_info.value.error_type == error_type
+    assert str(exc_info.value) == "client side failure"
+
+
+@pytest.mark.parametrize(
+    "status_code",
+    [
+        pytest.param(None, id="no-status"),
+        pytest.param(399, id="399-below-4xx"),
+        pytest.param(500, id="500-server-error"),
+        pytest.param(503, id="503-service-unavailable"),
+    ],
+)
+def test_non_fatal_non_4xx_status_raises_runtime(status_code: int | None):
+    base = _make_base()
+    msg = ErrorMessage(
+        error="server side failure",
+        status_code=status_code,
+        error_type="ShouldBeIgnoredForNon4xx",
+    )
+
+    with pytest.raises(RuntimeError) as exc_info:
+        base._handle_output_message(msg)
+
+    assert not isinstance(exc_info.value, OmniClientError)
+    assert str(exc_info.value) == "server side failure"
+
+
+def _make_enqueue_client_error(
+    status_code: int,
+    error_type: str,
+    error_text: str,
+) -> Callable[[FakeAsyncOmniEngine, dict[str, Any]], None]:
+    def _enqueue(engine: FakeAsyncOmniEngine, msg: dict[str, Any]) -> None:
+        engine.output_q.put_nowait(
+            ErrorMessage(
+                request_id=msg["request_id"],
+                stage_id=2,
+                error=error_text,
+                status_code=status_code,
+                error_type=error_type,
+            )
+        )
+
+    return _enqueue
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(("status_code", "error_type"), _NON_400_CLIENT_ERRORS)
+async def test_async_omni_propagates_non_400_client_error_status(
+    monkeypatch: pytest.MonkeyPatch,
+    status_code: int,
+    error_type: str,
+):
+    error_text = f"blocked with {status_code}"
+    engine = FakeAsyncOmniEngine(
+        stage_metadata=THREE_STAGE_META,
+        on_add_request=_make_enqueue_client_error(status_code, error_type, error_text),
+    )
+    _patch_engine(monkeypatch, engine)
+
+    app = AsyncOmni("dummy-model")
+    try:
+        with pytest.raises(OmniClientError) as exc_info:
+            async for _ in app.generate(prompt="blocked", request_id="req-1"):
+                pass
+    finally:
+        app.shutdown()
+
+    assert exc_info.value.status_code == status_code
+    assert exc_info.value.error_type == error_type
+    assert str(exc_info.value) == error_text
+
+
 def test_async_omni_errored_property_alive():
    omni = object.__new__(AsyncOmni)
    omni.engine = SimpleNamespace(
@@ -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).
@@ -1,5 +1,6 @@
 """Synthetic media generation and media/text utilities for tests."""

+import atexit
 import base64
 import concurrent.futures
 import gc
@@ -14,7 +15,6 @@ import re
 import subprocess
 import tempfile
 import time
-import uuid
 from contextlib import contextmanager
 from pathlib import Path
 from typing import Any
@@ -683,7 +683,10 @@ def convert_audio_file_to_text(output_path: str) -> str:


 def convert_audio_bytes_to_text(raw_bytes: bytes) -> str:
-    output_path = f"./test_{uuid.uuid4().hex}.wav"
+    output_fd, output_path = tempfile.mkstemp(prefix="test_", suffix=".wav")
+    os.close(output_fd)
+    if os.environ.get("VLLM_OMNI_KEEP_REQUEST_MEDIA", "").lower() not in ("1", "true", "yes"):
+        atexit.register(Path(output_path).unlink, missing_ok=True)
    data, samplerate = sf.read(io.BytesIO(raw_bytes))
    sf.write(output_path, data, samplerate, format="WAV", subtype="PCM_16")
    print(f"audio data is saved: {output_path}")
@@ -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
@@ -1473,6 +1477,8 @@ class OpenAIClientHandler:

        - ``send_frames``: optional ``str`` or sequence of ``str`` raw WebSocket text frames (omit when the server
          speaks first, e.g. ``/v1/realtime`` rejection path).
+        - ``ws_skip_types``: optional event ``type`` strings to ignore while waiting for the first matching frame
+          (e.g. ``["session.created"]`` on ``/v1/realtime``).
        - ``timeout``: seconds to wait for the first inbound text frame (default ``120``).
        - ``ws_max_size``: passed through as ``max_size`` to :func:`websockets.connect` when the key is present.
        """
@@ -1486,6 +1492,7 @@ class OpenAIClientHandler:

        timeout = float(cfg.get("timeout", 120.0))
        uri = self._build_ws_url(path)
+        skip_types = set(cfg.get("ws_skip_types") or [])

        connect_kw: dict[str, Any] = {}
        if "ws_max_size" in cfg:
@@ -1497,16 +1504,19 @@ class OpenAIClientHandler:
            async with websockets.connect(uri, **connect_kw) as ws:
                for frame in frames:
                    await ws.send(frame)
-                raw = await asyncio.wait_for(ws.recv(), timeout=timeout)
-                if not isinstance(raw, str):
-                    raise AssertionError(f"Expected JSON text frame from {uri}, got {type(raw).__name__}")
-                try:
-                    data = json.loads(raw)
-                except json.JSONDecodeError as exc:
-                    raise AssertionError(f"Expected JSON text frame from {uri}, body={raw[:500]!r}") from exc
-                if not isinstance(data, dict):
-                    raise AssertionError(f"Expected JSON object from {uri}, got {type(data).__name__}")
-                return WebSocketJsonResponse(json_body=data)
+                while True:
+                    raw = await asyncio.wait_for(ws.recv(), timeout=timeout)
+                    if not isinstance(raw, str):
+                        raise AssertionError(f"Expected JSON text frame from {uri}, got {type(raw).__name__}")
+                    try:
+                        data = json.loads(raw)
+                    except json.JSONDecodeError as exc:
+                        raise AssertionError(f"Expected JSON text frame from {uri}, body={raw[:500]!r}") from exc
+                    if not isinstance(data, dict):
+                        raise AssertionError(f"Expected JSON object from {uri}, got {type(data).__name__}")
+                    if skip_types and data.get("type") in skip_types:
+                        continue
+                    return WebSocketJsonResponse(json_body=data)

        resp = asyncio.run(_recv_first_json_object())
        _run_ws_expectations_from_request_config(cfg, resp)
@@ -0,0 +1,361 @@
+"""Tests for MossTTSCUDAGraphCodecWrapper numerical equivalence.
+
+Verifies that CUDA Graph-accelerated decoding produces results equivalent
+to eager mode, with special attention to the two-argument _decode interface
+(codes [NQ, 1, T] + lengths [1]) and the NQ-first input convention.
+"""
+
+import pytest
+import torch
+import torch.nn as nn
+
+from vllm_omni.model_executor.models.moss_tts.audio_tokenizer import (
+    MossAudioTokenizerDecoderOutput,
+)
+from vllm_omni.model_executor.models.moss_tts.moss_codec_cudagraph import (
+    MossTTSCUDAGraphCodecWrapper,
+)
+
+pytestmark = [pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")]
+
+DEVICE = torch.device("cuda:0")
+NUM_QUANTIZERS = 8
+DOWNSAMPLE_RATE = 4  # synthetic; real checkpoint uses 1920
+
+
+# ---------------------------------------------------------------------------
+# Synthetic codec model
+# ---------------------------------------------------------------------------
+
+
+class SyntheticCodecModel(nn.Module):
+    """Minimal stand-in for MossAudioTokenizerModel.
+
+    Exposes the same two-argument _decode(codes, lengths) interface and
+    returns a MossAudioTokenizerDecoderOutput, mirroring the real model.
+
+    Input:
+        codes:   [NQ, B, T]  long  — RVQ codes
+        lengths: [B]          long  — valid frame counts
+    Output:
+        MossAudioTokenizerDecoderOutput(audio=[B, 1, T*upsample], audio_lengths=[B])
+    """
+
+    def __init__(
+        self,
+        num_quantizers: int = NUM_QUANTIZERS,
+        downsample_rate: int = DOWNSAMPLE_RATE,
+    ):
+        super().__init__()
+        self.downsample_rate = downsample_rate
+        hidden = 32
+        # embed: treat NQ quantizers as input channels for Conv1d
+        self.embed = nn.Conv1d(num_quantizers, hidden, kernel_size=3, padding=1)
+        self.conv = nn.Conv1d(hidden, hidden, kernel_size=3, padding=1)
+        self.upsample = nn.ConvTranspose1d(hidden, 1, kernel_size=downsample_rate, stride=downsample_rate)
+
+    def _decode(self, codes: torch.Tensor, lengths: torch.Tensor) -> MossAudioTokenizerDecoderOutput:
+        """codes: [NQ, B, T], lengths: [B] → audio [B, 1, T*upsample]."""
+        nq, b, t = codes.shape
+        # treat NQ as channel dim → [B, NQ, T] for Conv1d (expects [B, C, T])
+        x = codes.permute(1, 0, 2).to(dtype=self.embed.weight.dtype)  # [B, NQ, T]
+        x = torch.relu(self.embed(x))  # [B, hidden, T]
+        x = torch.relu(self.conv(x))  # [B, hidden, T]
+        audio = self.upsample(x)  # [B, 1, T * upsample]
+        audio_lengths = lengths * self.downsample_rate
+        return MossAudioTokenizerDecoderOutput(audio=audio, audio_lengths=audio_lengths)
+
+    def batch_decode(
+        self,
+        codes_list: list[torch.Tensor],
+        num_quantizers: int | None = None,
+    ) -> MossAudioTokenizerDecoderOutput:
+        """Mirrors MossAudioTokenizerModel.batch_decode (fallback path)."""
+        device = codes_list[0].device
+        nq = num_quantizers or codes_list[0].shape[0]
+        max_t = max(c.shape[-1] for c in codes_list)
+        codes = torch.zeros(nq, len(codes_list), max_t, device=device, dtype=torch.long)
+        lengths = torch.zeros(len(codes_list), device=device, dtype=torch.long)
+        for i, c in enumerate(codes_list):
+            codes[:nq, i, : c.shape[-1]] = c[:nq]
+            lengths[i] = c.shape[-1]
+        return self._decode(codes, lengths)
+
+
+# ---------------------------------------------------------------------------
+# Fixtures
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture(scope="module")
+def model():
+    torch.manual_seed(0)
+    return SyntheticCodecModel().to(DEVICE).eval()
+
+
+@pytest.fixture(scope="module")
+def wrapper(model):
+    w = MossTTSCUDAGraphCodecWrapper(
+        model=model,
+        capture_sizes=[25, 50, 100],
+        num_quantizers=NUM_QUANTIZERS,
+        enabled=True,
+    )
+    w.warmup(DEVICE)
+    return w
+
+
+def _random_codes(t: int, device: torch.device = DEVICE) -> torch.Tensor:
+    """Return [NQ, T] long tensor (single-request convention)."""
+    return torch.randint(0, 64, (NUM_QUANTIZERS, t), dtype=torch.long, device=device)
+
+
+def _eager_decode(model: SyntheticCodecModel, codes_nq_t: torch.Tensor) -> MossAudioTokenizerDecoderOutput:
+    """Run the model directly in eager mode (reference output)."""
+    t = codes_nq_t.shape[-1]
+    codes_nq_1_t = codes_nq_t.unsqueeze(1)  # [NQ, 1, T]
+    lengths = torch.tensor([t], dtype=torch.long, device=codes_nq_t.device)
+    with torch.no_grad():
+        return model._decode(codes_nq_1_t, lengths)
+
+
+# ---------------------------------------------------------------------------
+# 1. Exact-size inputs — must be bit-identical
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("t", [25, 50, 100])
+def test_exact_size_bit_identical(model, wrapper, t):
+    codes = _random_codes(t)
+    ref = _eager_decode(model, codes)
+    with torch.no_grad():
+        out = wrapper.decode(codes)
+    torch.testing.assert_close(out.audio, ref.audio, atol=0, rtol=0)
+
+
+# ---------------------------------------------------------------------------
+# 2. Padded inputs — output trimmed to actual length
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("t", [10, 30, 47, 73, 99])
+def test_padded_output_shape(model, wrapper, t):
+    codes = _random_codes(t)
+    ref = _eager_decode(model, codes)
+    with torch.no_grad():
+        out = wrapper.decode(codes)
+    expected_len = t * DOWNSAMPLE_RATE
+    assert out.audio.shape[-1] == expected_len, f"got {out.audio.shape[-1]}, want {expected_len}"
+    assert out.audio.shape == ref.audio.shape
+
+
+@pytest.mark.parametrize("t", [10, 30, 47, 73, 99])
+def test_padded_interior_positions_close(model, wrapper, t):
+    """Positions well away from the zero-padding boundary must be numerically close."""
+    codes = _random_codes(t)
+    ref = _eager_decode(model, codes)
+    with torch.no_grad():
+        out = wrapper.decode(codes)
+    boundary = 2 * DOWNSAMPLE_RATE
+    if ref.audio.shape[-1] > boundary:
+        torch.testing.assert_close(
+            out.audio[..., :-boundary],
+            ref.audio[..., :-boundary],
+            atol=1e-5,
+            rtol=1e-5,
+        )
+
+
+# ---------------------------------------------------------------------------
+# 3. Fallback to eager (T > all capture sizes)
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("t", [101, 150, 200])
+def test_fallback_exact_match(model, wrapper, t):
+    codes = _random_codes(t)
+    ref = _eager_decode(model, codes)
+    with torch.no_grad():
+        out = wrapper.decode(codes)
+    torch.testing.assert_close(out.audio, ref.audio, atol=0, rtol=0)
+
+
+# ---------------------------------------------------------------------------
+# 4. Static buffer not aliased across calls
+# ---------------------------------------------------------------------------
+
+
+def test_output_not_aliased_after_later_replay(model, wrapper):
+    """clone() in decode() must prevent later replays from overwriting the result."""
+    codes_a = _random_codes(50)
+    codes_b = _random_codes(50)
+    with torch.no_grad():
+        out_a = wrapper.decode(codes_a)
+        saved = out_a.audio.clone()
+        _ = wrapper.decode(codes_b)  # replay overwrites static buffer
+    torch.testing.assert_close(out_a.audio, saved, atol=0, rtol=0)
+
+
+def test_deterministic_across_calls(model, wrapper):
+    codes = _random_codes(30)
+    with torch.no_grad():
+        out1 = wrapper.decode(codes)
+        out2 = wrapper.decode(codes)
+    torch.testing.assert_close(out1.audio, out2.audio, atol=0, rtol=0)
+
+
+# ---------------------------------------------------------------------------
+# 5. Disabled wrapper falls back to eager
+# ---------------------------------------------------------------------------
+
+
+def test_disabled_falls_back_to_eager(model, wrapper):
+    codes = _random_codes(30)
+    ref = _eager_decode(model, codes)
+    wrapper.enabled = False
+    with torch.no_grad():
+        out = wrapper.decode(codes)
+    wrapper.enabled = True
+    torch.testing.assert_close(out.audio, ref.audio, atol=0, rtol=0)
+
+
+# ---------------------------------------------------------------------------
+# 6. audio_lengths is consistent with audio shape
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("t", [25, 30, 50, 99, 100])
+def test_audio_lengths_consistent(model, wrapper, t):
+    codes = _random_codes(t)
+    with torch.no_grad():
+        out = wrapper.decode(codes)
+    assert out.audio_lengths is not None
+    assert int(out.audio_lengths[0].item()) == out.audio.shape[-1]
+
+
+# ---------------------------------------------------------------------------
+# 7. NQ-first input layout: codes_nq_t is [NQ, T], not [T, NQ]
+# ---------------------------------------------------------------------------
+
+
+def test_nq_first_layout_matches_eager(model, wrapper):
+    """Verify that wrapper correctly interprets [NQ, T] (NQ-first) layout."""
+    t = 25
+    codes = _random_codes(t)  # [NQ, T]
+    ref = _eager_decode(model, codes)
+    with torch.no_grad():
+        out = wrapper.decode(codes)
+    torch.testing.assert_close(out.audio, ref.audio, atol=0, rtol=0)
+
+
+# ---------------------------------------------------------------------------
+# 8. Non-contiguous input — copy_() must handle strided tensors correctly
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("t", [20, 47])
+def test_noncontiguous_input_matches_contiguous(model, wrapper, t):
+    """Padded path uses copy_() on a slice; must produce the same result
+    whether codes_nq_t is contiguous or a non-contiguous view (e.g. a column
+    slice of a larger buffer, which is the realistic production layout when the
+    caller slices out one request from a batched tensor)."""
+    codes_contiguous = _random_codes(t)
+    # Build a non-contiguous view: embed in a larger [NQ, 2, T] buffer and
+    # take column 0.  The resulting tensor has stride (2*T, 2, 1) → non-contiguous.
+    buf = torch.zeros(NUM_QUANTIZERS, 2, t, dtype=torch.long, device=DEVICE)
+    buf[:, 0, :] = codes_contiguous
+    codes_noncontig = buf[:, 0, :]  # [NQ, T], non-contiguous
+    assert not codes_noncontig.is_contiguous()
+
+    with torch.no_grad():
+        out_c = wrapper.decode(codes_contiguous)
+        out_nc = wrapper.decode(codes_noncontig)
+    torch.testing.assert_close(out_c.audio, out_nc.audio, atol=0, rtol=0)
+
+
+# ---------------------------------------------------------------------------
+# 9. Zero-padding uses token 0 — must be a valid codebook index (no OOB)
+# ---------------------------------------------------------------------------
+
+
+def test_zero_padding_no_cuda_error(model, wrapper):
+    """Padded inputs fill unused frames with code 0.  Verify that the zero-
+    padded region does not corrupt the output shape.
+
+    Note: SyntheticCodecModel uses Conv1d (not nn.Embedding), so it cannot
+    trigger a real CUDA index-out-of-bounds on large token ids.  This test
+    validates shape correctness only; OOB safety on the real codec requires
+    an integration test with MossAudioTokenizerModel."""
+    t = 7  # well below the 25-frame bucket → large zero-padded region
+    codes = _random_codes(t)
+    with torch.no_grad():
+        out = wrapper.decode(codes)
+    torch.accelerator.synchronize()  # surface any deferred CUDA errors
+    assert out.audio.shape[-1] == t * DOWNSAMPLE_RATE
+
+
+# ---------------------------------------------------------------------------
+# 10. Shape alignment: graph path and eager fallback return identical shapes
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("t", [25, 30, 101])
+def test_graph_and_eager_shape_identical(model, wrapper, t):
+    """Graph replay (t<=100) and eager fallback (t>100) must return tensors
+    with the same number of dimensions and the same batch/channel layout so
+    downstream consumers can treat both paths uniformly."""
+    codes = _random_codes(t)
+    ref = _eager_decode(model, codes)
+    with torch.no_grad():
+        out = wrapper.decode(codes)
+    assert out.audio.ndim == ref.audio.ndim, f"ndim mismatch: graph={out.audio.ndim} eager={ref.audio.ndim}"
+    assert out.audio.shape[:-1] == ref.audio.shape[:-1], (
+        f"batch/channel shape mismatch: graph={out.audio.shape} eager={ref.audio.shape}"
+    )
+
+
+# ---------------------------------------------------------------------------
+# 11. Half-precision weights: graph and eager must agree under fp16/bf16
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+def test_half_precision_graph_eager_agree(dtype):
+    """CUDA Graph must produce outputs consistent with eager when model weights
+    are in half precision.  The wrapper itself stores codes as long and lengths
+    as long; only the model's conv weights change dtype."""
+    torch.manual_seed(0)
+    model_hp = SyntheticCodecModel().to(device=DEVICE, dtype=dtype).eval()
+    wrapper_hp = MossTTSCUDAGraphCodecWrapper(
+        model=model_hp,
+        capture_sizes=[25, 50],
+        num_quantizers=NUM_QUANTIZERS,
+        enabled=True,
+    )
+    wrapper_hp.warmup(DEVICE)
+
+    codes = _random_codes(25)
+    ref = _eager_decode(model_hp, codes)
+    with torch.no_grad():
+        out = wrapper_hp.decode(codes)
+    # Half-precision accumulation introduces small numerical error; use
+    # tolerances that match typical fp16/bf16 rounding.
+    atol = 1e-2 if dtype == torch.float16 else 5e-2
+    torch.testing.assert_close(out.audio.float(), ref.audio.float(), atol=atol, rtol=1e-3)
+
+
+# ---------------------------------------------------------------------------
+# 12. Empty input (T=0): wrapper must not crash and return zero-length audio
+# ---------------------------------------------------------------------------
+
+
+def test_empty_input_t0(model, wrapper):
+    """T=0 input must be handled gracefully: no exception, audio length == 0.
+
+    The outer MossTTSCodecDecoder guards against empty segments, but the
+    wrapper should also handle T=0 without crashing so it is safe to call
+    directly."""
+    codes = _random_codes(0)  # [NQ, 0]
+    with torch.no_grad():
+        out = wrapper.decode(codes)
+    assert out.audio.shape[-1] == 0, f"expected empty audio, got shape {out.audio.shape}"
@@ -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]
@@ -0,0 +1,918 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Tests for StepAudioReasoningParser.
+
+Step-Audio supports two representations of thinking markers:
+
+1. **Special tokens**: ``<|THINK_START|>`` and ``<|THINK_END|>``
+   (single-token IDs in vocab).
+
+2. **Text markers**: ``<think>`` and ``</think>``
+   (multi-token sequences for </think>, e.g. [522, 26865]).
+
+These tests use standalone mock objects so they can run without
+installing the full vLLM/torch stack.
+"""
+
+import sys
+from types import ModuleType
+from unittest.mock import MagicMock
+
+import pytest
+
+# ---------------------------------------------------------------------------
+# Bootstrap minimal vLLM stubs so the parser can be imported
+# ---------------------------------------------------------------------------
+
+
+def _install_vllm_stubs() -> dict[str, ModuleType | None]:
+    """Temporarily install minimal stub modules for the vllm imports the
+    parser needs, so it can be imported without the full vLLM/torch stack.
+
+    Returns a snapshot of the original ``sys.modules`` entries that were
+    replaced. The caller MUST pass it to ``_restore_vllm_stubs`` once the
+    parser is imported. These stubs are process-global: if they leak, they
+    shadow the real vllm modules for every other test in the same worker
+    (and any forked server subprocess), e.g. replacing
+    ``vllm.entrypoints.openai.engine.protocol`` with a stub that lacks
+    ``ErrorResponse`` and breaking unrelated server-startup tests.
+    """
+    stubs: dict[str, ModuleType] = {}
+
+    # DeltaMessage stub
+    delta_mod = ModuleType("vllm.entrypoints.openai.engine.protocol")
+
+    class DeltaMessage:
+        __slots__ = ("reasoning", "content")
+
+        def __init__(self, reasoning=None, content=None):
+            self.reasoning = reasoning
+            self.content = content
+
+    delta_mod.DeltaMessage = DeltaMessage
+    stubs["vllm.entrypoints.openai.engine.protocol"] = delta_mod
+
+    # ReasoningParser stub
+    reasoning_mod = ModuleType("vllm.reasoning")
+
+    class ReasoningParser:
+        def __init__(self, tokenizer, *args, **kwargs):
+            self.model_tokenizer = tokenizer
+
+        @property
+        def vocab(self):
+            if not self.model_tokenizer:
+                return {}
+            return self.model_tokenizer.get_vocab()
+
+    class _ReasoningParserManager:
+        reasoning_parsers = {}
+        lazy_parsers = {}
+
+        @classmethod
+        def register_lazy_module(cls, name, module_path, class_name):
+            cls.lazy_parsers[name] = (module_path, class_name)
+
+    reasoning_mod.ReasoningParser = ReasoningParser
+    reasoning_mod.ReasoningParserManager = _ReasoningParserManager
+    stubs["vllm.reasoning"] = reasoning_mod
+    stubs["vllm.reasoning.abs_reasoning_parsers"] = reasoning_mod
+
+    # TokenizerLike stub
+    tokenizers_mod = ModuleType("vllm.tokenizers")
+    tokenizers_mod.TokenizerLike = object
+    stubs["vllm.tokenizers"] = tokenizers_mod
+
+    # ChatCompletionRequest / ResponsesRequest stubs
+    chat_mod = ModuleType("vllm.entrypoints.openai.chat_completion.protocol")
+    chat_mod.ChatCompletionRequest = MagicMock
+    stubs["vllm.entrypoints.openai.chat_completion.protocol"] = chat_mod
+    responses_mod = ModuleType("vllm.entrypoints.openai.responses.protocol")
+    responses_mod.ResponsesRequest = MagicMock
+    stubs["vllm.entrypoints.openai.responses.protocol"] = responses_mod
+
+    # Parent packages: only stub the ones not already importable, so we never
+    # clobber a real (installed) vllm package tree.
+    for mod_path in [
+        "vllm",
+        "vllm.entrypoints",
+        "vllm.entrypoints.openai",
+        "vllm.entrypoints.openai.engine",
+        "vllm.entrypoints.openai.chat_completion",
+        "vllm.entrypoints.openai.responses",
+    ]:
+        if mod_path not in sys.modules and mod_path not in stubs:
+            stubs[mod_path] = ModuleType(mod_path)
+
+    saved = {name: sys.modules.get(name) for name in stubs}
+    sys.modules.update(stubs)
+    return saved
+
+
+def _restore_vllm_stubs(saved: dict[str, ModuleType | None]) -> None:
+    """Undo ``_install_vllm_stubs`` so ``sys.modules`` is left untouched."""
+    for name, original in saved.items():
+        if original is None:
+            sys.modules.pop(name, None)
+        else:
+            sys.modules[name] = original
+
+
+# Install stubs only for the duration of the import. The parser binds these
+# stub classes into its own module namespace at import time, so it keeps using
+# them at runtime even after sys.modules is restored below.
+_saved_modules = _install_vllm_stubs()
+try:
+    from vllm_omni.reasoning.step_audio_reasoning_parser import (  # noqa: E402
+        StepAudioReasoningParser,
+    )
+finally:
+    _restore_vllm_stubs(_saved_modules)
+
+# Short aliases for marker constants
+TS = StepAudioReasoningParser.THINK_START_TEXT  # <think>
+TE = StepAudioReasoningParser.THINK_END_TEXT  # </think>
+TSS = StepAudioReasoningParser.THINK_START_SPECIAL  # <|THINK_START|>
+TES = StepAudioReasoningParser.THINK_END_SPECIAL  # <|THINK_END|>
+
+
+# ---------------------------------------------------------------------------
+# Mock tokenizers
+# ---------------------------------------------------------------------------
+
+
+class SingleTokenTokenizer:
+    """Tokenizer where both ``<think>`` and ``</think>`` are single tokens.
+
+    This exercises the fast (token-ID) path.
+    Also includes special-token forms ``<|THINK_START|>`` and ``<|THINK_END|>``.
+    """
+
+    def __init__(self):
+        self._vocab = {
+            TS: 100,  # <think>
+            TE: 101,  # </think>
+            TSS: 151669,  # <|THINK_START|>
+            TES: 151670,  # <|THINK_END|>
+            "Hello": 1,
+            " world": 2,
+            "reasoning": 3,
+            " content": 4,
+            "\n": 5,
+            "</": 6,
+            "think>": 7,
+            "some text": 8,
+        }
+        self._reverse = {v: k for k, v in self._vocab.items()}
+
+    def get_vocab(self):
+        return dict(self._vocab)
+
+    def decode(self, ids, skip_special_tokens=False):
+        return "".join(self._reverse.get(i, "") for i in ids)
+
+    def tokenize(self, text):
+        if not text:
+            return []
+        return [text]
+
+
+class MultiTokenTokenizer:
+    """Tokenizer where ``</think>`` is NOT a single token.
+
+    ``<think>`` is in the vocab (ID 100), but ``</think>`` is NOT.
+    Instead, ``</think>`` encodes as ["</", "think>"] → [522, 26865].
+
+    Both special tokens ``<|THINK_START|>`` and ``<|THINK_END|>`` ARE
+    present as single tokens.
+
+    This exercises the text-based path (with special-token fallback).
+    """
+
+    def __init__(self):
+        self._vocab = {
+            TS: 100,  # <think> - single token
+            # </think> is intentionally NOT in the vocab
+            "</": 522,
+            "think>": 26865,
+            TSS: 151669,  # <|THINK_START|>
+            TES: 151670,  # <|THINK_END|>
+            "Hello": 1,
+            " world": 2,
+            "reasoning": 3,
+            " content": 4,
+            "\n": 5,
+            "some text": 8,
+        }
+        self._reverse = {v: k for k, v in self._vocab.items()}
+
+    def get_vocab(self):
+        return dict(self._vocab)
+
+    def decode(self, ids, skip_special_tokens=False):
+        return "".join(self._reverse.get(i, "") for i in ids)
+
+    def tokenize(self, text):
+        if not text:
+            return []
+        return [text]
+
+
+class TextOnlyTokenizer:
+    """Tokenizer with only text-form markers (no special tokens at all).
+
+    Both ``<think>`` and ``</think>`` are single tokens.
+    ``<|THINK_START|>`` and ``<|THINK_END|>`` are NOT in the vocab.
+
+    This tests the pure text-based fallback paths.
+    """
+
+    def __init__(self):
+        self._vocab = {
+            TS: 100,  # <think>
+            TE: 101,  # </think>
+            "Hello": 1,
+            " world": 2,
+            "reasoning": 3,
+            " content": 4,
+            "\n": 5,
+            "</": 6,
+            "think>": 7,
+        }
+        self._reverse = {v: k for k, v in self._vocab.items()}
+
+    def get_vocab(self):
+        return dict(self._vocab)
+
+    def decode(self, ids, skip_special_tokens=False):
+        return "".join(self._reverse.get(i, "") for i in ids)
+
+    def tokenize(self, text):
+        if not text:
+            return []
+        return [text]
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _stream_extract(parser, deltas_and_ids):
+    """Run streaming extraction over a list of (delta_text, delta_ids) tuples.
+    Returns (reasoning_text, content_text)."""
+    prev_text = ""
+    prev_ids = []
+    all_r = []
+    all_c = []
+
+    for delta_text, delta_ids in deltas_and_ids:
+        cur_text = prev_text + delta_text
+        cur_ids = prev_ids + delta_ids
+        result = parser.extract_reasoning_streaming(prev_text, cur_text, delta_text, prev_ids, cur_ids, delta_ids)
+        if result is not None:
+            if result.reasoning:
+                all_r.append(result.reasoning)
+            if result.content:
+                all_c.append(result.content)
+        prev_text = cur_text
+        prev_ids = cur_ids
+
+    return "".join(all_r), "".join(all_c)
+
+
+# ---------------------------------------------------------------------------
+# Fixtures
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+def single_token_parser():
+    return StepAudioReasoningParser(SingleTokenTokenizer())
+
+
+@pytest.fixture
+def multi_token_parser():
+    return StepAudioReasoningParser(MultiTokenTokenizer())
+
+
+@pytest.fixture
+def text_only_parser():
+    return StepAudioReasoningParser(TextOnlyTokenizer())
+
+
+@pytest.fixture
+def request_obj():
+    return MagicMock(messages=[], model="test-model")
+
+
+# ---------------------------------------------------------------------------
+# Initialization tests
+# ---------------------------------------------------------------------------
+
+
+class TestInit:
+    def test_single_token_ids(self, single_token_parser):
+        """Text-form markers are in vocab as single tokens."""
+        assert single_token_parser.think_start_text_id == 100
+        assert single_token_parser.think_end_text_id == 101
+
+    def test_single_token_special_ids(self, single_token_parser):
+        """Special-token markers are also in vocab."""
+        assert single_token_parser.think_start_special_id == 151669
+        assert single_token_parser.think_end_special_id == 151670
+
+    def test_single_token_has_ids(self, single_token_parser):
+        assert single_token_parser._has_start_token_id is True
+        assert single_token_parser._has_end_token_id is True
+
+    def test_multi_token_text_ids(self, multi_token_parser):
+        """<think> is a single token; </think> is NOT."""
+        assert multi_token_parser.think_start_text_id == 100
+        assert multi_token_parser.think_end_text_id == -1
+
+    def test_multi_token_special_ids(self, multi_token_parser):
+        """Special tokens are still in vocab."""
+        assert multi_token_parser.think_start_special_id == 151669
+        assert multi_token_parser.think_end_special_id == 151670
+
+    def test_multi_token_has_ids(self, multi_token_parser):
+        """Even without </think> as single token, special tokens provide IDs."""
+        assert multi_token_parser._has_start_token_id is True
+        assert multi_token_parser._has_end_token_id is True
+
+    def test_text_only_ids(self, text_only_parser):
+        """Only text-form markers, no special tokens."""
+        assert text_only_parser.think_start_text_id == 100
+        assert text_only_parser.think_end_text_id == 101
+        assert text_only_parser.think_start_special_id == -1
+        assert text_only_parser.think_end_special_id == -1
+        assert text_only_parser._has_start_token_id is True
+        assert text_only_parser._has_end_token_id is True
+
+    def test_missing_tokenizer_raises(self):
+        with pytest.raises(ValueError, match="model tokenizer must be passed"):
+            StepAudioReasoningParser(None)
+
+
+# ---------------------------------------------------------------------------
+# extract_reasoning (non-streaming)
+# ---------------------------------------------------------------------------
+
+
+class TestExtractReasoning:
+    @pytest.mark.parametrize(
+        "parser_name",
+        [
+            "single_token_parser",
+            "multi_token_parser",
+            "text_only_parser",
+        ],
+    )
+    def test_basic_extraction(self, request, parser_name, request_obj):
+        parser = request.getfixturevalue(parser_name)
+        output = f"some reasoning{TE}Hello world"
+        reasoning, content = parser.extract_reasoning(output, request_obj)
+        assert reasoning == "some reasoning"
+        assert content == "Hello world"
+
+    @pytest.mark.parametrize(
+        "parser_name",
+        [
+            "single_token_parser",
+            "multi_token_parser",
+            "text_only_parser",
+        ],
+    )
+    def test_no_end_token(self, request, parser_name, request_obj):
+        parser = request.getfixturevalue(parser_name)
+        output = "still thinking..."
+        reasoning, content = parser.extract_reasoning(output, request_obj)
+        assert reasoning == "still thinking..."
+        assert content is None
+
+    @pytest.mark.parametrize(
+        "parser_name",
+        [
+            "single_token_parser",
+            "multi_token_parser",
+            "text_only_parser",
+        ],
+    )
+    def test_with_start_token_prefix(self, request, parser_name, request_obj):
+        parser = request.getfixturevalue(parser_name)
+        output = f"{TS}my reasoning{TE}answer"
+        reasoning, content = parser.extract_reasoning(output, request_obj)
+        assert reasoning == "my reasoning"
+        assert content == "answer"
+
+    @pytest.mark.parametrize(
+        "parser_name",
+        [
+            "single_token_parser",
+            "multi_token_parser",
+            "text_only_parser",
+        ],
+    )
+    def test_strips_newline_after_end(self, request, parser_name, request_obj):
+        parser = request.getfixturevalue(parser_name)
+        output = f"reasoning{TE}\nHello"
+        reasoning, content = parser.extract_reasoning(output, request_obj)
+        assert reasoning == "reasoning"
+        assert content == "Hello"
+
+    @pytest.mark.parametrize(
+        "parser_name",
+        [
+            "single_token_parser",
+            "multi_token_parser",
+            "text_only_parser",
+        ],
+    )
+    def test_empty_reasoning(self, request, parser_name, request_obj):
+        parser = request.getfixturevalue(parser_name)
+        output = f"{TE}Hello"
+        reasoning, content = parser.extract_reasoning(output, request_obj)
+        assert reasoning is None
+        assert content == "Hello"
+
+    @pytest.mark.parametrize(
+        "parser_name",
+        [
+            "single_token_parser",
+            "multi_token_parser",
+        ],
+    )
+    def test_special_end_token(self, request, parser_name, request_obj):
+        """Test with <|THINK_END|> special token form."""
+        parser = request.getfixturevalue(parser_name)
+        output = f"reasoning{TES}answer"
+        reasoning, content = parser.extract_reasoning(output, request_obj)
+        assert reasoning == "reasoning"
+        assert content == "answer"
+
+    @pytest.mark.parametrize(
+        "parser_name",
+        [
+            "single_token_parser",
+            "multi_token_parser",
+        ],
+    )
+    def test_special_start_token(self, request, parser_name, request_obj):
+        """Test with <|THINK_START|> special token form."""
+        parser = request.getfixturevalue(parser_name)
+        output = f"{TSS}reasoning{TE}answer"
+        reasoning, content = parser.extract_reasoning(output, request_obj)
+        assert reasoning == "reasoning"
+        assert content == "answer"
+
+
+# ---------------------------------------------------------------------------
+# is_reasoning_end
+# ---------------------------------------------------------------------------
+
+
+class TestIsReasoningEnd:
+    def test_single_token_text_end(self, single_token_parser):
+        assert single_token_parser.is_reasoning_end([1, 2, 101]) is True
+
+    def test_single_token_special_end(self, single_token_parser):
+        assert single_token_parser.is_reasoning_end([1, 2, 151670]) is True
+
+    def test_single_token_no_end(self, single_token_parser):
+        assert single_token_parser.is_reasoning_end([1, 2, 3]) is False
+
+    def test_multi_token_text_end_decoded(self, multi_token_parser):
+        """</think> = [522, 26865] detected via decode."""
+        assert multi_token_parser.is_reasoning_end([1, 522, 26865]) is True
+
+    def test_multi_token_special_end(self, multi_token_parser):
+        """<|THINK_END|> detected as single token ID."""
+        assert multi_token_parser.is_reasoning_end([1, 151670]) is True
+
+    def test_multi_token_no_end(self, multi_token_parser):
+        assert multi_token_parser.is_reasoning_end([1, 2, 3]) is False
+
+    def test_multi_token_partial_marker(self, multi_token_parser):
+        """Only </ without think> should NOT count."""
+        assert multi_token_parser.is_reasoning_end([1, 522, 3]) is False
+
+    def test_multi_token_empty_input(self, multi_token_parser):
+        assert multi_token_parser.is_reasoning_end([]) is False
+
+    def test_text_only_end(self, text_only_parser):
+        assert text_only_parser.is_reasoning_end([1, 2, 101]) is True
+
+    def test_text_only_no_end(self, text_only_parser):
+        assert text_only_parser.is_reasoning_end([1, 2, 3]) is False
+
+    def test_start_after_end_returns_false(self, single_token_parser):
+        """Multi-turn prompt: start marker AFTER end marker → reasoning
+        is still active (the generation prompt's start marker)."""
+        # IDs: [start=100, end=101, start=100] → last marker is start
+        assert single_token_parser.is_reasoning_end([100, 101, 100]) is False
+
+    def test_end_is_last_marker_returns_true(self, single_token_parser):
+        """End marker is the last marker → reasoning has ended."""
+        # IDs: [start=100, end=101] → last marker is end
+        assert single_token_parser.is_reasoning_end([100, 101]) is True
+
+    def test_start_only_no_end(self, single_token_parser):
+        """Only a start marker, no end marker → reasoning is active."""
+        assert single_token_parser.is_reasoning_end([100, 1, 2]) is False
+
+    def test_multi_turn_prompt_start_after_end(self, multi_token_parser):
+        """Simulated multi-turn prompt with previous assistant response.
+
+        Prompt contains: ... </think ... <think ...
+        The last marker is <think (start), so is_reasoning_end
+        should return False.
+        """
+        # Decode of these IDs: "reasoning" +  "</ " + "think>" + " " + "<think"
+        # The last marker is the start marker.
+        # Using: [3=reasoning, 522="</", 26865="think>", 100=TS]
+        ids = [3, 522, 26865, 100]
+        assert multi_token_parser.is_reasoning_end(ids) is False
+
+    def test_multi_turn_prompt_end_is_last(self, multi_token_parser):
+        """Previous turn ends with </think, no new start → ended."""
+        # [3=reasoning, 522="</", 26865="think>"]
+        ids = [3, 522, 26865]
+        assert multi_token_parser.is_reasoning_end(ids) is True
+
+    def test_special_start_after_end(self, multi_token_parser):
+        """Special token start after end marker → reasoning still active."""
+        # [151670=<|THINK_END|>, 151669=<|THINK_START|>]
+        ids = [151670, 151669]
+        assert multi_token_parser.is_reasoning_end(ids) is False
+
+
+# ---------------------------------------------------------------------------
+# is_reasoning_end_streaming
+# ---------------------------------------------------------------------------
+
+
+class TestIsReasoningEndStreaming:
+    def test_single_token_delta_has_end(self, single_token_parser):
+        assert single_token_parser.is_reasoning_end_streaming([1, 2], [101]) is True
+
+    def test_single_token_delta_no_end(self, single_token_parser):
+        assert single_token_parser.is_reasoning_end_streaming([1, 2], [3]) is False
+
+    def test_multi_token_delta_has_both_parts(self, multi_token_parser):
+        """Full </think> in delta → [522, 26865]."""
+        assert multi_token_parser.is_reasoning_end_streaming([1], [522, 26865]) is True
+
+    def test_multi_token_marker_spans_boundary(self, multi_token_parser):
+        """</ in previous, think> in delta."""
+        assert multi_token_parser.is_reasoning_end_streaming([1, 522], [26865]) is True
+
+    def test_multi_token_no_end(self, multi_token_parser):
+        assert multi_token_parser.is_reasoning_end_streaming([1, 2], [3]) is False
+
+
+# ---------------------------------------------------------------------------
+# extract_reasoning_streaming
+# ---------------------------------------------------------------------------
+
+
+class TestExtractReasoningStreaming:
+    """Test streaming extraction for both single-token and multi-token modes."""
+
+    @pytest.mark.parametrize(
+        "parser_name",
+        [
+            "single_token_parser",
+            "multi_token_parser",
+            "text_only_parser",
+        ],
+    )
+    def test_simple_streaming(self, request_, parser_name):
+        parser = request_.getfixturevalue(parser_name)
+        # Simulate: <think> reasoning </think>  final
+        if parser_name == "multi_token_parser":
+            deltas = [TS, "reasoning", TE, " content"]
+            token_ids_per_delta = [
+                [100],  # <think>
+                [3],  # reasoning
+                [522, 26865],  # </think>
+                [4],  #  content
+            ]
+        else:
+            deltas = [TS, "reasoning", TE, " content"]
+            token_ids_per_delta = [
+                [100],  # <think>
+                [3],  # reasoning
+                [101],  # </think>
+                [4],  #  content
+            ]
+
+        r, c = _stream_extract(parser, list(zip(deltas, token_ids_per_delta)))
+        assert r == "reasoning"
+        assert c == " content"
+
+    @pytest.mark.parametrize(
+        "parser_name",
+        [
+            "single_token_parser",
+            "multi_token_parser",
+            "text_only_parser",
+        ],
+    )
+    def test_no_end_token_streaming(self, request_, parser_name):
+        parser = request_.getfixturevalue(parser_name)
+        deltas_and_ids = [
+            (TS, [100]),
+            ("thinking", [3]),
+            (" more", [2]),
+        ]
+        r, c = _stream_extract(parser, deltas_and_ids)
+        assert r == "thinking more"
+        assert c == ""
+
+    def test_multi_token_end_split_across_chunks(self, multi_token_parser):
+        """The key test: </think> is split across streaming chunks.
+
+        Previous chunk ends with "reasoning</" and delta starts with
+        "think>\\nfinal answer".  This is the real-world scenario when
+        ``</think>`` is tokenized as two separate tokens.
+        """
+        deltas_and_ids = [
+            ("reasoning", [3]),
+            ("</", [522]),
+            ("think>\nHello", [26865, 1]),
+        ]
+        r, c = _stream_extract(multi_token_parser, deltas_and_ids)
+        assert r == "reasoning"
+        assert c == "Hello"
+
+    def test_multi_token_start_token_in_delta(self, multi_token_parser):
+        """Test that <think> start token is skipped in streaming output."""
+        deltas_and_ids = [
+            (TS, [100]),
+            ("reasoning", [3]),
+            (TE, [522, 26865]),
+            ("Hello", [1]),
+        ]
+        r, c = _stream_extract(multi_token_parser, deltas_and_ids)
+        assert r == "reasoning"
+        assert c == "Hello"
+
+    def test_multi_token_three_way_split(self, multi_token_parser):
+        """</think> split across 3 chunks: </ + thi + nk>."""
+        deltas_and_ids = [
+            ("some reasoning", [3]),
+            ("</", [522]),
+            ("thi", [8]),
+            ("nk>", [8]),
+            ("\nHello", [1]),
+        ]
+        r, c = _stream_extract(multi_token_parser, deltas_and_ids)
+        assert r == "some reasoning"
+        assert c == "Hello"
+
+    def test_special_end_token_in_stream(self, multi_token_parser):
+        """<|THINK_END|> special token in streaming (decoded as text)."""
+        deltas_and_ids = [
+            ("reasoning", [3]),
+            (TES, [151670]),
+            ("answer", [1]),
+        ]
+        r, c = _stream_extract(multi_token_parser, deltas_and_ids)
+        assert r == "reasoning"
+        assert c == "answer"
+
+    def test_no_start_token_in_stream(self, multi_token_parser):
+        """Model starts directly with reasoning (no <think> prefix)."""
+        deltas_and_ids = [
+            ("reasoning", [3]),
+            (TE, [522, 26865]),
+            ("answer", [1]),
+        ]
+        r, c = _stream_extract(multi_token_parser, deltas_and_ids)
+        assert r == "reasoning"
+        assert c == "answer"
+
+    def test_empty_reasoning(self, single_token_parser):
+        """Start immediately followed by end → empty reasoning."""
+        deltas_and_ids = [
+            (TS, [100]),
+            (TE, [101]),
+            ("Hello", [1]),
+        ]
+        r, c = _stream_extract(single_token_parser, deltas_and_ids)
+        assert r in ("", None)
+        assert c == "Hello"
+
+    def test_end_marker_no_content_after(self, multi_token_parser):
+        """End marker with no content following it."""
+        deltas_and_ids = [
+            ("reasoning", [3]),
+            (TE, [522, 26865]),
+        ]
+        r, c = _stream_extract(multi_token_parser, deltas_and_ids)
+        assert r == "reasoning"
+        assert c == ""
+
+    def test_delta_is_just_start_token(self, multi_token_parser):
+        """Delta consisting only of the start token should be skipped."""
+        deltas_and_ids = [
+            (TS, [100]),
+        ]
+        r, c = _stream_extract(multi_token_parser, deltas_and_ids)
+        assert r == ""
+        assert c == ""
+
+    @pytest.mark.parametrize(
+        "parser_name",
+        [
+            "single_token_parser",
+            "multi_token_parser",
+        ],
+    )
+    def test_only_first_newline_stripped_after_end(self, request_, parser_name):
+        """Only the first newline after </think> should be stripped."""
+        parser = request_.getfixturevalue(parser_name)
+        if parser_name == "multi_token_parser":
+            deltas_and_ids = [
+                ("reasoning", [3]),
+                (TE, [522, 26865]),
+                ("\n\nHello", [5, 1]),
+            ]
+        else:
+            deltas_and_ids = [
+                ("reasoning", [3]),
+                (TE, [101]),
+                ("\n\nHello", [5, 1]),
+            ]
+        r, c = _stream_extract(parser, deltas_and_ids)
+        assert r == "reasoning"
+        assert c == "\nHello"
+
+    @pytest.mark.parametrize(
+        "parser_name",
+        [
+            "single_token_parser",
+            "multi_token_parser",
+        ],
+    )
+    def test_content_after_end_multiple_deltas(self, request_, parser_name):
+        """Multiple content deltas after end marker."""
+        parser = request_.getfixturevalue(parser_name)
+        if parser_name == "multi_token_parser":
+            deltas_and_ids = [
+                (TS, [100]),
+                ("reasoning", [3]),
+                (TE, [522, 26865]),
+                ("Hello", [1]),
+                (" world", [2]),
+            ]
+        else:
+            deltas_and_ids = [
+                (TS, [100]),
+                ("reasoning", [3]),
+                (TE, [101]),
+                ("Hello", [1]),
+                (" world", [2]),
+            ]
+        r, c = _stream_extract(parser, deltas_and_ids)
+        assert r == "reasoning"
+        assert c == "Hello world"
+
+    def test_pending_buffer_resolved_on_non_marker(self, multi_token_parser):
+        """When buffered text turns out not to be a marker prefix, it is
+        emitted as reasoning on the next call."""
+        # "reason<" is buffered because "<" could start a marker.
+        # Next delta ">" doesn't form a marker, so the buffer is
+        # flushed as reasoning.
+        parser = multi_token_parser
+        deltas_and_ids = [
+            (TS, [100]),
+            ("reason", [3]),
+            ("<", [8]),  # "<" is buffered (prefix of markers)
+            ("> more", [8]),  # ">" doesn't complete any marker
+        ]
+        r, c = _stream_extract(parser, deltas_and_ids)
+        assert r == "reason> more"
+        assert c == ""
+
+    def test_multi_token_end_with_reasoning_before_buffer(self, multi_token_parser):
+        """Reasoning text before a buffered prefix should be emitted."""
+        # "reasoning<" → "reasoning" emitted, "<" buffered
+        # "</think>\nanswer" → buffer resolves as end marker
+        parser = multi_token_parser
+        deltas_and_ids = [
+            ("reasoning", [3]),
+            ("<", [8]),
+            (f"/{TE[2:]}\nanswer", [522, 26865, 1]),
+        ]
+        r, c = _stream_extract(parser, deltas_and_ids)
+        assert r == "reasoning"
+        assert c == "answer"
+
+    def test_multi_token_five_way_split(self, multi_token_parser):
+        """</think> split across 5 tiny chunks."""
+        # Split: each chunk is one character of </think>
+        chunks = list(TE)
+        deltas_and_ids = [
+            ("reasoning", [3]),
+        ]
+        for ch in chunks:
+            deltas_and_ids.append((ch, [8]))
+        deltas_and_ids.append(("\nHello", [1]))
+        r, c = _stream_extract(multi_token_parser, deltas_and_ids)
+        assert r == "reasoning"
+        assert c == "Hello"
+
+    def test_single_token_parser_with_pending(self, single_token_parser):
+        """Single-token parser: ambiguous '<' prefix buffered then resolved."""
+        # Even with single-token markers, the text-based path buffers
+        # '<' because it could be a prefix of a marker.
+        parser = single_token_parser
+        deltas_and_ids = [
+            (TS, [100]),
+            ("reason", [3]),
+            ("<", [8]),
+            ("> not a marker", [8]),
+        ]
+        r, c = _stream_extract(parser, deltas_and_ids)
+        assert r == "reason> not a marker"
+        assert c == ""
+
+
+# ---------------------------------------------------------------------------
+
+
+class TestCountReasoningTokens:
+    def test_single_token_basic(self, single_token_parser):
+        # IDs: [100=<think>, 3=reasoning, 101=</think>, 4=content]
+        assert single_token_parser.count_reasoning_tokens([100, 3, 101, 4]) == 1
+
+    def test_single_token_no_end(self, single_token_parser):
+        # All tokens after start are reasoning
+        assert single_token_parser.count_reasoning_tokens([100, 3, 2]) == 2
+
+    def test_multi_token_basic(self, multi_token_parser):
+        # </think> encodes as [522, 26865]
+        # Full sequence: <think> reasoning </think> content
+        # IDs: [100, 3, 522, 26865, 4]
+        result = multi_token_parser.count_reasoning_tokens([100, 3, 522, 26865, 4])
+        assert result > 0
+
+    def test_multi_token_no_end(self, multi_token_parser):
+        # No end token → all tokens inside <think> are counted.
+        result = multi_token_parser.count_reasoning_tokens([100, 3, 2])
+        assert result > 0
+
+    def test_text_only_basic(self, text_only_parser):
+        assert text_only_parser.count_reasoning_tokens([100, 3, 101, 4]) == 1
+
+    def test_text_only_no_end(self, text_only_parser):
+        assert text_only_parser.count_reasoning_tokens([100, 3, 2]) == 2
+
+
+# ---------------------------------------------------------------------------
+# extract_content_ids
+# ---------------------------------------------------------------------------
+
+
+class TestExtractContentIds:
+    def test_single_token_basic(self, single_token_parser):
+        # [1, 101=</think>, 2, 3] → content starts after ID 101
+        result = single_token_parser.extract_content_ids([1, 101, 2, 3])
+        assert result == [2, 3]
+
+    def test_single_token_no_end(self, single_token_parser):
+        result = single_token_parser.extract_content_ids([1, 2, 3])
+        assert result == []
+
+    def test_multi_token_basic(self, multi_token_parser):
+        # </think> = [522, 26865]
+        # [3, 522, 26865, 4] → content after </think>
+        result = multi_token_parser.extract_content_ids([3, 522, 26865, 4])
+        assert 4 in result
+
+    def test_multi_token_no_end(self, multi_token_parser):
+        result = multi_token_parser.extract_content_ids([1, 2, 3])
+        assert result == []
+
+    def test_text_only_basic(self, text_only_parser):
+        result = text_only_parser.extract_content_ids([1, 101, 2, 3])
+        assert result == [2, 3]
+
+    def test_text_only_no_end(self, text_only_parser):
+        result = text_only_parser.extract_content_ids([1, 2, 3])
+        assert result == []
+
+
+# ---------------------------------------------------------------------------
+# Fixture request workaround (pytest "request" fixture shadowed by parametrize)
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+def request_(request):
+    """Proxy for the pytest ``request`` fixture to avoid name collisions
+    with the ``request_obj`` fixture used in extract_reasoning tests."""
+    return request
@@ -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)
@@ -6,6 +6,11 @@ AC-5 (delay pattern), AC-7 (stage processor), and AC-10 (registry)
 without requiring the actual checkpoint or GPU.
 """

+import asyncio
+import time
+from collections import OrderedDict, defaultdict
+from types import SimpleNamespace
+
 import pytest
 import torch

@@ -232,7 +237,7 @@ class TestDelayPatternBehavior:


 class TestSamplerMethods:
-    """Test the actual _sample_audio_codes and _apply_delay_pattern_masking methods."""
+    """Test the actual sampler and batched delay masking methods."""

    def _make_minimal_talker(self):
        """Create a minimal talker-like object with sampler/masking methods."""
@@ -241,18 +246,65 @@ class TestSamplerMethods:
        class FakeTalker:
            num_codebooks = 8
            codebook_size = 1026
-            _audio_state = {}

        t = FakeTalker()
        t._sample_audio_codes = mod.HiggsAudioV3TalkerForConditionalGeneration._sample_audio_codes.__get__(t)
-        t._apply_delay_pattern_masking = (
-            mod.HiggsAudioV3TalkerForConditionalGeneration._apply_delay_pattern_masking.__get__(t)
+        return t
+
+    def _make_batched_sampler_talker(self, num_rows=4):
+        """Create a fake talker with GPU-resident sampler state helpers."""
+        from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_talker as mod
+
+        class FakeTalker:
+            num_codebooks = 8
+            codebook_size = 1026
+
+        t = FakeTalker()
+        t._decode_last_codes = torch.arange(num_rows * t.num_codebooks, dtype=torch.long).view(
+            num_rows, t.num_codebooks
        )
+        t._decode_has_codes = torch.tensor([True, False, True, True])[:num_rows].clone()
+        t._decode_delay_count = torch.tensor([0, 2, 7, 4], dtype=torch.int32)[:num_rows].clone()
+        t._decode_eoc_countdown = torch.tensor([-1, -1, 3, 1], dtype=torch.int32)[:num_rows].clone()
+        t._decode_generation_done = torch.tensor([False, False, False, True])[:num_rows].clone()
+        t._decode_active_audio_count = 0
+        t._codebook_index_cache = {}
+        t._row_index_cache = {}
+        t._boc_frame_cache = {}
+        t._last_audio_codes_buffer = None
+        t._last_audio_host_staging = None
+        t._last_audio_gpu_staging = None
+        t._last_audio_staging_event = None
+        t._audio_staging_events = []
+        t._audio_staging_event_cursor = 0
+        t._last_audio_codes = None
+        t._last_audio_code_valid = []
+        t._postprocess_cursor = 0
+        t._postprocess_audio_rows = 0
+        t._postprocess_audio_active_rows = 0
+
+        cls = mod.HiggsAudioV3TalkerForConditionalGeneration
+        for name in (
+            "_sample_audio_codes",
+            "_ensure_decode_state_capacity",
+            "_get_row_indices",
+            "_get_codebook_indices",
+            "_get_boc_frames",
+            "_get_audio_codes_buffer",
+            "_get_audio_gpu_staging_buffer",
+            "_get_audio_host_staging_buffer",
+            "_apply_delay_pattern_masking_batched",
+            "_update_delay_state_batched",
+            "_prefill_row_mask",
+            "_audio_seed_mask_from_step_input",
+        ):
+            setattr(t, name, getattr(cls, name).__get__(t))
+        t._device_cache_key = cls._device_cache_key
        return t

    def test_sample_respects_mask(self):
        """Tokens masked to -inf must never be sampled."""
-        t = self._make_minimal_talker()
+        t = self._make_batched_sampler_talker(num_rows=1)
        logits = torch.full((10, 1026), float("-inf"))
        # Only token 500 is allowed for each row
        logits[:, 500] = 0.0
@@ -262,7 +314,7 @@ class TestSamplerMethods:

    def test_sample_all_masked_falls_back_to_argmax(self):
        """All-masked row should fall back to argmax (least-negative logit)."""
-        t = self._make_minimal_talker()
+        t = self._make_batched_sampler_talker(num_rows=1)
        logits = torch.full((2, 1026), float("-inf"))
        # Row 0: all masked
        # Row 1: only token 42 allowed
@@ -275,11 +327,13 @@ class TestSamplerMethods:
        """During delay phase, codebooks beyond delay_count must have only BOC allowed."""
        from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import BOC_ID

-        t = self._make_minimal_talker()
+        t = self._make_batched_sampler_talker(num_rows=1)
        # Simulate delay_count=2 for batch row 0
-        t._audio_state[0] = {"num_delay": 2, "num_remaining_delays": None}
+        t._ensure_decode_state_capacity(1, torch.device("cpu"))
+        t._decode_delay_count[0] = 2
+        t._decode_eoc_countdown[0] = -1
        cb_logits = torch.zeros(1, 8, 1026)  # [1 audio row, 8 codebooks, 1026 vocab]
-        t._apply_delay_pattern_masking(cb_logits, [0])
+        t._apply_delay_pattern_masking_batched(cb_logits, torch.tensor([0], dtype=torch.long))
        # CBs 3-7 should have only BOC allowed (everything else -inf)
        for q in range(3, 8):
            row = cb_logits[0, q]
@@ -292,10 +346,12 @@ class TestSamplerMethods:
        """Active codebooks during delay should have BOC disallowed."""
        from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import BOC_ID

-        t = self._make_minimal_talker()
-        t._audio_state[0] = {"num_delay": 3, "num_remaining_delays": None}
+        t = self._make_batched_sampler_talker(num_rows=1)
+        t._ensure_decode_state_capacity(1, torch.device("cpu"))
+        t._decode_delay_count[0] = 3
+        t._decode_eoc_countdown[0] = -1
        cb_logits = torch.zeros(1, 8, 1026)
-        t._apply_delay_pattern_masking(cb_logits, [0])
+        t._apply_delay_pattern_masking_batched(cb_logits, torch.tensor([0], dtype=torch.long))
        # CBs 0-3 are active: BOC should be -inf
        for q in range(4):
            assert cb_logits[0, q, BOC_ID].item() == float("-inf")
@@ -304,10 +360,12 @@ class TestSamplerMethods:
        """Only codebook 0 should allow EOC during normal generation."""
        from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import EOC_ID

-        t = self._make_minimal_talker()
-        t._audio_state[0] = {"num_delay": 8, "num_remaining_delays": None}
+        t = self._make_batched_sampler_talker(num_rows=1)
+        t._ensure_decode_state_capacity(1, torch.device("cpu"))
+        t._decode_delay_count[0] = 8
+        t._decode_eoc_countdown[0] = -1
        cb_logits = torch.zeros(1, 8, 1026)
-        t._apply_delay_pattern_masking(cb_logits, [0])
+        t._apply_delay_pattern_masking_batched(cb_logits, torch.tensor([0], dtype=torch.long))
        # CB0 should keep EOC
        assert cb_logits[0, 0, EOC_ID].item() != float("-inf")
        # CB1-7 should have EOC masked
@@ -321,11 +379,13 @@ class TestSamplerMethods:
            EOC_ID,
        )

-        t = self._make_minimal_talker()
+        t = self._make_batched_sampler_talker(num_rows=1)
        # Ramp-down with 4 remaining delays: lock CBs 0-3 to EOC
-        t._audio_state[0] = {"num_delay": 8, "num_remaining_delays": 4}
+        t._ensure_decode_state_capacity(1, torch.device("cpu"))
+        t._decode_delay_count[0] = 8
+        t._decode_eoc_countdown[0] = 4
        cb_logits = torch.zeros(1, 8, 1026)
-        t._apply_delay_pattern_masking(cb_logits, [0])
+        t._apply_delay_pattern_masking_batched(cb_logits, torch.tensor([0], dtype=torch.long))
        # CBs 0-3 locked: only EOC allowed
        for q in range(4):
            row = cb_logits[0, q]
@@ -337,6 +397,195 @@ class TestSamplerMethods:
            assert cb_logits[0, q, BOC_ID].item() == float("-inf")
            assert cb_logits[0, q, EOC_ID].item() == float("-inf")

+    def test_batched_delay_masking_all_rows_matches_sparse_path(self):
+        """The all-row fast path must preserve sparse path masking semantics."""
+        torch.manual_seed(0)
+        sparse = self._make_batched_sampler_talker()
+        all_rows = self._make_batched_sampler_talker()
+        rows = torch.arange(4, dtype=torch.long)
+        logits_sparse = torch.randn(4, 8, 1026)
+        logits_all_rows = logits_sparse.clone()
+
+        sparse._apply_delay_pattern_masking_batched(logits_sparse, rows, all_rows=False)
+        all_rows._apply_delay_pattern_masking_batched(logits_all_rows, rows, all_rows=True)
+
+        assert torch.equal(logits_all_rows, logits_sparse)
+
+    def test_delay_state_update_all_rows_matches_sparse_path(self):
+        """The all-row fast path must update state and staging like index_copy_."""
+        from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import (
+            EOC_ID,
+        )
+
+        sparse = self._make_batched_sampler_talker()
+        all_rows = self._make_batched_sampler_talker()
+        rows = torch.arange(4, dtype=torch.long)
+        codes = torch.tensor(
+            [
+                [10, 11, 12, 13, 14, 15, 16, 17],
+                [20, 21, 22, 23, 24, 25, 26, 27],
+                [30, 31, EOC_ID, 33, 34, 35, 36, 37],
+                [40, 41, 42, 43, 44, 45, 46, 47],
+            ],
+            dtype=torch.long,
+        )
+        code_row_mask = torch.tensor([True, False, True, True])
+
+        sparse._update_delay_state_batched(
+            codes.clone(), rows, 4, torch.device("cpu"), code_row_mask=code_row_mask, all_rows=False
+        )
+        all_rows._update_delay_state_batched(
+            codes.clone(), rows, 4, torch.device("cpu"), code_row_mask=code_row_mask, all_rows=True
+        )
+
+        assert torch.equal(all_rows._decode_delay_count, sparse._decode_delay_count)
+        assert torch.equal(all_rows._decode_eoc_countdown, sparse._decode_eoc_countdown)
+        assert torch.equal(all_rows._decode_generation_done, sparse._decode_generation_done)
+        assert torch.equal(all_rows._decode_last_codes, sparse._decode_last_codes)
+        assert torch.equal(all_rows._decode_has_codes, sparse._decode_has_codes)
+        assert all_rows._last_audio_codes is None
+        assert sparse._last_audio_codes is None
+        assert torch.equal(all_rows._last_audio_host_staging, sparse._last_audio_host_staging)
+
+    def test_direct_audio_sampler_does_not_override_done_tokens(self):
+        """Done rows are represented directly in GPU sampled tokens."""
+        from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_talker as mod
+
+        t = self._make_batched_sampler_talker(num_rows=2)
+        t._resolve_token_ids = lambda: None
+        t._audio_continuation_id = 99999
+        t._eos_token_id = 151671
+        t._last_logits_hidden = torch.zeros(2, 16)
+        t._last_step_input_ids = torch.tensor([99999, 99999])
+        t._last_step_query_start_loc = None
+        t._decode_has_codes = torch.tensor([True, True])
+        t._decode_generation_done = torch.tensor([False, True])
+        t._decode_delay_count = torch.zeros(2, dtype=torch.long)
+        t._decode_eoc_countdown = torch.full((2,), -1, dtype=torch.long)
+        t._fast_audio_direct_rows = 2
+        t._last_audio_done_flags = None
+        t._fast_audio_sampler_gpu_fallback_reason = lambda **kwargs: None
+        t._audio_codebook_logits_from_rows = lambda hidden, rows, all_rows=False: torch.zeros(2, 8, 1026)
+        t._apply_delay_pattern_masking_batched = lambda cb_logits, audio_rows, all_rows=False: None
+        t._sample_audio_codes = lambda logits_2d: torch.zeros(int(logits_2d.shape[0]), dtype=torch.long)
+        t._update_delay_state_batched = lambda *args, **kwargs: None
+        t.sample = mod.HiggsAudioV3TalkerForConditionalGeneration.sample.__get__(t)
+
+        sampler_output = t.sample(torch.zeros(2, 200000), sampling_metadata=object())
+
+        assert not getattr(
+            mod.HiggsAudioV3TalkerForConditionalGeneration, "supports_sampled_token_ids_cpu_override", False
+        )
+        assert sampler_output.sampled_token_ids.tolist() == [[99999], [151671]]
+
+    def test_direct_audio_sampler_uses_gpu_tokens_for_active_rows(self):
+        """Direct audio sampling leaves scheduler-visible tokens to the GPU tensor."""
+        from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_talker as mod
+
+        t = self._make_batched_sampler_talker(num_rows=2)
+        t._resolve_token_ids = lambda: None
+        t._audio_continuation_id = 99999
+        t._eos_token_id = 151671
+        t._last_logits_hidden = torch.zeros(2, 16)
+        t._last_step_input_ids = torch.tensor([99999, 99999])
+        t._last_step_query_start_loc = None
+        t._decode_has_codes = torch.tensor([True, True])
+        t._decode_generation_done = torch.tensor([False, False])
+        t._decode_delay_count = torch.zeros(2, dtype=torch.long)
+        t._decode_eoc_countdown = torch.full((2,), -1, dtype=torch.long)
+        t._fast_audio_direct_rows = 2
+        t._last_audio_done_flags = [0, 0]
+        t._fast_audio_sampler_gpu_fallback_reason = lambda **kwargs: None
+        t._audio_codebook_logits_from_rows = lambda hidden, rows, all_rows=False: torch.zeros(2, 8, 1026)
+        t._apply_delay_pattern_masking_batched = lambda cb_logits, audio_rows, all_rows=False: None
+        t._sample_audio_codes = lambda logits_2d: torch.zeros(int(logits_2d.shape[0]), dtype=torch.long)
+        t._update_delay_state_batched = lambda *args, **kwargs: None
+        t.sample = mod.HiggsAudioV3TalkerForConditionalGeneration.sample.__get__(t)
+
+        sampler_output = t.sample(torch.zeros(2, 200000), sampling_metadata=object())
+
+        assert sampler_output.sampled_token_ids.tolist() == [[99999], [99999]]
+
+    def test_terminal_rampdown_frame_is_emitted_before_done(self):
+        """The final ramp-down frame must be emitted, not replaced by a done marker."""
+        from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import EOC_ID
+
+        t = self._make_batched_sampler_talker(num_rows=1)
+        t._decode_delay_count[0] = 8
+        t._decode_eoc_countdown[0] = 1
+        t._decode_generation_done[0] = False
+        t._decode_has_codes[0] = True
+        codes = torch.tensor([[10, 11, 12, 13, 14, 15, 16, 17]], dtype=torch.long)
+
+        t._update_delay_state_batched(
+            codes,
+            torch.tensor([0], dtype=torch.long),
+            1,
+            torch.device("cpu"),
+            code_row_mask=torch.tensor([True]),
+            all_rows=True,
+        )
+
+        staged = t._last_audio_host_staging[0]
+        assert staged[t.num_codebooks].item() == 1  # valid
+        assert staged[t.num_codebooks + 1].item() == 1  # done
+        assert not (staged[: t.num_codebooks] == -1).any()
+        assert staged[: t.num_codebooks - 1].eq(EOC_ID).all()
+        assert t._decode_generation_done[0].item() is True
+        assert t._decode_has_codes[0].item() is False
+
+    def test_decode_state_growth_preserves_inflight_rows(self):
+        """Growing decode buffers must copy old rows and initialize new rows inactive."""
+        t = self._make_batched_sampler_talker(num_rows=2)
+        old_codes = t._decode_last_codes.clone()
+        old_has = t._decode_has_codes.clone()
+        old_delay = t._decode_delay_count.clone()
+        old_rem = t._decode_eoc_countdown.clone()
+        old_done = t._decode_generation_done.clone()
+
+        t._ensure_decode_state_capacity(5, torch.device("cpu"))
+
+        assert t._decode_last_codes.shape[0] >= 5
+        assert torch.equal(t._decode_last_codes[:2], old_codes)
+        assert torch.equal(t._decode_has_codes[:2], old_has)
+        assert torch.equal(t._decode_delay_count[:2], old_delay)
+        assert torch.equal(t._decode_eoc_countdown[:2], old_rem)
+        assert torch.equal(t._decode_generation_done[:2], old_done)
+        assert not t._decode_has_codes[2:5].any()
+        assert not t._decode_generation_done[2:5].any()
+        assert t._decode_eoc_countdown[2:5].eq(-1).all()
+
+    def test_mixed_batch_prefill_mask_targets_request_rows(self):
+        """Mixed prefill/decode must reset only prefill request rows."""
+        t = self._make_batched_sampler_talker(num_rows=3)
+        t._last_step_query_start_loc = torch.tensor([0, 1, 4, 5])
+
+        mask = t._prefill_row_mask(3, torch.device("cpu"))
+
+        assert torch.equal(mask, torch.tensor([False, True, False]))
+
+    def test_mixed_batch_seed_mask_maps_token_tails_to_request_rows(self):
+        """Seed detection should return one flag per request, not per token."""
+        t = self._make_batched_sampler_talker(num_rows=3)
+        t._audio_continuation_id = 99999
+        t._last_step_query_start_loc = torch.tensor([0, 1, 4, 5])
+        t._last_step_input_ids = torch.tensor([99999, 11, 12, 13, 99999])
+
+        mask = t._audio_seed_mask_from_step_input(3, torch.device("cpu"))
+
+        assert torch.equal(mask, torch.tensor([True, False, True]))
+
+    def test_mixed_batch_seed_mask_allows_prefill_tail_audio_token(self):
+        """A prefill span ending in <|audio|> starts audio-code generation."""
+        t = self._make_batched_sampler_talker(num_rows=3)
+        t._audio_continuation_id = 99999
+        t._last_step_query_start_loc = torch.tensor([0, 1, 4, 5])
+        t._last_step_input_ids = torch.tensor([99999, 11, 12, 99999, 99999])
+
+        mask = t._audio_seed_mask_from_step_input(3, torch.device("cpu"))
+
+        assert torch.equal(mask, torch.tensor([True, True, True]))
+

 # ---- AC-6: Feedback Method Tests ----

@@ -348,10 +597,14 @@ class TestFeedbackMethods:

        class FakeTalker:
            _last_audio_codes = torch.tensor([[100, 200, 300, 400, 500, 600, 700, 800]])
+            _last_audio_code_valid = [True]
+            _last_audio_host_staging = None
+            _last_audio_staging_event = None
            _postprocess_cursor = 0

        t = FakeTalker()
        t.postprocess = mod.HiggsAudioV3TalkerForConditionalGeneration.postprocess.__get__(t)
+        t._postprocess_impl = mod.HiggsAudioV3TalkerForConditionalGeneration._postprocess_impl.__get__(t)
        result = t.postprocess(torch.zeros(1, 64))
        assert "codes" in result
        assert "audio" in result["codes"]
@@ -363,10 +616,14 @@ class TestFeedbackMethods:

        class FakeTalker:
            _last_audio_codes = torch.tensor([[-1, -1, -1, -1, -1, -1, -1, -1]])
+            _last_audio_code_valid = [False]
+            _last_audio_host_staging = None
+            _last_audio_staging_event = None
            _postprocess_cursor = 0

        t = FakeTalker()
        t.postprocess = mod.HiggsAudioV3TalkerForConditionalGeneration.postprocess.__get__(t)
+        t._postprocess_impl = mod.HiggsAudioV3TalkerForConditionalGeneration._postprocess_impl.__get__(t)
        result = t.postprocess(torch.zeros(1, 64))
        assert result == {}

@@ -378,10 +635,14 @@ class TestFeedbackMethods:
            _last_audio_codes = torch.tensor(
                [[100, 200, 300, 400, 500, 600, 700, 800], [-1, -1, -1, -1, -1, -1, -1, -1]]
            )
+            _last_audio_code_valid = [True, False]
+            _last_audio_host_staging = None
+            _last_audio_staging_event = None
            _postprocess_cursor = 0

        t = FakeTalker()
        t.postprocess = mod.HiggsAudioV3TalkerForConditionalGeneration.postprocess.__get__(t)
+        t._postprocess_impl = mod.HiggsAudioV3TalkerForConditionalGeneration._postprocess_impl.__get__(t)
        r1 = t.postprocess(torch.zeros(1, 64))
        assert "codes" in r1
        r2 = t.postprocess(torch.zeros(1, 64))
@@ -405,14 +666,23 @@ class TestAudioFeedback:
            codebook_size = 1026
            _audio_continuation_id = 99999  # fake audio token
            _last_step_query_start_loc = None
-            _audio_state = {}
+            _decode_active_audio_count = 0
+            _decode_last_codes = torch.zeros(4, 8, dtype=torch.long)
+            _decode_has_codes = torch.zeros(4, dtype=torch.bool)
+            _use_external_decode_cudagraph = False
            multimodal_embedding = embed
            model = type("M", (), {"embed_tokens": lambda self, ids: torch.zeros(ids.shape[0], 16)})()

        t = FakeTalker()
-        t._decode_last_codes = torch.zeros(256, 8, dtype=torch.long)
-        t._decode_has_codes = torch.zeros(256, dtype=torch.bool)
        t._apply_audio_feedback = mod.HiggsAudioV3TalkerForConditionalGeneration._apply_audio_feedback.__get__(t)
+        t._is_single_token_decode_step = (
+            mod.HiggsAudioV3TalkerForConditionalGeneration._is_single_token_decode_step.__get__(t)
+        )
+        t._step_request_count = mod.HiggsAudioV3TalkerForConditionalGeneration._step_request_count.__get__(t)
+        t._decode_request_token_positions = (
+            mod.HiggsAudioV3TalkerForConditionalGeneration._decode_request_token_positions.__get__(t)
+        )
+        t._ensure_decode_state_capacity = lambda min_bs, device: None
        return t

    def test_text_positions_unchanged(self):
@@ -428,9 +698,11 @@ class TestAudioFeedback:
        """Audio position with state should have its embedding replaced."""
        t = self._make_feedback_talker()
        audio_id = t._audio_continuation_id
-        t._audio_state[1] = {"last_codes": torch.zeros(8, dtype=torch.long)}
        t._decode_last_codes[1] = torch.zeros(8, dtype=torch.long)
        t._decode_has_codes[1] = True
+        t._decode_active_audio_count = 1
+        t._ensure_decode_state_capacity = lambda min_bs, device: None
+        t._last_step_query_start_loc = torch.tensor([0, 1, 2, 3])
        input_ids = torch.tensor([1, audio_id, 3])
        hidden = torch.zeros(3, 16)
        result = t._apply_audio_feedback(hidden, input_ids)
@@ -444,12 +716,80 @@ class TestAudioFeedback:
        """Audio position without state should not be replaced."""
        t = self._make_feedback_talker()
        audio_id = t._audio_continuation_id
-        # No state for row 0 — _decode_has_codes[0] defaults to False
+        # No state for row 0
        input_ids = torch.tensor([audio_id])
        hidden = torch.zeros(1, 16)
        result = t._apply_audio_feedback(hidden, input_ids)
        assert torch.equal(result, hidden)

+    def test_prefill_span_does_not_receive_audio_feedback(self):
+        """A prefill span can have numel == hidden rows but must not be treated as decode."""
+        t = self._make_feedback_talker()
+        t._last_step_input_ids = torch.tensor([1, 2, 3, 4])
+        t._last_step_query_start_loc = torch.tensor([0, 4])
+        t._decode_last_codes[0] = torch.zeros(8, dtype=torch.long)
+        t._decode_has_codes[0] = True
+        t._decode_active_audio_count = 1
+        hidden = torch.zeros(4, 16)
+
+        result = t._apply_audio_feedback(hidden, t._last_step_input_ids)
+
+        assert torch.equal(result, hidden)
+
+    def test_multi_request_decode_receives_audio_feedback(self):
+        """Concurrent decode has numel > 1 but each request span is exactly one token."""
+        t = self._make_feedback_talker()
+        t._last_step_input_ids = torch.tensor([10, 11, 12, 13])
+        t._last_step_query_start_loc = torch.tensor([0, 1, 2, 3, 4])
+        t._decode_last_codes[2] = torch.zeros(8, dtype=torch.long)
+        t._decode_has_codes[2] = True
+        t._decode_active_audio_count = 1
+        hidden = torch.zeros(4, 16)
+
+        result = t._apply_audio_feedback(hidden, t._last_step_input_ids)
+
+        assert result[2].abs().sum() > 0
+        assert result[0].abs().sum() == 0
+        assert result[1].abs().sum() == 0
+        assert result[3].abs().sum() == 0
+
+    def test_external_decode_graph_capture_keeps_feedback_ops_without_query_start_loc(self):
+        """FULL_DECODE capture has no request qsl but still needs dense decode feedback ops."""
+        t = self._make_feedback_talker()
+        t._use_external_decode_cudagraph = True
+        t._last_step_input_ids = torch.tensor([10, 11, 12, 13])
+        t._last_step_query_start_loc = None
+        t._decode_last_codes[2] = torch.zeros(8, dtype=torch.long)
+        t._decode_has_codes[2] = True
+        hidden = torch.zeros(4, 16)
+
+        result = t._apply_audio_feedback(hidden, t._last_step_input_ids)
+
+        assert result[2].abs().sum() > 0
+        assert result[0].abs().sum() == 0
+        assert result[1].abs().sum() == 0
+        assert result[3].abs().sum() == 0
+
+    def test_mixed_prefill_decode_receives_feedback_only_on_decode_rows(self):
+        """Mixed batches should skip prefill spans but keep audio feedback for decode spans."""
+        t = self._make_feedback_talker()
+        t._last_step_input_ids = torch.tensor([99999, 21, 22, 23, 99999])
+        t._last_step_query_start_loc = torch.tensor([0, 1, 4, 5])
+        t._decode_last_codes[0] = torch.zeros(8, dtype=torch.long)
+        t._decode_last_codes[2] = torch.zeros(8, dtype=torch.long)
+        t._decode_has_codes[0] = True
+        t._decode_has_codes[2] = True
+        t._decode_active_audio_count = 2
+        hidden = torch.zeros(5, 16)
+
+        result = t._apply_audio_feedback(hidden, t._last_step_input_ids)
+
+        assert result[0].abs().sum() > 0
+        assert result[4].abs().sum() > 0
+        assert result[1].abs().sum() == 0
+        assert result[2].abs().sum() == 0
+        assert result[3].abs().sum() == 0
+

 # ---- AC-8: Codec Strictness ----

@@ -775,6 +1115,23 @@ class TestStageInputProcessor:
        with pytest.raises(ValueError, match="Not enough frames"):
            _revert_delay_pattern(codes)

+    def test_talker2code2wav_skips_too_few_frames_without_crashing(self):
+        from vllm_omni.model_executor.stage_input_processors.higgs_audio_v3 import (
+            talker2code2wav,
+        )
+
+        class Output:
+            multimodal_output = {"codes": {"audio": torch.zeros(5, 8, dtype=torch.long)}}
+
+        class TalkerOutput:
+            finished = True
+            outputs = [Output()]
+
+        result = talker2code2wav([TalkerOutput()])
+
+        assert len(result) == 1
+        assert result[0]["prompt_token_ids"] == []
+
    def test_filter_real_code_frames(self):
        from vllm_omni.model_executor.stage_input_processors.higgs_audio_v3 import (
            _filter_real_code_frames,
@@ -800,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 ----

@@ -826,8 +1208,14 @@ class TestRegistry:

        # __file__ = tests/unit/higgs_audio_v3/test_*.py → 4 dirnames to repo root
        repo_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
-        yaml_path = os.path.join(repo_root, "vllm_omni", "deploy", "higgs_multimodal_qwen3.yaml")
-        assert os.path.isfile(yaml_path), f"Deploy YAML not found at {yaml_path}"
+        deploy_dir = os.path.join(repo_root, "vllm_omni", "deploy")
+        for name in (
+            "higgs_multimodal_qwen3.yaml",
+            "higgs_multimodal_qwen3_high_throughput.yaml",
+            "higgs_multimodal_qwen3_low_latency.yaml",
+        ):
+            yaml_path = os.path.join(deploy_dir, name)
+            assert os.path.isfile(yaml_path), f"Deploy YAML not found at {yaml_path}"


 # ---- AC-3: Prompt Builder ----
@@ -914,298 +1302,89 @@ class TestPromptBuilder:
        assert 151704 not in ids  # <|ref_text|>


-# ---- Regression: ramp-down done-before-emit off-by-one ----
+class TestVoiceCloneReferenceCache:
+    def test_audio_tokenizer_dir_env_accepts_parent_or_subdir(self, tmp_path, monkeypatch):
+        from vllm_omni.model_executor.models.higgs_audio_v2 import higgs_audio_v2_tokenizer as tok

+        parent = tmp_path / "OmniVoice"
+        subdir = parent / "audio_tokenizer"
+        subdir.mkdir(parents=True)
+        (subdir / "config.json").write_text("{}", encoding="utf-8")

-def _simulate_delay_state_machine(
-    sampled_codes_per_step: list[torch.Tensor],
-) -> list[torch.Tensor]:
-    """Reproduce the delay-pattern state machine from sample().
+        monkeypatch.setenv("HIGGS_AUDIO_TOKENIZER_PATH", str(parent))
+        assert tok._resolve_audio_tokenizer_dir() == str(subdir)

-    Feeds ``sampled_codes_per_step`` (each [Q]) through the exact same
-    logic as HiggsAudioV3TalkerForConditionalGeneration.sample() and
-    returns the list of emitted frames (non -1 codes).
-    """
-    from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import (
-        BOC_ID,
-        EOC_ID,
-    )
+        monkeypatch.setenv("HIGGS_AUDIO_TOKENIZER_PATH", str(subdir))
+        assert tok._resolve_audio_tokenizer_dir() == str(subdir)

-    num_codebooks = 8
-    bos = BOC_ID
-    eos_stream = EOC_ID
+    def test_higgs_v3_ref_code_cache_returns_clone(self):
+        from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech

-    state: dict = {
-        "num_delay": 0,
-        "num_remaining_delays": None,
-        "audio_out_ids": None,
-        "last_codes": torch.full((num_codebooks,), bos, dtype=torch.long),
-        "should_terminate": False,
-    }
+        serving = object.__new__(OmniOpenAIServingSpeech)
+        serving._higgs_audio_v3_ref_code_cache = OrderedDict()
+        serving._higgs_audio_v3_ref_code_cache_bytes = 0
+        serving._higgs_audio_v3_ref_code_inflight = {}

-    emitted: list[torch.Tensor] = []
-    for step_codes in sampled_codes_per_step:
-        if state["should_terminate"]:
-            break
+        codes = torch.arange(16, dtype=torch.long).reshape(2, 8)
+        serving._put_higgs_audio_v3_ref_codes("ref-a", codes)

-        num_delay = state["num_delay"]
-        num_remaining_delays = state["num_remaining_delays"]
-        this_codes = step_codes.clone()
+        cached = serving._get_higgs_audio_v3_ref_codes("ref-a")
+        assert cached is not None
+        cached.fill_(0)

-        if num_delay + 1 < num_codebooks:
-            this_codes[num_delay + 1 :] = bos
-            num_delay += 1
+        cached_again = serving._get_higgs_audio_v3_ref_codes("ref-a")
+        assert cached_again is not None
+        assert torch.equal(cached_again, codes)

-        if num_remaining_delays is not None:
-            this_codes[: num_codebooks - num_remaining_delays] = eos_stream
-            num_remaining_delays -= 1
-        else:
-            eos_positions = (this_codes == eos_stream).nonzero(as_tuple=False).reshape(-1)
-            if eos_positions.numel() > 0:
-                last_eos_idx = int(eos_positions[-1].item())
-                this_codes[: last_eos_idx + 1] = eos_stream
-                num_remaining_delays = num_codebooks - last_eos_idx - 1
+    def test_higgs_v3_ref_code_inflight_deduplicates_concurrent_encode(self):
+        from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech

-        if num_remaining_delays is not None and num_remaining_delays <= 0:
-            state["num_delay"] = 0
-            state["num_remaining_delays"] = None
-            state["should_terminate"] = True
-            continue
+        serving = object.__new__(OmniOpenAIServingSpeech)
+        serving._higgs_audio_v3_ref_code_cache = OrderedDict()
+        serving._higgs_audio_v3_ref_code_cache_bytes = 0
+        serving._higgs_audio_v3_ref_code_inflight = {}
+        calls = 0

-        state["num_delay"] = num_delay
-        state["num_remaining_delays"] = num_remaining_delays
-        state["last_codes"] = this_codes.clone()
-        emitted.append(this_codes)
+        def encode_reference_audio(_wav, _sr):
+            nonlocal calls
+            calls += 1
+            time.sleep(0.05)
+            return torch.arange(16, dtype=torch.long).reshape(2, 8)

-    return emitted
+        def apply_delay_pattern(codes):
+            return codes + 1

+        async def run():
+            return await asyncio.gather(
+                *[
+                    serving._resolve_higgs_audio_v3_ref_codes(
+                        "ref-a",
+                        object(),
+                        24000,
+                        encode_reference_audio,
+                        apply_delay_pattern,
+                    )
+                    for _ in range(3)
+                ]
+            )

-class TestRampDownOffByOne:
-    """Regression tests for the done-before-emit off-by-one bug.
+        results = asyncio.run(run())
+        assert calls == 1
+        assert sum(int(inflight_wait) for _, _, inflight_wait in results) == 2
+        for codes, cache_hit, _ in results:
+            assert cache_hit is False
+            assert torch.equal(codes, torch.arange(16, dtype=torch.long).reshape(2, 8) + 1)

-    When EOC fires early, the ramp-down must emit enough frames so that
-    Stage1's _revert_delay_pattern receives T >= Q = 8.
-    """
-
-    def test_eoc_at_step0_terminates(self):
-        """EOC on cb0 at step 0 terminates; T < Q is handled by Stage1 defensive try/except."""
-        from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import EOC_ID
-
-        Q = 8
-        steps = []
-        for s in range(20):
-            codes = torch.randint(0, 1024, (Q,))
-            if s == 0:
-                codes[0] = EOC_ID
-            steps.append(codes)
-
-        emitted = _simulate_delay_state_machine(steps)
-        assert len(emitted) == Q - 1, f"EOC at step 0: expected {Q - 1} emitted frames, got {len(emitted)}"
-
-    def test_eoc_at_step1_emits_enough_frames(self):
-        """EOC on cb0 at step 1 must produce >= Q emitted frames."""
-        from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import EOC_ID
-
-        Q = 8
-        steps = []
-        for s in range(20):
-            codes = torch.randint(0, 1024, (Q,))
-            if s == 1:
-                codes[0] = EOC_ID
-            steps.append(codes)
-
-        emitted = _simulate_delay_state_machine(steps)
-        assert len(emitted) >= Q, f"EOC at step 1: expected >= {Q} emitted frames, got {len(emitted)}"
-
-    def test_eoc_after_full_rampup_emits_enough(self):
-        """EOC after all codebooks are active produces T >> Q."""
-        from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import EOC_ID
-
-        Q = 8
-        eoc_step = 10
-        steps = []
-        for s in range(30):
-            codes = torch.randint(0, 1024, (Q,))
-            if s == eoc_step:
-                codes[0] = EOC_ID
-            steps.append(codes)
-
-        emitted = _simulate_delay_state_machine(steps)
-        assert len(emitted) >= Q, f"EOC at step {eoc_step}: expected >= {Q} emitted frames, got {len(emitted)}"
-
-    def test_last_rampdown_frame_not_all_negative(self):
-        """The last emitted ramp-down frame must contain real codes, not -1."""
-        from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import EOC_ID
-
-        Q = 8
-        steps = []
-        for s in range(20):
-            codes = torch.randint(0, 1024, (Q,))
-            if s == 1:
-                codes[0] = EOC_ID
-            steps.append(codes)
-
-        emitted = _simulate_delay_state_machine(steps)
-        last = emitted[-1]
-        assert (last >= 0).all(), f"Last ramp-down frame contains negative values: {last.tolist()}"
-
-    def test_emitted_frames_revertible_by_stage1(self):
-        """Emitted frames from EOC-at-step-1 must be revertible by _revert_delay_pattern."""
-        from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import EOC_ID
-        from vllm_omni.model_executor.stage_input_processors.higgs_audio_v3 import (
-            _revert_delay_pattern,
+        cached, cache_hit, inflight_wait = asyncio.run(
+            serving._resolve_higgs_audio_v3_ref_codes(
+                "ref-a",
+                object(),
+                24000,
+                encode_reference_audio,
+                apply_delay_pattern,
+            )
        )
-
-        Q = 8
-        steps = []
-        for s in range(20):
-            codes = torch.randint(0, 1024, (Q,))
-            if s == 1:
-                codes[0] = EOC_ID
-            steps.append(codes)
-
-        emitted = _simulate_delay_state_machine(steps)
-        # Stack to [T, Q], transpose to [Q, T]
-        codes_qt = torch.stack(emitted, dim=0).t().contiguous()
-        assert codes_qt.shape[0] == Q
-        assert codes_qt.shape[1] >= Q
-
-        # Must not raise
-        reverted = _revert_delay_pattern(codes_qt)
-        assert reverted.shape[0] == Q
-        assert reverted.shape[1] == codes_qt.shape[1] - Q + 1
-
-
-class TestStage1DefensiveHandling:
-    """Stage1 must not crash the server when receiving T < Q frames."""
-
-    def test_revert_delay_pattern_raises_on_insufficient_frames(self):
-        """_revert_delay_pattern correctly raises ValueError for T < Q."""
-        from vllm_omni.model_executor.stage_input_processors.higgs_audio_v3 import (
-            _revert_delay_pattern,
-        )
-
-        codes = torch.zeros(8, 7)  # T=7 < Q=8
-        with pytest.raises(ValueError, match="Not enough frames"):
-            _revert_delay_pattern(codes)
-
-    def test_talker2code2wav_survives_insufficient_frames(self):
-        """talker2code2wav must not crash when audio_codes has T < Q frames."""
-        from vllm_omni.model_executor.stage_input_processors.higgs_audio_v3 import (
-            talker2code2wav,
-        )
-
-        # Build a fake talker output with only 7 frames (T=7 < Q=8)
-        class FakeOutput:
-            def __init__(self):
-                self.multimodal_output = {
-                    "codes": {"audio": torch.zeros(7, 8, dtype=torch.long)},
-                }
-
-        class FakeTalkerOutput:
-            finished = True
-            outputs = [FakeOutput()]
-
-        # Must not raise — should return gracefully
-        result = talker2code2wav([FakeTalkerOutput()])
-        assert len(result) == 1
-        # Result is an OmniTokensPrompt with empty prompt_token_ids
-        out = result[0]
-        token_ids = out.prompt_token_ids if hasattr(out, "prompt_token_ids") else out.get("prompt_token_ids", [])
-        assert token_ids == []
-
-
-# ---- Regression: buffer growth + prefill guard ----
-
-
-class TestBufferGrowthPreservesState:
-    """Bug #1: buffer growth in _apply_audio_feedback must not wipe existing state."""
-
-    def _make_feedback_talker_with_buffers(self):
-        from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_talker as mod
-
-        embed = mod.HiggsFusedMultiTextEmbedding(num_codebooks=8, vocab_size=1026, hidden_size=16)
-        torch.nn.init.ones_(embed.weight)
-
-        class FakeTalker:
-            num_codebooks = 8
-            codebook_size = 1026
-            multimodal_embedding = embed
-
-        t = FakeTalker()
-        t._decode_last_codes = torch.zeros(4, 8, dtype=torch.long)
-        t._decode_has_codes = torch.zeros(4, dtype=torch.bool)
-        t._apply_audio_feedback = mod.HiggsAudioV3TalkerForConditionalGeneration._apply_audio_feedback.__get__(t)
-        return t
-
-    def test_buffer_growth_preserves_existing_codes(self):
-        """Growing the buffer must copy old _decode_last_codes data."""
-        t = self._make_feedback_talker_with_buffers()
-        # Set state for slot 2
-        t._decode_last_codes[2] = torch.tensor([100, 200, 300, 400, 500, 600, 700, 800])
-        t._decode_has_codes[2] = True
-
-        # Trigger growth by calling with bs=8 > current capacity 4
-        hidden = torch.randn(8, 16)
-        t._apply_audio_feedback(hidden, torch.zeros(8, dtype=torch.long))
-
-        # After growth, slot 2 data must still be present
-        assert t._decode_has_codes[2].item() is True
-        assert t._decode_last_codes[2, 0].item() == 100
-
-    def test_buffer_growth_new_slots_are_inactive(self):
-        """New slots from buffer growth must have has_codes=False."""
-        t = self._make_feedback_talker_with_buffers()
-        t._decode_has_codes[1] = True
-
-        hidden = torch.randn(8, 16)
-        t._apply_audio_feedback(hidden, torch.zeros(8, dtype=torch.long))
-
-        # Slots 4-7 (new) must be inactive
-        assert not t._decode_has_codes[4].item()
-        assert not t._decode_has_codes[7].item()
-        # Slot 1 (old) must still be active
-        assert t._decode_has_codes[1].item()
-
-
-class TestPrefillGuard:
-    """Bug #2: _apply_audio_feedback must NOT run during prefill.
-
-    The prefill detection uses attn_metadata.max_query_len (== 1 for decode,
-    > 1 for prefill). This correctly handles concurrent decode (N requests,
-    each contributing 1 token → numel=N but max_query_len=1).
-    """
-
-    @staticmethod
-    def _is_prefill_via_max_query_len(max_query_len: int | None, input_ids: torch.Tensor) -> bool:
-        """Mirror the actual detection logic in forward()."""
-        if max_query_len is not None:
-            return int(max_query_len) > 1
-        return input_ids is not None and int(input_ids.numel()) > 1
-
-    def test_prefill_detected_via_max_query_len(self):
-        """Prefill: max_query_len > 1 → is_prefill=True."""
-        input_ids = torch.tensor([1, 2, 3, 4, 5])
-        assert self._is_prefill_via_max_query_len(5, input_ids) is True
-
-    def test_single_decode_detected_via_max_query_len(self):
-        """Single-request decode: max_query_len=1 → is_prefill=False."""
-        input_ids = torch.tensor([42])
-        assert self._is_prefill_via_max_query_len(1, input_ids) is False
-
-    def test_concurrent_decode_not_misclassified(self):
-        """Multi-request decode: numel=4 but max_query_len=1 → is_prefill=False.
-
-        This is the regression caught by reviewer: with c=4 concurrent decode
-        requests, input_ids.numel()==4 but each request contributes exactly 1
-        token, so max_query_len==1 and audio feedback must still apply.
-        """
-        input_ids = torch.tensor([42, 43, 44, 45])  # 4 concurrent decode
-        assert self._is_prefill_via_max_query_len(1, input_ids) is False
-
-    def test_fallback_when_max_query_len_unavailable(self):
-        """When attn_metadata is unavailable, fall back to numel heuristic."""
-        prefill_ids = torch.tensor([1, 2, 3])
-        assert self._is_prefill_via_max_query_len(None, prefill_ids) is True
-        decode_ids = torch.tensor([42])
-        assert self._is_prefill_via_max_query_len(None, decode_ids) is False
+        assert calls == 1
+        assert cache_hit is True
+        assert inflight_wait is False
+        assert torch.equal(cached, torch.arange(16, dtype=torch.long).reshape(2, 8) + 1)
Author	SHA1	Message	Date
Zhou Taichang	7b837944a1	[Rebase] Rebase to vllm 0.23.0 (#4286 ) Signed-off-by: tzhouam <tzhouam@connect.ust.hk>	2026-06-13 16:43:02 +00:00
Chendi.Xue	a9ee92ddec	[XPU] update Dreamzero to support any HW (#4399 ) Signed-off-by: Chendi Xue <chendi.xue@intel.com>	2026-06-13 15:32:39 +00:00
Yuanheng Zhao	31886d6ece	[Doc][Recipe] Update CUDA verifications for inclusionAI/Ming-omni-tts-0.5B (#4324 ) Signed-off-by: Yuanheng <jonathan.zhaoyh@gmail.com>	2026-06-13 18:12:37 +08:00
Hongsheng Liu	fa94f41807	[Bugfix] Surface all-rank diffusion RPC failures (#4403 ) Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>	2026-06-13 15:42:16 +08:00
Yueqian Lin	d75a6433b4	[Refactor] TTS serving adapter framework + migrate TTS models (excl. ming_flash_omni) (#4330 ) Signed-off-by: Yueqian Lin <linyueqian@outlook.com>	2026-06-13 04:22:54 +00:00
Chendi.Xue	6506ed3d4d	[XPU][CI] Fix docker build slowness (#4402 ) Signed-off-by: Chendi Xue <chendi.xue@intel.com>	2026-06-13 10:59:15 +08:00
akshatvishu	d4f2a49a68	[Bugfix] Ming-omni-tts: generalize dict checks to Mapping (#4397 ) Signed-off-by: akshatvishu <akshatnayak197@gmail.com>	2026-06-13 09:37:57 +08:00
wuli666 linyueqian wuli666	9bc8705c98	[New Model] Step-Audio2 (#464 ) Signed-off-by: linyueqian <linyueqian@outlook.com> Signed-off-by: wuli666 <djjpro975@gmail.com> Co-authored-by: linyueqian <linyueqian@outlook.com> Co-authored-by: wuli666 <djjpro975@gmail.com>	2026-06-13 01:00:33 +00:00
Henry J jjw Claude Fable 5 Yueqian Lin	1733831f78	[Bugfix] Fix cross-request codes.ref leak in Qwen3-TTS make_omni_output (#4373 ) Signed-off-by: jjw <m@jjw.in> Co-authored-by: jjw <m@jjw.in> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>	2026-06-13 00:41:49 +00:00
Joni JuvonenandHongsheng Liu	20b7a67b0d	[Perf][Wan2.2] Skip attention mask for zero-padded SP sequences to avoid varlen path (#3763 ) Signed-off-by: jjuvonen-amd <joni.juvonen@amd.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-06-13 07:50:39 +08:00
wangyu	9577ccba50	[CI] skip MOSS-TTS-Nano E2E tests pending issue#4361 (#4391 )	2026-06-13 00:07:57 +08:00
wangyuandHongsheng Liu	54a9d241c0	[CI] Diff-gate L2/L3 E2E jobs and migrate single-GPU tests to gpu_1_queue (#4365 ) Signed-off-by: wangyu <410167048@qq.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-06-12 19:05:43 +08:00
zyz111222andCanlin Guo	6de0ca36bb	[Hardware][Ascend] Adapt Qwen3 TTS for 310P (#4283 ) Signed-off-by: zouyizhou <zouyizhou@huawei.com> Co-authored-by: Canlin Guo <canlinguosdu@gmail.com>	2026-06-12 09:39:21 +00:00
ChenWenjing	293396737d	[ci]skip voxcpm2 pcm hnr test (#4375 ) Signed-off-by: CHEN <116010019@link.cuhk.edu.cn>	2026-06-12 17:14:33 +08:00
tanhaoan333	aa59dc50c7	[bugfix]VoxCPM2 audio encoder adapt other than CUDA (#4374 ) Signed-off-by: tanhaoan333 <tanhaoan@huawei.com>	2026-06-12 08:54:29 +00:00
Samit	f44f3ef9f7	[platform] fix: set UnspecifiedOmniPlatform device_type to cpu (#4357 ) Signed-off-by: samithuang <285365963@qq.com>	2026-06-12 08:40:00 +00:00
Megha Agarwal Copilot Hongsheng Liu meghaagr13 Boao Shi Taichang Zhou Cursor	8f869c8ea5	[Refactor] Output Processor Phase 2: separate multimodal output channel (#1601 ) (#2744 ) Signed-off-by: Megha Agarwal <agarwalmegha1308@gmail.com> Signed-off-by: meghaagr13 <divyanshsinghvi@gmail.com> Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com> Co-authored-by: meghaagr13 <divyanshsinghvi@gmail.com> Co-authored-by: Boao Shi <aoibosh@connect.hku.hk> Co-authored-by: Taichang Zhou <tzhouam@connect.ust.hk> Co-authored-by: Cursor <cursoragent@cursor.com>	2026-06-12 15:19:14 +08:00
Ting FU	1be1a71156	[Model] Add more resolution support for HunyuanImage3.0 (#4004 ) Signed-off-by: Semmer <semmer@live.cn>	2026-06-12 14:44:00 +08:00
Chendi.Xue	3fe3aa359b	[XPU][COSMOS3] removing cuda hardcode and make VLLM_VIDEO_SYNC_TIMEOUT a tunable config (#4360 ) Signed-off-by: Chendi Xue <chendi.xue@intel.com>	2026-06-12 04:55:07 +00:00
YU XINYUAN YU Xinyuan Yueqian Lin	541f06c60d	[Bugfix] Fix Qwen3-TTS gradio streaming TTFP by using audio/pcm (#4346 ) Signed-off-by: YU Xinyuan <yuxinyuan02@corp.netease.com> Co-authored-by: YU Xinyuan <yuxinyuan02@corp.netease.com> Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>	2026-06-12 03:19:48 +00:00
yangyonggit	a3834e32cf	feat(moss-tts): add CUDA Graph support for codec decoder (#4157 ) Signed-off-by: leo.yang <leo.yang.engineer@gmail.com>	2026-06-11 22:17:31 +00:00
hurukawa Yueqian Lin Sy03	152343d5fb	Feat: `non_streaming_mode` for Qwen3-TTS Base Models During Online Inference (#4198 ) Signed-off-by: nagisa-kun <1434936049@qq.com> Signed-off-by: Yueqian Lin <linyueqian@outlook.com> Co-authored-by: Yueqian Lin <linyueqian@outlook.com> Co-authored-by: Sy03 <1370724210@qq.com>	2026-06-11 22:15:05 +00:00
MaciejBalaNVandHongsheng Liu	03cc888485	[refactor] Refactor guardrail error handling - add 400 error code (#4297 ) Signed-off-by: Maciej Bala <mbala@nvidia.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-06-11 21:09:24 +00:00
Sy03 Yueqian Lin Hongsheng Liu	e27fd46ab9	[Perf] Optimize Higgs Audio v3 serving (#4204 ) Signed-off-by: Sy03 <1370724210@qq.com> Signed-off-by: Yueqian Lin <linyueqian@outlook.com> Co-authored-by: Yueqian Lin <linyueqian@outlook.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-06-11 13:29:31 -07:00
dengyunyang	eb964dcec5	[HunyuanImage][Perf] opt prepare_attention_mask for e2e latency 6% reduction (#4333 ) Signed-off-by: dengyunyang <584797741@qq.com>	2026-06-11 19:06:22 +08:00
Chendi.XueandHongsheng Liu	5414f78e5f	[WAN2.2-S2V] Add server API for image + audio (#3394 ) Signed-off-by: Chendi Xue <chendi.xue@intel.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-06-11 18:50:03 +08:00
WeiQing ChenandHongsheng Liu	f3d40b5402	[Bugfix] Add Cosmos3-Nano baselines and fix USP gather (#4301 ) Signed-off-by: david6666666 <530634352@qq.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-06-11 18:36:26 +08:00
Hongsheng Liu	3529b8280d	[Doc] Clean up PR template (#4336 )	2026-06-11 15:49:22 +08:00
tanhaoan333	e26f5cb918	[NPU] Support VoxCPM2 model (#4310 ) Signed-off-by: tanhaoan333 <tanhaoan@huawei.com>	2026-06-11 14:12:06 +08:00
wangyu	0d3c259a76	[CI] skip unrelated L2 ready tests via diff-aware upload, aligned L3 tweak and fix issue 4334 (#4313 ) Signed-off-by: wangyu <410167048@qq.com> Signed-off-by: wangyu <53896905+yenuo26@users.noreply.github.com>	2026-06-11 12:58:40 +08:00
Kushanam	74b5ae5ebe	[Bugfix] Fix SP denoise indentation bug on BAGEL (#4328 ) Signed-off-by: Kushan A. <kahmadian@nvidia.com>	2026-06-11 03:32:22 +00:00
ChenWenjingandSy03	1fe267d6c9	[Bugfix][VoxCPM2] Fix VoxCPM2 concurrent speech quality (#4319 ) Signed-off-by: CHEN <116010019@link.cuhk.edu.cn> Co-authored-by: Sy03 <1370724210@qq.com>	2026-06-11 11:12:07 +08:00
dengyunyangandHongsheng Liu	14879d6619	[Bugfix][HunyuanImage] fix accuracy in stream mode (#4265 ) Signed-off-by: dengyunyang <584797741@qq.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-06-11 10:39:00 +08:00
zhumingjue138	18e083e512	[Docs] add doc for failure mode (#3926 ) Signed-off-by: zhumingjue <zhumingjue@huawei.com>	2026-06-11 10:36:34 +08:00
WeiQing Chen baonudesifeizhai lishunyang12	a1a53e5f94	[Skills] Add quantization Claude skill (#4252 ) Signed-off-by: david6666666 <530634352@qq.com> Co-authored-by: baonudesifeizhai <85092850+baonudesifeizhai@users.noreply.github.com> Co-authored-by: lishunyang12 <125541396+lishunyang12@users.noreply.github.com>	2026-06-11 06:41:31 +08:00
wangyu	1b2ab23586	[Test] Automatically clean up audio files generated from requests, and realtime invalid-param coverage (#4294 ) Signed-off-by: wangyu <410167048@qq.com> Signed-off-by: [Your Name] <your.email@example.com>	2026-06-11 05:57:57 +08:00
Mike Qiu Mike_Qiu Claude Opus 4.6 Yueqian Lin Hongsheng Liu	57ea047acf	Step audio R1 reasoning parser (#2846 ) Signed-off-by: Mike_Qiu <qiudayu.qdy@antgroup.com> Signed-off-by: Yueqian Lin <linyueqian@outlook.com> Signed-off-by: Mike Qiu <qdy220091330@gmail.com> Co-authored-by: Mike_Qiu <qiudayu.qdy@antgroup.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> Co-authored-by: Yueqian Lin <linyueqian@outlook.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-06-10 21:03:19 +00:00