Remove mm_prefix_lm patch because vllm==0.18.0 already support (#2047)

Signed-off-by: princepride <wangzhipeng628@gmail.com>

.buildkite/bootstrap-amd-omni.sh

@@ -90,10 +90,18 @@ upload_pipeline() {
     FAIL_FAST=$(fail_fast)
 
     cd .buildkite
+
+    # Select test definition file: merge suite for main, ready suite for PRs
+    if [[ $BUILDKITE_BRANCH == "main" ]]; then
+        TEST_YAML="test-amd-merge.yml"
+    else
+        TEST_YAML="test-amd-ready.yaml"
+    fi
+
     (
         set -x
         # Output pipeline.yaml with all blank lines removed
-        minijinja-cli test-template.j2 test-amd.yaml \
+        minijinja-cli test-template.j2 "$TEST_YAML" \
             -D branch="$BUILDKITE_BRANCH" \
             -D list_file_diff="$LIST_FILE_DIFF" \
             -D run_all="$RUN_ALL" \

.buildkite/bootstrap-intel-omni.sh

@@ -0,0 +1,51 @@
+#!/bin/bash
+# vllm-omni Intel bootstrap
+# Uses static pipeline-intel.yaml for Intel XPU tests
+# Based on: bootstrap-amd-omni.sh
+
+set -euo pipefail
+
+if [[ -z "${DOCS_ONLY_DISABLE:-}" ]]; then
+    DOCS_ONLY_DISABLE=0
+fi
+
+upload_pipeline() {
+    echo "--- 🛠 Preparing Intel pipeline"
+    buildkite-agent pipeline upload .buildkite/pipeline-intel.yaml
+}
+
+get_diff() {
+    $(git add .)
+    echo $(git diff --name-only --diff-filter=ACMDR $(git merge-base origin/main HEAD))
+}
+
+get_diff_main() {
+    $(git add .)
+    echo $(git diff --name-only --diff-filter=ACMDR HEAD~1)
+}
+
+file_diff=$(get_diff)
+if [[ $BUILDKITE_BRANCH == "main" ]]; then
+    file_diff=$(get_diff_main)
+fi
+
+# ----------------------------------------------------------------------
+# Early exit: skip pipeline if all changed files are under docs/
+# ----------------------------------------------------------------------
+if [[ "${DOCS_ONLY_DISABLE}" != "1" ]] && [[ -n "${file_diff:-}" ]]; then
+    docs_only=1
+    while IFS= read -r f; do
+        [[ -z "$f" ]] && continue
+        if [[ "$f" != docs/* ]]; then
+            docs_only=0
+            break
+        fi
+    done < <(printf '%s\n' "$file_diff" | tr ' ' '\n' | tr -d '\r')
+
+    if [[ "$docs_only" -eq 1 ]]; then
+        buildkite-agent annotate ":memo: CI skipped — docs only" --style "info"
+        exit 0
+    fi
+fi
+
+upload_pipeline

.buildkite/pipeline-intel.yaml

@@ -0,0 +1,16 @@
+# .buildkite/pipeline-intel.yaml
+steps:
+  - group: "INTEL Tests"
+    steps:
+      - label: "Intel: XPU Test"
+        command: "bash .buildkite/scripts/hardware_ci/run-xpu-test.sh"
+        agents:
+          queue: "intel-gpu-omni"
+        env:
+          DOCKER_BUILDKIT: "1"
+          # Buildkite will automatically replace this with the actual commit hash
+          VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}"
+          VLLM_VERSION: "v0.17.0"
+        priority: 100
+        timeout_in_minutes: 60
+        soft_fail: true

.buildkite/scripts/hardware_ci/run-amd-test.sh

@@ -87,7 +87,11 @@ HF_CACHE="$(realpath ~)/huggingface"
 mkdir -p "${HF_CACHE}"
 HF_MOUNT="/root/.cache/huggingface"
 
-commands=$@
+if [[ -n "${TEST_COMMAND:-}" ]]; then
+    commands="$TEST_COMMAND"
+else
+    commands="$@"
+fi
 echo "Commands:$commands"
 
 PARALLEL_JOB_COUNT=8
@@ -102,6 +106,7 @@ if [[ -z "$render_gid" ]]; then
 fi
 
 # check if the command contains shard flag, we will run all shards in parallel because the host have 8 GPUs.
+# TODO: @tjtanaa reenable to run VLLM_ROCM_USE_AITER=1 when AITER is shipped with prebuilt kernels.
 if [[ $commands == *"--shard-id="* ]]; then
   # assign job count as the number of shards used
   commands=$(echo "$commands" | sed -E "s/--num-shards[[:blank:]]*=[[:blank:]]*[0-9]*/--num-shards=${PARALLEL_JOB_COUNT} /g" | sed 's/ \\ / /g')
@@ -118,7 +123,7 @@ if [[ $commands == *"--shard-id="* ]]; then
         --rm \
         -e MIOPEN_DEBUG_CONV_DIRECT=0 \
         -e MIOPEN_DEBUG_CONV_GEMM=0 \
-        -e VLLM_ROCM_USE_AITER=1 \
+        -e VLLM_ROCM_USE_AITER=0 \
         -e HIP_VISIBLE_DEVICES="${GPU}" \
         -e HF_TOKEN \
         -e AWS_ACCESS_KEY_ID \
@@ -153,7 +158,7 @@ else
           --rm \
           -e MIOPEN_DEBUG_CONV_DIRECT=0 \
           -e MIOPEN_DEBUG_CONV_GEMM=0 \
-          -e VLLM_ROCM_USE_AITER=1 \
+          -e VLLM_ROCM_USE_AITER=0 \
           -e HF_TOKEN \
           -e AWS_ACCESS_KEY_ID \
           -e AWS_SECRET_ACCESS_KEY \

.buildkite/test-amd-merge.yml

@@ -0,0 +1,197 @@
+steps:
+
+- label: "Simple Unit Test"
+  agent_pool: mi325_1
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export VLLM_ROCM_USE_AITER=0
+    - "timeout 20m pytest -v -s -m 'core_model and cpu' --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
+
+- label: "Diffusion Model Test"
+  agent_pool: mi325_2
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export GPU_ARCHS=gfx942
+    - timeout 20m pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "advanced_model and diffusion" --run-level "advanced_model"
+
+- label: "Diffusion Images API LoRA E2E"
+  agent_pool: mi325_1
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - timeout 20m pytest -s -v tests/e2e/online_serving/test_images_generations_lora.py
+
+- label: "Diffusion Model CPU offloading Test"
+  agent_pool: mi325_1
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export GPU_ARCHS=gfx942
+    - export VLLM_LOGGING_LEVEL=DEBUG
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - |
+      timeout 20m bash -c '
+        set +e
+        pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
+        EXIT1=\$?
+        pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
+        EXIT2=\$?
+        exit \$((EXIT1 | EXIT2))
+      '
+
+## ISSUE depends on `diffusers` package: https://github.com/huggingface/diffusers/issues/13274
+# - label: "Audio Generation Model Test"
+#   agent_pool: mi325_1
+#   depends_on: amd-build
+#   mirror_hardwares: [amdproduction]
+#   grade: Blocking
+#   commands:
+#     - export GPU_ARCHS=gfx942
+#     - export VLLM_LOGGING_LEVEL=DEBUG
+#     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+#     - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_model.py
+
+- label: "Diffusion Cache Backend Test"
+  agent_pool: mi325_1
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export GPU_ARCHS=gfx942
+    - export VLLM_LOGGING_LEVEL=DEBUG
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - timeout 15m pytest -s -v -m "core_model and cache and diffusion and not distributed_cuda and L4"
+
+- label: "Diffusion Sequence Parallelism Test"
+  agent_pool: mi325_2
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export VLLM_LOGGING_LEVEL=DEBUG
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - timeout 20m pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py
+
+# merge-only tests
+- label: "Diffusion Tensor Parallelism Test"
+  agent_pool: mi325_2
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - pytest -s -v tests/e2e/offline_inference/test_zimage_parallelism.py
+
+- label: "Diffusion GPU Worker Test"
+  agent_pool: mi325_2
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py
+
+- label: "Benchmark & Engine Test"
+  agent_pool: mi325_2
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - |
+      timeout 20m bash -c '
+        set +e
+        pytest -s -v tests/benchmarks/test_serve_cli.py
+        EXIT1=\$?
+        pytest -s -v tests/engine/test_async_omni_engine_abort.py
+        EXIT2=\$?
+        exit \$((EXIT1 | EXIT2))
+      '
+
+- label: "Omni Model Test Qwen2-5-Omni"
+  agent_pool: mi325_2
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export VLLM_LOGGING_LEVEL=DEBUG
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - timeout 20m pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
+
+- label: "Omni Model Test Qwen3-Omni"
+  agent_pool: mi325_2
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export VLLM_LOGGING_LEVEL=DEBUG
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - export VLLM_TEST_CLEAN_GPU_MEMORY=1
+    - timeout 10m pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
+    - timeout 20m pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "advanced_model" --run-level "advanced_model"
+
+- label: "Qwen3-TTS E2E Test"
+  agent_pool: mi325_2
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export VLLM_LOGGING_LEVEL=DEBUG
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - timeout 20m pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py
+
+- label: "Diffusion Image Edit Test"
+  agent_pool: mi325_1
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export GPU_ARCHS=gfx942
+    - export VLLM_LOGGING_LEVEL=DEBUG
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - timeout 20m pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py
+
+# split Bagel Model Test with H100 (Real Weights) into three tests
+- label: "Bagel Text2Img Model Test"
+  agent_pool: mi325_1
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export GPU_ARCHS=gfx942
+    - export VLLM_TEST_CLEAN_GPU_MEMORY=1
+    - export VLLM_LOGGING_LEVEL=DEBUG
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - export VLLM_ROCM_USE_AITER_RMSNORM=0
+    - timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py -m "advanced_model" --run-level "advanced_model" -k "shared_memory" -k "rocm"
+
+- label: "Bagel Img2Img Model Test"
+  agent_pool: mi325_1
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export GPU_ARCHS=gfx942
+    - export VLLM_TEST_CLEAN_GPU_MEMORY=1
+    - export VLLM_LOGGING_LEVEL=DEBUG
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - export VLLM_ROCM_USE_AITER_RMSNORM=0
+    - timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py -m "advanced_model" --run-level "advanced_model" -k "rocm"
+
+- label: "Bagel Online Serving Test"
+  agent_pool: mi325_1
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export GPU_ARCHS=gfx942
+    - export VLLM_TEST_CLEAN_GPU_MEMORY=1
+    - export VLLM_IMAGE_FETCH_TIMEOUT=60
+    - export VLLM_LOGGING_LEVEL=DEBUG
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - export VLLM_ROCM_USE_AITER_RMSNORM=0
+    - timeout 40m pytest -s -v tests/e2e/online_serving/test_bagel_online.py -m "advanced_model" --run-level "advanced_model" -k "rocm"

.buildkite/test-amd-ready.yaml

@@ -0,0 +1,178 @@
+steps:
+
+- label: "Simple Unit Test"
+  agent_pool: mi325_1
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export VLLM_ROCM_USE_AITER=0
+    - "timeout 20m pytest -v -s -m 'core_model and cpu' --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
+
+- label: "Diffusion Model Test"
+  agent_pool: mi325_2
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - timeout 20m pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "core_model and diffusion" --run-level "core_model"
+
+- label: "Diffusion Model CPU offloading Test"
+  agent_pool: mi325_1
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export GPU_ARCHS=gfx942
+    - export VLLM_LOGGING_LEVEL=DEBUG
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - |
+      timeout 20m bash -c '
+        set +e
+        pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
+        EXIT1=\$?
+        pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
+        EXIT2=\$?
+        exit \$((EXIT1 | EXIT2))
+      '
+
+## ISSUE depends on `diffusers` package: https://github.com/huggingface/diffusers/issues/13274
+# - label: "Audio Generation Model Test"
+#   agent_pool: mi325_1
+#   depends_on: amd-build
+#   mirror_hardwares: [amdproduction]
+#   grade: Blocking
+#   commands:
+#     - export GPU_ARCHS=gfx942
+#     - export VLLM_LOGGING_LEVEL=DEBUG
+#     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+#     - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_model.py
+
+- label: "Diffusion Cache Backend Test"
+  agent_pool: mi325_1
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export GPU_ARCHS=gfx942
+    - export VLLM_LOGGING_LEVEL=DEBUG
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - timeout 15m pytest -s -v -m "core_model and cache and diffusion and not distributed_cuda and L4"
+
+- label: "Diffusion Sequence Parallelism Test"
+  agent_pool: mi325_2
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export VLLM_LOGGING_LEVEL=DEBUG
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - timeout 20m pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py -m core_model
+
+- label: "Diffusion GPU Worker Test"
+  agent_pool: mi325_2
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py
+
+- label: "Benchmark & Engine Test"
+  agent_pool: mi325_2
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - |
+      timeout 30m bash -c '
+        set +e
+        pytest -s -v tests/benchmarks/test_serve_cli.py
+        EXIT1=\$?
+        pytest -s -v tests/engine/test_async_omni_engine_abort.py
+        EXIT2=\$?
+        exit \$((EXIT1 | EXIT2))
+      '
+
+- label: "Omni Model Test Qwen2-5-Omni"
+  agent_pool: mi325_2
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export VLLM_LOGGING_LEVEL=DEBUG
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - timeout 17m pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
+
+- label: "Omni Model Test Qwen3-Omni"
+  agent_pool: mi325_2
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export VLLM_LOGGING_LEVEL=DEBUG
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - export VLLM_TEST_CLEAN_GPU_MEMORY=1
+    - timeout 10m pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
+    - timeout 10m pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "core_model" --run-level "core_model"
+
+- label: "Qwen3-TTS E2E Test"
+  agent_pool: mi325_2
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export VLLM_LOGGING_LEVEL=DEBUG
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - timeout 20m pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py
+
+- label: "Diffusion Image Edit Test"
+  agent_pool: mi325_1
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export GPU_ARCHS=gfx942
+    - export VLLM_LOGGING_LEVEL=DEBUG
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - timeout 20m pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py
+
+- label: "Bagel Text2Img Model Test"
+  agent_pool: mi325_1
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export GPU_ARCHS=gfx942
+    - export VLLM_TEST_CLEAN_GPU_MEMORY=1
+    - export VLLM_LOGGING_LEVEL=DEBUG
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - export VLLM_ROCM_USE_AITER_RMSNORM=0
+    - timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py -m "core_model" --run-level "core_model" -k "rocm"
+
+- label: "Bagel Img2Img Model Test"
+  agent_pool: mi325_1
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export GPU_ARCHS=gfx942
+    - export VLLM_TEST_CLEAN_GPU_MEMORY=1
+    - export VLLM_LOGGING_LEVEL=DEBUG
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - export VLLM_ROCM_USE_AITER_RMSNORM=0
+    - timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py -m "core_model" --run-level "core_model" -k "rocm"
+
+- label: "Bagel Online Serving Test"
+  agent_pool: mi325_1
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export GPU_ARCHS=gfx942
+    - export VLLM_TEST_CLEAN_GPU_MEMORY=1
+    - export VLLM_IMAGE_FETCH_TIMEOUT=60
+    - export VLLM_LOGGING_LEVEL=DEBUG
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - export VLLM_ROCM_USE_AITER_RMSNORM=0
+    - timeout 40m pytest -s -v tests/e2e/online_serving/test_bagel_online.py -m "core_model" --run-level "core_model" -k "rocm"

.buildkite/test-amd.yaml

@@ -1,7 +1,7 @@
 steps:
 
 - label: "Diffusion Model Test"
-  timeout_in_minutes: 20
+  timeout_in_minutes: 30
   agent_pool: mi325_2
   depends_on: amd-build
   mirror_hardwares: [amdproduction]
@@ -11,7 +11,7 @@ steps:
     - pytest -s -v tests/e2e/offline_inference/test_t2i_model.py
 
 - label: "Diffusion Images API LoRA E2E"
-  timeout_in_minutes: 20
+  timeout_in_minutes: 30
   agent_pool: mi325_1
   depends_on: amd-build
   mirror_hardwares: [amdproduction]
@@ -101,6 +101,7 @@ steps:
     - export VLLM_LOGGING_LEVEL=DEBUG
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
     - export VLLM_TEST_CLEAN_GPU_MEMORY="1"
+    - export VLLM_ROCM_USE_AITER=0
     - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
     - pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py
 

.buildkite/test-merge.yml

@@ -12,11 +12,12 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
   - label: "Diffusion Model Test"
-    timeout_in_minutes: 20
+    timeout_in_minutes: 30
     depends_on: upload-merge-pipeline
     commands:
       - pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "advanced_model and diffusion" --run-level "advanced_model"
@@ -29,11 +30,12 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
   - label: "Diffusion Images API LoRA E2E"
-    timeout_in_minutes: 20
+    timeout_in_minutes: 30
     depends_on: upload-merge-pipeline
     commands:
       - pytest -s -v tests/e2e/online_serving/test_images_generations_lora.py
@@ -46,6 +48,7 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -64,6 +67,7 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -81,6 +85,7 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -98,6 +103,7 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -116,6 +122,7 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -134,6 +141,7 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -152,21 +160,17 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
-  - label: "Benchmark & Engine Test"
+  - label: "Engine Test"
     depends_on: upload-merge-pipeline
     commands:
       - |
         timeout 15m bash -c '
                 export VLLM_WORKER_MULTIPROC_METHOD=spawn
-                set +e
-                pytest -s -v tests/benchmarks/test_serve_cli.py
-                EXIT1=$$?
                 pytest -s -v tests/engine/test_async_omni_engine_abort.py
-                EXIT2=$$?
-                exit $$((EXIT1 | EXIT2))
         '
     agents:
       queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
@@ -178,6 +182,7 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -197,6 +202,7 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -226,6 +232,11 @@ steps:
                 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:
@@ -261,6 +272,11 @@ steps:
                 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:
@@ -272,12 +288,64 @@ steps:
                   path: /mnt/hf-cache
                   type: DirectoryOrCreate
 
-  # - label: "Bagel Text2Img Model Test with H100"
-  #   timeout_in_minutes: 30
+  - label: "Bagel Model Test with H100 (Real Weights)"
+    timeout_in_minutes: 60
+    depends_on: upload-merge-pipeline
+    commands:
+      - |
+        timeout 55m bash -c '
+          set -e
+          export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          export VLLM_TEST_CLEAN_GPU_MEMORY=1
+          export VLLM_IMAGE_FETCH_TIMEOUT=60
+          pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py -m "advanced_model" --run-level "advanced_model" -k "shared_memory"
+          pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py -m "advanced_model" --run-level "advanced_model"
+          pytest -s -v tests/e2e/online_serving/test_bagel_online.py -m "advanced_model" --run-level "advanced_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: 1
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: hf-cache
+                    mountPath: /root/.cache/huggingface
+                env:
+                  - name: HF_HOME
+                    value: /root/.cache/huggingface
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
+            nodeSelector:
+              node.kubernetes.io/instance-type: gpu-h100-sxm
+            volumes:
+              - name: devshm
+                emptyDir:
+                  medium: Memory
+              - name: hf-cache
+                hostPath:
+                  path: /mnt/hf-cache
+                  type: DirectoryOrCreate
+
+  # - label: "Voxtral TTS E2E Test"
+  #   timeout_in_minutes: 20
   #   depends_on: upload-merge-pipeline
   #   commands:
-  #     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  #     - pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py
+  #     - |
+  #       timeout 20m bash -c '
+  #         export VLLM_LOGGING_LEVEL=DEBUG
+  #         export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  #         pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
+  #       '
   #   agents:
   #     queue: "mithril-h100-pool"
   #   plugins:

.buildkite/test-nightly.yml

@@ -11,6 +11,9 @@ steps:
         EXIT1=$$?
         pytest -s -v tests/examples/online_serving/test_qwen3_omni.py -m "advanced_model" --run-level "advanced_model"
         EXIT2=$$?
+        # TODO: uncomment when mistralai/tts-model is public
+        # pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
+        # EXIT3=$$?
         exit $$((EXIT1 | EXIT2))
     agents:
       queue: "mithril-h100-pool"
@@ -30,6 +33,11 @@ steps:
                 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:
@@ -58,13 +66,14 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
   - label: ":full_moon: Diffusion Model Test with H100"
     timeout_in_minutes: 60
     depends_on: upload-nightly-pipeline
-    # if: build.env("NIGHTLY") == "1"
+    if: build.env("NIGHTLY") == "1"
     commands:
       - export VLLM_WORKER_MULTIPROC_METHOD=spawn
       - pytest -s -v tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and diffusion and H100" --run-level "advanced_model"
@@ -86,6 +95,11 @@ steps:
                 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:
@@ -116,6 +130,7 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -149,6 +164,13 @@ steps:
                 env:
                   - name: HF_HOME
                     value: /root/.cache/huggingface
+                  - name: VLLM_HTTP_TIMEOUT_KEEP_ALIVE
+                    value: "120"
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
             nodeSelector:
               node.kubernetes.io/instance-type: gpu-h100-sxm
             volumes:
@@ -163,7 +185,7 @@ steps:
   - label: ":email: Nightly Collection & Email"
     key: nightly-perf-distribution
     depends_on:
-    - nightly-performance
+      - nightly-performance
     if: build.env("NIGHTLY") == "1"
     commands:
       - pip install openpyxl

.buildkite/test-ready.yml

@@ -5,6 +5,24 @@ steps:
       - "timeout 20m pytest -v -s -m 'core_model and cpu' --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
     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
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Voxtral TTS CUDA Unit Test"
+    timeout_in_minutes: 10
+    depends_on: upload-ready-pipeline
+    commands:
+      - "timeout 10m pytest -s -v tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py"
+    agents:
+      queue: "gpu_1_queue"
     plugins:
       - docker#v5.2.0:
           image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
@@ -18,7 +36,7 @@ steps:
   - label: "Diffusion Model Test"
     depends_on: upload-ready-pipeline
     commands:
-      - timeout 20m pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "core_model and diffusion" --run-level "core_model"
+      - timeout 30m pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "core_model and diffusion" --run-level "core_model"
     agents:
       queue: "gpu_1_queue"
     plugins:
@@ -28,6 +46,7 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -44,6 +63,7 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -60,6 +80,7 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -76,6 +97,7 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -93,6 +115,7 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -110,22 +133,18 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
 
-  - label: "Benchmark & Engine Test"
+  - label: "Engine Test"
     depends_on: upload-ready-pipeline
     commands:
       - |
         timeout 15m bash -c '
                 export VLLM_WORKER_MULTIPROC_METHOD=spawn
-                set +e
-                pytest -s -v tests/benchmarks/test_serve_cli.py
-                EXIT1=$$?
                 pytest -s -v tests/engine/test_async_omni_engine_abort.py
-                EXIT2=$$?
-                exit $$((EXIT1 | EXIT2))
         '
     agents:
       queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
@@ -137,6 +156,7 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -159,6 +179,7 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -221,9 +242,49 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
+  # TODO: uncomment when mistralai/tts-model is public
+  # - label: "Voxtral-TTS E2E Test"
+  #   depends_on: upload-ready-pipeline
+  #   commands:
+  #     - |
+  #       timeout 20m bash -c '
+  #         export VLLM_LOGGING_LEVEL=DEBUG
+  #         export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  #         pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_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: 1
+  #               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: "Diffusion Image Edit Test with H100 (1 GPU)"
   #   depends_on: upload-ready-pipeline
   #   commands:
@@ -261,39 +322,132 @@ steps:
   #                 path: /mnt/hf-cache
   #                 type: DirectoryOrCreate
 
-  # - label: "Bagel Text2Img Model Test with H100"
-  #   depends_on: upload-ready-pipeline
-  #   commands:
-  #     - |
-  #       timeout 30m bash -c '
-  #         export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  #         pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py
-  #       '
-  #   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
-  #           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: "Bagel Text2Img Model Test with H100"
+    depends_on: upload-ready-pipeline
+    commands:
+      - |
+        timeout 30m bash -c '
+          export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          export VLLM_TEST_CLEAN_GPU_MEMORY=1
+          pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py -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: 1
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: hf-cache
+                    mountPath: /root/.cache/huggingface
+                env:
+                  - name: HF_HOME
+                    value: /root/.cache/huggingface
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
+            nodeSelector:
+              node.kubernetes.io/instance-type: gpu-h100-sxm
+            volumes:
+              - name: devshm
+                emptyDir:
+                  medium: Memory
+              - name: hf-cache
+                hostPath:
+                  path: /mnt/hf-cache
+                  type: DirectoryOrCreate
+
+  - label: "Bagel Img2Img Model Test with H100"
+    depends_on: upload-ready-pipeline
+    commands:
+      - |
+        timeout 30m bash -c '
+          export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          export VLLM_TEST_CLEAN_GPU_MEMORY=1
+          pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py -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: 1
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: hf-cache
+                    mountPath: /root/.cache/huggingface
+                env:
+                  - name: HF_HOME
+                    value: /root/.cache/huggingface
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
+            nodeSelector:
+              node.kubernetes.io/instance-type: gpu-h100-sxm
+            volumes:
+              - name: devshm
+                emptyDir:
+                  medium: Memory
+              - name: hf-cache
+                hostPath:
+                  path: /mnt/hf-cache
+                  type: DirectoryOrCreate
+
+  - label: "Bagel Online Serving Test with H100"
+    depends_on: upload-ready-pipeline
+    commands:
+      - |
+        timeout 40m bash -c '
+          export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          export VLLM_TEST_CLEAN_GPU_MEMORY=1
+          export VLLM_IMAGE_FETCH_TIMEOUT=60
+          pytest -s -v tests/e2e/online_serving/test_bagel_online.py -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: 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

.buildkite/test-template-amd-omni.j2

@@ -40,9 +40,14 @@
           {% else %}
           queue: amd_mi325_1
           {% endif %}
-        command: bash .buildkite/scripts/hardware_ci/run-amd-test.sh "(command rocm-smi || true) && cd {{ (step.working_dir or default_working_dir) | safe  }} ; {{ step.command  or (step.commands | join(" && ")) | safe }}"
+{% set cmd_body = (step.command or (step.commands | join("\n"))) | trim %}
+{% set indented_cmd = cmd_body | replace("\n", "\n            ") %}
+        command: bash .buildkite/scripts/hardware_ci/run-amd-test.sh
         env:
           DOCKER_BUILDKIT: "1"
+          TEST_COMMAND: |-
+            (command rocm-smi || true) && cd {{ (step.working_dir or default_working_dir) | safe }}
+            {{ indented_cmd | safe }}
         priority: 100
         {% if step.grade and step.grade == "Blocking" %}
         soft_fail: false

.github/ISSUE_TEMPLATE/400-bug-report.yml

@@ -74,28 +74,21 @@ body:
       If relevant, add a minimal example so that we can reproduce the error by running the code. It is very important for the snippet to be as succinct (minimal) as possible, so please take time to trim down any irrelevant code to help us debug efficiently. We are going to copy-paste your code and we expect to get the same result as you did: avoid any external data, and include the relevant imports, etc. For example:
 
       ```python
-      from vllm_omni import OmniLLM, create_ar_stage_config, create_dit_stage_config
+      from vllm_omni.entrypoints.omni import Omni
+      from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+      from vllm import SamplingParams
 
-      # Create stage configurations
-      ar_config = create_ar_stage_config(
-          stage_id=0,
-          model_path="Qwen/Qwen3-0.6B",
-          input_modalities=["text"],
-          output_modalities=["text"]
+      omni = Omni(
+          model="Qwen/Qwen-Image",
+          stage_configs_path="/path/to/stage_configs.yaml",
       )
 
-      dit_config = create_dit_stage_config(
-          stage_id=1,
-          model_path="stabilityai/stable-diffusion-2-1",
-          input_modalities=["text"],
-          output_modalities=["image"]
-      )
-
-      # Initialize OmniLLM
-      omni = OmniLLM([ar_config, dit_config])
-
-      # Generate
-      outputs = omni.generate(prompt="A scenic watercolor painting of a lighthouse at sunset")
+      prompts = [{"prompt": "A scenic watercolor painting of a lighthouse at sunset"}]
+      sampling_params_list = [
+          SamplingParams(max_tokens=1),
+          OmniDiffusionSamplingParams(num_outputs_per_prompt=1),
+      ]
+      outputs = omni.generate(prompts=prompts, sampling_params_list=sampling_params_list)
       ```
 
       If the code is too long (hopefully, it isn't), feel free to put it in a public gist and link it in the issue: https://gist.github.com.

README.md

@@ -17,7 +17,7 @@ Easy, fast, and cheap omni-modality model serving for everyone
 
 *Latest News* 🔥
 
-- [2026/03] Joined vLLM HK First Meetup! Video on [youtube](https://youtu.be/sgwNfsNnR9I)
+- [2026/03] Check out our first public [project deepdive](https://youtu.be/sgwNfsNnR9I) at the vLLM Hong Kong Meetup!
 - [2026/03] **[vllm-omni-skills](https://github.com/hsliuustc0106/vllm-omni-skills)** is a community-driven collection of AI assistant skills that help developers work with vLLM-Omni more effectively. These skills can be used with popular agentic AI coding assistants like **Cursor IDE**, **Claude**, **Codex**, and more.
 - [2026/02] We released [0.16.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.16.0) - A major alignment + capability release that rebases onto **upstream vLLM v0.16.0** and significantly expands performance, distributed execution, and production readiness across **Qwen3-Omni / Qwen3-TTS**, **Bagel**, **MiMo-Audio**, **GLM-Image** and the **Diffusion (DiT) image/video stack**—while also improving platform coverage (CUDA / ROCm / NPU / XPU), CI quality, and documentation.
 - [2026/02] We released [0.14.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.14.0) - This is the first **stable release** of vLLM-Omni that expands Omni’s diffusion / image-video generation and audio / TTS stack, improves distributed execution and memory efficiency, and broadens platform/backend coverage (GPU/ROCm/NPU/XPU). It also brings meaningful upgrades to serving APIs, profiling & benchmarking, and overall stability. Please check our latest [paper](https://arxiv.org/abs/2602.02204) for architecture design and performance results.

benchmarks/diffusion/backends.py

@@ -1,3 +1,4 @@
+import asyncio
 import base64
 import mimetypes
 import os
@@ -35,6 +36,7 @@ class RequestFuncOutput:
     error: str = ""
     start_time: float = 0.0
     response_body: dict[str, Any] = field(default_factory=dict)
+    stage_durations: dict[str, float] = field(default_factory=dict)
     peak_memory_mb: float = 0.0
     slo_achieved: bool | None = None
 
@@ -55,6 +57,7 @@ async def async_request_chat_completions(
     input: RequestFuncInput,
     session: aiohttp.ClientSession,
     pbar: tqdm | None = None,
+    enable_diffusion_pipeline_profiler: bool = False,
 ) -> RequestFuncOutput:
     output = RequestFuncOutput()
     output.start_time = time.perf_counter()
@@ -108,6 +111,18 @@ async def async_request_chat_completions(
                 output.success = True
                 if "peak_memory_mb" in resp_json:
                     output.peak_memory_mb = resp_json["peak_memory_mb"]
+                try:
+                    choices = resp_json.get("choices", [])
+                    if choices and isinstance(choices, list):
+                        msg = choices[0].get("message", {})
+                        if isinstance(msg, dict):
+                            content = msg.get("content", [])
+                            if content and isinstance(content, list) and len(content) > 0:
+                                first_item = content[0]
+                                if isinstance(first_item, dict):
+                                    output.stage_durations = first_item.get("stage_durations")
+                except (IndexError, TypeError, AttributeError):
+                    pass
             else:
                 output.error = f"HTTP {response.status}: {await response.text()}"
                 output.success = False
@@ -192,7 +207,140 @@ async def async_request_openai_images(
     return output
 
 
+async def async_request_v1_videos(
+    input: RequestFuncInput,
+    session: aiohttp.ClientSession,
+    pbar: tqdm | None = None,
+) -> RequestFuncOutput:
+    output = RequestFuncOutput()
+    output.start_time = time.perf_counter()
+
+    files = dict(input.extra_body)
+    if input.prompt:
+        files.setdefault("prompt", input.prompt)
+    if input.width and input.height:
+        files.setdefault("height", input.height)
+        files.setdefault("width", input.width)
+    if input.num_frames:
+        files.setdefault("num_frames", input.num_frames)
+    if input.num_inference_steps:
+        files.setdefault("num_inference_steps", input.num_inference_steps)
+    if input.seed is not None:
+        files.setdefault("seed", input.seed)
+    if input.fps:
+        files.setdefault("fps", input.fps)
+
+    form = aiohttp.FormData()
+    for k, v in files.items():
+        form.add_field(k, str(v))
+
+    image_file = None
+    if input.image_paths and len(input.image_paths) > 0:
+        image_path = input.image_paths[0]
+        image_file = open(image_path, "rb")
+        form.add_field(
+            "input_reference",
+            image_file,
+            filename=os.path.basename(image_path),
+            content_type="application/octet-stream",
+        )
+
+    job_id = None
+    job_status = None
+    poll_json = {}
+    resp_json = {}
+
+    try:
+        # invoke a post request (POST /v1/videos)
+        async with session.post(input.api_url, data=form) as response:
+            if response.status == 200:
+                resp_json = await response.json()
+                job_id = resp_json.get("id")
+                job_status = resp_json.get("status")
+                if not job_id or not job_status:
+                    output.error = "API response missing job 'id' or 'status' field."
+                    output.success = False
+                    return output
+            else:
+                output.error = f"HTTP {response.status}: {await response.text()}"
+                output.success = False
+                return output
+
+        # invoke a poll request (GET /v1/videos/{video_id})
+        poll_interval = 2.0  # Unit(s)
+        timeout_seconds = 600.0
+        deadline = time.perf_counter() + timeout_seconds
+        job_url = f"{input.api_url}/{job_id}"
+
+        while job_status not in {"completed", "failed"}:
+            await asyncio.sleep(poll_interval)
+
+            async with session.get(job_url) as poll_response:
+                if poll_response.status != 200:
+                    output.error = f"Polling failed HTTP {poll_response.status}: {await poll_response.text()}"
+                    output.success = False
+                    return output
+
+                poll_json = await poll_response.json()
+                job_status = poll_json.get("status")
+
+                if time.perf_counter() >= deadline:
+                    output.error = f"Timed out waiting for video job {job_id} to complete."
+                    output.success = False
+                    return output
+
+        if job_status == "failed":
+            output.error = f"Video job failed: {poll_json}"
+            output.success = False
+            return output
+
+        # invoke a get request (GET /v1/videos/{video_id}/content)
+        content_url = f"{job_url}/content"
+        async with session.get(content_url) as content_response:
+            if content_response.status != 200:
+                output.error = (
+                    f"Content retrieval failed HTTP {content_response.status}: {await content_response.text()}"
+                )
+                output.success = False
+                return output
+
+            video_bytes = await content_response.read()
+            output.response_body = video_bytes
+            output.success = True
+            if "peak_memory_mb" in poll_json:
+                output.peak_memory_mb = poll_json["peak_memory_mb"]
+            elif "peak_memory_mb" in resp_json:
+                output.peak_memory_mb = resp_json["peak_memory_mb"]
+    except Exception as e:
+        output.error = str(e)
+        output.success = False
+    finally:
+        if image_file is not None:
+            image_file.close()
+
+        if job_id is not None:
+            try:
+                async with session.delete(f"{input.api_url}/{job_id}") as _:
+                    pass
+            except Exception as e:
+                print(f"Failed to clean up video job {job_id}: {e}")
+
+    output.latency = time.perf_counter() - output.start_time
+
+    if output.success and input.slo_ms is not None:
+        output.slo_achieved = (output.latency * 1000.0) <= float(input.slo_ms)
+
+    if pbar:
+        pbar.update(1)
+    return output
+
+
 backends_function_mapping = {
-    "vllm-omni": (async_request_chat_completions, "/v1/chat/completions"),
-    "openai": (async_request_openai_images, "/v1/images/generations"),
+    "2i": {
+        "vllm-omni": (async_request_chat_completions, "/v1/chat/completions"),
+        "openai": (async_request_openai_images, "/v1/images/generations"),
+    },
+    "2v": {
+        "v1/videos": (async_request_v1_videos, "/v1/videos"),
+    },
 }

benchmarks/diffusion/diffusion_benchmark_serving.py

@@ -9,6 +9,7 @@ If you want to use i2v, i2i dataset, you should `uv pip install gdown` first
 Supports multiple backends:
     - vllm-omni: Uses /v1/chat/completions endpoint (default)
     - openai: Uses /v1/images/generations endpoint
+    - v1/videos: Use /v1/videos endpoint
 
 Usage:
     # Video (vllm-omni backend)
@@ -28,6 +29,16 @@ Usage:
         --backend vllm-omni --dataset vbench --task t2i --num-prompts 10 \
         --height 1024 --width 1024
 
+    python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
+        --backend vllm-omni --dataset random --task t2i --num-prompts 1 \
+        --max-concurrency 1 --enable-negative-prompt \
+        --random-request-config '[
+            {"width":512,"height":512,"num_inference_steps":20,"weight":0.15},
+            {"width":768,"height":768,"num_inference_steps":20,"weight":0.25},
+            {"width":1024,"height":1024,"num_inference_steps":25,"weight":0.45},
+            {"width":1536,"height":1536,"num_inference_steps":35,"weight":0.15}
+        ]'
+
     i2i:
     python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
         --backend vllm-omni --dataset vbench --task i2i --num-prompts 10
@@ -38,6 +49,16 @@ Usage:
         --backend openai --dataset vbench --task t2i --num-prompts 10 \
         --height 1024 --width 1024 --port 3000
 
+    # Video (v1/vedeos)
+    t2v:
+    python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
+        --backend v1/videos --dataset random --task t2v --num-prompts 1 \
+        --max-concurrency 1 --enable-negative-prompt \
+        --random-request-config '[
+            {"width":854,"height":480,"num_inference_steps":18,"num_frames":120,"fps":24,"weight":1}
+        ]'
+
+
 """
 
 import argparse
@@ -45,7 +66,10 @@ import ast
 import asyncio
 import glob
 import json
+import logging
 import os
+import random
+import tempfile
 import time
 import uuid
 from abc import ABC, abstractmethod
@@ -57,8 +81,11 @@ import aiohttp
 import numpy as np
 import requests
 from backends import RequestFuncInput, RequestFuncOutput, backends_function_mapping
+from PIL import Image
 from tqdm.asyncio import tqdm
 
+logger = logging.getLogger(__name__)
+
 
 class BaseDataset(ABC):
     def __init__(self, args, api_url: str, model: str):
@@ -516,26 +543,66 @@ class TraceDataset(BaseDataset):
 
 
 class RandomDataset(BaseDataset):
-    def __init__(self, args, api_url: str, model: str):
-        self.args = args
-        self.api_url = api_url
-        self.model = model
+    def __init__(self, args, api_url: str, model: str, enable_negative_prompt: bool = False):
+        super().__init__(args, api_url, model)
         self.num_prompts = args.num_prompts
+        self.enable_negative_prompt = enable_negative_prompt
+        self.random_request_config = getattr(args, "random_request_config", None)
+        if self.random_request_config:
+            self.random_request_config = json.loads(self.random_request_config)
+            self._weights = [p["weight"] for p in self.random_request_config]
+
+            self.random_request_config = [
+                {k: v for k, v in p.items() if k != "weight"} for p in self.random_request_config
+            ]
+
+            seed = getattr(args, "random_request_seed", 42)
+            self._rng = random.Random(seed)
+
+            self._sampled_requests = self._rng.choices(
+                self.random_request_config,
+                weights=self._weights,
+                k=self.num_prompts,
+            )
+        else:
+            self._sampled_requests = None
+
+        # Random image generate
+        if self.args.task in ["i2v", "ti2v", "ti2i", "i2i"]:
+            img = Image.new("RGB", (512, 512), (255, 255, 255))
+
+            image_path = os.path.join(tempfile.gettempdir(), "diffusion_benchmark_random_image.png")
+            self._random_image_path = [image_path]
+            img.save(image_path)
+        else:
+            self._random_image_path = None
 
     def __len__(self) -> int:
         return self.num_prompts
 
     def __getitem__(self, idx: int) -> RequestFuncInput:
+        extra_body = {}
+        if self.enable_negative_prompt:
+            extra_body["negative_prompt"] = f"Negative prompt {idx} for benchmarking diffusion models"
+
+        params = {
+            "width": self.args.width,
+            "height": self.args.height,
+            "num_frames": self.args.num_frames,
+            "num_inference_steps": self.args.num_inference_steps,
+            "fps": self.args.fps,
+        }
+        if self._sampled_requests:
+            profile = self._sampled_requests[idx]
+            params.update(profile)
         return RequestFuncInput(
             prompt=f"Random prompt {idx} for benchmarking diffusion models",
             api_url=self.api_url,
             model=self.model,
-            width=self.args.width,
-            height=self.args.height,
-            num_frames=self.args.num_frames,
-            num_inference_steps=self.args.num_inference_steps,
             seed=self.args.seed,
-            fps=self.args.fps,
+            extra_body=extra_body,
+            image_paths=self._random_image_path,
+            **params,
         )
 
     def get_requests(self) -> list[RequestFuncInput]:
@@ -686,6 +753,7 @@ def calculate_metrics(
         "latency_mean": np.mean(latencies) if latencies else 0,
         "latency_median": np.median(latencies) if latencies else 0,
         "latency_p99": np.percentile(latencies, 99) if latencies else 0,
+        "latency_p95": np.percentile(latencies, 95) if latencies else 0,
         "latency_p50": np.percentile(latencies, 50) if latencies else 0,
         "peak_memory_mb_max": max(peak_memories) if peak_memories else 0,
         "peak_memory_mb_mean": np.mean(peak_memories) if peak_memories else 0,
@@ -742,8 +810,32 @@ async def benchmark(args):
     if args.base_url is None:
         args.base_url = f"http://{args.host}:{args.port}"
 
+    VIDEO_TASKS = {"t2v", "i2v", "ti2v"}
+    IMAGE_TASKS = {"t2i", "i2i", "ti2i"}
+
+    if args.task in VIDEO_TASKS:
+        task_type = "2v"
+    elif args.task in IMAGE_TASKS:
+        task_type = "2i"
+    else:
+        raise ValueError(
+            f"Unsupported task: '{args.task}'. "
+            f"Valid video tasks: {sorted(VIDEO_TASKS)}, "
+            f"Valid image tasks: {sorted(IMAGE_TASKS)}"
+        )
+
+    valid_backends = sorted(backends_function_mapping[task_type].keys())
+
+    if args.backend not in valid_backends:
+        logger.error(
+            f"Invalid backend '{args.backend}' for task '{args.task}' (task type: '{task_type}').\n"
+            f"Valid backends for this task type: {valid_backends}\n"
+            f"Example usage: --task {args.task} --backend {valid_backends[0]}"
+        )
+        raise ValueError("Backend validation failed. See log above for valid options.")
+
     # Setup API URL and request function based on backend
-    request_func, api_url = backends_function_mapping[args.backend]
+    request_func, api_url = backends_function_mapping[task_type][args.backend]
     api_url = f"{args.base_url}{api_url}"
 
     if args.dataset == "vbench":
@@ -751,7 +843,7 @@ async def benchmark(args):
     elif args.dataset == "trace":
         dataset = TraceDataset(args, api_url, args.model)
     elif args.dataset == "random":
-        dataset = RandomDataset(args, api_url, args.model)
+        dataset = RandomDataset(args, api_url, args.model, args.enable_negative_prompt)
     else:
         raise ValueError(f"Unknown dataset: {args.dataset}")
 
@@ -847,6 +939,7 @@ async def benchmark(args):
     print("{:<40} {:<15.4f}".format("Latency Mean (s):", metrics["latency_mean"]))
     print("{:<40} {:<15.4f}".format("Latency Median (s):", metrics["latency_median"]))
     print("{:<40} {:<15.4f}".format("Latency P99 (s):", metrics["latency_p99"]))
+    print("{:<40} {:<15.4f}".format("Latency P95 (s):", metrics["latency_p95"]))
 
     if args.slo:
         print(f"{'-' * 50}")
@@ -883,7 +976,7 @@ if __name__ == "__main__":
         "--backend",
         type=str,
         default="vllm-omni",
-        choices=["vllm-omni", "openai"],
+        choices=["vllm-omni", "openai", "v1/videos"],
         help="Backend to target the benchmark to.",
     )
     parser.add_argument(
@@ -972,6 +1065,25 @@ if __name__ == "__main__":
         help="SLO target multiplier: slo_ms = estimated_exec_time_ms * slo_scale (default: 3).",
     )
     parser.add_argument("--disable-tqdm", action="store_true", help="Disable progress bar.")
+    parser.add_argument(
+        "--enable-negative-prompt",
+        action="store_true",
+        default=False,
+        help="Generate negative prompts when using the random dataset.",
+    )
+    parser.add_argument(
+        "--random-request-config",
+        type=str,
+        default=None,
+        help=(
+            "JSON string defining random request profiles. "
+            "Each profile may contain: width, height, num_inference_steps, etc. "
+            "The 'weight' field controls sampling probability (relative weight). "
+            "Example: "
+            '[{"width":512,"height":512,"num_inference_steps":20,"weight":0.15},'
+            '{"width":768,"height":768,"num_inference_steps":20,"weight":0.85}]'
+        ),
+    )
 
     args = parser.parse_args()
 

benchmarks/diffusion/performance_dashboard/qwen_image_serving_performance.md

@@ -0,0 +1,169 @@
+# Qwen-Image Serving Performance Dashboard
+
+This document describes how to deploy and benchmark **Qwen-Image** using vLLM-Omni. It includes service startup configuration, acceleration-related options, benchmark methodology, dataset settings, and performance results.
+
+---
+
+# 1. Overview
+
+Qwen-Image is a multimodal text-to-image generation model served through the vLLM-Omni infrastructure.
+
+This document covers:
+
+* Service launch configuration (including acceleration options)
+* Benchmark scripts and usage
+* Dataset and workload settings
+* Performance measurement results
+* Reproducibility guidelines
+
+---
+
+# 2. Test Environment
+| Component | Specification |
+|------------|----------------|
+| GPU | NVIDIA A100-SXM4-80GB |
+| Diffusion Attention Backend | FlashAttention |
+
+# 3. Service Launch Configuration
+
+## 3.1 Basic Serving Command
+
+```bash
+vllm serve Qwen/Qwen-Image --omni \
+    --port 8091
+```
+
+## 3.2 Key Parameters
+
+| Parameter             | Description              |
+| --------------------- | ------------------------ |
+| `--cfg-parallel-size` | CFG parallelism degree   |
+| `--ulysses-degree`    | Ulysses parallel degree  |
+| `--vae-patch-parallel-size`    | VAE parallel degree  |
+| `--tensor-parallel-size` | Tensor parallelism degree |
+
+Record these parameters when reporting performance results.
+
+---
+
+# 4. Benchmark Script
+
+## 4.1 Benchmark Entry
+
+```bash
+python benchmarks/diffusion/diffusion_benchmark_serving.py \
+    --backend vllm-omni \
+    --dataset <DATASET_NAME> \
+    --task t2i \
+    --num-prompts <N> \
+    --max-concurrency <C> \
+    --enable-negative-prompt \
+    --random-request-config <CFG>
+```
+
+## 4.2 Key Benchmark Arguments
+
+| Parameter              | Description                       |
+| ---------------------- | --------------------------------- |
+| `--backend`            | Serving backend (use `vllm-omni`) |
+| `--dataset`            | Dataset name (`random` or custom) |
+| `--task`               | Task type (e.g., `t2i`)           |
+| `--num-prompts`        | Total number of requests          |
+| `--max-concurrency`    | Client-side concurrency           |
+| `--random-request-config`| JSON string defining random request |
+
+---
+
+# 5. Dataset & Workload Settings
+
+## 5.1 Recommended Evaluation Configurations
+
+### Dataset A ( 512 Resolution)
+
+* Dataset: `random`
+* Task: t2i
+* Concurrency: 1
+* Mix Resolution
+```
+[
+ {"width":512,"height":512,"num_inference_steps":20,"weight":1}
+]
+```
+
+### Dataset B (1536 Resolution)
+
+* Dataset: `random`
+* Task: t2i
+* Concurrency: 1
+* Mix Resolution
+```
+[
+ {"width":1536,"height":1536,"num_inference_steps":35,"weight":1}
+]
+```
+
+### Dataset C (Mix Resolution)
+
+* Dataset: `random`
+* Task: t2i
+* Concurrency: 1
+* Mix Resolution
+```
+[
+ {"width":512,"height":512,"num_inference_steps":20,"weight":0.15},
+ {"width":768,"height":768,"num_inference_steps":20,"weight":0.25},
+ {"width":1024,"height":1024,"num_inference_steps":25,"weight":0.45},
+ {"width":1536,"height":1536,"num_inference_steps":35,"weight":0.15}
+]
+```
+---
+
+## 5.2 Example Benchmark Command
+
+```bash
+python benchmarks/diffusion/diffusion_benchmark_serving.py \
+    --backend vllm-omni \
+    --dataset random \
+    --task t2i \
+    --num-prompts 1 \
+    --max-concurrency 1 \
+    --enable-negative-prompt \
+    --random-request-config '[
+        {"width":512,"height":512,"num_inference_steps":20,"weight":1}
+    ]'
+```
+
+---
+
+# 6. Performance Metrics
+
+The following metrics are collected during benchmarking:
+
+| Metric             | Description                   | Unit    |
+| ------------------ | ----------------------------- | ------- |
+| Mean Latency        | Mean of latency       | seconds |
+| P99 Latency        | P99 of latency             | seconds |
+
+---
+
+# 7. Performance Results
+
+| Dataset Configuration | Max Concur. | CFG | Usp | Tp | VAE Parallel | Mean Latency (s) | P99 Latency (s) |
+|-----------------------|-----|-----|-----|----|--------------|------------------|------------------|
+| Dataset A | 1 | 2 | 2 | Off | Off          | 2.2087          | 2.2087          |
+| Dataset B | 1 | 2 | 2 | Off | Off          | 19.6739          | 19.6739          |
+| Dataset C | 1 | 2 | 2 | Off | Off          | 5.67259        | 18.6234 |
+---
+
+# 8. Reproducibility Checklist
+
+To ensure consistent and comparable benchmark results:
+
+* Record GPU type
+* Record parallel configuration
+* Record benchmark parameters (resolution, concurrency, number of prompts)
+* Ensure no background workload on GPUs during testing
+
+---
+
+This document serves as the official Qwen-Image serving performance reference under vLLM-Omni.

benchmarks/diffusion/performance_dashboard/wan_2_2_serving_performance.md

@@ -0,0 +1,170 @@
+# Wan2.2 Serving Performance Dashboard
+
+This document describes how to deploy and benchmark **Wan-AI/Wan2.2-T2V-A14B-Diffusers** using vLLM-Omni. It includes service startup configuration, acceleration-related options, benchmark methodology, dataset settings, and performance results.
+
+---
+
+# 1. Overview
+
+Wan-AI/Wan2.2-T2V-A14B-Diffusers is a multimodal text-to-video generation model served through the vLLM-Omni infrastructure.
+
+This document covers:
+
+* Service launch configuration (including acceleration options)
+* Benchmark scripts and usage
+* Dataset and workload settings
+* Performance measurement results
+* Reproducibility guidelines
+
+---
+
+# 2. Test Environment
+| Component | Specification |
+|------------|----------------|
+| GPU | NVIDIA A100-SXM4-80GB |
+| Diffusion Attention Backend | FlashAttention |
+
+# 3. Service Launch Configuration
+
+## 3.1 Basic Serving Command
+
+```bash
+vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni \
+    --port 8091
+```
+
+## 3.2 Key Parameters
+
+| Parameter             | Description              |
+| --------------------- | ------------------------ |
+| `--cfg-parallel-size` | CFG parallelism degree   |
+| `--ulysses-degree`    | Ulysses parallel degree  |
+| `--vae-patch-parallel-size`    | VAE parallel degree  |
+| `--tensor-parallel-size` | Tensor parallelism degree |
+| `--use-hsdp` | Enable Hybrid Sharded Data Parallel to shard model weights across GPUs |
+
+Record these parameters when reporting performance results.
+
+---
+
+# 4. Benchmark Script
+
+## 4.1 Benchmark Entry
+
+```bash
+python benchmarks/diffusion/diffusion_benchmark_serving.py \
+    --backend v1/videos \
+    --dataset <DATASET_NAME> \
+    --task t2v \
+    --num-prompts <N> \
+    --max-concurrency <C> \
+    --enable-negative-prompt \
+    --random-request-config <CFG>
+```
+
+## 4.2 Key Benchmark Arguments
+
+| Parameter              | Description                       |
+| ---------------------- | --------------------------------- |
+| `--backend`            | Serving backend (use `v1/videos`) |
+| `--dataset`            | Dataset name (`random` or custom) |
+| `--task`               | Task type (e.g., `t2v`)           |
+| `--num-prompts`        | Total number of requests          |
+| `--max-concurrency`    | Client-side concurrency           |
+| `--random-request-config`| JSON string defining random request |
+
+---
+
+# 5. Dataset & Workload Settings
+
+## 5.1 Recommended Evaluation Configurations
+
+### Dataset A (480p)
+
+* Dataset: `random`
+* Task: t2v
+* Concurrency: 1
+* Mix Resolution
+```
+[
+    {"width":854,"height":480,"num_inference_steps":3,"num_frames":80,"fps":16,"weight":1}
+]
+```
+### Dataset B (720p)
+
+* Dataset: `random`
+* Task: t2v
+* Concurrency: 1
+* Mix Resolution
+```
+[
+    {"width":1280,"height":720,"num_inference_steps":6,"num_frames":80,"fps":16,"weight":1}
+]
+```
+### Dataset C (Mix Resolution)
+
+* Dataset: `random`
+* Task: t2v
+* Concurrency: 1
+* Mix Resolution
+```
+[
+ {"width":854,"height":480,"num_inference_steps":3,"num_frames":80,"fps":16,"weight":0.15},
+ {"width":854,"height":480,"num_inference_steps":4,"num_frames":120,"fps":24,"weight":0.25},
+ {"width":1280,"height":720,"num_inference_steps":6,"num_frames":80,"fps":16,"weight":0.6}
+]
+```
+---
+
+## 5.2 Example Benchmark Command
+
+```bash
+python benchmarks/diffusion/diffusion_benchmark_serving.py \
+    --backend v1/videos \
+    --dataset random \
+    --task t2v \
+    --num-prompts 1 \
+    --max-concurrency 1 \
+    --enable-negative-prompt \
+    --random-request-config '[
+        {"width":854,"height":480,"num_inference_steps":18,"num_frames": 33,"fps":16",weight":1}
+    ]'
+```
+
+---
+
+# 6. Performance Metrics
+
+The following metrics are collected during benchmarking:
+
+| Metric             | Description                   | Unit    |
+| ------------------ | ----------------------------- | ------- |
+| Mean Latency        | Mean of latency       | seconds |
+| P99 Latency        | P99 of latency             | seconds |
+
+---
+
+# 7. Performance Results
+
+| Dataset Configuration | Max Concur. | CFG | Usp | Tp | Hsdp | VAE Parallel | Mean Latency (s) | P99 Latency (s) |
+|-----------------------|-----|-----|-----|-----|----|--------------|------------------|------------------|
+| Dataset A | 1 | 2 | 2 | 1 | On | 1          | 24.6766          | 24.6766          |
+| Dataset A | 1 | 2 | 2 | 1 | On | 4          | 21.6810          | 21.6810          |
+| Dataset B | 1 | 2 | 2 | 1 | On | 1          | 124.6639         | 124.6639          |
+| Dataset B | 1 | 2 | 2 | 1 | On | 4          | 117.44          | 117.44          |
+| Dataset C | 1 | 2 | 2 | 1 | On | 1          | 79.2175        | 124.2565 |
+| Dataset C | 1 | 2 | 2 | 1 | On | 4          | 74.4977        | 117.710 |
+---
+
+# 8. Reproducibility Checklist
+
+To ensure consistent and comparable benchmark results:
+
+* Record GPU type
+* Record parallel configuration
+* Record benchmark parameters (resolution, concurrency, number of prompts)
+* Ensure no background workload on GPUs during testing
+
+---
+
+This document serves as the official Wan2.2 serving performance reference under vLLM-Omni.

benchmarks/qwen3-omni/README.md

@@ -56,9 +56,9 @@ What it does:
 - Runs `examples/offline_inference/qwen3_omni/end2end.py` with `--log-stats`.
 - Uses `benchmarks/build_dataset/top100.txt` and writes to:
   - Logs: `benchmarks/qwen3-omni/vllm_omni/logs/`
-    - `omni_llm_pipeline_text.orchestrator.stats.jsonl` — per-stage latency stats.
-    - `omni_llm_pipeline_text.overall.stats.jsonl` — end-to-end latency/TPS.
-    - `omni_llm_pipeline_text.stage{0,1,2}.log` — per-stage detailed logs/errors.
+    - `omni_pipeline_text.orchestrator.stats.jsonl` — per-stage latency stats.
+    - `omni_pipeline_text.overall.stats.jsonl` — end-to-end latency/TPS.
+    - `omni_pipeline_text.stage{0,1,2}.log` — per-stage detailed logs/errors.
   - Outputs: `benchmarks/qwen3-omni/vllm_omni/outputs/` — ~100 text and `.wav` files.
 
 Key checks:

benchmarks/qwen3-omni/vllm_omni/eval_qwen3_moe_omni.sh

@@ -26,12 +26,12 @@ else
                       --log-stats \
                       --log-dir $log_dir
     echo "Logs and outputs are saved in ${log_dir} and ${outputs_dir} respectively:"
-    echo "  - omni_llm_pipeline_text                       run dir/base name"
-    echo "  - omni_llm_pipeline_text.orchestrator.stats.jsonl  orchestrator-stage latency stats"
-    echo "  - omni_llm_pipeline_text.overall.stats.jsonl       overall latency/TPS stats"
-    echo "  - omni_llm_pipeline_text.stage0.log                per-stage detailed logs"
-    echo "  - omni_llm_pipeline_text.stage1.log"
-    echo "  - omni_llm_pipeline_text.stage2.log"
+    echo "  - omni_pipeline_text                       run dir/base name"
+    echo "  - omni_pipeline_text.orchestrator.stats.jsonl  orchestrator-stage latency stats"
+    echo "  - omni_pipeline_text.overall.stats.jsonl       overall latency/TPS stats"
+    echo "  - omni_pipeline_text.stage0.log                per-stage detailed logs"
+    echo "  - omni_pipeline_text.stage1.log"
+    echo "  - omni_pipeline_text.stage2.log"
     echo "Key checks: overall.stats.jsonl for end-to-end latency/TPS; orchestrator.stats.jsonl for stable per-stage latency; stage*.log for errors or long tails."
     echo "  - outputs/             Generated txt and wav files, there should be 100 text and wav files generated respectively"
 fi

benchmarks/qwen3-tts/README.md

@@ -32,8 +32,8 @@ bash run_benchmark.sh --hf-only
 # Use a different model (e.g. 1.7B)
 MODEL=Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice bash run_benchmark.sh --async-only
 
-# Use batch_size=4 config for higher throughput
-STAGE_CONFIG=vllm_omni/configs/qwen3_tts_bs4.yaml bash run_benchmark.sh --async-only
+# Use bs16 config for higher throughput
+STAGE_CONFIG=vllm_omni/configs/qwen3_tts_bs16.yaml bash run_benchmark.sh --async-only
 
 # Custom GPU, prompt count, concurrency levels
 GPU_DEVICE=1 NUM_PROMPTS=20 CONCURRENCY="1 4" bash run_benchmark.sh
@@ -54,7 +54,7 @@ CUDA_VISIBLE_DEVICES=0 python -m vllm_omni.entrypoints.cli.main serve \
 ### 2) Run online serving benchmark
 
 ```bash
-python vllm_omni/bench_tts_serve.py \
+python benchmarks/qwen3-tts/vllm_omni/bench_tts_serve.py \
     --port 8000 \
     --num-prompts 50 \
     --max-concurrency 1 4 10 \
@@ -65,7 +65,7 @@ python vllm_omni/bench_tts_serve.py \
 ### 3) Run HuggingFace baseline
 
 ```bash
-python transformers/bench_tts_hf.py \
+python benchmarks/qwen3-tts/transformers/bench_tts_hf.py \
     --model "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice" \
     --num-prompts 50 \
     --gpu-device 0 \
@@ -75,7 +75,7 @@ python transformers/bench_tts_hf.py \
 ### 4) Generate comparison plots
 
 ```bash
-python plot_results.py \
+python benchmarks/qwen3-tts/plot_results.py \
     --results results/bench_async_chunk_*.json results/bench_hf_transformers_*.json \
     --labels "vllm-omni" "hf_transformers" \
     --output results/comparison.png
@@ -83,14 +83,14 @@ python plot_results.py \
 
 ## Stage Configs
 
-| Config | Batch Size | Description |
-|--------|:----------:|-------------|
-| `vllm_omni/configs/qwen3_tts_bs1.yaml` | 1 | Single-request processing (lower latency) |
-| `vllm_omni/configs/qwen3_tts_bs4.yaml` | 4 | Concurrent request processing (higher throughput) |
+| Config | max_num_seqs | Description |
+|--------|:------------:|-------------|
+| `vllm_omni/configs/qwen3_tts_bs1.yaml` | 1 | Single-request processing (lowest latency) |
+| `vllm_omni/configs/qwen3_tts_bs16.yaml` | 16 | High-throughput concurrent processing |
 
-Both configs use a 2-stage pipeline (Talker -> Code2Wav) with `async_chunk` streaming enabled. The `SharedMemoryConnector` streams codec frames (25-frame chunks with 25-frame context overlap) between stages.
+All configs use a 2-stage pipeline (Talker -> Code2Wav) with `async_chunk` streaming enabled. The `SharedMemoryConnector` streams codec frames (25-frame chunks with 25-frame context overlap) between stages.
 
-The model is specified via the CLI `--model` flag (or `MODEL` env var in `run_benchmark.sh`), so the same configs work for both the 0.6B and 1.7B model variants.
+The model is specified via the CLI `--model` flag (or `MODEL` env var), so the same configs work for both the 0.6B and 1.7B model variants.
 
 ## Metrics
 

benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs1.yaml

@@ -4,19 +4,18 @@ async_chunk: true
 stage_args:
   - stage_id: 0
     stage_type: llm
+    is_comprehension: true
     runtime:
       devices: "0"
-      max_batch_size: 1
     engine_args:
+      max_num_seqs: 1
       model_stage: qwen3_tts
       model_arch: Qwen3TTSTalkerForConditionalGeneration
-      hf_overrides:
-        architectures: [Qwen3TTSTalkerForConditionalGeneration]
       worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       enforce_eager: false
       trust_remote_code: true
-      async_scheduling: false
+      async_scheduling: true
       enable_prefix_caching: false
       engine_output_type: latent
       gpu_memory_utilization: 0.3
@@ -39,20 +38,18 @@ stage_args:
     stage_type: llm
     runtime:
       devices: "0"
-      max_batch_size: 1
     engine_args:
+      max_num_seqs: 1
       model_stage: code2wav
       model_arch: Qwen3TTSCode2Wav
-      hf_overrides:
-        architectures: [Qwen3TTSCode2Wav]
       worker_type: generation
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       enforce_eager: true
       trust_remote_code: true
-      async_scheduling: false
+      async_scheduling: true
       enable_prefix_caching: false
       engine_output_type: audio
-      gpu_memory_utilization: 0.2
+      gpu_memory_utilization: 0.3
       distributed_executor_backend: "mp"
       max_num_batched_tokens: 8192
       max_model_len: 32768

benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs16.yaml

@@ -0,0 +1,94 @@
+# Qwen3-TTS max_num_seqs=16 config (streaming with async_chunk)
+# High-throughput concurrent request processing
+# 2-stage pipeline: Talker -> Code2Wav
+async_chunk: true
+stage_args:
+  - stage_id: 0
+    stage_type: llm
+    is_comprehension: true
+    runtime:
+      devices: "0"
+    engine_args:
+      max_num_seqs: 16
+      model_stage: qwen3_tts
+      model_arch: Qwen3TTSTalkerForConditionalGeneration
+      worker_type: ar
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      enforce_eager: false
+      trust_remote_code: true
+      async_scheduling: true
+      enable_prefix_caching: false
+      engine_output_type: latent
+      gpu_memory_utilization: 0.3
+      distributed_executor_backend: "mp"
+      max_num_batched_tokens: 4096
+      max_model_len: 4096
+      custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_tts.talker2code2wav_async_chunk
+    output_connectors:
+      to_stage_1: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 0.9
+      top_k: 50
+      max_tokens: 4096
+      seed: 42
+      detokenize: false
+      repetition_penalty: 1.05
+      stop_token_ids: [2150]
+
+  - stage_id: 1
+    stage_type: llm
+    runtime:
+      devices: "0"
+    engine_args:
+      max_num_seqs: 16
+      model_stage: code2wav
+      model_arch: Qwen3TTSCode2Wav
+      worker_type: generation
+      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
+      enforce_eager: true
+      trust_remote_code: true
+      async_scheduling: true
+      enable_prefix_caching: false
+      engine_output_type: audio
+      gpu_memory_utilization: 0.2
+      distributed_executor_backend: "mp"
+      max_num_batched_tokens: 16384
+      max_model_len: 32768
+    engine_input_source: [0]
+    final_output: true
+    final_output_type: audio
+    input_connectors:
+      from_stage_0: connector_of_shared_memory
+    tts_args:
+      max_instructions_length: 500
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 65536
+      seed: 42
+      detokenize: true
+      repetition_penalty: 1.0
+
+runtime:
+  enabled: true
+  defaults:
+    window_size: -1
+    max_inflight: 16
+
+  connectors:
+    connector_of_shared_memory:
+      name: SharedMemoryConnector
+      extra:
+        shm_threshold_bytes: 65536
+        codec_streaming: true
+        connector_get_sleep_s: 0.01
+        connector_get_max_wait_first_chunk: 3000
+        connector_get_max_wait: 300
+        codec_chunk_frames: 25
+        codec_left_context_frames: 25
+
+  edges:
+    - from: 0
+      to: 1
+      window_size: -1

benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs4.yaml

@@ -1,23 +1,22 @@
 # Qwen3-TTS batch_size=4 config (streaming with async_chunk)
-# Enables concurrent request processing with max_inflight=4
+# Enables concurrent request processing
 # 2-stage pipeline: Talker -> Code2Wav
 async_chunk: true
 stage_args:
   - stage_id: 0
     stage_type: llm
+    is_comprehension: true
     runtime:
       devices: "0"
-      max_batch_size: 4
     engine_args:
+      max_num_seqs: 4
       model_stage: qwen3_tts
       model_arch: Qwen3TTSTalkerForConditionalGeneration
-      hf_overrides:
-        architectures: [Qwen3TTSTalkerForConditionalGeneration]
       worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       enforce_eager: false
       trust_remote_code: true
-      async_scheduling: false
+      async_scheduling: true
       enable_prefix_caching: false
       engine_output_type: latent
       gpu_memory_utilization: 0.3
@@ -40,17 +39,15 @@ stage_args:
     stage_type: llm
     runtime:
       devices: "0"
-      max_batch_size: 4
     engine_args:
+      max_num_seqs: 4
       model_stage: code2wav
       model_arch: Qwen3TTSCode2Wav
-      hf_overrides:
-        architectures: [Qwen3TTSCode2Wav]
       worker_type: generation
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       enforce_eager: true
       trust_remote_code: true
-      async_scheduling: false
+      async_scheduling: true
       enable_prefix_caching: false
       engine_output_type: audio
       gpu_memory_utilization: 0.2

docker/Dockerfile.ci

@@ -11,10 +11,29 @@ RUN apt-get update && \
     apt-get clean && \
     rm -rf /var/lib/apt/lists/*
 
-# Install vllm-omni into the same uv-managed Python environment used by the base image.
-# Use bash -c so that $(python3 -c ...) is expanded inside the container.
-RUN uv pip install --system --no-cache-dir ".[dev]"
+RUN uv pip uninstall --system -y vllm || true
 
+# Install vLLM from precompiled wheel at the selected commit.
+# Must use direct URL because the wheel has a PEP 440 local version identifier
+# (e.g. +g0a0a1a198) which pip/uv refuse to install from a PEP 503 package index.
+ENV VLLM_PRECOMPILED_WHEEL_COMMIT=89138b21cc246ae944c741d5c399c148e2b770ab
+RUN VLLM_WHEEL_URL=$(python3 -c "import urllib.request,re; \
+    html=urllib.request.urlopen('https://wheels.vllm.ai/${VLLM_PRECOMPILED_WHEEL_COMMIT}/vllm/').read().decode(); \
+    m=re.search(r'>(\S+x86_64\.whl)<',html); \
+    print('https://wheels.vllm.ai/${VLLM_PRECOMPILED_WHEEL_COMMIT}/'+m.group(1).replace('+','%2B'))") && \
+    echo "Installing vLLM from: ${VLLM_WHEEL_URL}" && \
+    uv pip install --system --force-reinstall "${VLLM_WHEEL_URL}"
+
+RUN uv pip install --system ".[dev]"
+
+RUN uv pip install --system --upgrade \
+        "flashinfer-cubin==0.6.6" \
+        "nvidia-cublas-cu12==12.9.1.4" \
+        "numpy==2.2.6"
+
+RUN uv pip install --system --upgrade \
+    "flashinfer-jit-cache==0.6.6" \
+    --index-url https://flashinfer.ai/whl/cu129
 RUN ln -sf /usr/bin/python3 /usr/bin/python
 
 ENTRYPOINT []

docker/Dockerfile.npu

@@ -1,17 +1,17 @@
 ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
-ARG VLLM_ASCEND_TAG=v0.14.0rc1
+ARG VLLM_ASCEND_TAG=v0.17.0rc1
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 
-WORKDIR /vllm-workspace/vllm-ascend
-RUN git checkout e2175d9c7e62b437391dfee996b1375674ba7c18
-RUN pip install -v -e .
-
 ARG APP_DIR=/vllm-workspace/vllm-omni
 WORKDIR ${APP_DIR}
 
 COPY . .
 
-RUN pip install -v -e .
+RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
+    source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
+    source /usr/local/Ascend/nnal/atb/set_env.sh && \
+    export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
+    python3 -m pip install -v -e /vllm-workspace/vllm-omni/ --no-build-isolation
 
 ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
 

docker/Dockerfile.npu.a3

@@ -1,17 +1,17 @@
 ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
-ARG VLLM_ASCEND_TAG=v0.14.0rc1-a3
+ARG VLLM_ASCEND_TAG=v0.17.0rc1-a3
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 
-WORKDIR /vllm-workspace/vllm-ascend
-RUN git checkout e2175d9c7e62b437391dfee996b1375674ba7c18
-RUN pip install -v -e .
-
 ARG APP_DIR=/vllm-workspace/vllm-omni
 WORKDIR ${APP_DIR}
 
 COPY . .
 
-RUN pip install -v -e .
+RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
+    source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
+    source /usr/local/Ascend/nnal/atb/set_env.sh && \
+    export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
+    python3 -m pip install -v -e /vllm-workspace/vllm-omni/ --no-build-isolation
 
 ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
 

docs/.nav.yml

@@ -58,7 +58,9 @@ nav:
       - Quantization:
         - Overview: user_guide/diffusion/quantization/overview.md
         - FP8: user_guide/diffusion/quantization/fp8.md
+        - Int8: user_guide/diffusion/quantization/int8.md
         - GGUF: user_guide/diffusion/quantization/gguf.md
+      - Step Execution: user_guide/diffusion/step_execution.md
       - Parallelism Acceleration: user_guide/diffusion/parallelism_acceleration.md
       - CPU Offloading: user_guide/diffusion/cpu_offload_diffusion.md
       - LoRA: user_guide/diffusion/lora.md
@@ -90,6 +92,7 @@ nav:
       - design/feature/cache_dit.md
       - design/feature/teacache.md
       - design/feature/async_chunk_design.md
+      - design/feature/diffusion_step_execution.md
     - Module Design:
       - design/module/ar_module.md
       - design/module/dit_module.md

docs/api/README.md

@@ -6,19 +6,13 @@ Main entry points for vLLM-Omni inference and serving.
 
 - [vllm_omni.entrypoints.async_omni.AsyncOmni][]
 - [vllm_omni.entrypoints.async_omni_diffusion.AsyncOmniDiffusion][]
-- [vllm_omni.entrypoints.async_omni_llm.AsyncOmniLLM][]
 - [vllm_omni.entrypoints.cli.benchmark.base.OmniBenchmarkSubcommandBase][]
 - [vllm_omni.entrypoints.cli.benchmark.main.OmniBenchmarkSubcommand][]
 - [vllm_omni.entrypoints.cli.benchmark.serve.OmniBenchmarkServingSubcommand][]
 - [vllm_omni.entrypoints.cli.serve.OmniServeCommand][]
 - [vllm_omni.entrypoints.client_request_state.ClientRequestState][]
 - [vllm_omni.entrypoints.omni.Omni][]
-- [vllm_omni.entrypoints.omni.OmniBase][]
-- [vllm_omni.entrypoints.omni_diffusion.OmniDiffusion][]
-- [vllm_omni.entrypoints.omni_llm.OmniLLM][]
-- [vllm_omni.entrypoints.omni_stage.OmniStage][]
-- [vllm_omni.entrypoints.stage_utils.OmniStageTaskType][]
-- [vllm_omni.entrypoints.zmq_utils.ZmqQueue][]
+- [vllm_omni.entrypoints.omni_base.OmniBase][]
 
 ## Inputs
 
@@ -48,9 +42,7 @@ Engine classes for offline and online inference.
 - [vllm_omni.engine.OmniEngineCoreOutputs][]
 - [vllm_omni.engine.OmniEngineCoreRequest][]
 - [vllm_omni.engine.PromptEmbedsPayload][]
-- [vllm_omni.engine.arg_utils.AsyncOmniEngineArgs][]
 - [vllm_omni.engine.arg_utils.OmniEngineArgs][]
-- [vllm_omni.engine.input_processor.OmniInputProcessor][]
 - [vllm_omni.engine.output_processor.MultimodalOutputProcessor][]
 - [vllm_omni.engine.output_processor.OmniRequestState][]
 

docs/configuration/README.md

@@ -14,6 +14,10 @@ For introduction, please check [Introduction for stage config](./stage_configs.m
 
 - **[GPU Memory Calculation and Configuration](./gpu_memory_utilization.md)** - Guide on how to calculate memory requirements and set up `gpu_memory_utilization` for optimal performance
 
+## Multi-Stage Recipes
+
+- **[Prefill-Decode Disaggregation](./pd_disaggregation.md)** - How to derive a PD-aware Qwen3-Omni stage config from the default config without introducing another bundled YAML
+
 ## Optimization Features
 
 - **[TeaCache Configuration](../user_guide/diffusion/teacache.md)** - Enable TeaCache adaptive caching for DiT models to achieve 1.5x-2.0x speedup with minimal quality loss

docs/configuration/gpu_memory_utilization.md

@@ -158,7 +158,7 @@ The stage initialized but ran out of memory during processing.
 
 **Solutions:**
 1. Reduce `max_num_batched_tokens`
-2. Reduce `max_batch_size` in runtime config
+2. Reduce `max_num_seqs` in engine_args
 3. Lower `gpu_memory_utilization` slightly
 4. Enable quantization if supported
 

docs/configuration/pd_disaggregation.md

@@ -0,0 +1,177 @@
+# Prefill-Decode (PD) Disaggregation
+
+PD disaggregation splits the Qwen3-Omni thinker into separate prefill and decode
+stages so prompt processing and token generation can run on different workers.
+
+This is documented as a stage-config recipe instead of a bundled YAML because the
+deployment-specific values usually change per environment:
+
+- GPU placement
+- `tensor_parallel_size`
+- connector backend and connector ports
+- connector IPs or bootstrap addresses
+
+Start from the [default Qwen3-Omni stage config](gh-file:vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml)
+and copy it to your own file, for example `qwen3_omni_pd.yaml`. Then apply the
+changes below.
+
+## Requirements
+
+- 3+ GPUs for a basic layout: prefill, decode, and talker+code2wav
+- A KV connector supported by vLLM, such as `MooncakeConnector`
+- Matching `tensor_parallel_size` on the prefill and decode thinker stages
+
+## 1. Split the thinker into prefill and decode stages
+
+Replace the original thinker stage with two stages:
+
+```yaml
+stage_args:
+  - stage_id: 0
+    stage_type: llm
+    is_prefill_only: true
+    runtime:
+      devices: "0"
+    engine_args:
+      max_num_seqs: 16
+      model_stage: thinker
+      model_arch: Qwen3OmniMoeForConditionalGeneration
+      worker_type: ar
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      gpu_memory_utilization: 0.9
+      enforce_eager: true
+      trust_remote_code: true
+      engine_output_type: latent
+      distributed_executor_backend: "mp"
+      enable_prefix_caching: false
+      max_num_batched_tokens: 32768
+      hf_config_name: thinker_config
+      tensor_parallel_size: 1
+      kv_transfer_config:
+        kv_connector: "MooncakeConnector"
+        kv_role: "kv_producer"
+        kv_rank: 0
+        kv_parallel_size: 2
+        kv_connector_extra_config:
+          mooncake_bootstrap_port: 25201
+    final_output: false
+    is_comprehension: true
+    default_sampling_params:
+      temperature: 0.4
+      top_p: 0.9
+      top_k: 1
+      max_tokens: 2048
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.05
+
+  - stage_id: 1
+    stage_type: llm
+    is_decode_only: true
+    runtime:
+      devices: "1"
+    engine_args:
+      max_num_seqs: 64
+      model_stage: thinker
+      model_arch: Qwen3OmniMoeForConditionalGeneration
+      worker_type: ar
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      gpu_memory_utilization: 0.9
+      enforce_eager: true
+      trust_remote_code: true
+      engine_output_type: latent
+      distributed_executor_backend: "mp"
+      enable_prefix_caching: false
+      max_num_batched_tokens: 32768
+      hf_config_name: thinker_config
+      tensor_parallel_size: 1
+      kv_transfer_config:
+        kv_connector: "MooncakeConnector"
+        kv_role: "kv_consumer"
+        kv_rank: 1
+        kv_parallel_size: 2
+        kv_connector_extra_config:
+          mooncake_bootstrap_port: 25202
+    engine_input_source: [0]
+    final_output: true
+    final_output_type: text
+    is_comprehension: true
+    default_sampling_params:
+      temperature: 0.4
+      top_p: 0.9
+      top_k: 1
+      max_tokens: 2048
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.05
+```
+
+Notes:
+
+- `is_prefill_only: true` marks the thinker stage that only saves KV.
+- `is_decode_only: true` marks the thinker stage that resumes from remote KV.
+- `kv_transfer_config` is required on both stages.
+- The orchestrator forces the prefill stage to run with `max_tokens=1`, so the
+  prefill side only processes the prompt and exports KV.
+
+## 2. Shift the downstream stages by one index
+
+After inserting the extra thinker stage, renumber the remaining stages:
+
+```yaml
+  - stage_id: 2
+    runtime:
+      devices: "2"
+    engine_input_source: [1]
+    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker
+
+  - stage_id: 3
+    runtime:
+      devices: "2"
+    engine_args:
+      max_num_seqs: 1
+    engine_input_source: [2]
+    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav
+```
+
+Compared with the default Qwen3-Omni config:
+
+- the talker becomes stage `2` instead of stage `1`
+- the code2wav stage becomes stage `3` instead of stage `2`
+- the talker now reads from decode stage `1`
+
+## 3. Add runtime edges for the four-stage pipeline
+
+```yaml
+runtime:
+  enabled: true
+  defaults:
+    window_size: -1
+    max_inflight: 1
+  edges:
+    - from: 0
+      to: 1
+      window_size: -1
+    - from: 1
+      to: 2
+      window_size: -1
+    - from: 2
+      to: 3
+      window_size: -1
+```
+
+## 4. Launch with your custom config
+
+```bash
+vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
+    --stage-configs-path /path/to/qwen3_omni_pd.yaml
+```
+
+## Operational Notes
+
+- `MooncakeConnector` does not support heterogeneous TP sizes across the PD
+  pair. Keep prefill and decode at the same `tensor_parallel_size`.
+- If the thinker requires TP=2, both thinker stages must use TP=2 and be given
+  separate GPU sets, for example `"0,1"` for prefill and `"2,3"` for decode.
+- Choose connector ports and addresses that match your deployment. The values
+  shown above are examples only.

docs/configuration/stage_configs.md

@@ -18,7 +18,7 @@ If users want to modify some part of it. The custom stage_configs file can be in
 For offline (Assume necessary dependencies have ben imported):
 ```python
 model_name = "Qwen/Qwen2.5-Omni-7B"
-omni_llm = Omni(model=model_name, stage_configs_path="/path/to/custom_stage_configs.yaml")
+omni = Omni(model=model_name, stage_configs_path="/path/to/custom_stage_configs.yaml")
 ```
 
 For online serving:
@@ -30,15 +30,15 @@ vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --stage-configs-path /path/to
 
 Below is a specific example of stage_configs.yaml in Qwen2.5-omni.
 ```python
-# stage config for running qwen2.5-omni with architecture of OmniLLM.
+# stage config for running qwen2.5-omni with AsyncOmniEngine + Orchestrator runtime.
 stage_args:
   - stage_id: 0 # mark the unique id for each stage
     runtime: # The disaggregated configuration
       process: true  # Run this stage in a separate process
       devices: "0" # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device)
-      max_batch_size: 1 # the batch_size for offline inference
     engine_args: # Engine arguments for a certain engine
       model_stage: thinker
+      max_num_seqs: 1
       model_arch: Qwen2_5OmniForConditionalGeneration # The model implementation registered in model_executor/models/registry.py
       worker_type: ar # The specific worker used
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler # The specific scehduler used
@@ -62,9 +62,9 @@ stage_args:
     runtime:
       process: true
       devices: "1"
-      max_batch_size: 3
     engine_args:
       model_stage: talker
+      max_num_seqs: 3
       model_arch: Qwen2_5OmniForConditionalGeneration
       worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -88,9 +88,9 @@ stage_args:
     runtime:
       process: true
       devices: "0"            # Example: use a different GPU than the previous stage; use "0" if single GPU
-      max_batch_size: 1
     engine_args:
       model_stage: code2wav
+      max_num_seqs: 1
       model_arch: Qwen2_5OmniForConditionalGeneration
       worker_type: generation
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
@@ -159,9 +159,9 @@ Visible devices for this stage, specified as a string. This controls which GPU d
 
 Default: `"0"`
 
-#### `runtime.max_batch_size`
+#### `engine_args.max_num_seqs`
 
-The maximum batch size for offline inference in this stage. This limits how many sequences can be processed together in a single batch during offline inference operations.
+The maximum number of sequences for concurrent processing in this stage. For LLM stages, this controls the vLLM scheduler's maximum concurrent sequences. For all stage types, this also controls how many tasks can be batched together in the task processing loop.
 
 Default: `1`
 

docs/configuration/stage_configs/qwen2_5_omni.yaml

@@ -1,12 +1,12 @@
-# stage config for running qwen2.5-omni with architecture of OmniLLM.
+# stage config for running qwen2.5-omni with AsyncOmniEngine + Orchestrator runtime.
 stage_args:
   - stage_id: 0
     runtime:
       process: true            # Run this stage in a separate process
       devices: "0"            # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device)
-      max_batch_size: 1
     engine_args:
       model_stage: thinker
+      max_num_seqs: 1
       model_arch: Qwen2_5OmniForConditionalGeneration
       worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -30,9 +30,9 @@ stage_args:
     runtime:
       process: true
       devices: "1"
-      max_batch_size: 1
     engine_args:
       model_stage: talker
+      max_num_seqs: 1
       model_arch: Qwen2_5OmniForConditionalGeneration
       worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -56,9 +56,9 @@ stage_args:
     runtime:
       process: true
       devices: "0"            # Example: use a different GPU than the previous stage; use "0" if single GPU
-      max_batch_size: 1
     engine_args:
       model_stage: code2wav
+      max_num_seqs: 1
       model_arch: Qwen2_5OmniForConditionalGeneration
       worker_type: generation
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler

docs/contributing/README.md

@@ -107,7 +107,7 @@ Only specific types of PRs will be reviewed. The PR title is prefixed appropriat
 - `[CI/Build]` for build or continuous integration improvements.
 - `[Doc]` for documentation fixes and improvements.
 - `[Model]` for adding a new model or improving an existing model. Model name should appear in the title.
-- `[Frontend]` For changes on the vLLM-Omni frontend (e.g., OpenAI API server, `OmniLLM` class, etc.)
+- `[Frontend]` For changes on the vLLM-Omni frontend (e.g., OpenAI API server, `Omni`/`AsyncOmni`, etc.)
 - `[Kernel]` for changes affecting CUDA kernels or other compute kernels.
 - `[Core]` for changes in the core vLLM-Omni logic (e.g., `OmniProcessor`, `OmniARScheduler`, etc.)
 - `[Hardware][Vendor]` for hardware-specific changes. Vendor name should appear in the prefix, such as [Ascend] for Ascend NPUs.

docs/contributing/ci/CI_5levels.md

@@ -168,11 +168,6 @@ vllm_omni/                                    tests/
 │   └── arg_utils.py                            │   └── test_arg_utils.py               ⬜
 │
 ├── entrypoints/                        →     ├── entrypoints/
-│   ├── omni.py                                 │   ├── test_omni.py                    ⬜  (E2E covered by e2e/offline, e2e/online)
-│   ├── omni_llm.py                             │   ├── test_omni_llm.py                ✅
-│   ├── omni_stage.py                            │   ├── test_omni_stage.py              ⬜  (partial in test_omni_stage_diffusion_config.py)
-│   ├── omni_diffusion.py                       │   ├── test_omni_diffusion.py          ✅
-│   ├── async_omni.py                            │   ├── test_async_omni.py              ✅ actually in e2e/online_serving/test_async_omni.py
 │   ├── async_omni_diffusion.py                 │   ├── test_async_omni_diffusion_config.py  ✅
 │   ├── stage_utils.py                          │   ├── test_stage_utils.py            ✅
 │   ├── cli/                                     │   ├── cli/                           (benchmarks/test_serve_cli.py covers CLI serve)

docs/contributing/ci/tests_style.md

@@ -24,8 +24,6 @@ End-to-end tests verify the complete functionality of a system or component. For
 
 - **`tests/e2e/online_serving/`**: Tests for online serving scenarios (e.g., API server tests)
 
-**Example:** The test file for `vllm_omni/entrypoints/omni_llm.py` should be located at `tests/entrypoints/test_omni_llm.py`.
-
 ## Test Directory Structure
 
 The ideal directory structure mirrors the source code organization. Legend: `✅` = test exists, `⬜` = suggested to add.
@@ -75,11 +73,6 @@ vllm_omni/                                    tests/
 │   └── arg_utils.py                            │   └── test_arg_utils.py               ⬜
 │
 ├── entrypoints/                        →     ├── entrypoints/
-│   ├── omni.py                                 │   ├── test_omni.py                    ⬜  (E2E covered by e2e/offline, e2e/online)
-│   ├── omni_llm.py                             │   ├── test_omni_llm.py                ✅
-│   ├── omni_stage.py                            │   ├── test_omni_stage.py              ⬜  (partial in test_omni_stage_diffusion_config.py)
-│   ├── omni_diffusion.py                       │   ├── test_omni_diffusion.py          ✅
-│   ├── async_omni.py                            │   ├── test_async_omni.py              ✅ actually in e2e/online_serving/test_async_omni.py
 │   ├── async_omni_diffusion.py                 │   ├── test_async_omni_diffusion_config.py  ✅
 │   ├── stage_utils.py                          │   ├── test_stage_utils.py            ✅
 │   ├── cli/                                     │   ├── cli/                           (benchmarks/test_serve_cli.py covers CLI serve)
@@ -170,7 +163,7 @@ vllm_omni/                                    tests/
 
 ### Naming Conventions
 
-- **Unit Tests**: Use `test_<module_name>.py` format. Example: `omni_llm.py` → `test_omni_llm.py`
+- **Unit Tests**: Use `test_<module_name>.py` format. Example: `stage_utils.py` → `test_stage_utils.py`
 
 - **E2E Tests**: Place in `tests/e2e/offline_inference/` or `tests/e2e/online_serving/` with descriptive names. Example: `tests/e2e/offline_inference/test_qwen3_omni.py`, `tests/e2e/offline_inference/test_diffusion_model.py`
 

docs/contributing/model/adding_diffusion_model.md

@@ -739,6 +739,22 @@ See detailed guide: [How to add Sequence Parallel support](../../design/feature/
 omni = Omni(model="your-model", ulysses_degree=2, ring_degree=2)
 ```
 
+### Step Execution
+
+See detailed design guide: [How to add step execution support](../../design/feature/diffusion_step_execution.md)
+
+Use this only when your pipeline can be split into stable request-scoped and
+step-scoped phases. The reference implementation is
+`QwenImagePipeline`, which maps its request-level `forward()` into:
+
+1. `prepare_encode()` for prompt encoding, latent init, timestep prep, and per-request scheduler setup.
+2. `denoise_step()` for one transformer/noise prediction.
+3. `step_scheduler()` for one scheduler update and `step_index` advance.
+4. `post_decode()` for the final VAE decode.
+
+Do not enable `step_execution=True` until those four methods are implemented
+and validated against the request-level path.
+
 ### Cache Acceleration
 
 #### TeaCache
@@ -815,6 +831,107 @@ class WanTransformer3DModel(nn.Module):
 
 ---
 
+### Diffusion Timing (Performance Profiling)
+When adapting a new diffusion model, it is often useful to analyze the latency of key components such as text encoding, diffusion denoising, and VAE decoding.
+vLLM-Omni provides a timing utility via `DiffusionPipelineProfilerMixin` to help developers quickly identify performance bottlenecks.
+
+!!! info
+      `DiffusionPipelineProfilerMixin` is different from using `torch.profiler` for diffusion models, as introduced in this [tutorial](https://github.com/vllm-project/vllm-omni/blob/main/docs/contributing/profiling.md#3-profiling-diffusion-models). `DiffusionPipelineProfilerMixin` only prints the timing information of multiple functions (such as `vae.decode`), while `torch.profiler` saves detailed GPU/CPU computation time, call/execution steps.
+
+This tool automatically measures the execution time of selected pipeline modules and prints the results in the logs.
+
+**Enabling Diffusion Timing**
+
+
+Enable timing by setting:
+```
+vllm serve Qwen/Qwen-Image --omni --port 8091 --enable-diffusion-pipeline-profiler
+```
+You can optionally specify which modules to profile:
+```
+class YourPipeline(xxx, DiffusionPipelineProfilerMixin):
+    def __init__(self, xxx):
+        ...
+        self.setup_diffusion_pipeline_profiler(profiler_targets=["diffuse"], enable_diffusion_pipeline_profiler)
+```
+If not specified, the default targets are used:
+```
+["vae.encode", "vae.decode", "diffuse", "text_encoder.forward", "tokenizer.forward"]
+```
+**Adding DiffusionPipelineProfilerMixin to a Pipeline**
+To enable timing support in your pipeline, inherit from DiffusionPipelineProfilerMixin.
+```python
+from vllm_omni.diffusion.utils.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
+
+class YourModelPipeline(nn.Module, DiffusionPipelineProfilerMixin):
+    # Optional: Specify custom timing targets
+    _PROFILER_TARGETS = ["vae.encode", "vae.decode", "diffuse", "text_encoder.forward", "tokenizer.forward"]
+
+    def __init__(
+        self,
+        *,
+        od_config: OmniDiffusionConfig,
+        prefix: str = "",
+    ):
+        super().__init__()
+        self.od_config = od_config
+        self.parallel_config = od_config.parallel_config
+        # initialize pipeline components
+        ...
+
+        # initialize timing profiler
+        self.setup_diffusion_pipeline_profiler(enable_diffusion_pipeline_profiler)
+```
+The mixin dynamically wraps selected methods and records their execution time during inference.
+
+If you need to fetch the execution time of different modules, you will need to pass `self.stage_durations` to `DiffusionOutput`, as shown below:
+
+```diff
+-   return DiffusionOutput(output=img)
++   return DiffusionOutput(
+            output=image, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None
+        )
+```
+
+**Pipeline Design for Timing**
+The current diffusion timing utility is function-based, meaning it measures the execution time of individual methods.
+
+When implementing a new pipeline, avoid putting all logic inside a single function (e.g., forward). Instead, structure the pipeline in a modular way by separating key stages into independent methods, such as the diffusion loop.
+
+For example:
+```
+def forward(self, req: OmniDiffusionRequest):
+    prompt_embeds = self.encode_prompt(req)
+    latents = self.diffuse(prompt_embeds, req)
+    images = self.vae.decode(latents)
+    return DiffusionOutput(output=images)
+```
+This allows the timing utility to measure each stage (e.g., encode_prompt, diffuse, vae.decode) separately and helps identify performance bottlenecks more easily.
+
+
+**Default Profiled Modules**
+
+By default, the following pipeline modules are timed:
+```
+vae.encode
+vae.decode
+diffuse
+text_encoder.forward
+tokenizer.forward
+```
+
+**Example Output**
+
+When enabled, timing logs appear like this:
+```
+[DiffusionTiming] text_encoder.forward took 0.018s
+[DiffusionTiming] diffuse took 2.412s
+[DiffusionTiming] vae.decode took 0.063s
+```
+These measurements help identify bottlenecks during model adaptation and optimization
+
+
+
 ## Troubleshooting
 
 
@@ -866,7 +983,7 @@ hidden_states = hidden_states.reshape(batch_size, seq_len, -1)
 
 1. **Reduce batch size:**
    ```python
-   omni.generate(prompts=[...], max_batch_size=2)
+   omni.generate(prompts=[...], max_num_seqs=2)
    ```
 
 2. **Use smaller image size:**

docs/contributing/model/adding_omni_model.md

@@ -330,54 +330,46 @@ Stage transitions are the mechanism by which outputs from one stage are converte
 
 ### Where Stage Transitions Are Called
 
-Stage transitions happen automatically in the orchestrator (`OmniLLM` class) during the generation loop. Here's the detailed flow:
+Stage transitions happen automatically in the runtime orchestrator. Here's the detailed flow:
 
-1. **Location**: `vllm_omni/entrypoints/omni_llm.py` in the `_run_generation()` method
+1. **Location**: `vllm_omni/engine/orchestrator.py` in `_forward_to_next_stage()`
 2. **Trigger**: When a stage completes processing and produces outputs
 3. **Execution Flow**:
    ```python
-   # In omni_llm.py, _run_generation() method (around line 345-460)
+   # In orchestrator.py
+   next_stage_id = stage_id + 1
+   next_client = self.stage_clients[next_stage_id]
+   params = req_state.sampling_params_list[next_stage_id]
 
-   # Main orchestrator loop polls each stage for completed requests
-   for stage_id, stage in enumerate(self.stage_list):
-       result = stage.try_collect()  # Get completed request
-       if result is None:
-           continue
+   # Save current stage outputs so stage_input_processors can consume them.
+   self.stage_clients[stage_id].set_engine_outputs([output])
 
-       # Store outputs from this stage
-       engine_outputs = _load(result, obj_key="engine_outputs", shm_key="engine_outputs_shm")
-       stage.set_engine_outputs(engine_outputs)
+   # THIS IS WHERE STAGE TRANSITION HAPPENS
+   next_inputs = next_client.process_engine_inputs(
+       stage_list=self.stage_clients,
+       prompt=req_state.prompt,
+   )
 
-       # Check if there's a next stage to forward to
-       next_stage_id = stage_id + 1
-       if next_stage_id < len(self.stage_list):
-           next_stage: OmniStage = self.stage_list[next_stage_id]
-
-           # THIS IS WHERE STAGE TRANSITION HAPPENS
-           next_inputs = next_stage.process_engine_inputs(
-               self.stage_list,
-               [request_id_to_prompt[req_id]]
-           )
-
-           # Submit to next stage
-           task = {
-               "type": OmniStageTaskType.GENERATE,
-               "request_id": req_id,
-               "engine_inputs": next_inputs[0],
-               "sampling_params": sampling_params_list[next_stage_id],
-           }
-           next_stage.submit(task)
+   # Build and submit request(s) to the next stage.
+   for next_input in next_inputs:
+       request = build_engine_core_request_from_tokens(
+           request_id=req_id,
+           prompt=next_input,
+           params=params,
+           model_config=self.stage_vllm_configs[next_stage_id].model_config,
+       )
+       await next_client.add_request_async(request)
    ```
 
 ### How Stage Transitions Work
 
 The stage transition process follows these steps:
 
-1. **Stage Completion**: When a stage finishes processing a request, it stores outputs via `stage.set_engine_outputs(engine_outputs)`
+1. **Stage Completion**: When a stage finishes processing a request, the orchestrator stores outputs via `stage_client.set_engine_outputs(...)`
 
 2. **Transition Detection**: The orchestrator checks if there's a next stage and calls `process_engine_inputs()` on it
 
-3. **Input Processing**: The `process_engine_inputs()` method in `OmniStage` (`omni_stage.py`) handles the transition:
+3. **Input Processing**: The stage input processor configured in stage YAML (under `vllm_omni/model_executor/stage_input_processors/`) handles the transition:
    ```python
    def process_engine_inputs(
        self, stage_list: list[Any], prompt: OmniTokensPrompt | TextPrompt = None

docs/contributing/model/adding_tts_model.md

@@ -361,9 +361,9 @@ stage_args:
     stage_type: llm
     runtime:
       devices: "0"
-      max_batch_size: 64
     engine_args:
       model_stage: ar_stage
+      max_num_seqs: 64
       model_arch: YourTTSModelForConditionalGeneration
       worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -403,9 +403,9 @@ stage_args:
     stage_type: llm
     runtime:
       devices: "0"
-      max_batch_size: 64
     engine_args:
       model_stage: ar_stage
+      max_num_seqs: 64
       model_arch: YourTTSModelForConditionalGeneration
       worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler

docs/contributing/profiling.md

@@ -23,30 +23,32 @@ export VLLM_PROFILER_MAX_ITERS=1
 The profiler is default to function across all stages. But It is highly recommended to profile specific stages by passing the stages list, preventing from producing too large trace files:
 ```python
 # Profile all stages
-omni_llm.start_profile()
+omni.start_profile()
 
 # Only profile Stage 1
-omni_llm.start_profile(stages=[1])
+omni.start_profile(stages=[1])
 ```
 
 ```python
 # Stage 0 (Thinker) and Stage 2 (Audio Decoder) for qwen omni
-omni_llm.start_profile(stages=[0, 2])
+omni.start_profile(stages=[0, 2])
 ```
 
 **Python Usage**: Wrap your generation logic with `start_profile()` and `stop_profile()`.
 
 ```python
-from vllm_omni import omni_llm
+from vllm_omni.entrypoints.omni import Omni
+
+omni = Omni(model="Qwen/Qwen3-Omni-30B-A3B-Instruct")
 
 profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
 
 # 1. Start profiling if enabled
 if profiler_enabled:
-    omni_llm.start_profile(stages=[0])
+    omni.start_profile(stages=[0])
 
 # Initialize generator
-omni_generator = omni_llm.generate(prompts, sampling_params_list, py_generator=args.py_generator)
+omni_generator = omni.generate(prompts, sampling_params_list, py_generator=args.py_generator)
 
 total_requests = len(prompts)
 processed_count = 0
@@ -57,21 +59,21 @@ for stage_outputs in omni_generator:
     # ... [Output processing logic for text/audio would go here] ...
 
     # Update count to track when to stop profiling
-    processed_count += len(stage_outputs.request_output)
+    processed_count += 1
 
     # 2. Check if all requests are done to stop the profiler safely
     if profiler_enabled and processed_count >= total_requests:
         print(f"[Info] Processed {processed_count}/{total_requests}. Stopping profiler inside active loop...")
 
         # Stop the profiler while workers are still active
-        omni_llm.stop_profile()
+        omni.stop_profile()
 
         # Wait for traces to flush to disk
         print("[Info] Waiting 30s for workers to write trace files to disk...")
         time.sleep(30)
         print("[Info] Trace export wait time finished.")
 
-omni_llm.close()
+omni.close()
 ```
 
 

docs/design/architecture_overview.md

@@ -67,12 +67,12 @@ According to analysis for current popular open-source models, most of them have
 | Component         | Description                                                                                                                              |
 | ----------------- | ---------------------------------------------------------------------------------------------------------------------------------------- |
 | **OmniRouter**    | provide an intelligent router for Omni-modality requests dispatch                                                                        |
-| **EntryPoints**   | define the APIs for offline/online serving (APIServer, Omni/AsyncOmni) and provide the OmniStage abstraction for different AR/DiT stages |
+| **EntryPoints**   | define the APIs for offline/online serving (APIServer, Omni/AsyncOmni), while `AsyncOmniEngine` and `Orchestrator` coordinate multi-stage AR/DiT execution |
 | **AR**            | adapted for omni-modality models while inheriting efficient features from vLLM, such as cache management                                 |
 | **Diffusion**     | natively implemented and optimized using acceleration components                                                                         |
 | **OmniConnector** | supports fully disaggregation based on E/P/D/G (Encoding/Processing/Decoding/Generation) disaggregation across stages                    |
 
-Disaggregated stages are managed through configuration, such as in the Qwen3-Omni example, where stages like Thinker, Talker, and Code2wav are defined as separate OmniStage instances with specific resources and input/output type.
+Disaggregated stages are managed through stage configuration. In Qwen3-Omni, Thinker/Talker/Code2wav are declared as separate configured stages, and runtime routing is handled by `Orchestrator` over `StageEngineCoreClient` / `StageDiffusionClient`.
 
 ## Main features
 
@@ -127,10 +127,10 @@ Taking **Qwen3-Omni** as an example:
 The **Omni** class provides a Python interface for offline batched inference. Users initialize the Omni class with a Hugging Face model name and use the generate method, passing inputs that include both text prompts and multi-modal data:
 
 ```
-# Create an omni_lm with HF model name.
+# Create an omni runtime with HF model name.
 from vllm_omni.entrypoints.omni import Omni
 
-omni_lm = Omni(model="Qwen/Qwen3-Omni-30B-A3B-Instruct")
+omni = Omni(model="Qwen/Qwen3-Omni-30B-A3B-Instruct")
 
 # Example prompts.
 om_inputs = {"prompt": prompt,
@@ -140,7 +140,7 @@ om_inputs = {"prompt": prompt,
              }}
 
 # Generate texts and audio from the multi-modality inputs.
-outputs = omni_lm.generate(om_inputs, sampling_params_list)
+outputs = omni.generate(om_inputs, sampling_params_list)
 ```
 
 ## Online Serving

docs/design/feature/async_chunk_design.md

@@ -171,7 +171,7 @@ stage_args:
 - `custom_process_next_stage_input_func: str`: Path to custom chunk processing function; receives `(transfer_manager, pooling_output, request)`. For qwen3-omni: `thinker2talker_async_chunk`, `talker2code2wav_async_chunk`
 - `stage_connector_config: dict`: Connector configuration
 - `worker_type: str`: Model type, e.g. `"ar"` or `"generation"` (used by OmniChunkTransferAdapter for mode-specific payload handling)
-- `max_batch_size: int`: Maximum batch size for the stage
+- `max_num_seqs: int`: Maximum number of sequences for concurrent processing in the stage
 
 
 ### Connector Configuration
@@ -195,9 +195,9 @@ stage_args:
   - stage_id: 2  # code2wav stage
     runtime:
       devices: "1"
-      max_batch_size: 64  # Enables batched audio generation
     engine_args:
       model_stage: code2wav
+      max_num_seqs: 64  # Enables batched audio generation
 ```
 
 ## Related Files

docs/design/feature/diffusion_step_execution.md

@@ -0,0 +1,121 @@
+# Diffusion Step Execution
+
+This guide documents vLLM-Omni's stepwise diffusion contract for model authors
+and contributors implementing `step_execution=True` support for a diffusion
+pipeline.
+
+For end-user enablement, supported models, and current limitations, see
+[Step Execution](../../user_guide/diffusion/step_execution.md).
+
+## Current Support Scope
+
+`step_execution` is **not** a generic diffusion toggle. It only works for
+pipelines that implement the segmented stateful contract in
+[`vllm_omni/diffusion/models/interface.py`](gh-file:vllm_omni/diffusion/models/interface.py).
+
+Current in-tree support:
+
+| Pipeline | Example models | Step execution |
+|----------|----------------|----------------|
+| `QwenImagePipeline` | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Yes |
+| All other diffusion pipelines | `QwenImageEditPipeline`, `QwenImageEditPlusPipeline`, `QwenImageLayeredPipeline`, GLM-Image, Wan, Flux, etc. | No |
+
+Current engine/runtime limitations:
+
+- `StepScheduler` only schedules `batch_size=1`.
+- `cache_backend` is not supported in step mode.
+- Request-mode extras such as KV transfer are not wired into step mode yet.
+- Unsupported pipelines now fail early during model loading instead of failing on the first request.
+
+## Execution Contract
+
+Step mode is driven by four pipeline methods plus the shared mutable request
+state object:
+
+- `prepare_encode(state)`: one-time request preparation.
+- `denoise_step(state)`: compute the noise prediction for the current step.
+- `step_scheduler(state, noise_pred)`: mutate latents and advance step state.
+- `post_decode(state)`: decode the final output after denoising is complete.
+
+The state lives in
+[`vllm_omni/diffusion/worker/utils.py`](gh-file:vllm_omni/diffusion/worker/utils.py)
+as `DiffusionRequestState`. Store request-scoped tensors there, or use
+`state.extra` for model-specific fields that do not justify extending the core
+dataclass.
+
+The worker-side step loop lives in
+[`vllm_omni/diffusion/worker/diffusion_model_runner.py`](gh-file:vllm_omni/diffusion/worker/diffusion_model_runner.py):
+
+1. `prepare_encode()` runs once for a new request.
+2. `denoise_step()` runs every scheduler tick.
+3. `step_scheduler()` mutates `state.latents` and advances `state.step_index`.
+4. `post_decode()` runs exactly once after `state.denoise_completed` becomes true.
+
+## Recommended Split
+
+When converting an existing request-level `forward()` pipeline, keep the split
+strict and mechanical:
+
+| Request-level phase | Stepwise method | What belongs there |
+|---------------------|-----------------|--------------------|
+| Input validation, prompt encoding, latent init, timestep prep, per-request scheduler creation | `prepare_encode()` | Anything that should happen once per request |
+| Transformer forward / noise prediction | `denoise_step()` | Pure denoise computation for the current timestep |
+| `scheduler.step(...)` and `step_index += 1` | `step_scheduler()` | Only latent/state mutation for one step |
+| VAE decode / postprocess | `post_decode()` | Final decode only |
+
+Keep the stepwise path reusing the same helpers as the request-level path
+whenever possible. Reimplementing the denoise loop from scratch is the easiest
+way to introduce behavioral drift.
+
+## Qwen-Image Reference
+
+[`pipeline_qwen_image.py`](gh-file:vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py)
+is the reference implementation and is split correctly for the current
+contract:
+
+- `prepare_encode()` reuses `_prepare_generation_context()` so prompt encoding,
+  latent init, timestep creation, CFG setup, and shape bookkeeping stay aligned
+  with `forward()`.
+- `prepare_encode()` deep-copies `self.scheduler` **after**
+  `prepare_timesteps()` so request-specific scheduler state is isolated.
+- `denoise_step()` reuses `_build_denoise_kwargs()` plus
+  `predict_noise_maybe_with_cfg()`, so sequential CFG, CFG-parallel, and
+  non-CFG behavior stay identical to the request-level path.
+- `step_scheduler()` only calls
+  `scheduler_step_maybe_with_cfg(..., per_request_scheduler=state.scheduler)`
+  and increments `state.step_index`.
+- `post_decode()` reuses `_decode_latents()`, so the final image decode matches
+  the normal `forward()` path.
+
+That decomposition is the target pattern for future models.
+
+## Rules For New Pipelines
+
+- Do not keep request-scoped scheduler state on `self.scheduler`. Copy it into
+  `state.scheduler` during `prepare_encode()`.
+- Do not mutate `state.step_index` inside `denoise_step()`. Only
+  `step_scheduler()` should advance the step.
+- Do not decode partial outputs in `denoise_step()` or `step_scheduler()`.
+- If the request-level pipeline has condition latents, masks, or edit-specific
+  tensors, store them in `state` or `state.extra`, not in global pipeline
+  attributes.
+- Preserve CFG behavior by sharing the same helper path used by `forward()`.
+- Keep `post_decode()` equivalent to the tail of `forward()`.
+
+## Validation Checklist
+
+Before marking a pipeline as `supports_step_execution = True`, verify:
+
+- Stepwise output matches request-level output for the same seed and sampling params.
+- Per-request scheduler state is isolated across concurrent requests.
+- Abort during denoise does not leak cached state.
+- `step_index` reported by `RunnerOutput` matches the scheduler progress.
+- CFG-parallel and non-CFG paths both work if the request-level pipeline supports them.
+
+## Related Files
+
+- Contract: [`vllm_omni/diffusion/models/interface.py`](gh-file:vllm_omni/diffusion/models/interface.py)
+- State: [`vllm_omni/diffusion/worker/utils.py`](gh-file:vllm_omni/diffusion/worker/utils.py)
+- Runner loop: [`vllm_omni/diffusion/worker/diffusion_model_runner.py`](gh-file:vllm_omni/diffusion/worker/diffusion_model_runner.py)
+- Scheduler transport: [`vllm_omni/diffusion/sched/interface.py`](gh-file:vllm_omni/diffusion/sched/interface.py)
+- Reference pipeline: [`vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py`](gh-file:vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py)

docs/design/module/ar_module.md

@@ -85,9 +85,6 @@ classDiagram
     class InputProcessor {
         +process_inputs() EngineCoreRequest
     }
-    class OmniInputProcessor {
-        +process_inputs() OmniEngineCoreRequest
-    }
 
     class VLLMOutputProcessor {
         +process_outputs() OutputProcessorOutput
@@ -96,7 +93,6 @@ classDiagram
         +process_outputs() OutputProcessorOutput
         +_route_and_normalize()
     }
-    InputProcessor <|-- OmniInputProcessor
     VLLMOutputProcessor <|-- MultimodalOutputProcessor
 ```
 
@@ -105,7 +101,7 @@ classDiagram
 - **Scheduler**: `OmniARScheduler` extends `vllm.v1.core.sched.scheduler.Scheduler` to enrich scheduled requests with omni-specific payloads
 - **Worker**: `GPUARWorker` extends `vllm.v1.worker.gpu_worker.Worker` to initialize AR-specific model runners
 - **ModelRunner**: `GPUARModelRunner` extends `OmniGPUModelRunner` → `vllm.v1.worker.gpu_model_runner.GPUModelRunner` to expose hidden states and handle multimodal outputs
-- **InputProcessor**: `OmniInputProcessor` extends `vllm.v1.engine.input_processor.InputProcessor` to serialize prompt embeddings and additional information
+- **InputProcessor**: Stage-0 uses upstream `vllm.v1.engine.input_processor.InputProcessor`; `AsyncOmniEngine` then restores omni-specific payloads (for example `additional_information` and `prompt_embeds`) when building `OmniEngineCoreRequest`
 - **OutputProcessor**: `MultimodalOutputProcessor` extends `vllm.v1.engine.output_processor.OutputProcessor` to route and accumulate multimodal outputs
 
 ## 3. Scheduler Design
@@ -118,7 +114,7 @@ The following diagram illustrates the request flow through the AR module compone
 
 ```mermaid
 flowchart TD
-    A[OmniInputProcessor] -->|OmniEngineCoreRequest| B[OmniARScheduler]
+    A[InputProcessor stage-0 in AsyncOmniEngine] -->|EngineCoreRequest then upgraded to OmniEngineCoreRequest| B[OmniARScheduler]
     B -->|schedule: OmniNewRequestData| C[GPUARWorker]
     C -->|SchedulerOutput| D[GPUARModelRunner]
     D -->|execute_model: None| E[Model Forward Pass]
@@ -297,15 +293,17 @@ The input/output processing pipeline handles serialization, routing, and accumul
 ```mermaid
 sequenceDiagram
     participant Client
-    participant OmniInputProcessor
+    participant AsyncOmniEngine
+    participant InputProcessor
     participant Scheduler
     participant ModelRunner
     participant MultimodalOutputProcessor
     participant Client
 
-    Client->>OmniInputProcessor: prompt + prompt_embeds + additional_info
-    OmniInputProcessor->>OmniInputProcessor: Serialize tensors to payloads
-    OmniInputProcessor->>Scheduler: OmniEngineCoreRequest (with payloads)
+    Client->>AsyncOmniEngine: prompt + prompt_embeds + additional_info
+    AsyncOmniEngine->>InputProcessor: process_inputs()
+    InputProcessor->>Scheduler: EngineCoreRequest
+    AsyncOmniEngine->>AsyncOmniEngine: _upgrade_to_omni_request() + serialize_additional_information()
     Scheduler->>ModelRunner: OmniNewRequestData (with payloads)
     ModelRunner->>ModelRunner: Decode payloads → CPU tensors
     ModelRunner->>ModelRunner: Overlay prompt_embeds on inputs_embeds
@@ -318,43 +316,18 @@ sequenceDiagram
     MultimodalOutputProcessor->>Client: RequestOutput (with multimodal_output)
 ```
 
-### OmniInputProcessor
+### Stage-0 Input Processing
 
-`OmniInputProcessor` extends the base input processor to serialize prompt embeddings and additional information for inter-stage transfer.
-
-#### Payload Serialization
-
-Converts PyTorch tensors to serialized payloads:
+Stage-0 now uses upstream `InputProcessor` directly, and `AsyncOmniEngine` upgrades the request to `OmniEngineCoreRequest` while restoring omni-specific payloads.
 
 ```python
-def process_inputs(self, ...) -> OmniEngineCoreRequest:
-    # Serialize prompt_embeds
-    if "prompt_embeds" in decoder_inputs:
-        pe_cpu = decoder_inputs["prompt_embeds"].detach().to("cpu").contiguous()
-        prompt_embeds_payload = PromptEmbedsPayload(
-            data=pe_cpu.numpy().tobytes(),
-            shape=[seq_len, hidden_size],
-            dtype=dtype_str,
-        )
-
-    # Serialize additional_information
-    if "additional_information" in decoder_inputs:
-        entries = {}
-        for key, value in raw_info.items():
-            if isinstance(value, torch.Tensor):
-                entry = AdditionalInformationEntry(
-                    tensor_data=value.numpy().tobytes(),
-                    tensor_shape=list(value.shape),
-                    tensor_dtype=dtype_str,
-                )
-            entries[key] = entry
-        additional_information_payload = AdditionalInformationPayload(entries=entries)
-
-    return OmniEngineCoreRequest(
-        # ... standard fields ...
-        prompt_embeds=prompt_embeds_payload,
-        additional_information=additional_information_payload,
-    )
+request = self.input_processor.process_inputs(
+    request_id=request_id,
+    prompt=prompt,
+    params=params,
+    supported_tasks=self.supported_tasks,
+)
+request = _upgrade_to_omni_request(request, prompt)
 ```
 
 ### MultimodalOutputProcessor
@@ -400,13 +373,13 @@ The AR module of vLLM-Omni extends vLLM through strategic inheritance and minima
 ### Key Design Patterns
 
 1. **Inheritance over composition**: Extends vLLM classes to preserve compatibility with existing scheduling, batching, and execution mechanisms
-2. **Payload serialization**: Uses `PromptEmbedsPayload` and `AdditionalInformationPayload` for efficient inter-stage data transfer
+2. **Payload serialization**: Uses serialized `additional_information` payloads together with prompt-embedding handoff for efficient inter-stage transfer
 3. **Two-phase execution**: Maintains vLLM's execute/sample separation for AR models while supporting single-phase execution for generation models
 4. **Multimodal routing**: Routes outputs by `output_type` and accumulates tensors incrementally to support streaming
 
 ### Differences from vLLM
 
-- **Payload support**: Serialized prompt embeddings and additional information enable direct transfer between pipeline stages
+- **Payload support**: Serialized additional information and prompt embeddings enable direct transfer between pipeline stages
 - **Multimodal handling**: Extended input/output processors support images, audio, and other modalities alongside text
 - **Hidden state exposure**: AR model runners expose per-request hidden states via `pooler_output` for downstream consumption
 - **Generation scheduler**: Fast-path scheduling for basic heterogeneous architectures that complete in one step

docs/design/module/async_omni_architecture.md

@@ -0,0 +1,203 @@
+# AsyncOmni Architecture (Qwen3-Omni Example)
+
+## 1. System Architecture
+
+```text
+• ┌─────────────────────────────────────────────────────────────────────────────────┐
+  │                                    API Layer                                    │
+  │  ┌─────────────────────────────────────┐  ┌──────────────────────────────────┐  │
+  │  │ AsyncOmni (EngineClient)            │  │ Omni                             │  │
+  │  │ • generate() / abort() / shutdown() │  │ • generate()                     │  │
+  │  │ • _final_output_handler()           │  │                                  |  │
+  │  └─────────────────────────────────────┘  └──────────────────────────────────┘  │
+  ├─────────────────────────────────────────────────────────────────────────────────┤
+  │                              Engine Layer (Proxy)                               │
+  │  ┌───────────────────────────────────────────────────────────────────────────┐  │
+  │  │ AsyncOmniEngine                                                           │  │
+  │  │ • _bootstrap_orchestrator() & _initialize_stages()                        │  │
+  │  │ • add_request() / add_request_async() -> input_processor.process_inputs() │  │
+  │  │ • try_get_output() / try_get_output_async()                               │  │
+  │  └───────────────────┬─────────────────────────────────▲─────────────────────┘  │
+  │         request_queue (janus.Queue)        output_queue (janus.Queue)           │
+  ├──────────────────────┼─────────────────────────────────┼────────────────────────┤
+  │                      ▼        Orchestration Layer      │                        │
+  │  ┌───────────────────────────────────────────────────────────────────────────┐  │
+  │  │ Orchestrator [background thread]                                          │  │
+  │  │ • _request_handler()                                                      │  │
+  │  │     -  stage_client.add_request_async() & _prewarm_async_chunk_stages()   │  │
+  │  │ • _orchestration_output_handler()                                         │  │
+  │  │     -  _process_stage_outputs() -> output_processors[i].process_outputs() │  │
+  │  │     -  _route_output() & _forward_to_next_stage()                         │  │
+  │  └──────────┬─────────────────────────┬────────────────────────┬─────────────┘  │
+  ├─────────────┼─────────────────────────┼────────────────────────┼────────────────┤
+  │             │                 Communication Layer              │                │
+  │  ┌───────────────────────┐ ┌───────────────────────┐ ┌───────────────────────┐  │
+  │  │ StageEngineCoreClient │ │ StageEngineCoreClient │ │ StageDiffusionClient  │  │
+  │  │ • ZMQ ROUTER / PULL   │ │ • ZMQ ROUTER / PULL   │ │ • ZMQ ROUTER / PULL   │  │
+  │  │ • Msgpack codec       │ │ • Msgpack codec       │ │ • Msgpack codec       │  │
+  │  └──────────┬────────────┘ └──────────┬────────────┘ └──────────┬────────────┘  │
+  │             ▼ ZMQ IPC                 ▼ ZMQ IPC                 ▼ ZMQ IPC       │
+  ├─────────────────────────────────────────────────────────────────────────────────┤
+  │                                 Execution Layer                                 │
+  │  ┌───────────────────────┐ ┌───────────────────────┐ ┌───────────────────────┐  │
+  │  │ StageCoreProc         │ │ StageCoreProc         │ │ DiffusionEngine       │  │
+  │  │ [background process]  │ │ [background process]  │ │ [background process]  │  │
+  │  └───────────────────────┘ └───────────────────────┘ └───────────────────────┘  │
+  └─────────────────────────────────────────────────────────────────────────────────┘
+```
+
+## 2. Execution Flow (Arrow Steps, one generate request)
+
+```text
+[1] App
+    -> AsyncOmni.generate(prompt, request_id)
+
+[2] AsyncOmni
+    -> _final_output_handler()   (started on first request)
+    -> AsyncOmniEngine.add_request(stage_id=0, ...)
+
+[3] AsyncOmniEngine.add_request
+    -> (if stage-0 is llm and input is not EngineCoreRequest)
+       InputProcessor.process_inputs()
+       OutputProcessor[0].add_request()
+    -> request_queue.put(add_request_msg)
+
+[4] Orchestrator._request_handler
+    -> _handle_add_request(msg)
+    -> stage_clients[0].add_request_async(...)
+
+[5] Orchestrator._orchestration_loop (loop)
+    -> poll stage output
+       - llm stage: await get_output_async()
+       - diffusion stage: get_diffusion_output_async()
+    -> (llm stage) output_processors[i].process_outputs(...)
+    -> _route_output(...)
+    -> if finished and not final_stage and non-async-chunk:
+         _forward_to_next_stage(...)
+         -> next_stage.add_request_async(...)
+    -> output_queue.put(output)
+
+[6] AsyncOmni._final_output_loop (background coroutine)
+    -> AsyncOmniEngine.try_get_output_async()
+    -> route by request_id to ClientRequestState.queue
+
+[7] AsyncOmni._process_orchestrator_results
+    -> read from ClientRequestState.queue
+    -> _process_single_result(...)
+    -> yield OmniRequestOutput
+
+[8] Exit condition
+    -> receive result["finished"] == True
+    -> generate() ends
+```
+
+## 3. Runtime Sequence (one generate request)
+
+```mermaid
+sequenceDiagram
+    participant APP as App
+    participant AO as AsyncOmni
+    participant ENG as AsyncOmniEngine
+    participant ORCH as Orchestrator
+    participant S0 as Stage-0 Client
+    participant SN as Next Stage Client
+
+    APP->>AO: generate
+    AO->>AO: start output_handler once
+    AO->>ENG: add_request(stage_id=0, ...)
+    ENG->>ENG: input_processor.process_inputs()
+    ENG->>ORCH: request_queue.put(add_request)
+
+    ORCH->>ORCH: _handle_add_request
+    ORCH->>S0: add_request_async
+
+    loop poll route forward
+        ORCH->>S0: get_output_async / get_diffusion_output_async
+        ORCH->>ORCH: _route_output
+        alt need forward to next stage
+            ORCH->>SN: add_request_async
+        end
+        ORCH-->>ENG: output_queue.put
+    end
+
+    AO->>ENG: try_get_output_async
+    ENG-->>AO: message
+    AO-->>APP: yield OmniRequestOutput
+```
+
+## 4. Comparison
+
+Previous topology (reference):
+
+```text
+┌────────────────────────────────────────────────────────────────────────────┐
+│ Main Process                                                               │
+│  ┌──────────────────────┐   ┌────────────────────────────────────────────┐ │
+│  │ generate()           │   │ final_output_handler()                     │ │
+│  └──────────────────────┘   └────────────────────────────────────────────┘ │
+└──────────┬─────────────────────────┬─────────────────────────┬─────────────┘
+  mp.Queue (in_q/out_q)    mp.Queue (in_q/out_q)    mp.Queue (in_q/out_q)
+           ▼▲                        ▼▲                        ▼▲
+┌───────────────────────┐  ┌───────────────────────┐  ┌──────────────────────┐
+│ Worker Proc-0         │  │ Worker Proc-1         │  │ Worker Proc-2        │
+│ (Thinker LLM)         │  │ (Talker LLM)          │  │ (Vocoder)            │
+│  ┌────────────────┐   │  │  ┌────────────────┐   │  │  ┌────────────────┐  │
+│  │_stage_worker   │   │  │  │_stage_worker   │   │  │  │_stage_worker   │  │
+│  │_async()        │   │  │  │_async()        │   │  │  │_async()        │  │
+│  └────────────────┘   │  │  └────────────────┘   │  │  └────────────────┘  │
+│  ┌────────────────┐   │  │  ┌────────────────┐   │  │  ┌────────────────┐  │
+│  │output_handler()│   │  │  │output_handler()│   │  │  │output_handler()│  │
+│  └────────────────┘   │  │  └────────────────┘   │  │  └────────────────┘  │
+└──────────┬────────────┘  └──────────┬────────────┘  └──────────┬───────────┘
+       ZMQ ▼ ▲ ZMQ               ZMQ ▼ ▲ ZMQ               ZMQ ▼ ▲ ZMQ
+┌──────────────────────┐   ┌──────────────────────┐   ┌──────────────────────┐
+│ EngineCore Proc-0    │   │ EngineCore Proc-1    │   │ EngineCore Proc-2    │
+│ (Thinker)            │   │ (Talker)             │   │ (Vocoder)            │
+└──────────────────────┘   └──────────────────────┘   └──────────────────────┘
+```
+
+Current topology:
+
+```text
+┌────────────────────────────────────────────────────────────────────────────┐
+│ Main Process                                                               │
+│  ┌──────────────────────────────────────────────────────────────────────┐  │
+│  │ Main Thread                                                          │  │
+│  │  ┌──────────────────────┐   ┌─────────────────────────────────────┐  │  │
+│  │  │ generate()           │   │ final_output_handler()              │  │  │
+│  │  └──────────────────────┘   └─────────────────────────────────────┘  │  │
+│  └──────────────────────────────────────────────────────────────────────┘  │
+│         janus.Queue (request_queue) ▼  ▲ janus.Queue (output_queue)        │
+│  ┌──────────────────────────────────────────────────────────────────────┐  │
+│  │ Orchestrator Thread                                                  │  │
+│  │  ┌──────────────────────┐  ┌──────────────────────────────────────┐  │  │
+│  │  │ _request_handler()   │  │ _orchestration_output_handler()      │  │  │
+│  │  └──────────────────────┘  └──────────────────────────────────────┘  │  │
+│  │  ┌────────────────────────────────────────────────────────────────┐  │  │
+│  │  │ _orchestration_loop(): poll/process/route outputs for all stages│  │  │
+│  │  └────────────────────────────────────────────────────────────────┘  │  │
+│  └───────┬─────────────────────────┬─────────────────────────┬──────────┘  │
+└──────────┬─────────────────────────┬─────────────────────────┬─────────────┘
+       ZMQ ▼ ▲ ZMQ               ZMQ ▼ ▲ ZMQ               ZMQ ▼ ▲ ZMQ  
+  ┌──────────────────────┐  ┌──────────────────────┐  ┌──────────────────────┐
+  │ EngineCore Proc-0    │  │ EngineCore Proc-1    │  │ EngineCore Proc-2    │
+  │ (Thinker)            │  │ (Talker)             │  │ (Vocoder)            │
+  └──────────────────────┘  └──────────────────────┘  └──────────────────────┘
+```
+
+
+Test scripts:
+```bash
+# enter offline inference folder.
+cd examples/offline_inference/qwen2_5_omni
+python end2end.py --output-dir output_audio --query-type use_mixed_modalities
+
+cd ../qwen3_omni
+python end2end.py --output-dir output_audio --query-type text --async-chunk --enable-stats
+
+cd ../bagel
+python end2end.py --prompts "A cute cat"
+
+cd ../text_to_image
+python text_to_image.py --prompt "a cup of coffee on the table" --output output.png
+```

docs/features/custom_pipeline.md

@@ -97,7 +97,7 @@ outputs = omni.generate(
 )
 
 # Access custom trajectory data
-output = outputs[0].request_output[0]
+output = outputs[0].request_output
 print(f"Trajectory timesteps shape: {output.metrics['trajectory_timesteps'].shape}")
 print(f"Trajectory latents shape: {output.latents.shape}")
 ```

docs/getting_started/installation/npu/npu.inc.md

@@ -33,18 +33,25 @@ docker run --rm \
     -p 8000:8000 \
     -it $IMAGE bash
 
+cd /vllm-workspace/vllm
+git pull origin main
+git fetch origin --tags
+git checkout v0.16.0
+
 # Because vllm-ascend will release v0.16.0rc1 after vllm-omni 0.16.0,
 # we have to pin vllm-ascend at the current commit.
 cd /vllm-workspace/vllm-ascend
+git pull origin main
 git checkout e2175d9c7e62b437391dfee996b1375674ba7c18
 pip install -v -e .
 
 # Inside the container, install vLLM-Omni from source
 cd /vllm-workspace
 git clone -b v0.16.0 https://github.com/vllm-project/vllm-omni.git
-
 cd vllm-omni
-pip install -v -e .
+pip install -v -e . --no-build-isolation
+# or VLLM_OMNI_TARGET_DEVICE=npu pip install -v -e .
+
 export VLLM_WORKER_MULTIPROC_METHOD=spawn
 ```
 
@@ -61,22 +68,22 @@ We are keeping [issue #886](https://github.com/vllm-project/vllm-omni/issues/886
 You can also build vLLM-Omni from the latest main branch if you want to use the latest features or bug fixes. (But sometimes it will break for a while. You can check [issue #886](https://github.com/vllm-project/vllm-omni/issues/886) for the status of the latest commit of vLLM-Omni main branch on NPU.)
 
 ```bash
-# Pin vLLM version to 0.16.0
+# Pin vLLM version to 0.18.0
 cd /vllm-workspace/vllm
 git pull origin main
 git fetch origin --tags
-git checkout v0.16.0
+git checkout v0.18.0
 VLLM_TARGET_DEVICE=empty pip install -v -e .
 
 # Because vllm-ascend has not yet entered continuous development and has not been officially released, we need to pin it to a specific commit. Please note that this commit may change over time.
-cd ../vllm-ascend
+cd /vllm-workspace/vllm-ascend
 git pull origin main
 git fetch origin --tags
-git checkout e2175d9c7e62b437391dfee996b1375674ba7c18
+git checkout 1e05c4908f31737bc4eef865a9f351d030a77c9d
 pip install -v -e .
 
 # Install vLLM-Omni from the latest main branch
-cd ../vllm-omni
+cd /vllm-workspace/vllm-omni
 git clone https://github.com/vllm-project/vllm-omni.git
 pip install -v -e . --no-build-isolation
 # or VLLM_OMNI_TARGET_DEVICE=npu pip install -v -e .

docs/getting_started/quickstart.md

@@ -42,7 +42,7 @@ if __name__ == "__main__":
     omni = Omni(model="Tongyi-MAI/Z-Image-Turbo")
     prompt = "a cup of coffee on the table"
     outputs = omni.generate(prompt)
-    images = outputs[0].request_output[0].images
+    images = outputs[0].request_output.images
     images[0].save("coffee.png")
 ```
 
@@ -70,7 +70,7 @@ if __name__ == "__main__":
     ]
     omni_outputs = omni.generate(prompts)
     for i_prompt, prompt_output in enumerate(omni_outputs):
-        this_request_output = prompt_output.request_output[0]
+        this_request_output = prompt_output.request_output
         this_images = this_request_output.images
         for i_image, image in enumerate(this_images):
             image.save(f"p{i_prompt}-img{i_image}.jpg")
@@ -82,7 +82,7 @@ if __name__ == "__main__":
 
 !!! info
 
-    For diffusion pipelines, the stage config field `stage_args.[].runtime.max_batch_size` is 1 by default, and the input
+    For diffusion pipelines, the stage config field `stage_args.[].engine_args.max_num_seqs` is 1 by default, and the input
     list is sliced into single-item requests before feeding into the diffusion pipeline. For models that do internally support
     batched inputs, you can [modify this configuration](../configuration/stage_configs.md) to let the model accept a longer batch of prompts.
 

docs/mkdocs/hooks/generate_api_readme.py

@@ -141,6 +141,7 @@ def scan_package(package_name: str = "vllm_omni") -> dict[str, list[str]]:
                 "vllm_omni.diffusion.models.qwen_image",
                 "vllm_omni.entrypoints.async_diffusion",
                 "vllm_omni.entrypoints.openai",
+                "vllm_omni.model_executor.models.voxtral_tts.configuration_voxtral_tts",
             ]
             if any(module_path.startswith(prefix) for prefix in excluded_prefixes):
                 continue

docs/models/supported_models.md

@@ -28,7 +28,7 @@ th {
 | `QwenImageLayeredPipeline` | Qwen-Image-Layered | `Qwen/Qwen-Image-Layered` |
 | `GlmImagePipeline` | GLM-Image | `zai-org/GLM-Image` |
 |`ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` |
-| `WanPipeline` | Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` |
+| `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` |
 | `LTX2Pipeline` | LTX-2-T2V | `Lightricks/LTX-2` |
 | `LTX2ImageToVideoPipeline` | LTX-2-I2V | `Lightricks/LTX-2` |
@@ -52,6 +52,7 @@ th {
 |`Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` |
 |`FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` |
 |`DreamIDOmniPipeline`| DreamID-Omni | `XuGuo699/DreamID-Omni` |
+|`VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` |
 
 
 ## List of Supported Models for NPU
@@ -75,7 +76,7 @@ th {
 | `QwenImageLayeredPipeline` | Qwen-Image-Layered | `Qwen/Qwen-Image-Layered` |
 | `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2511 | `Qwen/Qwen-Image-Edit-2511` |
 |`ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` |
-| `WanPipeline` | Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` |
+| `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` |
 |`LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` |
 |`LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` |

docs/serving/speech_api.md

@@ -4,6 +4,7 @@ vLLM-Omni provides an OpenAI-compatible API for text-to-speech (TTS) generation.
 
 - **Qwen3-TTS** (`Qwen/Qwen3-TTS-12Hz-*`) -- Qwen3-based TTS with CustomVoice, VoiceDesign, and Base (voice cloning) task types. Output: 24 kHz.
 - **Fish Speech S2 Pro** (`fishaudio/s2-pro`) -- Dual-AR TTS with DAC codec. Supports text-to-speech and voice cloning via reference audio. Output: 44.1 kHz.
+- **Voxtral TTS** (`mistralai/Voxtral-4B-TTS-2603`) -- AR + FlowMatching TTS with preset voices. Output: 24 kHz.
 
 Each server instance runs a single model (specified at startup via `vllm serve <model> --omni`).
 
@@ -28,6 +29,14 @@ vllm-omni serve fishaudio/s2-pro \
     --trust-remote-code \
     --enforce-eager \
     --gpu-memory-utilization 0.9
+
+# Voxtral TTS
+vllm serve mistralai/Voxtral-4B-TTS-2603 \
+    --stage-configs-path vllm_omni/model_executor/stage_configs/voxtral_tts.yaml \
+    --omni \
+    --port 8091 \
+    --trust-remote-code \
+    --enforce-eager
 ```
 
 ### Generate Speech
@@ -347,6 +356,12 @@ curl -X POST http://localhost:8091/v1/audio/speech \
 
 Fish Speech uses `ref_audio` and `ref_text` for voice cloning (no `task_type` needed). The `voice` field should be set to `"default"`. See the [Fish Speech online serving example](../user_guide/examples/online_serving/fish_speech.md) for details.
 
+### Voxtral TTS
+
+| Model | Description |
+|-------|-------------|
+| `mistralai/Voxtral-4B-TTS-2603` | 3B AR + FlowMatching TTS. Supports text-to-speech with preset voices. |
+
 ## Error Responses
 
 ### 400 Bad Request

docs/user_guide/diffusion/parallelism_acceleration.md

@@ -24,22 +24,22 @@ The following table shows which models are currently supported by parallelism me
 
 ### ImageGen
 
-| Model                    | Model Identifier                     | Ulysses-SP | Ring-SP | CFG-Parallel | Tensor-Parallel | VAE-Patch-Parallel | Expert-Parallel |
-|--------------------------|--------------------------------------|:----------:|:-------:|:------------:|:---------------:|:------------------:|:---------------:|
-| **LongCat-Image**        | `meituan-longcat/LongCat-Image`      |     ✅      |    ✅    |      ❌       |        ✅        |         ❌          |        N/A        |
-| **LongCat-Image-Edit**   | `meituan-longcat/LongCat-Image-Edit` |     ✅      |    ✅    |      ❌       |        ✅        |         ❌          |        N/A        |
-| **Ovis-Image**           | `OvisAI/Ovis-Image`                  |     ❌      |    ❌    |      ❌       |        ❌        |         ❌          |        N/A        |
-| **Qwen-Image**           | `Qwen/Qwen-Image`                    |     ✅      |    ✅    |      ✅       |        ✅        |         ✅          |        N/A        |
-| **Qwen-Image-Edit**      | `Qwen/Qwen-Image-Edit`               |     ✅      |    ✅    |      ✅       |        ✅        |         ❌          |        N/A        |
-| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509`          |     ✅      |    ✅    |      ✅       |        ✅        |         ❌          |        N/A        |
-| **Qwen-Image-Layered**   | `Qwen/Qwen-Image-Layered`            |     ✅      |    ✅    |      ✅       |        ✅        |         ❌          |        N/A        |
-| **Z-Image**              | `Tongyi-MAI/Z-Image-Turbo`           |     ✅      |    ✅    |      ❌       |  ✅ (TP=2 only)  |         ✅          |        N/A        |
-| **Stable-Diffusion3.5**  | `stabilityai/stable-diffusion-3.5`   |     ❌      |    ❌    |      ❌       |        ✅        |         ✅          |        N/A        |
-| **FLUX.2-klein**         | `black-forest-labs/FLUX.2-klein-4B`  |     ❌      |    ❌    |      ❌       |        ✅        |         ❌          |        N/A        |
-| **FLUX.1-dev**           | `black-forest-labs/FLUX.1-dev`       |     ❌      |    ❌    |      ✅       |        ✅        |         ❌          |        N/A        |
-| **FLUX.2-dev**           | `black-forest-labs/FLUX.2-dev`       |     ❌      |    ❌    |      ❌       |        ✅        |         ❌          |        N/A        |
-| **HunyuanImage3.0**      | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` |     ❌      |    ❌    |      ❌       |        ✅        |         ❌          |        ✅        |
-| **DreamID-Omni**           | `XuGuo699/DreamID-Omni`       |     ❌      |    ❌    |      ✅       |        ❌        |         ❌          |        N/A        |
+| Model                    | Model Identifier                     | Ulysses-SP | Ring-SP | CFG-Parallel | Tensor-Parallel | VAE-Patch-Parallel | Expert-Parallel | HSDP |
+|--------------------------|--------------------------------------|:----------:|:-------:|:------------:|:---------------:|:------------------:|:---------------:|:----:|
+| **LongCat-Image**        | `meituan-longcat/LongCat-Image`      |     ✅      |    ✅    |      ❌       |        ✅        |         ❌          |      N/A        | ❌   |
+| **LongCat-Image-Edit**   | `meituan-longcat/LongCat-Image-Edit` |     ✅      |    ✅    |      ❌       |        ✅        |         ❌          |      N/A        | ❌   |
+| **Ovis-Image**           | `OvisAI/Ovis-Image`                  |     ❌      |    ❌    |      ❌       |        ❌        |         ❌          |      N/A        | ❌   |
+| **Qwen-Image**           | `Qwen/Qwen-Image`                    |     ✅      |    ✅    |      ✅       |        ✅        |         ✅          |      N/A        | ❌   |
+| **Qwen-Image-Edit**      | `Qwen/Qwen-Image-Edit`               |     ✅      |    ✅    |      ✅       |        ✅        |         ❌          |      N/A        | ❌   |
+| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509`          |     ✅      |    ✅    |      ✅       |        ✅        |         ❌          |      N/A        | ❌   |
+| **Qwen-Image-Layered**   | `Qwen/Qwen-Image-Layered`            |     ✅      |    ✅    |      ✅       |        ✅        |         ❌          |      N/A        | ❌   |
+| **Z-Image**              | `Tongyi-MAI/Z-Image-Turbo`           |     ✅      |    ✅    |      ❌       |  ✅ (TP=2 only)  |         ✅          |      N/A        | ❌   |
+| **Stable-Diffusion3.5**  | `stabilityai/stable-diffusion-3.5`   |     ❌      |    ❌    |      ❌       |        ✅        |         ✅          |      N/A        | ❌   |
+| **FLUX.2-klein**         | `black-forest-labs/FLUX.2-klein-4B`  |     ✅      |    ✅    |      ❌       |        ✅        |         ❌          |      N/A        | ✅   |
+| **FLUX.1-dev**           | `black-forest-labs/FLUX.1-dev`       |     ❌      |    ❌    |      ✅       |        ✅        |         ❌          |      N/A        | ✅   |
+| **FLUX.2-dev**           | `black-forest-labs/FLUX.2-dev`       |     ❌      |    ❌    |      ❌       |        ✅        |         ❌          |      N/A        | ✅   |
+| **HunyuanImage3.0**      | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` |     ❌      |    ❌    |      ❌       |        ✅        |         ❌          |      ✅        | ❌   |
+| **DreamID-Omni**           | `XuGuo699/DreamID-Omni`       |     ❌      |    ❌    |      ✅       |        ❌        |         ❌          |      N/A        | ❌   |
 
 !!! note "TP Limitations for Diffusion Models"
     We currently implement Tensor Parallelism (TP) only for the DiT (Diffusion Transformer) blocks. This is because the `text_encoder` component in vLLM-Omni uses the original Transformers implementation, which does not yet support TP.
@@ -58,6 +58,8 @@ The following table shows which models are currently supported by parallelism me
 
 | Model | Model Identifier | Ulysses-SP | Ring-Attention | Tensor-Parallel | HSDP | VAE-Patch-Parallel |
 |-------|------------------|:----------:|:--------------:|:---------------:|:----:| :----:|
+| **Wan2.1** | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers` | ✅ | ✅ | ✅ | ✅ | ✅ |
+| **Wan2.1** | `Wan-AI/Wan2.1-T2V-14B-Diffusers` | ✅ | ✅ | ✅ | ✅ | ✅ |
 | **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ✅ | ✅ | ✅ | ✅ | ✅ |
 | **LTX-2** | `Lightricks/LTX-2` | ✅ | ✅ | ✅ | ❌ | ❌ |
 

docs/user_guide/diffusion/quantization/gguf.md

@@ -22,10 +22,16 @@
 CLI (examples/offline_inference/text_to_image/text_to_image.py)
   |
   v
-Omni (vllm_omni/entrypoints/omni.py)
+Omni (vllm_omni/entrypoints/__init__.py)
   |
   v
-OmniStage (diffusion)
+AsyncOmniEngine
+  |
+  v
+Orchestrator
+  |
+  v
+StageDiffusionClient
   |
   v
 DiffusionWorker
@@ -64,7 +70,13 @@ _generate_with_async_omni
 AsyncOmni
   |
   v
-DiffusionEngine
+AsyncOmniEngine
+  |
+  v
+Orchestrator
+  |
+  v
+StageDiffusionClient
   |
   v
 OmniRequestOutput

docs/user_guide/diffusion/quantization/int8.md

@@ -0,0 +1,75 @@
+# Int8 Quantization
+
+## Overview
+
+Int8 quantization converts BF16/FP16 weights to Int8 at model load time. No calibration or pre-quantized checkpoint needed.
+
+Depending on the model, either all layers can be quantized, or some sensitive layers should stay in BF16/FP16. See the [per-model table](#supported-models) for which case applies.
+
+## Configuration
+
+1. **Python API**: set `quantization="int8"`. To skip sensitive layers, use `quantization_config` with `ignored_layers`.
+
+```python
+from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+# All layers quantized
+omni = Omni(model="<your-model>", quantization="int8")
+
+# Skip sensitive layers
+omni = Omni(
+    model="<your-model>",
+    quantization_config={
+        "method": "int8",
+        "ignored_layers": ["<layer-name>"],
+    },
+)
+
+outputs = omni.generate(
+    "A cat sitting on a windowsill",
+    OmniDiffusionSamplingParams(num_inference_steps=50),
+)
+```
+
+2. **CLI**: pass `--quantization int8` and optionally `--ignored-layers`.
+
+```bash
+# All layers
+python text_to_image.py --model <your-model> --quantization int8
+
+# Skip sensitive layers
+python text_to_image.py --model <your-model> --quantization int8 --ignored-layers "img_mlp"
+
+# Online serving
+vllm serve <your-model> --omni --quantization int8
+```
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `method` | str | — | Quantization method (`"int8"`) |
+| `ignored_layers` | list[str] | `[]` | Layer name patterns to keep in BF16/FP16 |
+| `activation_scheme` | str | `"dynamic"` | `"dynamic"` (no calibration) |
+
+
+The available `ignored_layers` names depend on the model architecture (e.g., `to_qkv`, `to_out`, `img_mlp`, `txt_mlp`). Consult the transformer source for your target model.
+
+## Supported Models
+
+| Model | HF Models | Recommendation | `ignored_layers` |
+|-------|-----------|---------------|------------------|
+| Z-Image | `Tongyi-MAI/Z-Image-Turbo` | All layers | None |
+| Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | All layers | None |
+
+## Combining with Other Features
+
+Int8 quantization can be combined with cache acceleration:
+
+```python
+omni = Omni(
+    model="<your-model>",
+    quantization="int8",
+    cache_backend="tea_cache",
+    cache_config={"rel_l1_thresh": 0.2},
+)
+```

docs/user_guide/diffusion/quantization/overview.md

@@ -7,12 +7,20 @@ vLLM-Omni supports quantization of DiT linear layers to reduce memory usage and
 | Method | Guide |
 |--------|-------|
 | FP8 | [FP8](fp8.md) |
+| Int8 | [Int8](int8.md) |
 | GGUF | [GGUF](gguf.md) |
 
-## Device Compatibility
+## Device Compatibility for FP8
 
 | GPU Generation | Example GPUs | FP8 Mode |
 |---------------|-------------------|----------|
 | Ada/Hopper (SM 89+) | RTX 4090, H100, H200 | Full W8A8 with native hardware |
 
 Kernel selection is automatic.
+
+## Device Compatibility for Int8
+
+| Device Type | Generation | Example | Int8 Mode |
+|-------------|---------------|-------------------|----------|
+| NVIDIA GPU | Ada/Hopper (SM 89+) | RTX 4090, H100, H200 | Full W8A8 with native hardware |
+| Ascend NPU | Atlas A2/Atlas A3 | Atlas 800T A2/Atlas 900 A3 | Full W8A8 with native hardware |

docs/user_guide/diffusion/step_execution.md

@@ -0,0 +1,61 @@
+# Step Execution
+
+Step execution is an opt-in diffusion execution mode enabled with
+`step_execution=True` when constructing `Omni`.
+
+It is not a generic diffusion toggle for every pipeline. Only pipelines that
+implement the stepwise contract support it today.
+
+## Quick Start
+
+```python
+from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+omni = Omni(
+    model="Qwen/Qwen-Image",
+    step_execution=True,
+)
+
+outputs = omni.generate(
+    "A cat sitting on a windowsill",
+    OmniDiffusionSamplingParams(
+        num_inference_steps=50,
+    ),
+)
+```
+
+## Supported Pipelines
+
+| Pipeline | Example models | Step execution |
+|----------|----------------|----------------|
+| `QwenImagePipeline` | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Yes |
+| All other diffusion pipelines | `QwenImageEditPipeline`, `QwenImageEditPlusPipeline`, `QwenImageLayeredPipeline`, GLM-Image, Wan, Flux, etc. | No |
+
+## Current Limitations
+
+- `step_execution` currently supports `batch_size=1` only.
+- `cache_backend` is not supported together with step execution.
+- Unsupported pipelines fail early during model loading.
+- Request-mode extras such as KV transfer are not wired into step mode yet.
+
+## When To Use It
+
+Use step execution only when you specifically need the pipeline to run through
+its stepwise request state machine. For normal diffusion inference, leave it
+disabled unless your workflow depends on this mode.
+
+If you are looking for general diffusion speedups, see
+[Diffusion Acceleration Overview](../diffusion_acceleration.md).
+
+## Troubleshooting
+
+If model loading fails with a message mentioning `prepare_encode()`,
+`denoise_step()`, `step_scheduler()`, and `post_decode()`, the selected
+pipeline does not support step execution.
+
+## For Model Authors
+
+If you want to add step execution support to a new diffusion pipeline, see the
+implementation guide:
+[Diffusion Step Execution Design](../../design/feature/diffusion_step_execution.md).

docs/user_guide/diffusion_acceleration.md

@@ -16,7 +16,7 @@ Both methods can provide significant speedups (typically **1.5x-2.0x**) while ma
 
 vLLM-Omni also supports quantization methods:
 
-3. **[FP8 Quantization](diffusion/quantization/overview.md)** - Reduces DiT linear layers from BF16 to FP8, providing ~1.28x speedup with minimal quality loss. Supports per-layer skip for sensitive layers.
+3. **[Quantization](diffusion/quantization/overview.md)** - Reduces DiT linear layers from BF16 to FP8 or Int8, providing ~1.28x speedup with minimal quality loss. Supports per-layer skip for sensitive layers.
 
 vLLM-Omni also supports parallelism methods for diffusion models, including:
 
@@ -46,6 +46,7 @@ vLLM-Omni also supports parallelism methods for diffusion models, including:
 | Method | Configuration | Description | Best For |
 |--------|--------------|-------------|----------|
 | **FP8** | `quantization="fp8"` | FP8 W8A8 on Ada/Hopper, weight-only on older GPUs | Memory reduction, inference speedup |
+| **Int8** | `quantization="int8"` | Int8 W8A8 | Memory reduction, inference speedup |
 
 ## Supported Models
 
@@ -75,17 +76,19 @@ The following table shows which models are currently supported by each accelerat
 
 | Model | Model Identifier | TeaCache | Cache-DiT | Ulysses-SP | Ring-Attention | CFG-Parallel | HSDP | VAE-Patch-Parallel |
 |-------|------------------|:--------:|:---------:|:----------:|:--------------:|:------------:|:----:|:----:|
+| **Wan2.1-T2V** | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers` | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
+| **Wan2.1-T2V** | `Wan-AI/Wan2.1-T2V-14B-Diffusers` | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
 | **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
 | **LTX-2** | `Lightricks/LTX-2` | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 | **DreamID-Omni** | `XuGuo699/DreamID-Omni` | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ |
 
 ### Quantization
 
-| Model | Model Identifier | FP8 |
-|-------|------------------|:---:|
-| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ |
-| **Qwen-Image-2512** | `Qwen/Qwen-Image-2512` | ✅ |
-| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ |
+| Model | Model Identifier | FP8 | Int8 |
+|-------|------------------|:---:|:---:|
+| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ |
+| **Qwen-Image-2512** | `Qwen/Qwen-Image-2512` | ✅ | ✅ |
+| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ |
 
 
 ## Performance Benchmarks
@@ -338,13 +341,30 @@ outputs = omni.generate(
 )
 ```
 
+### Using Int8 Quantization
+
+```python
+from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+omni = Omni(
+    model="<your-model>",
+    quantization="int8",
+)
+
+outputs = omni.generate(
+    "A cat sitting on a windowsill",
+    OmniDiffusionSamplingParams(num_inference_steps=50),
+)
+```
+
 ## Documentation
 
 For detailed information on each acceleration method:
 
 - **[TeaCache Guide](diffusion/teacache.md)** - Complete TeaCache documentation, configuration options, and best practices
 - **[Cache-DiT Acceleration Guide](diffusion/cache_dit_acceleration.md)** - Comprehensive Cache-DiT guide covering DBCache, TaylorSeer, SCM, and configuration parameters
-- **[FP8 Quantization Guide](diffusion/quantization/overview.md)** - FP8 quantization for DiT models with per-layer control
+- **[Quantization Guide](diffusion/quantization/overview.md)** - Quantization for DiT models with per-layer control
 - **[Tensor Parallelism](diffusion/parallelism_acceleration.md#tensor-parallelism)** - Guidance on how to enable TP for diffusion models.
 - **[Sequence Parallelism](diffusion/parallelism_acceleration.md#sequence-parallelism)** - Guidance on how to set sequence parallelism with configuration.
 - **[CFG-Parallel](diffusion/parallelism_acceleration.md#cfg-parallel)** - Guidance on how to set CFG-Parallel to run positive/negative branches across ranks.

docs/user_guide/examples/offline_inference/bagel.md

@@ -130,7 +130,7 @@ The default yaml configuration deploys Thinker and DiT on the same GPU. You can
 | :------------------------------- | :------------------------------ | :----------------------- |
 | `stage_type`                     | `llm`                           | Stage type               |
 | `devices`                        | `"0"`                           | GPU device ID            |
-| `max_batch_size`                 | `1`                             | Maximum batch size       |
+| `max_num_seqs`                   | `1`                             | Maximum batch size       |
 | `model_stage`                    | `thinker`                       | Model stage identifier   |
 | `model_arch`                     | `BagelForConditionalGeneration` | Model architecture       |
 | `gpu_memory_utilization`         | `0.4`                           | GPU memory utilization   |
@@ -146,7 +146,7 @@ The default yaml configuration deploys Thinker and DiT on the same GPU. You can
 | :------------------------------- | :---------- | :-------------------------- |
 | `stage_type`                     | `diffusion` | Stage type                  |
 | `devices`                        | `"0"`       | GPU device ID               |
-| `max_batch_size`                 | `1`         | Maximum batch size          |
+| `max_num_seqs`                   | `1`         | Maximum batch size          |
 | `model_stage`                    | `dit`       | Model stage identifier      |
 | `gpu_memory_utilization`         | `0.4`       | GPU memory utilization      |
 | `omni_kv_config.need_recv_cache` | `true`      | Whether to receive KV cache |

docs/user_guide/examples/offline_inference/qwen3_tts.md

@@ -107,7 +107,7 @@ completes. This demonstrates that audio data is available progressively rather t
 
 ## Batched Decoding
 
-The Code2Wav stage (stage 1) supports batched decoding, where multiple requests are decoded in a single forward pass through the SpeechTokenizer. To use it, provide a stage config with `max_batch_size > 1` and pass multiple prompts via `--txt-prompts` with a matching `--batch-size`.
+The Code2Wav stage (stage 1) supports batched decoding, where multiple requests are decoded in a single forward pass through the SpeechTokenizer. To use it, provide a stage config with `max_num_seqs > 1` and pass multiple prompts via `--txt-prompts` with a matching `--batch-size`.
 
 ```
 python end2end.py --query-type CustomVoice \
@@ -116,7 +116,7 @@ python end2end.py --query-type CustomVoice \
     --stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts_batch.yaml
 ```
 
-**Important:** `--batch-size` must match a CUDA graph capture size (1, 2, 4, 8, 16...) because the Talker's code predictor KV cache is sized to `max_num_seqs`, and CUDA graphs pad the batch to the next capture size. Both stages need `max_batch_size >= batch_size` in the stage config for batching to take effect. If only stage 1 has a higher `max_batch_size`, it won't help — stage 1 can only batch chunks from requests that are in-flight simultaneously, which requires stage 0 to also process multiple requests concurrently.
+**Important:** `--batch-size` must match a CUDA graph capture size (1, 2, 4, 8, 16...) because the Talker's code predictor KV cache is sized to `max_num_seqs`, and CUDA graphs pad the batch to the next capture size. Both stages need `max_num_seqs >= batch_size` in the stage config for batching to take effect. If only stage 1 has a higher `max_num_seqs`, it won't help — stage 1 can only batch chunks from requests that are in-flight simultaneously, which requires stage 0 to also process multiple requests concurrently.
 
 ## Notes
 

docs/user_guide/examples/offline_inference/text_to_image.md

@@ -19,7 +19,7 @@ if __name__ == "__main__":
     omni = Omni(model="Qwen/Qwen-Image")
     prompt = "a cup of coffee on the table"
     outputs = omni.generate(prompt)
-    images = outputs[0].request_output[0].images
+    images = outputs[0].request_output.images
     images[0].save("coffee.png")
 ```
 
@@ -37,7 +37,7 @@ if __name__ == "__main__":
     ]
     outputs = omni.generate(prompts)
     for i, output in enumerate(outputs):
-      image = output.request_output[0].images[0].save(f"{i}.jpg")
+      image = output.request_output.images[0].save(f"{i}.jpg")
 ```
 
 !!! info
@@ -49,7 +49,7 @@ if __name__ == "__main__":
 
 !!! info
 
-    For diffusion pipelines, the stage config field `stage_args.[].runtime.max_batch_size` is 1 by default, and the input
+    For diffusion pipelines, the stage config field `stage_args.[].engine_args.max_num_seqs` is 1 by default, and the input
     list is sliced into single-item requests before feeding into the diffusion pipeline. For models that do internally support
     batched inputs, you can [modify this configuration](https://github.com/vllm-project/vllm-omni/tree/main/configuration/stage_configs.md) to let the model accept a longer batch of prompts.
 
@@ -72,7 +72,7 @@ if __name__ == "__main__":
       }
     ])
     for i, output in enumerate(outputs):
-      image = output.request_output[0].images[0].save(f"{i}.jpg")
+      image = output.request_output.images[0].save(f"{i}.jpg")
 ```
 
 ## Local CLI Usage

docs/user_guide/examples/offline_inference/text_to_video.md

@@ -3,7 +3,10 @@
 Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/text_to_video>.
 
 
-The `Wan-AI/Wan2.2-T2V-A14B-Diffusers` pipeline generates short videos from text prompts.
+The Wan T2V pipeline generates short videos from text prompts. Supported models include:
+
+- `Wan-AI/Wan2.2-T2V-A14B-Diffusers` (14B, requires ~48GB VRAM)
+- `Wan-AI/Wan2.1-T2V-1.3B-Diffusers` (1.3B, requires ~16GB VRAM, use `--flow-shift 3.0 --boundary-ratio 0.0`)
 
 ## Local CLI Usage
 

examples/offline_inference/bagel/README.md

@@ -127,7 +127,7 @@ The default yaml configuration deploys Thinker and DiT on the same GPU. You can
 | :------------------------------- | :------------------------------ | :----------------------- |
 | `stage_type`                     | `llm`                           | Stage type               |
 | `devices`                        | `"0"`                           | GPU device ID            |
-| `max_batch_size`                 | `1`                             | Maximum batch size       |
+| `max_num_seqs`                   | `1`                             | Maximum batch size       |
 | `model_stage`                    | `thinker`                       | Model stage identifier   |
 | `model_arch`                     | `BagelForConditionalGeneration` | Model architecture       |
 | `gpu_memory_utilization`         | `0.4`                           | GPU memory utilization   |
@@ -143,7 +143,7 @@ The default yaml configuration deploys Thinker and DiT on the same GPU. You can
 | :------------------------------- | :---------- | :-------------------------- |
 | `stage_type`                     | `diffusion` | Stage type                  |
 | `devices`                        | `"0"`       | GPU device ID               |
-| `max_batch_size`                 | `1`         | Maximum batch size          |
+| `max_num_seqs`                   | `1`         | Maximum batch size          |
 | `model_stage`                    | `dit`       | Model stage identifier      |
 | `gpu_memory_utilization`         | `0.4`       | GPU memory utilization      |
 | `omni_kv_config.need_recv_cache` | `true`      | Whether to receive KV cache |

examples/offline_inference/bagel/end2end.py

@@ -34,6 +34,13 @@ def parse_args():
         help="Path to input image for img2img.",
     )
 
+    parser.add_argument(
+        "--output",
+        type=str,
+        default=".",
+        help="Output directory to save images.",
+    )
+
     # OmniLLM init args
     parser.add_argument("--log-stats", action="store_true", default=False)
     parser.add_argument("--init-sleep-seconds", type=int, default=20)
@@ -58,6 +65,11 @@ def parse_args():
         help="CFG parallel size: 1=batched (single GPU), 2=parallel with 2 branches (text CFG only), 3=parallel (3 GPUs).",
     )
     parser.add_argument("--seed", type=int, default=None, help="Random seed for generation.")
+    parser.add_argument(
+        "--enable-diffusion-pipeline-profiler",
+        action="store_true",
+        help="Enable diffusion pipeline profiler to display stage durations.",
+    )
 
     args = parser.parse_args()
     return args
@@ -65,6 +77,7 @@ def parse_args():
 
 def main():
     args = parse_args()
+    os.makedirs(args.output, exist_ok=True)
     model_name = args.model
     prompts: list[OmniPromptType] = []
     try:
@@ -103,6 +116,7 @@ def main():
             "shm_threshold_bytes": args.shm_threshold_bytes,
             "worker_backend": args.worker_backend,
             "ray_address": args.ray_address,
+            "enable_diffusion_pipeline_profiler": args.enable_diffusion_pipeline_profiler,
         }
     )
 
@@ -163,25 +177,17 @@ def main():
 
     omni_outputs = list(omni.generate(prompts=formatted_prompts, sampling_params_list=params_list))
 
-    for i, req_output in enumerate(omni_outputs):
+    img_idx = 0
+    for req_output in omni_outputs:
         images = getattr(req_output, "images", None)
-        if not images and hasattr(req_output, "output"):
-            if isinstance(req_output.output, list):
-                images = req_output.output
-            else:
-                images = [req_output.output]
 
-        if images:
-            for j, img in enumerate(images):
-                img.save(f"output_{i}_{j}.png")
+        if not images:
+            continue
 
-        if hasattr(req_output, "request_output") and req_output.request_output:
-            for stage_out in req_output.request_output:
-                if hasattr(stage_out, "images") and stage_out.images:
-                    for k, img in enumerate(stage_out.images):
-                        save_path = f"output_{i}_stage_{getattr(stage_out, 'stage_id', '?')}_{k}.png"
-                        img.save(save_path)
-                        print(f"[Info] Saved stage output image to {save_path}")
+        for j, img in enumerate(images):
+            save_path = os.path.join(args.output, f"output_{img_idx}_{j}.png")
+            img.save(save_path)
+        img_idx += 1
 
     print(omni_outputs)
 

examples/offline_inference/cosyvoice3/verify_e2e_cosyvoice.py

@@ -167,27 +167,26 @@ def run_e2e():
     print(f"Received {len(outputs)} outputs.")
     for i, output in enumerate(outputs):
         try:
-            ro_list = output.request_output or []
-            if not ro_list:
+            ro = output.request_output
+            if ro is None:
                 print("No request_output found.")
                 continue
 
-            for ro in ro_list:
-                # Multimodal output may be attached to RequestOutput or CompletionOutput.
-                mm = getattr(ro, "multimodal_output", None)
-                if not mm and ro.outputs:
-                    mm = getattr(ro.outputs[0], "multimodal_output", None)
+            # Multimodal output may be attached to RequestOutput or CompletionOutput.
+            mm = getattr(ro, "multimodal_output", None)
+            if not mm and ro.outputs:
+                mm = getattr(ro.outputs[0], "multimodal_output", None)
 
-                if mm:
-                    print(f"Multimodal output keys: {mm.keys()}")
-                    if "audio" in mm:
-                        audio_out = mm["audio"]
-                        print(f"Generated Audio Shape: {audio_out.shape}")
-                        out_path = f"output_{i}.wav"
-                        sf.write(out_path, audio_out.cpu().numpy().squeeze(), 22050)
-                        print(f"Saved audio to {out_path}")
-                else:
-                    print("No multimodal output found.")
+            if mm:
+                print(f"Multimodal output keys: {mm.keys()}")
+                if "audio" in mm:
+                    audio_out = mm["audio"]
+                    print(f"Generated Audio Shape: {audio_out.shape}")
+                    out_path = f"output_{i}.wav"
+                    sf.write(out_path, audio_out.cpu().numpy().squeeze(), 22050)
+                    print(f"Saved audio to {out_path}")
+            else:
+                print("No multimodal output found.")
         except Exception as e:
             print(f"Error inspecting output: {e}")
     omni.close()

examples/offline_inference/custom_pipeline/image_to_image/image_edit.py

@@ -209,7 +209,7 @@ async def main():
     if not outputs:
         raise ValueError("No output produced from omni.generate()")
 
-    first_out = outputs[0].request_output[0]
+    first_out = outputs[0].request_output
     req_out: OmniRequestOutput = first_out
 
     # Verify trajectory data (from custom pipeline)

examples/offline_inference/fish_speech/end2end.py

@@ -129,8 +129,14 @@ def main(args):
 
     t_start = time.perf_counter()
     for stage_outputs in omni.generate(inputs):
-        for output in stage_outputs.request_output:
-            _save_wav(output_dir, output.request_id, output.outputs[0].multimodal_output)
+        request_output = stage_outputs.request_output
+        if request_output is None or not request_output.outputs:
+            continue
+        _save_wav(
+            output_dir,
+            request_output.request_id,
+            request_output.outputs[0].multimodal_output,
+        )
     t_end = time.perf_counter()
     logger.info("Total inference time: %.1f ms", (t_end - t_start) * 1000)
 

examples/offline_inference/glm_image/end2end.py

@@ -247,6 +247,7 @@ def main(args: argparse.Namespace) -> None:
         stage_configs_path=config_path,
         log_stats=args.enable_stats,
         stage_init_timeout=args.stage_init_timeout,
+        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
     )
 
     init_time = time.time() - start_time
@@ -310,36 +311,35 @@ def main(args: argparse.Namespace) -> None:
 
     output_count = 0
     for stage_outputs in omni.generate(prompts, sampling_params_list, py_generator=True):
+        output = stage_outputs.request_output
         if stage_outputs.final_output_type == "image":
-            for output in stage_outputs.request_output:
-                request_id = output.request_id
+            request_id = output.request_id
 
-                # Get generated images
-                images = output.images if hasattr(output, "images") else []
-                if not images and hasattr(output, "multimodal_output"):
-                    images = output.multimodal_output.get("images", [])
+            # Get generated images
+            images = output.images if hasattr(output, "images") else []
+            if not images and hasattr(output, "multimodal_output"):
+                images = output.multimodal_output.get("images", [])
 
-                # Save each generated image
-                for idx, img in enumerate(images):
-                    if args.num_prompts == 1 and len(images) == 1:
-                        output_path = args.output
-                    else:
-                        base, ext = os.path.splitext(args.output)
-                        output_path = f"{base}_{request_id}_{idx}{ext}"
+            # Save each generated image
+            for idx, img in enumerate(images):
+                if args.num_prompts == 1 and len(images) == 1:
+                    output_path = args.output
+                else:
+                    base, ext = os.path.splitext(args.output)
+                    output_path = f"{base}_{request_id}_{idx}{ext}"
 
-                    if isinstance(img, Image.Image):
-                        save_image(img, output_path)
-                    else:
-                        print(f"Warning: Unexpected image type for request {request_id}: {type(img)}")
+                if isinstance(img, Image.Image):
+                    save_image(img, output_path)
+                else:
+                    print(f"Warning: Unexpected image type for request {request_id}: {type(img)}")
 
-                    output_count += 1
+                output_count += 1
 
         elif stage_outputs.final_output_type == "text":
             # AR stage output (intermediate, for debugging)
             if args.verbose:
-                for output in stage_outputs.request_output:
-                    print(f"AR output for request {output.request_id}:")
-                    print(f"  Token count: {len(output.outputs[0].token_ids)}")
+                print(f"AR output for request {output.request_id}:")
+                print(f"  Token count: {len(output.outputs[0].token_ids)}")
 
     gen_time = time.time() - gen_start_time
     print(f"\nGeneration completed in {gen_time:.2f}s")
@@ -464,6 +464,11 @@ def parse_args() -> argparse.Namespace:
         default=False,
         help="Enable verbose output",
     )
+    parser.add_argument(
+        "--enable-diffusion-pipeline-profiler",
+        action="store_true",
+        help="Enable diffusion pipeline profiler to display stage durations.",
+    )
 
     return parser.parse_args()
 

examples/offline_inference/helios/end2end.py

@@ -173,6 +173,11 @@ def parse_args() -> argparse.Namespace:
         default=1,
         help="Number of initial denoising steps using zero prediction (default: 1).",
     )
+    parser.add_argument(
+        "--enable-diffusion-pipeline-profiler",
+        action="store_true",
+        help="Enable diffusion pipeline profiler to display stage durations.",
+    )
 
     # Memory & parallelism
     parser.add_argument("--vae-use-slicing", action="store_true", help="Enable VAE slicing.")
@@ -207,6 +212,7 @@ def main():
         enable_cpu_offload=args.enable_cpu_offload,
         parallel_config=parallel_config,
         enforce_eager=args.enforce_eager,
+        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
     )
 
     # Validate I2V / V2V arguments
@@ -293,12 +299,11 @@ def main():
                 )
 
             if hasattr(first_item, "is_pipeline_output") and first_item.is_pipeline_output:
-                if isinstance(first_item.request_output, list) and len(first_item.request_output) > 0:
-                    inner_output = first_item.request_output[0]
-                    if isinstance(inner_output, OmniRequestOutput) and hasattr(inner_output, "images"):
-                        frames = inner_output.images[0] if inner_output.images else None
-                        if frames is None:
-                            raise ValueError("No video frames found in output.")
+                inner_output = first_item.request_output
+                if isinstance(inner_output, OmniRequestOutput) and hasattr(inner_output, "images"):
+                    frames = inner_output.images[0] if inner_output.images else None
+                    if frames is None:
+                        raise ValueError("No video frames found in output.")
             elif hasattr(first_item, "images") and first_item.images:
                 frames = first_item.images
             else:

examples/offline_inference/hunyuan_image3/image_to_text.py

@@ -9,7 +9,7 @@ from PIL import Image
 from vllm_omni.entrypoints.omni import Omni
 
 """
-The tencent/HunyuanImage-3.0-Instruct base model is built on the Hunyuan v1 architecture, specifically the tencent/Hunyuan-A13B-Instruct model. It utilizes two tokenizer delimiter templates:
+The tencent/HunyuanImage-3.0-Instruct base model uses the tencent/Hunyuan-A13B-Instruct backbone. It utilizes two tokenizer delimiter templates:
 
 1) Pretrained template (default for gen_text mode), which concatenates system, image
    tokens, and user question WITHOUT role delimiters:
@@ -41,6 +41,11 @@ def parse_args() -> argparse.Namespace:
         required=True,
         help="Pretrain template prompt: <|startoftext|>{system}<img>{question}",
     )
+    parser.add_argument(
+        "--enable-diffusion-pipeline-profiler",
+        action="store_true",
+        help="Enable diffusion pipeline profiler to display stage durations.",
+    )
     return parser.parse_args()
 
 
@@ -52,7 +57,10 @@ def load_image(image_path: str) -> Image.Image:
 
 
 def main(args: argparse.Namespace) -> None:
-    omni = Omni(model=args.model)
+    omni = Omni(
+        model=args.model,
+        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
+    )
 
     prompt_dict = {
         "prompt": args.prompt,
@@ -70,8 +78,8 @@ def main(args: argparse.Namespace) -> None:
     prompts = [prompt_dict]
     omni_outputs = omni.generate(prompts=prompts)
 
-    prompt_text = omni_outputs[0].request_output[0].prompt
-    generated_text = omni_outputs[0].request_output[0].outputs[0].text
+    prompt_text = omni_outputs[0].request_output.prompt
+    generated_text = omni_outputs[0].request_output.outputs[0].text
     print(f"Prompt: {prompt_text}")
     print(f"Text: {generated_text}")
 

examples/offline_inference/image_to_image/image_edit.py

@@ -325,6 +325,11 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Enable layerwise (blockwise) offloading on DiT modules.",
     )
+    parser.add_argument(
+        "--enable-diffusion-pipeline-profiler",
+        action="store_true",
+        help="Enable diffusion pipeline profiler to display stage durations.",
+    )
     return parser.parse_args()
 
 
@@ -389,6 +394,7 @@ def main():
         parallel_config=parallel_config,
         enforce_eager=args.enforce_eager,
         enable_cpu_offload=args.enable_cpu_offload,
+        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
     )
     print("Pipeline loaded")
 
@@ -463,12 +469,12 @@ def main():
         raise ValueError("No output generated from omni.generate()")
 
     # Extract images from OmniRequestOutput
-    # omni.generate() returns list[OmniRequestOutput], extract images from request_output[0].images
+    # omni.generate() returns list[OmniRequestOutput], extract images from request_output.images
     first_output = outputs[0]
     if not hasattr(first_output, "request_output") or not first_output.request_output:
         raise ValueError("No request_output found in OmniRequestOutput")
 
-    req_out = first_output.request_output[0]
+    req_out = first_output.request_output
     if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
         raise ValueError("Invalid request_output structure or missing 'images' key")
 

examples/offline_inference/image_to_video/image_to_video.py

@@ -176,6 +176,11 @@ def parse_args() -> argparse.Namespace:
             "Default 1 means pure sharding (no replication). "
         ),
     )
+    parser.add_argument(
+        "--enable-diffusion-pipeline-profiler",
+        action="store_true",
+        help="Enable diffusion pipeline profiler to display stage durations.",
+    )
     return parser.parse_args()
 
 
@@ -281,6 +286,7 @@ def main():
         model_class_name=model_class_name,
         cache_backend=args.cache_backend,
         cache_config=cache_config,
+        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
     )
 
     if profiler_enabled:
@@ -338,8 +344,6 @@ def main():
             audio = frames.multimodal_output["audio"]
         if frames.is_pipeline_output and frames.request_output is not None:
             inner_output = frames.request_output
-            if isinstance(inner_output, list):
-                inner_output = inner_output[0] if inner_output else None
             if isinstance(inner_output, OmniRequestOutput):
                 if inner_output.multimodal_output and "audio" in inner_output.multimodal_output:
                     audio = inner_output.multimodal_output["audio"]

examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py

@@ -43,6 +43,11 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--top-p", type=float, default=0.9)
     parser.add_argument("--seed", type=int, default=42)
     parser.add_argument("--trust-remote-code", action="store_true")
+    parser.add_argument(
+        "--enable-diffusion-pipeline-profiler",
+        action="store_true",
+        help="Enable diffusion pipeline profiler to display stage durations.",
+    )
     return parser.parse_args()
 
 
@@ -70,6 +75,7 @@ def main() -> None:
         model=args.model,
         stage_configs_path=args.stage_config,
         trust_remote_code=args.trust_remote_code,
+        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
     )
     try:
         sp = SamplingParams(
@@ -99,14 +105,12 @@ def main() -> None:
 
     lines: list[str] = []
     for stage_outputs in outputs:
-        req_outputs = getattr(stage_outputs, "request_output", stage_outputs)
-        req_outputs = req_outputs if isinstance(req_outputs, list) else [req_outputs]
-        for ro in req_outputs:
-            text = ro.outputs[0].text if getattr(ro, "outputs", None) else str(ro)
-            lines.append(f"request_id: {getattr(ro, 'request_id', 'unknown')}\n")
-            lines.append("answer:\n")
-            lines.append(text.strip() + "\n")
-            lines.append("\n")
+        ro = getattr(stage_outputs, "request_output", stage_outputs)
+        text = ro.outputs[0].text if getattr(ro, "outputs", None) else str(ro)
+        lines.append(f"request_id: {getattr(ro, 'request_id', 'unknown')}\n")
+        lines.append("answer:\n")
+        lines.append(text.strip() + "\n")
+        lines.append("\n")
 
     print("\n".join(lines))
 

examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_t2i.py

@@ -148,19 +148,16 @@ def _collect_images(outputs: list) -> list[torch.Tensor]:
     """Extract all image tensors produced by the final (DiT) stage."""
     images: list[torch.Tensor] = []
     for out in outputs:
-        ro_list = getattr(out, "request_output", out)
-        if not isinstance(ro_list, list):
-            ro_list = [ro_list]
-        for ro_item in ro_list:
-            for completion in getattr(ro_item, "outputs", None) or []:
-                mm = getattr(completion, "multimodal_output", None)
-                if not isinstance(mm, dict) or "image" not in mm:
-                    raise RuntimeError(f"Missing image in multimodal output: {mm}")
-                payload = mm["image"]
-                for tensor in payload if isinstance(payload, list) else [payload]:
-                    if not isinstance(tensor, torch.Tensor):
-                        raise TypeError(f"Expected image tensor, got {type(tensor)}")
-                    images.append(tensor)
+        ro_item = getattr(out, "request_output", out)
+        for completion in getattr(ro_item, "outputs", None) or []:
+            mm = getattr(completion, "multimodal_output", None)
+            if not isinstance(mm, dict) or "image" not in mm:
+                raise RuntimeError(f"Missing image in multimodal output: {mm}")
+            payload = mm["image"]
+            for tensor in payload if isinstance(payload, list) else [payload]:
+                if not isinstance(tensor, torch.Tensor):
+                    raise TypeError(f"Expected image tensor, got {type(tensor)}")
+                images.append(tensor)
     return images
 
 

examples/offline_inference/mimo_audio/end2end.py

@@ -180,15 +180,16 @@ def main(args):
     # Get the query function and call it with appropriate parameters
     query_func = query_map[args.query_type]
 
-    omni_llm = Omni(
+    omni = Omni(
         model=model_name,
         stage_configs_path=args.stage_configs_path,
         log_stats=args.enable_stats,
-        log_file=("omni_llm_pipeline.log" if args.enable_stats else None),
+        log_file=("omni_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,
+        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
     )
 
     thinker_sampling_params = SamplingParams(
@@ -290,7 +291,7 @@ def main(args):
     prompts = [copy.deepcopy(query_result) for _ in range(args.num_prompts)]
 
     print("prompts", prompts)
-    omni_outputs = omni_llm.generate(prompts, sampling_params_list)
+    omni_outputs = omni.generate(prompts, sampling_params_list)
 
     output_dir = args.output_dir if getattr(args, "output_dir", None) else args.output_wav
     if args.query_type is not None:
@@ -298,45 +299,44 @@ def main(args):
     os.makedirs(output_dir, exist_ok=True)
 
     for stage_outputs in omni_outputs:
+        output = stage_outputs.request_output
         if stage_outputs.final_output_type == "text":
-            for output in stage_outputs.request_output:
-                request_id = output.request_id
-                text_output = output.outputs[0].text
-                # Save aligned text file per request
-                prompt_text = output.prompt
-                out_txt = os.path.join(output_dir, f"{request_id}.txt")
-                lines = []
-                lines.append("Prompt:\n")
-                lines.append(str(prompt_text) + "\n")
-                lines.append("vllm_text_output:\n")
-                lines.append(str(text_output).strip() + "\n")
-                try:
-                    with open(out_txt, "w", encoding="utf-8") as f:
-                        print("lines", lines)
-                        f.writelines(lines)
-                except Exception as e:
-                    print(f"[Warn] Failed writing text file {out_txt}: {e}")
-                print(f"Request ID: {request_id}, Text saved to {out_txt}\n")
+            request_id = output.request_id
+            text_output = output.outputs[0].text
+            # Save aligned text file per request
+            prompt_text = output.prompt
+            out_txt = os.path.join(output_dir, f"{request_id}.txt")
+            lines = []
+            lines.append("Prompt:\n")
+            lines.append(str(prompt_text) + "\n")
+            lines.append("vllm_text_output:\n")
+            lines.append(str(text_output).strip() + "\n")
+            try:
+                with open(out_txt, "w", encoding="utf-8") as f:
+                    print("lines", lines)
+                    f.writelines(lines)
+            except Exception as e:
+                print(f"[Warn] Failed writing text file {out_txt}: {e}")
+            print(f"Request ID: {request_id}, Text saved to {out_txt}\n")
         elif stage_outputs.final_output_type == "audio":
-            for output in stage_outputs.request_output:
-                request_id = output.request_id
-                audio_tensor = output.outputs[0].multimodal_output.get("audio")
+            request_id = output.request_id
+            audio_tensor = output.outputs[0].multimodal_output.get("audio")
 
-                if audio_tensor is None:
-                    continue
+            if audio_tensor is None:
+                continue
 
-                output_wav = os.path.join(output_dir, f"{request_id}.wav")
+            output_wav = os.path.join(output_dir, f"{request_id}.wav")
 
-                # Convert to numpy array and ensure correct format
-                audio_numpy = audio_tensor.float().detach().cpu().numpy()
+            # Convert to numpy array and ensure correct format
+            audio_numpy = audio_tensor.float().detach().cpu().numpy()
 
-                # Ensure audio is 1D (flatten if needed)
-                if audio_numpy.ndim > 1:
-                    audio_numpy = audio_numpy.flatten()
+            # Ensure audio is 1D (flatten if needed)
+            if audio_numpy.ndim > 1:
+                audio_numpy = audio_numpy.flatten()
 
-                # Save audio file with explicit WAV format
-                sf.write(output_wav, audio_numpy, samplerate=24000, format="WAV")
-                print(f"Request ID: {request_id}, Audio saved to {output_wav}")
+            # Save audio file with explicit WAV format
+            sf.write(output_wav, audio_numpy, samplerate=24000, format="WAV")
+            print(f"Request ID: {request_id}, Audio saved to {output_wav}")
 
 
 def parse_args():
@@ -434,6 +434,11 @@ def parse_args():
         default="../../../model_executor/stage_configs/mimo_audio.yaml",
         help="Path to a stage configs file.",
     )
+    parser.add_argument(
+        "--enable-diffusion-pipeline-profiler",
+        action="store_true",
+        help="Enable diffusion pipeline profiler to display stage durations.",
+    )
 
     return parser.parse_args()
 

examples/offline_inference/qwen2_5_omni/end2end.py

@@ -320,13 +320,14 @@ def main(args):
         query_result = query_func(audio_path=audio_path, sampling_rate=sampling_rate)
     else:
         query_result = query_func()
-    omni_llm = Omni(
+    omni = Omni(
         model=model_name,
         log_stats=args.log_stats,
         stage_init_timeout=args.stage_init_timeout,
         batch_timeout=args.batch_timeout,
         init_timeout=args.init_timeout,
         shm_threshold_bytes=args.shm_threshold_bytes,
+        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
     )
     thinker_sampling_params = SamplingParams(
         temperature=0.0,  # Deterministic - no randomness
@@ -378,9 +379,11 @@ def main(args):
             prompt["modalities"] = output_modalities
 
     profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
-    if profiler_enabled:
-        omni_llm.start_profile(stages=[0])
-    omni_generator = omni_llm.generate(prompts, sampling_params_list, py_generator=args.py_generator)
+    if profiler_enabled and hasattr(omni, "start_profile"):
+        omni.start_profile(stages=[0])
+    elif profiler_enabled:
+        print("[Warn] VLLM_TORCH_PROFILER_DIR is set, but current engine does not support profiler controls.")
+    omni_generator = omni.generate(prompts, sampling_params_list, py_generator=args.py_generator)
 
     # Determine output directory: prefer --output-dir; fallback to --output-wav
     output_dir = args.output_dir if getattr(args, "output_dir", None) else args.output_wav
@@ -389,43 +392,42 @@ def main(args):
     total_requests = len(prompts)
     processed_count = 0
     for stage_outputs in omni_generator:
+        output = stage_outputs.request_output
         if stage_outputs.final_output_type == "text":
-            for output in stage_outputs.request_output:
-                request_id = output.request_id
-                text_output = output.outputs[0].text
-                # Save aligned text file per request
-                prompt_text = output.prompt
-                out_txt = os.path.join(output_dir, f"{request_id}.txt")
-                lines = []
-                lines.append("Prompt:\n")
-                lines.append(str(prompt_text) + "\n")
-                lines.append("vllm_text_output:\n")
-                lines.append(str(text_output).strip() + "\n")
-                try:
-                    with open(out_txt, "w", encoding="utf-8") as f:
-                        f.writelines(lines)
-                except Exception as e:
-                    print(f"[Warn] Failed writing text file {out_txt}: {e}")
-                print(f"Request ID: {request_id}, Text saved to {out_txt}")
+            request_id = output.request_id
+            text_output = output.outputs[0].text
+            # Save aligned text file per request
+            prompt_text = output.prompt
+            out_txt = os.path.join(output_dir, f"{request_id}.txt")
+            lines = []
+            lines.append("Prompt:\n")
+            lines.append(str(prompt_text) + "\n")
+            lines.append("vllm_text_output:\n")
+            lines.append(str(text_output).strip() + "\n")
+            try:
+                with open(out_txt, "w", encoding="utf-8") as f:
+                    f.writelines(lines)
+            except Exception as e:
+                print(f"[Warn] Failed writing text file {out_txt}: {e}")
+            print(f"Request ID: {request_id}, Text saved to {out_txt}")
         elif stage_outputs.final_output_type == "audio":
-            for output in stage_outputs.request_output:
-                request_id = output.request_id
-                audio_tensor = output.outputs[0].multimodal_output["audio"]
-                output_wav = os.path.join(output_dir, f"output_{request_id}.wav")
-                sf.write(output_wav, audio_tensor.detach().cpu().numpy(), samplerate=24000)
-                print(f"Request ID: {request_id}, Saved audio to {output_wav}")
+            request_id = output.request_id
+            audio_tensor = output.outputs[0].multimodal_output["audio"]
+            output_wav = os.path.join(output_dir, f"output_{request_id}.wav")
+            sf.write(output_wav, audio_tensor.detach().cpu().numpy(), samplerate=24000)
+            print(f"Request ID: {request_id}, Saved audio to {output_wav}")
 
-        processed_count += len(stage_outputs.request_output)
-        if profiler_enabled and processed_count >= total_requests:
+        processed_count += 1
+        if profiler_enabled and hasattr(omni, "stop_profile") and processed_count >= total_requests:
             print(f"[Info] Processed {processed_count}/{total_requests}. Stopping profiler inside active loop...")
             # Stop the profiler while workers are still alive
-            omni_llm.stop_profile()
+            omni.stop_profile()
 
             print("[Info] Waiting 30s for workers to write massive trace files to disk...")
             time.sleep(30)
             print("[Info] Trace export wait finished.")
 
-    omni_llm.close()
+    omni.close()
 
 
 def parse_args():
@@ -539,6 +541,11 @@ def parse_args():
         default=False,
         help="Use py_generator mode. The returned type of Omni.generate() is a Python Generator object.",
     )
+    parser.add_argument(
+        "--enable-diffusion-pipeline-profiler",
+        action="store_true",
+        help="Enable diffusion pipeline profiler to display stage durations.",
+    )
     return parser.parse_args()
 
 

examples/offline_inference/qwen3_omni/end2end.py

@@ -294,11 +294,12 @@ def main(args):
     else:
         query_result = query_func()
 
-    omni_llm = Omni(
+    omni = Omni(
         model=model_name,
         stage_configs_path=args.stage_configs_path,
         log_stats=args.log_stats,
         stage_init_timeout=args.stage_init_timeout,
+        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
     )
 
     thinker_sampling_params = SamplingParams(
@@ -354,8 +355,8 @@ def main(args):
 
     profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
     if profiler_enabled:
-        omni_llm.start_profile(stages=[0])
-    omni_generator = omni_llm.generate(prompts, sampling_params_list, py_generator=args.py_generator)
+        omni.start_profile(stages=[0])
+    omni_generator = omni.generate(prompts, sampling_params_list, py_generator=args.py_generator)
     # Determine output directory: prefer --output-dir; fallback to --output-wav
     output_dir = args.output_dir if getattr(args, "output_dir", None) else args.output_wav
     os.makedirs(output_dir, exist_ok=True)
@@ -366,51 +367,50 @@ def main(args):
     print(f"query type: {args.query_type}")
 
     for stage_outputs in omni_generator:
+        output = stage_outputs.request_output
         if stage_outputs.final_output_type == "text":
-            for output in stage_outputs.request_output:
-                request_id = output.request_id
-                text_output = output.outputs[0].text
-                # Save aligned text file per request
-                prompt_text = output.prompt
-                out_txt = os.path.join(output_dir, f"{request_id}.txt")
-                lines = []
-                lines.append("Prompt:\n")
-                lines.append(str(prompt_text) + "\n")
-                lines.append("vllm_text_output:\n")
-                lines.append(str(text_output).strip() + "\n")
-                try:
-                    with open(out_txt, "w", encoding="utf-8") as f:
-                        f.writelines(lines)
-                except Exception as e:
-                    print(f"[Warn] Failed writing text file {out_txt}: {e}")
-                print(f"Request ID: {request_id}, Text saved to {out_txt}")
+            request_id = output.request_id
+            text_output = output.outputs[0].text
+            # Save aligned text file per request
+            prompt_text = output.prompt
+            out_txt = os.path.join(output_dir, f"{request_id}.txt")
+            lines = []
+            lines.append("Prompt:\n")
+            lines.append(str(prompt_text) + "\n")
+            lines.append("vllm_text_output:\n")
+            lines.append(str(text_output).strip() + "\n")
+            try:
+                with open(out_txt, "w", encoding="utf-8") as f:
+                    f.writelines(lines)
+            except Exception as e:
+                print(f"[Warn] Failed writing text file {out_txt}: {e}")
+            print(f"Request ID: {request_id}, Text saved to {out_txt}")
         elif stage_outputs.final_output_type == "audio":
-            for output in stage_outputs.request_output:
-                request_id = output.request_id
-                audio_tensor = output.outputs[0].multimodal_output["audio"]
-                output_wav = os.path.join(output_dir, f"output_{request_id}.wav")
+            request_id = output.request_id
+            audio_tensor = output.outputs[0].multimodal_output["audio"]
+            output_wav = os.path.join(output_dir, f"output_{request_id}.wav")
 
-                # Convert to numpy array and ensure correct format
-                audio_numpy = audio_tensor.float().detach().cpu().numpy()
+            # Convert to numpy array and ensure correct format
+            audio_numpy = audio_tensor.float().detach().cpu().numpy()
 
-                # Ensure audio is 1D (flatten if needed)
-                if audio_numpy.ndim > 1:
-                    audio_numpy = audio_numpy.flatten()
+            # Ensure audio is 1D (flatten if needed)
+            if audio_numpy.ndim > 1:
+                audio_numpy = audio_numpy.flatten()
 
-                # Save audio file with explicit WAV format
-                sf.write(output_wav, audio_numpy, samplerate=24000, format="WAV")
-                print(f"Request ID: {request_id}, Saved audio to {output_wav}")
+            # Save audio file with explicit WAV format
+            sf.write(output_wav, audio_numpy, samplerate=24000, format="WAV")
+            print(f"Request ID: {request_id}, Saved audio to {output_wav}")
 
-        processed_count += len(stage_outputs.request_output)
+        processed_count += 1
         if profiler_enabled and processed_count >= total_requests:
             print(f"[Info] Processed {processed_count}/{total_requests}. Stopping profiler inside active loop...")
             # Stop the profiler while workers are still alive
-            omni_llm.stop_profile()
+            omni.stop_profile()
 
             print("[Info] Waiting 30s for workers to write trace files to disk...")
             time.sleep(30)
             print("[Info] Trace export wait time finished.")
-    omni_llm.close()
+    omni.close()
 
 
 def parse_args():
@@ -527,6 +527,11 @@ def parse_args():
         default=False,
         help="Use py_generator mode. The returned type of Omni.generate() is a Python Generator object.",
     )
+    parser.add_argument(
+        "--enable-diffusion-pipeline-profiler",
+        action="store_true",
+        help="Enable diffusion pipeline profiler to display stage durations.",
+    )
 
     return parser.parse_args()
 

examples/offline_inference/qwen3_omni/end2end_async_chunk.py

@@ -231,52 +231,47 @@ async def run_single_request(
             sampling_params_list=sampling_params_list,
             output_modalities=output_modalities,
         ):
-            if not isinstance(omni_output.request_output, list):
-                outputs_list = [omni_output.request_output]
-            else:
-                outputs_list = omni_output.request_output
+            output = omni_output.request_output
+            if omni_output.final_output_type == "text":
+                if stage_0_first_output_ts is None:
+                    stage_0_first_output_ts = time.perf_counter()
+                text_output = output.outputs[0].text
+                if output.finished:
+                    text_parts.append(text_output)
+            elif omni_output.final_output_type == "audio":
+                mm_out = output.outputs[0].multimodal_output
+                if mm_out and "audio" in mm_out:
+                    if first_audio_ts is None:
+                        first_audio_ts = time.perf_counter()
+                    if audio_sr is None and "sr" in mm_out:
+                        sr_val = mm_out["sr"]
+                        audio_sr = sr_val.item() if hasattr(sr_val, "item") else int(sr_val)
+                        samplerate = audio_sr
+                    audio_data = mm_out["audio"]
+                    if isinstance(audio_data, list):
+                        new_chunks = audio_data[audio_list_consumed:]
+                        audio_list_consumed = len(audio_data)
+                    elif isinstance(audio_data, torch.Tensor):
+                        new_chunks = [audio_data]
+                        audio_last_tensor = audio_data
+                    else:
+                        new_chunks = []
 
-            for output in outputs_list:
-                if omni_output.final_output_type == "text":
-                    if stage_0_first_output_ts is None:
-                        stage_0_first_output_ts = time.perf_counter()
-                    text_output = output.outputs[0].text
-                    if output.finished:
-                        text_parts.append(text_output)
-                elif omni_output.final_output_type == "audio":
-                    mm_out = output.outputs[0].multimodal_output
-                    if mm_out and "audio" in mm_out:
-                        if first_audio_ts is None:
-                            first_audio_ts = time.perf_counter()
-                        if audio_sr is None and "sr" in mm_out:
-                            sr_val = mm_out["sr"]
-                            audio_sr = sr_val.item() if hasattr(sr_val, "item") else int(sr_val)
-                            samplerate = audio_sr
-                        audio_data = mm_out["audio"]
-                        if isinstance(audio_data, list):
-                            new_chunks = audio_data[audio_list_consumed:]
-                            audio_list_consumed = len(audio_data)
-                        elif isinstance(audio_data, torch.Tensor):
-                            new_chunks = [audio_data]
-                            audio_last_tensor = audio_data
-                        else:
-                            new_chunks = []
-
-                        if stream_audio_to_disk and new_chunks:
-                            if sf_writer is None:
-                                sf_writer = sf.SoundFile(
-                                    wav_file,
-                                    mode="w",
-                                    samplerate=samplerate,
-                                    channels=1,
-                                    subtype="FLOAT",
-                                )
-                            for chunk in new_chunks:
-                                chunk_np = chunk.float().detach().cpu().numpy().flatten()
-                                sf_writer.write(chunk_np)
-                                audio_samples_written += len(chunk_np)
-                        else:
-                            audio_chunks.extend(new_chunks)
+                    if stream_audio_to_disk and new_chunks:
+                        if sf_writer is None:
+                            sf_writer = sf.SoundFile(
+                                wav_file,
+                                mode="w",
+                                samplerate=samplerate,
+                                channels=1,
+                                subtype="FLOAT",
+                            )
+                        for chunk in new_chunks:
+                            chunk_np = chunk.float().detach().cpu().numpy().flatten()
+                            sf_writer.write(chunk_np)
+                            audio_samples_written += len(chunk_np)
+                    else:
+                        audio_chunks.extend(new_chunks)
     finally:
         if sf_writer is not None:
             sf_writer.close()
@@ -387,6 +382,7 @@ async def run_all(args):
             stage_configs_path=args.stage_configs_path,
             log_stats=args.log_stats,
             stage_init_timeout=args.stage_init_timeout,
+            enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
         )
 
         # Use default sampling params from stage config (they are pre-configured
@@ -589,6 +585,11 @@ def parse_args():
         default=16000,
         help="Sampling rate for audio loading.",
     )
+    parser.add_argument(
+        "--enable-diffusion-pipeline-profiler",
+        action="store_true",
+        help="Enable diffusion pipeline profiler to display stage durations.",
+    )
     return parser.parse_args()
 
 

examples/offline_inference/qwen3_tts/README.md

@@ -104,7 +104,7 @@ completes. This demonstrates that audio data is available progressively rather t
 
 ## Batched Decoding
 
-The Code2Wav stage (stage 1) supports batched decoding, where multiple requests are decoded in a single forward pass through the SpeechTokenizer. To use it, provide a stage config with `max_batch_size > 1` and pass multiple prompts via `--txt-prompts` with a matching `--batch-size`.
+The Code2Wav stage (stage 1) supports batched decoding, where multiple requests are decoded in a single forward pass through the SpeechTokenizer. To use it, provide a stage config with `max_num_seqs > 1` and pass multiple prompts via `--txt-prompts` with a matching `--batch-size`.
 
 ```
 python end2end.py --query-type CustomVoice \
@@ -113,7 +113,7 @@ python end2end.py --query-type CustomVoice \
     --stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts_batch.yaml
 ```
 
-**Important:** `--batch-size` must match a CUDA graph capture size (1, 2, 4, 8, 16...) because the Talker's code predictor KV cache is sized to `max_num_seqs`, and CUDA graphs pad the batch to the next capture size. Both stages need `max_batch_size >= batch_size` in the stage config for batching to take effect. If only stage 1 has a higher `max_batch_size`, it won't help — stage 1 can only batch chunks from requests that are in-flight simultaneously, which requires stage 0 to also process multiple requests concurrently.
+**Important:** `--batch-size` must match a CUDA graph capture size (1, 2, 4, 8, 16...) because the Talker's code predictor KV cache is sized to `max_num_seqs`, and CUDA graphs pad the batch to the next capture size. Both stages need `max_num_seqs >= batch_size` in the stage config for batching to take effect. If only stage 1 has a higher `max_num_seqs`, it won't help — stage 1 can only batch chunks from requests that are in-flight simultaneously, which requires stage 0 to also process multiple requests concurrently.
 
 ## Notes
 

examples/offline_inference/qwen3_tts/end2end.py

@@ -51,16 +51,68 @@ def _estimate_prompt_len(
 
             tok = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True, padding_side="left")
             cfg = Qwen3TTSConfig.from_pretrained(model_name, trust_remote_code=True)
-            _cache[model_name] = (tok, getattr(cfg, "talker_config", None))
 
-        tok, tcfg = _cache[model_name]
+            # Load speech tokenizer (codec encoder) for exact ref_code_len.
+            speech_tok = None
+            try:
+                import os
+
+                from transformers.utils import cached_file
+
+                from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_tokenizer import Qwen3TTSTokenizer
+
+                st_cfg_path = cached_file(model_name, "speech_tokenizer/config.json")
+                if st_cfg_path:
+                    speech_tok = Qwen3TTSTokenizer.from_pretrained(
+                        os.path.dirname(st_cfg_path), torch_dtype=torch.bfloat16
+                    )
+                    logger.info("Loaded speech tokenizer for exact ref_code_len estimation")
+            except Exception as e:
+                logger.debug("Could not load speech tokenizer: %s", e)
+
+            _cache[model_name] = (tok, getattr(cfg, "talker_config", None), speech_tok)
+
+        tok, tcfg, speech_tok = _cache[model_name]
         task_type = (additional_information.get("task_type") or ["CustomVoice"])[0]
+
+        def _estimate_ref_code_len(ref_audio: object) -> int | None:
+            """Encode ref_audio with the actual codec to get exact frame count."""
+            if not isinstance(ref_audio, (str, list)):
+                return None
+            audio_path = ref_audio[0] if isinstance(ref_audio, list) else ref_audio
+            if not isinstance(audio_path, str) or not audio_path.strip():
+                return None
+            try:
+                from vllm.multimodal.media import MediaConnector
+
+                connector = MediaConnector(allowed_local_media_path="/")
+                audio, sr = connector.fetch_audio(audio_path)
+                import numpy as np
+
+                wav_np = np.asarray(audio, dtype=np.float32)
+
+                if speech_tok is not None:
+                    enc = speech_tok.encode(wav_np, sr=int(sr), return_dict=True)
+                    ref_code = getattr(enc, "audio_codes", None)
+                    if isinstance(ref_code, list):
+                        ref_code = ref_code[0] if ref_code else None
+                    if ref_code is not None and hasattr(ref_code, "shape"):
+                        shape = ref_code.shape
+                        return int(shape[0]) if len(shape) == 2 else int(shape[1]) if len(shape) == 3 else None
+
+                # Fallback: estimate from duration
+                codec_hz = getattr(tcfg, "codec_frame_rate", None) or 12
+                return int(len(audio) / sr * codec_hz)
+            except Exception:
+                return None
+
         return Qwen3TTSTalkerForConditionalGeneration.estimate_prompt_len_from_additional_information(
             additional_information=additional_information,
             task_type=task_type,
             tokenize_prompt=lambda t: tok(t, padding=False)["input_ids"],
             codec_language_id=getattr(tcfg, "codec_language_id", None),
             spk_is_dialect=getattr(tcfg, "spk_is_dialect", None),
+            estimate_ref_code_len=_estimate_ref_code_len,
         )
     except Exception as exc:
         logger.warning("Failed to estimate prompt length, using fallback 2048: %s", exc)
@@ -319,14 +371,15 @@ def main(args):
         stage_configs_path=args.stage_configs_path,
         log_stats=args.log_stats,
         stage_init_timeout=args.stage_init_timeout,
+        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
     )
 
     batch_size = args.batch_size
     for batch_start in range(0, len(inputs), batch_size):
         batch = inputs[batch_start : batch_start + batch_size]
         for stage_outputs in omni.generate(batch):
-            for output in stage_outputs.request_output:
-                _save_wav(output_dir, output.request_id, output.outputs[0].multimodal_output)
+            output = stage_outputs.request_output
+            _save_wav(output_dir, output.request_id, output.outputs[0].multimodal_output)
 
 
 async def main_streaming(args):
@@ -340,6 +393,7 @@ async def main_streaming(args):
         stage_configs_path=args.stage_configs_path,
         log_stats=args.log_stats,
         stage_init_timeout=args.stage_init_timeout,
+        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
     )
 
     for i, prompt in enumerate(inputs):
@@ -481,6 +535,11 @@ def parse_args():
         default=1,
         help="Number of prompts per batch (default: 1, sequential).",
     )
+    parser.add_argument(
+        "--enable-diffusion-pipeline-profiler",
+        action="store_true",
+        help="Enable diffusion pipeline profiler to display stage durations.",
+    )
 
     return parser.parse_args()
 

examples/offline_inference/text_to_audio/text_to_audio.py

@@ -90,6 +90,11 @@ def parse_args() -> argparse.Namespace:
         default=44100,
         help="Sample rate for output audio (Stable Audio uses 44100 Hz).",
     )
+    parser.add_argument(
+        "--enable-diffusion-pipeline-profiler",
+        action="store_true",
+        help="Enable diffusion pipeline profiler to display stage durations.",
+    )
     return parser.parse_args()
 
 
@@ -133,7 +138,10 @@ def main():
     print(f"{'=' * 60}\n")
 
     # Initialize Omni with Stable Audio model
-    omni = Omni(model=args.model)
+    omni = Omni(
+        model=args.model,
+        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
+    )
 
     # Calculate audio end time
     audio_end_in_s = args.audio_start + args.audio_length
@@ -177,7 +185,7 @@ def main():
     output = outputs[0]
     if not hasattr(output, "request_output") or not output.request_output:
         raise ValueError("No request_output found in OmniRequestOutput")
-    request_output = output.request_output[0]
+    request_output = output.request_output
     if not hasattr(request_output, "multimodal_output"):
         raise ValueError("No multimodal_output found in request_output")
 

examples/offline_inference/text_to_image/README.md

@@ -16,7 +16,7 @@ if __name__ == "__main__":
     omni = Omni(model="Qwen/Qwen-Image")
     prompt = "a cup of coffee on the table"
     outputs = omni.generate(prompt)
-    images = outputs[0].request_output[0].images
+    images = outputs[0].request_output.images
     images[0].save("coffee.png")
 ```
 
@@ -34,7 +34,7 @@ if __name__ == "__main__":
     ]
     outputs = omni.generate(prompts)
     for i, output in enumerate(outputs):
-      image = output.request_output[0].images[0].save(f"{i}.jpg")
+      image = output.request_output.images[0].save(f"{i}.jpg")
 ```
 
 !!! info
@@ -46,7 +46,7 @@ if __name__ == "__main__":
 
 !!! info
 
-    For diffusion pipelines, the stage config field `stage_args.[].runtime.max_batch_size` is 1 by default, and the input
+    For diffusion pipelines, the stage config field `stage_args.[].engine_args.max_num_seqs` is 1 by default, and the input
     list is sliced into single-item requests before feeding into the diffusion pipeline. For models that do internally support
     batched inputs, you can [modify this configuration](../../../configuration/stage_configs.md) to let the model accept a longer batch of prompts.
 
@@ -69,7 +69,7 @@ if __name__ == "__main__":
       }
     ])
     for i, output in enumerate(outputs):
-      image = output.request_output[0].images[0].save(f"{i}.jpg")
+      image = output.request_output.images[0].save(f"{i}.jpg")
 ```
 
 ## Local CLI Usage

examples/offline_inference/text_to_image/gradio_demo.py

@@ -112,7 +112,7 @@ def build_demo(args: argparse.Namespace) -> gr.Blocks:
         )
         images_outputs = []
         for output in outputs:
-            req_out = output.request_output[0]
+            req_out = output.request_output
             if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
                 raise ValueError("Invalid request_output structure or missing 'images' key")
             images = req_out.images

examples/offline_inference/text_to_image/text_to_image.py

@@ -14,7 +14,6 @@ from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.lora.request import LoRARequest
 from vllm_omni.lora.utils import stable_lora_int_id
-from vllm_omni.outputs import OmniRequestOutput
 from vllm_omni.platforms import current_omni_platform
 
 
@@ -132,12 +131,10 @@ def parse_args() -> argparse.Namespace:
         "--quantization",
         type=str,
         default=None,
-        choices=["fp8", "gguf"],
-        help=(
-            "Quantization method for the transformer. "
-            "Options: 'fp8' (FP8 W8A8), 'gguf' (GGUF quantized weights). "
-            "Default: None (no quantization, uses BF16)."
-        ),
+        choices=["fp8", "int8", "gguf"],
+        help="Quantization method for the transformer. "
+        "Options: 'fp8' (FP8 W8A8 on Ada/Hopper, weight-only on older GPUs), 'int8' (Int8 W8A8), 'gguf' (GGUF quantized weights). "
+        "Default: None (no quantization, uses BF16).",
     )
     parser.add_argument(
         "--gguf-model",
@@ -218,6 +215,11 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="[NextStep-1.1 only] Apply layer normalization to sampled tokens.",
     )
+    parser.add_argument(
+        "--enable-diffusion-pipeline-profiler",
+        action="store_true",
+        help="Enable diffusion pipeline profiler to display stage durations.",
+    )
     return parser.parse_args()
 
 
@@ -305,6 +307,7 @@ def main():
         "parallel_config": parallel_config,
         "enforce_eager": args.enforce_eager,
         "enable_cpu_offload": args.enable_cpu_offload,
+        "enable_diffusion_pipeline_profiler": args.enable_diffusion_pipeline_profiler,
         **lora_args,
         **quant_kwargs,
     }
@@ -400,20 +403,18 @@ def main():
         else:
             print("[Profiler] No valid profiling data returned.")
 
-    # Extract images from OmniRequestOutput
-    # omni.generate() returns list[OmniRequestOutput], extract images from the first output
+    # omni.generate() returns list[OmniRequestOutput]
     if not outputs or len(outputs) == 0:
         raise ValueError("No output generated from omni.generate()")
     logger.info(f"Outputs: {outputs}")
 
-    # Extract images from request_output[0]['images']
     first_output = outputs[0]
     if not hasattr(first_output, "request_output") or not first_output.request_output:
         raise ValueError("No request_output found in OmniRequestOutput")
 
-    req_out = first_output.request_output[0]
-    if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
-        raise ValueError("Invalid request_output structure or missing 'images' key")
+    req_out = first_output.request_output
+    if not hasattr(req_out, "images"):
+        raise ValueError("Invalid request_output structure or missing 'images'.")
 
     images = req_out.images
     if not images:

examples/offline_inference/text_to_video/text_to_video.py

@@ -132,6 +132,11 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Enable expert parallelism for MoE layers.",
     )
+    parser.add_argument(
+        "--enable-diffusion-pipeline-profiler",
+        action="store_true",
+        help="Enable diffusion pipeline profiler to display stage durations.",
+    )
     return parser.parse_args()
 
 
@@ -184,6 +189,7 @@ def main():
         enforce_eager=args.enforce_eager,
         cache_backend=args.cache_backend,
         cache_config=cache_config,
+        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
     )
 
     if profiler_enabled:
@@ -240,8 +246,6 @@ def main():
             audio = frames.multimodal_output["audio"]
         if frames.is_pipeline_output and frames.request_output is not None:
             inner_output = frames.request_output
-            if isinstance(inner_output, list):
-                inner_output = inner_output[0] if inner_output else None
             if isinstance(inner_output, OmniRequestOutput):
                 if inner_output.multimodal_output and "audio" in inner_output.multimodal_output:
                     audio = inner_output.multimodal_output["audio"]

examples/offline_inference/voxtral_tts/README.md

@@ -0,0 +1,64 @@
+# Voxtral TTS Offline Inference
+
+`end2end.py` runs Voxtral TTS end-to-end offline inference using vLLM. It supports both blocking (`Omni`) and streaming (`AsyncOmni`) generation, batched prompts with configurable concurrency, and voice selection via preset name or reference audio file.
+
+When `mistral_common` has `SpeechRequest` support, prompt token IDs are built via `encode_speech_request`. Otherwise, the script falls back to manual token construction.
+
+## Usage Examples
+
+
+```bash
+# Basic single-prompt with cheerful_female voice preset
+python3 examples/offline_inference/voxtral_tts/end2end.py \
+    --stage-configs-path vllm_omni/model_executor/stage_configs/voxtral_tts.yaml \
+    --write-audio --voice cheerful_female \
+    --model mistralai/Voxtral-4B-TTS-2603 \
+    --text "That eerie silence after the first storm was just the calm before another round of chaos, wasn't it?"
+
+# 32 replicate prompts with cheerful_female voice preset
+python3 examples/offline_inference/voxtral_tts/end2end.py \
+    --stage-configs-path vllm_omni/model_executor/stage_configs/voxtral_tts.yaml \
+    --num-prompts 32 --write-audio --voice cheerful_female \
+    --model mistralai/Voxtral-4B-TTS-2603 \
+    --text "That eerie silence after the first storm was just the calm before another round of chaos, wasn't it?"
+
+# Streaming with neutral_female voice preset
+python3 examples/offline_inference/voxtral_tts/end2end.py \
+    --stage-configs-path vllm_omni/model_executor/stage_configs/voxtral_tts.yaml \
+    --streaming --write-audio --voice neutral_female \
+    --model mistralai/Voxtral-4B-TTS-2603 \
+    --text "That eerie silence after the first storm was just the calm before another round of chaos, wasn't it?"
+
+# 32 prompts, 8 concurrent requests per wave, streaming with neutral_female voice
+python3 examples/offline_inference/voxtral_tts/end2end.py \
+    --stage-configs-path vllm_omni/model_executor/stage_configs/voxtral_tts.yaml \
+    --num-prompts 32 --concurrency 8 --streaming --write-audio --voice neutral_female \
+    --model mistralai/Voxtral-4B-TTS-2603 \
+    --text "That eerie silence after the first storm was just the calm before another round of chaos, wasn't it?"
+
+# Short debug prompt with reference audio
+# Note: Reference audio capability is not yet released.
+python3 examples/offline_inference/voxtral_tts/end2end.py \
+    --stage-configs-path vllm_omni/model_executor/stage_configs/voxtral_tts.yaml \
+    --write-audio \
+    --model mistralai/Voxtral-4B-TTS-2603 \
+    --text "This is a test message." \
+    --audio-path path/to/reference_audio.wav
+```
+
+## Arguments
+
+| Argument | Description |
+|---|---|
+| `--model PATH` | HuggingFace repo ID or local directory path (default: `mistralai/Voxtral-4B-TTS-2603`) |
+| `--text TEXT` | Text to synthesize (default: `"This is a test message."`) |
+| `--audio-path PATH` | Path to reference audio file for voice cloning |
+| `--output-dir DIR` | Directory to write output WAV files (default: `output_audio`) |
+| `--stage-configs-path PATH` | Path to stage configs YAML (currently it must be set for VoxtralTTS) |
+| `--num-prompts N` | Number of replicate prompts to run for measuring performance (default: 1) |
+| `--streaming` | Use streaming generation via `AsyncOmni` (default: blocking `Omni`) |
+| `--concurrency N` | Max concurrent requests per wave (must be used with `--streaming`, must evenly divide `--num-prompts`) |
+| `--voice NAME` | Voice preset to use instead of reference audio file (e.g., casual_female, casual_male, cheerful_female, neutral_female, neutral_male) |
+| `--write-audio` | Write generated audio to WAV files |
+| `--profiling-mode` | Enable profiling mode (reduces max tokens to 50) |
+| `--log-stats` | Enable detailed statistics logging |

examples/offline_inference/voxtral_tts/end2end.py

@@ -0,0 +1,377 @@
+"""
+This example shows how to use vLLM for running Voxtral TTS
+"""
+
+import asyncio
+import gc
+import logging
+import os
+import time
+import uuid
+from argparse import Namespace
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import soundfile as sf
+import torch
+from mistral_common.protocol.instruct.chunk import TextChunk
+
+try:
+    from mistral_common.protocol.speech.request import SpeechRequest
+    from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
+except ImportError:
+    raise ImportError(
+        "Could not import SpeechRequest or MistralTokenizer from mistral_common. "
+        "Please pull the latest mistral-common code and install it with: "
+        "pip install -e ."
+    )
+from vllm import SamplingParams
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+from vllm_omni import AsyncOmni
+from vllm_omni.entrypoints.omni import Omni
+
+logger = logging.getLogger(__name__)
+
+
+# ---- Streaming version ----
+async def run_streaming(inputs, sampling_params_list, model_name, args, output_dir):
+    async_omni = AsyncOmni(
+        model=model_name,
+        stage_configs_path=args.stage_configs_path,
+        log_stats=args.log_stats,
+    )
+
+    # Normalize to a list so batch and single-input share the same code path
+    if isinstance(inputs, list):
+        inputs_list = inputs
+    else:
+        inputs_list = [inputs]
+
+    total_audio_dur = 0.0
+    total_gen_time = 0.0
+    total_ttfa = 0.0
+    ttfa_count = 0
+    total_waits = 0
+    total_chunks = 0
+    results_lock = asyncio.Lock()
+
+    async def _generate_one(batch_idx, single_input):
+        nonlocal total_audio_dur, total_gen_time, total_ttfa, ttfa_count
+        nonlocal total_waits, total_chunks
+        request_id = str(uuid.uuid4())
+        all_audio_chunks = []
+        chunk_arrival_times = []
+        chunk_durations = []
+        chunk_idx = 0
+        gen_start = time.time()
+        ttfa = None
+        accumulated_sample = 0
+
+        async for stage_output in async_omni.generate(single_input, request_id, sampling_params_list):
+            mm_output = stage_output.multimodal_output
+            finished = stage_output.finished
+            if not mm_output or "audio" not in mm_output:
+                continue
+
+            now = time.time()
+            if ttfa is None:
+                ttfa = now - gen_start
+
+            audio_chunk = mm_output["audio"]
+            # audio_chunk is a 1-D tensor from decode_helper_async
+            if isinstance(audio_chunk, torch.Tensor):
+                if finished:
+                    # Last chunk may return whole audio instead of chunk delta.
+                    # Cut accumulated samples from previous chunks.
+                    audio_numpy = audio_chunk[accumulated_sample:].float().detach().cpu().numpy()
+                else:
+                    audio_numpy = audio_chunk.float().detach().cpu().numpy()
+            elif isinstance(audio_chunk, list):
+                # Audio_chunk list contain all previous chunks. Use chunk_idx to index
+                audio_numpy = audio_chunk[chunk_idx].float().detach().cpu().numpy()
+            else:
+                audio_numpy = audio_chunk
+
+            accumulated_sample += len(audio_numpy)
+            all_audio_chunks.append(audio_numpy)
+            chunk_arrival_times.append(now)
+            chunk_durations.append(len(audio_numpy) / 24000)
+            chunk_idx += 1
+
+        gen_elapsed = time.time() - gen_start
+
+        # Analyze wait/no-wait per chunk:
+        # A "wait" means the client has finished playing all previous audio
+        # before this chunk arrived, so there is a playback gap.
+        # buffer_time > 0: client still has audio to play (no wait)
+        # buffer_time <= 0: client ran out of audio and had to wait
+        chunk_labels = []
+        accumulated_audio_dur = 0.0
+        if chunk_arrival_times:
+            first_arrival = chunk_arrival_times[0]
+            for i in range(len(chunk_arrival_times)):
+                if i == 0:
+                    chunk_labels.append("no_wait")
+                else:
+                    playback_elapsed = chunk_arrival_times[i] - first_arrival
+                    buffer_time = accumulated_audio_dur - playback_elapsed
+                    if buffer_time <= 0:
+                        chunk_labels.append("wait")
+                    else:
+                        chunk_labels.append("no_wait")
+                accumulated_audio_dur += chunk_durations[i]
+
+        req_wait_count = sum(1 for chunk_label in chunk_labels if chunk_label == "wait")
+        req_wait_rate = req_wait_count / len(chunk_labels) if chunk_labels else 0.0
+
+        # Concatenate all chunks for this request's audio
+        if all_audio_chunks:
+            full_audio = np.concatenate(all_audio_chunks)
+            output_audio_dur = len(full_audio) / 24000
+            output_path = os.path.join(output_dir, f"tts_output_{batch_idx}.wav")
+            if args.write_audio:
+                sf.write(output_path, full_audio, 24000)
+                print(
+                    f"Request {batch_idx}: saved {len(full_audio)} samples ({output_audio_dur:.2f}s) to {output_path}"
+                )
+            # Per-chunk wait details
+            for i, label in enumerate(chunk_labels):
+                dur_ms = chunk_durations[i] * 1000
+                arrival_ms = (chunk_arrival_times[i] - first_arrival) * 1000 if i > 0 else 0.0
+                print(
+                    f"  Request {batch_idx} chunk {i}: {label} | arrived={arrival_ms:.1f}ms | chunk_dur={dur_ms:.1f}ms"
+                )
+            print(
+                f"Request {batch_idx}: TTFA={ttfa:.4f}s | "
+                f"Generation={gen_elapsed:.4f}s | "
+                f"Audio={output_audio_dur:.2f}s | "
+                f"RTF={output_audio_dur / gen_elapsed:.4f} | "
+                f"WaitRate={req_wait_rate:.2%} ({req_wait_count}/{len(chunk_labels)})"
+            )
+            async with results_lock:
+                total_audio_dur += output_audio_dur
+                total_gen_time += gen_elapsed
+                total_waits += req_wait_count
+                total_chunks += len(chunk_labels)
+                if ttfa is not None:
+                    total_ttfa += ttfa
+                    ttfa_count += 1
+
+    # Launch requests in waves of `concurrency` size
+    concurrency = args.concurrency or len(inputs_list)
+    gen_start_all = time.time()
+    for wave_start in range(0, len(inputs_list), concurrency):
+        wave = inputs_list[wave_start : wave_start + concurrency]
+        wave_idx_offset = wave_start
+        print(f"\n--- Wave {wave_start // concurrency + 1} (requests {wave_start}-{wave_start + len(wave) - 1}) ---")
+        await asyncio.gather(*[_generate_one(wave_idx_offset + i, inp) for i, inp in enumerate(wave)])
+    generation_time = time.time() - gen_start_all
+    avg_ttfa = total_ttfa / ttfa_count if ttfa_count else float("nan")
+    overall_wait_rate = total_waits / total_chunks if total_chunks else float("nan")
+    print(
+        f"\nAll requests: Generation={generation_time:.4f}s | "
+        f"TotalAudio={total_audio_dur:.2f}s | "
+        f"Concurrency={concurrency} | "
+        f"AvgTTFA={avg_ttfa:.4f}s | "
+        f"RTF(total)={total_audio_dur / generation_time:.4f} | "
+        f"RTF(per-request)={total_audio_dur / total_gen_time:.4f} | "
+        f"WaitRate={overall_wait_rate:.2%} ({total_waits}/{total_chunks})"
+    )
+
+    async_omni.shutdown()
+    torch.cuda.empty_cache()
+    gc.collect()
+
+
+# ---- Non-streaming version ----
+def run_non_streaming(inputs, sampling_params_list, model_name, args, output_dir):
+    llm = Omni(
+        model=model_name,
+        log_stats=args.log_stats,
+        stage_configs_path=args.stage_configs_path,
+    )
+
+    if args.profiling_mode:
+        llm.start_profile()
+
+    vllm_start = time.time()
+    outputs = llm.generate(inputs, sampling_params_list)
+    vllm_elapsed = time.time() - vllm_start
+
+    if args.profiling_mode:
+        llm.stop_profile()
+        time.sleep(10)
+
+    print(f"vLLM run time: {vllm_elapsed:.4f}s")
+    output_audio_dur = 0.0
+
+    for batch_idx, o in enumerate(outputs):
+        audio_tensor = torch.cat(o.multimodal_output["audio"])
+        audio_array = audio_tensor.tolist()
+        output_audio_dur += float(len(audio_array)) / 24000
+        if args.write_audio:
+            output_path = os.path.join(output_dir, f"tts_output_{batch_idx}.wav")
+            sf.write(output_path, audio_array, 24000)
+            print(f"Audio saved to {output_path}")
+
+    print(f"Total audio duration: {output_audio_dur:.2f}s")
+    print(f"RTF: {output_audio_dur / vllm_elapsed:.4f}")
+
+    del llm
+    torch.cuda.empty_cache()
+    gc.collect()
+
+
+def parse_args() -> Namespace:
+    parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with Voxtral TTS")
+    parser.add_argument(
+        "--model",
+        type=str,
+        default="mistralai/Voxtral-4B-TTS-2603",
+        help="Model name or path.",
+    )
+    parser.add_argument(
+        "--text",
+        type=str,
+        default="This is a test message.",
+        help="Text to synthesize.",
+    )
+    parser.add_argument(
+        "--audio-path",
+        type=str,
+        default=None,
+        help="Path to reference audio file for voice cloning.",
+    )
+    parser.add_argument(
+        "--output-dir",
+        type=str,
+        default="output_audio",
+        help="Directory to write output wav files.",
+    )
+    parser.add_argument(
+        "--stage-configs-path",
+        type=str,
+        default=None,
+        help="Path to stage configs YAML. Auto-resolved from model if not set.",
+    )
+    parser.add_argument(
+        "--num-prompts", type=int, default=1, help="Number of replicate prompts to run for measuring performance"
+    )
+    parser.add_argument(
+        "--profiling-mode",
+        default=False,
+        action="store_true",
+        help="Set max_num_tokens=2 to reduce profiling time",
+    )
+    parser.add_argument(
+        "--log-stats",
+        action="store_true",
+        default=False,
+        help="Enable writing detailed statistics (default: disabled)",
+    )
+    parser.add_argument(
+        "--write-audio",
+        action="store_true",
+        default=False,
+        help="Write audio output to files.",
+    )
+    parser.add_argument(
+        "--streaming",
+        action="store_true",
+        default=False,
+        help="Use streaming generation via AsyncOmni instead of blocking Omni.",
+    )
+    parser.add_argument(
+        "--concurrency",
+        type=int,
+        default=None,
+        help="Max concurrent requests per wave (default: all at once). Must evenly divide --num-prompts.",
+    )
+    parser.add_argument(
+        "--voice",
+        type=str,
+        default=None,
+        help="Voice to use instead of audio file.",
+    )
+    return parser.parse_args()
+
+
+def compose_request(
+    model_name: str,
+    text_chunk: TextChunk,
+    audio_prompt_file: str,
+    args: Any,
+) -> dict:
+    """Build the full TTS input dict (prompt_token_ids, multi_modal_data or additional_information)."""
+    inputs: dict[str, Any] = {}
+    if Path(model_name).is_dir():
+        mistral_tokenizer = MistralTokenizer.from_file(str(Path(model_name) / "tekken.json"))
+    else:
+        mistral_tokenizer = MistralTokenizer.from_hf_hub(model_name)
+    instruct_tokenizer = mistral_tokenizer.instruct_tokenizer
+
+    if args.voice is not None:
+        tokenized = instruct_tokenizer.encode_speech_request(SpeechRequest(input=text_chunk.text, voice=args.voice))
+        inputs["additional_information"] = {"voice": [args.voice]}
+        inputs["prompt_token_ids"] = tokenized.tokens
+    else:
+        with open(audio_prompt_file, "rb") as f:
+            ref_audio_bytes = f.read()
+        tokenized = instruct_tokenizer.encode_speech_request(
+            SpeechRequest(input=text_chunk.text, ref_audio=ref_audio_bytes)
+        )
+        audio = tokenized.audios[0]
+        inputs["multi_modal_data"] = {"audio": [(audio.audio_array, audio.sampling_rate)]}
+        inputs["prompt_token_ids"] = tokenized.tokens
+
+    return inputs
+
+
+def main(args: Any) -> None:
+    max_num_tokens = 50 if args.profiling_mode else 2500
+
+    model_name = args.model
+    output_dir = args.output_dir
+
+    if args.voice is None and args.audio_path is None:
+        raise ValueError("Either --voice or --audio-path must be provided.")
+
+    audio_prompt_file = args.audio_path
+    text_chunk = TextChunk(text=args.text)
+
+    if args.write_audio:
+        os.makedirs(output_dir, exist_ok=True)
+
+    inputs = compose_request(model_name, text_chunk, audio_prompt_file, args)
+
+    sampling_params = SamplingParams(
+        max_tokens=max_num_tokens,
+    )
+    sampling_params_list = [
+        sampling_params,
+        sampling_params,
+    ]
+
+    if args.num_prompts > 1:
+        inputs = [inputs] * args.num_prompts
+
+    if args.concurrency is not None:
+        assert args.streaming, "--concurrency must be used with --streaming on AsyncOmni"
+        assert args.num_prompts % args.concurrency == 0, (
+            f"--num-prompts ({args.num_prompts}) must be divisible by --concurrency ({args.concurrency})"
+        )
+
+    torch.cuda.empty_cache()
+    gc.collect()
+
+    if args.streaming:
+        asyncio.run(run_streaming(inputs, sampling_params_list, model_name, args, output_dir))
+    else:
+        run_non_streaming(inputs, sampling_params_list, model_name, args, output_dir)
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    main(args)

examples/online_serving/qwen3_omni/qwen3_omni_moe_thinking.yaml

@@ -9,9 +9,9 @@ stage_args:
   - stage_id: 0
     runtime:
       devices: "0,1"
-      max_batch_size: 1
     engine_args:
       model_stage: thinker
+      max_num_seqs: 1
       model_arch: Qwen3OmniMoeForConditionalGeneration
       worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler

examples/online_serving/voxtral_tts/gradio_demo.py

@@ -0,0 +1,269 @@
+"""
+- Make sure to install the following for this example to function correctly:
+- `pip install -e .`
+- `pip install gradio==5.50 mistral_common=1.10.0`
+
+Example use case:
+
+python examples/online_serving/voxtral_tts/gradio_demo.py --host slurm-199-077 --port 8000
+
+"""
+
+import argparse
+import io
+import json
+import logging
+import socket
+import uuid
+from datetime import datetime
+from pathlib import Path
+from typing import Any
+
+import gradio as gr
+import httpx
+import numpy as np
+import soundfile as sf
+
+logger = logging.getLogger()
+
+LOGFORMAT = "%(asctime)s - %(levelname)s - %(message)s"
+TIMEFORMAT = "%Y-%m-%d %H:%M:%S"
+
+logging.basicConfig(level=logging.INFO, force=True, format=LOGFORMAT, datefmt=TIMEFORMAT)
+logger.setLevel(logging.INFO)
+
+
+_BASE_URL = f"{socket.gethostname()}:7860"
+
+
+def run_inference(
+    voice_name: str,
+    text_prompt: str,
+    base_url: str,
+    model: str,
+) -> tuple[int, np.ndarray]:
+    """Call /v1/audio/speech and return (sample_rate, audio_array)."""
+    text_prompt = text_prompt.strip()
+    if not text_prompt:
+        raise gr.Error("Please enter a text prompt.")
+
+    payload: dict[str, Any] = {
+        "input": text_prompt,
+        "model": model,
+        "response_format": "wav",
+        "voice": voice_name,
+    }
+
+    response = httpx.post(
+        f"{base_url}/audio/speech",
+        json=payload,
+        timeout=120.0,
+    )
+    response.raise_for_status()
+
+    audio_array, sr = sf.read(io.BytesIO(response.content), dtype="float32")
+    return sr, audio_array
+
+
+def _save_example(
+    outputs_dir: Path,
+    voice_name: str,
+    text_prompt: str,
+    sr: int,
+    audio_array: np.ndarray,
+) -> tuple[str, str]:
+    """
+    Save inputs/outputs for sharing.
+    Returns (share_id, saved_audio_path)
+    """
+    share_id = uuid.uuid4().hex
+
+    # Save generated audio
+    saved_audio_path = outputs_dir / f"{share_id}_gen.wav"
+    sf.write(str(saved_audio_path), audio_array, sr)
+
+    meta = {
+        "id": share_id,
+        "created_at": datetime.utcnow().isoformat(),
+        "voice_name": voice_name,
+        "text_prompt": text_prompt,
+        "generated_audio_path": str(saved_audio_path),
+    }
+
+    with open(outputs_dir / f"{share_id}.json", "w") as f:
+        json.dump(meta, f, ensure_ascii=False)
+
+    return share_id, str(saved_audio_path)
+
+
+def _load_from_share(
+    outputs_dir: Path | None,
+    request: gr.Request,
+) -> tuple[str | None, str, str | None, dict[str, Any], str]:
+    """
+    Called on page load. If ?share_id=... is present, load stored example.
+    Returns: (voice_name, text_prompt, output_audio, submit_btn_update, share_link_text)
+    """
+    if outputs_dir is None:
+        return "Neutral_Male", "", None, gr.update(interactive=False), ""
+
+    share_id = None
+    if request and request.query_params:
+        share_id = request.query_params.get("share_id")
+
+    if not share_id:
+        return "Neutral_Male", "", None, gr.update(interactive=False), ""
+
+    meta_path = outputs_dir / f"{share_id}.json"
+    if not meta_path.exists():
+        logger.warning("No stored example for share_id=%s", share_id)
+        return "Neutral_Male", "", None, gr.update(interactive=False), ""
+
+    with open(meta_path) as f:
+        meta = json.load(f)
+
+    voice = meta.get("voice_name", "Neutral_Male")
+    text_prompt = meta.get("text_prompt", "")
+    gen_path = meta.get("generated_audio_path")
+
+    if _BASE_URL:
+        share_link = f"http://{_BASE_URL}?share_id={share_id}"
+    else:
+        share_link = f"Use this query on your current URL: ?share_id={share_id}"
+
+    return voice, text_prompt, gen_path, gr.update(interactive=True), share_link
+
+
+def main(
+    model: str,
+    host: str,
+    port: str,
+    output_dir: str | None = None,
+) -> None:
+    base_url = f"http://{host}:{port}/v1"
+    logger.info(f"Using speech API at: {base_url}/audio/speech")
+
+    outputs_dir: Path | None = None
+    if output_dir is not None:
+        outputs_dir = Path(output_dir)
+        outputs_dir.mkdir(parents=True, exist_ok=True)
+
+    with gr.Blocks(title="Voxtral TTS", fill_height=True) as demo:
+        gr.Markdown("## Voxtral TTS")
+
+        with gr.Row():
+            with gr.Column():
+                voice_name = gr.Dropdown(
+                    ["Neutral_Male", "Neutral_Female", "Cheerful_Female", "Casual_Male", "Casual_Female"],
+                    label="Voice",
+                    value="Neutral_Male",
+                )
+                text_prompt = gr.Textbox(
+                    label="Text prompt",
+                    placeholder="Enter the text you want to synthesize...",
+                    lines=4,
+                )
+                with gr.Row():
+                    reset_btn = gr.Button("Clear")
+                    submit_btn = gr.Button("Generate audio", interactive=False)
+
+            with gr.Column():
+                output_audio = gr.Audio(
+                    label="Generated audio",
+                    show_download_button=True,
+                    interactive=False,
+                    autoplay=True,
+                    type="filepath",
+                )
+                share_link_box = gr.Textbox(
+                    label="Shareable link",
+                    interactive=False,
+                    show_copy_button=True,
+                    visible=outputs_dir is not None,
+                )
+
+        # --- UI logic: disable submit until text is non-empty ---
+        def _toggle_submit(text: str):
+            enabled = bool(text.strip())
+            return gr.update(interactive=enabled)
+
+        text_prompt.change(
+            fn=_toggle_submit,
+            inputs=[text_prompt],
+            outputs=submit_btn,
+        )
+
+        # --- Wiring inference + persistence to the button ---
+        def _on_submit(voice: str, text: str):
+            assert text.strip() != ""
+            sr, audio_array = run_inference(voice, text, base_url, model)
+            if outputs_dir is not None:
+                share_id, saved_audio_path = _save_example(
+                    outputs_dir,
+                    voice_name=voice,
+                    text_prompt=text,
+                    sr=sr,
+                    audio_array=audio_array,
+                )
+                share_link = f"{_BASE_URL}?share_id={share_id}"
+                return saved_audio_path, share_link
+            return (sr, audio_array), ""
+
+        submit_btn.click(
+            fn=_on_submit,
+            inputs=[voice_name, text_prompt],
+            outputs=[output_audio, share_link_box],
+        )
+
+        # --- Clear everything and disable submit again ---
+        def _on_reset():
+            return (
+                "Neutral_Male",  # voice_name
+                "",  # text_prompt
+                None,  # output_audio
+                gr.update(interactive=False),  # submit_btn
+                "",  # share_link_box
+            )
+
+        reset_btn.click(
+            fn=_on_reset,
+            inputs=[],
+            outputs=[voice_name, text_prompt, output_audio, submit_btn, share_link_box],
+        )
+
+        demo.load(
+            fn=lambda request: _load_from_share(outputs_dir, request),
+            inputs=[],
+            outputs=[voice_name, text_prompt, output_audio, submit_btn, share_link_box],
+        )
+
+    launch_kwargs: dict[str, Any] = {
+        "server_name": "0.0.0.0",
+        "share": True,
+    }
+    if outputs_dir is not None:
+        launch_kwargs["allowed_paths"] = [str(outputs_dir)]
+    demo.launch(**launch_kwargs)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Voxtral TTS Gradio Demo")
+    parser.add_argument("--model", type=str, default="mistralai/Voxtral-4B-TTS-2603", help="Name of model repo on HF")
+    parser.add_argument("--host", type=str, default="localhost", help="Name of host")
+    parser.add_argument("--port", type=str, default="8091", help="port number")
+    parser.add_argument(
+        "--output-dir",
+        type=str,
+        default=None,
+        help="Directory to save generated audio and share links. "
+        "If not provided, save/share functionality is disabled.",
+    )
+
+    args = parser.parse_args()
+
+    main(
+        model=args.model,
+        host=args.host,
+        port=args.port,
+        output_dir=args.output_dir,
+    )

requirements/common.txt

@@ -17,3 +17,4 @@ einops>=0.8.1
 prettytable>=3.8.0
 aenum==3.1.16
 pyzmq>=25.0.0
+janus>=1.0.0

tests/benchmarks/test_serve_cli.py

@@ -1,38 +1,101 @@
+import json
+import os
 import subprocess
+import textwrap
 from pathlib import Path
 
 import pytest
 
-from tests.conftest import OmniServerParams
-from tests.utils import hardware_test
-from vllm_omni.platforms import current_omni_platform
-
-models = ["Qwen/Qwen2.5-Omni-7B"]
-stage_configs = [str(Path(__file__).parent.parent / "e2e" / "stage_configs" / "qwen2_5_omni_ci.yaml")]
-
-if current_omni_platform.is_xpu():
-    stage_configs = [str(Path(__file__).parent.parent / "e2e" / "stage_configs" / "xpu" / "qwen2_5_omni_ci.yaml")]
-
-# Create parameter combinations for model and stage config
-test_params = [
-    OmniServerParams(model=model, stage_config_path=stage_config) for model in models for stage_config in stage_configs
-]
-
 
 @pytest.mark.core_model
 @pytest.mark.benchmark
-@hardware_test(res={"cuda": "L4", "xpu": "B60"}, num_cards=3)
-@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_bench_serve_chat(omni_server):
-    command = [
+@pytest.mark.cpu
+def test_bench_serve_cli_mocks_http_request(tmp_path: Path):
+    num_prompts = 5
+    port = 18000
+    result_filename = "bench-result.json"
+    calls_filename = "http-post-calls.json"
+    result_path = tmp_path / result_filename
+    calls_path = tmp_path / calls_filename
+
+    sitecustomize_path = tmp_path / "sitecustomize.py"
+    sitecustomize_path.write_text(
+        textwrap.dedent(
+            """
+            import atexit
+            import json
+            import os
+
+            POST_CALLS = []
+            SSE_CHUNKS = [
+                b'data: {"choices":[{"delta":{"content":"hi"}}],"modality":"text"}\\n\\n',
+                b'data: {"choices":[{"delta":{"content":" there"}}],"modality":"text","metrics":{"num_tokens_out":4,"num_tokens_in":5}}\\n\\n',
+                b"data: [DONE]\\n\\n",
+            ]
+
+            class _Content:
+                def __init__(self, chunks):
+                    self._chunks = chunks
+
+                async def iter_any(self):
+                    for chunk in self._chunks:
+                        yield chunk
+
+            class MockResponse:
+                def __init__(self):
+                    self.status = 200
+                    self.reason = "OK"
+                    self.content = _Content(SSE_CHUNKS)
+
+                async def __aenter__(self):
+                    return self
+
+                async def __aexit__(self, exc_type, exc, tb):
+                    return False
+
+            class MockClientSession:
+                def __init__(self, *args, **kwargs):
+                    pass
+
+                def post(self, url=None, *args, **kwargs):
+                    if url is not None:
+                        POST_CALLS.append(url)
+                    return MockResponse()
+
+                async def close(self):
+                    return None
+
+            import vllm_omni.benchmarks.patch.patch as patch_mod
+
+            patch_mod.aiohttp.ClientSession = MockClientSession
+            patch_mod.aiohttp.TCPConnector = lambda *args, **kwargs: object()
+
+            calls_file = os.environ.get("VLLM_OMNI_TEST_POST_CALLS_FILE")
+
+            if calls_file:
+                def _write_calls():
+                    with open(calls_file, "w", encoding="utf-8") as f:
+                        json.dump({"requested_urls": POST_CALLS}, f)
+
+                atexit.register(_write_calls)
+            """
+        ),
+        encoding="utf-8",
+    )
+
+    env = os.environ.copy()
+    env["PYTHONPATH"] = str(tmp_path) + os.pathsep + env.get("PYTHONPATH", "")
+    env["VLLM_OMNI_TEST_POST_CALLS_FILE"] = str(calls_path)
+
+    cmd = [
         "vllm",
         "bench",
         "serve",
         "--omni",
         "--model",
-        omni_server.model,
+        "Qwen/Qwen2.5-Omni-7B",
         "--port",
-        str(omni_server.port),
+        str(port),
         "--dataset-name",
         "random",
         "--random-input-len",
@@ -40,14 +103,41 @@ def test_bench_serve_chat(omni_server):
         "--random-output-len",
         "4",
         "--num-prompts",
-        "5",
+        str(num_prompts),
         "--endpoint",
         "/v1/chat/completions",
         "--backend",
         "openai-chat-omni",
+        "--disable-tqdm",
+        "--num-warmups",
+        "0",
+        "--ready-check-timeout-sec",
+        "0",
+        "--save-result",
+        "--result-dir",
+        str(tmp_path),
+        "--result-filename",
+        result_filename,
     ]
-    result = subprocess.run(command, capture_output=True, text=True)
-    print(result.stdout)
-    print(result.stderr)
+    proc = subprocess.run(
+        cmd,
+        cwd=str(Path(__file__).resolve().parents[2]),
+        env=env,
+        capture_output=True,
+        text=True,
+    )
 
-    assert result.returncode == 0, f"Benchmark failed: {result.stderr}"
+    assert proc.returncode == 0, f"CLI failed: stdout={proc.stdout}\nstderr={proc.stderr}"
+    print(f"CLI output: {proc.stdout}")
+    assert result_path.exists()
+    result = json.loads(result_path.read_text(encoding="utf-8"))
+    assert calls_path.exists()
+    calls = json.loads(calls_path.read_text(encoding="utf-8"))
+
+    expected_url = f"http://127.0.0.1:{port}/v1/chat/completions"
+    requested_urls = calls["requested_urls"]
+    sent_requests = len(requested_urls)
+
+    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}"

tests/comfyui/test_comfyui_integration.py

@@ -11,6 +11,7 @@ import time
 import traceback
 from collections.abc import Iterable, Sequence
 from enum import StrEnum, auto
+from types import SimpleNamespace
 from typing import Any, NamedTuple
 from unittest.mock import AsyncMock, MagicMock, patch
 
@@ -31,6 +32,7 @@ from vllm import SamplingParams
 from vllm.outputs import CompletionOutput, RequestOutput
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 
+from vllm_omni.entrypoints.async_omni import AsyncOmni as RealAsyncOmni
 from vllm_omni.entrypoints.cli.serve import OmniServeCommand
 from vllm_omni.inputs.data import OmniSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
@@ -43,7 +45,7 @@ class ServerCase(NamedTuple):
 
     served_model: str
     stage_list: list
-    stage_configs: list[dict]
+    stage_configs: list[Any]
     outputs: list[OmniRequestOutput]
 
 
@@ -262,6 +264,40 @@ def _assert_model_param_values(received: OmniSamplingParams, expected: dict):
         )
 
 
+def _make_stage_config(
+    stage_type: str,
+    *,
+    is_comprehension: bool = False,
+    model_stage: str | None = None,
+):
+    engine_args = SimpleNamespace()
+    if model_stage is not None:
+        engine_args.model_stage = model_stage
+    return SimpleNamespace(
+        stage_type=stage_type,
+        is_comprehension=is_comprehension,
+        engine_args=engine_args,
+    )
+
+
+def _stage_type(stage: Any) -> str | None:
+    return getattr(stage, "stage_type", None)
+
+
+def _build_output_modalities(stage_configs: list[Any]) -> list[str]:
+    modalities: list[str] = []
+    for stage in stage_configs:
+        final_output = getattr(stage, "final_output", False)
+        final_output_type = getattr(stage, "final_output_type", None)
+        if final_output and isinstance(final_output_type, str):
+            modalities.append(final_output_type)
+    if modalities:
+        return modalities
+    if any(_stage_type(stage) == "diffusion" for stage in stage_configs):
+        return ["image"]
+    return ["text", "audio"]
+
+
 def _build_mock_outputs(outputs: Iterable[OmniRequestOutput], sampling_case: SamplingCase, server_case: ServerCase):
     async def _mock_generate(*args, **kwargs):
         received_sampling_params_list: Sequence[OmniSamplingParams] | None = (
@@ -367,18 +403,28 @@ def mock_async_omni(server_case: ServerCase, sampling_case: SamplingCase):
             new=_mock_preprocess_chat,
         ),
     ):
-        mock_instance = AsyncMock()
+        mock_instance = AsyncMock(spec=RealAsyncOmni)
         mock_instance.generate = _build_mock_outputs(server_case.outputs, sampling_case, server_case)
 
         mock_instance.stage_list = server_case.stage_list
         mock_instance.stage_configs = server_case.stage_configs
+        mock_instance.output_modalities = _build_output_modalities(server_case.stage_configs)
         mock_instance.default_sampling_params_list = [
-            SamplingParams() if stage.get("stage_type") != "diffusion" else MagicMock()
+            SamplingParams() if _stage_type(stage) != "diffusion" else MagicMock()
             for stage in server_case.stage_configs
         ]
         mock_instance.errored = False
         mock_instance.dead_error = RuntimeError("Mock engine error")
-        mock_instance.model_config = MagicMock(max_model_len=4096, io_processor_plugin=None)
+        mock_instance.model_config = MagicMock(
+            max_model_len=4096,
+            io_processor_plugin=None,
+            allowed_local_media_path=None,
+            allowed_media_domains=None,
+        )
+        # Mimic Qwen3-TTS talker speaker config so CustomVoice validation passes.
+        mock_instance.model_config.hf_config = MagicMock()
+        mock_instance.model_config.hf_config.talker_config = MagicMock()
+        mock_instance.model_config.hf_config.talker_config.speaker_id = {"Vivian": 0}
         mock_instance.io_processor = MagicMock()
         mock_instance.input_processor = MagicMock()
         mock_instance.shutdown = MagicMock()
@@ -444,7 +490,7 @@ def api_server(unused_tcp_port_factory, server_case: ServerCase, mock_async_omni
             ServerCase(
                 served_model="Tongyi-MAI/Z-Image-Turbo",
                 stage_list=["diffusion"],
-                stage_configs=[{"stage_type": "diffusion"}],
+                stage_configs=[_make_stage_config("diffusion")],
                 outputs=[_build_diffusion_image_output_for_images_endpoint()],
             ),
             "Tongyi-MAI/Z-Image-Turbo",
@@ -455,7 +501,7 @@ def api_server(unused_tcp_port_factory, server_case: ServerCase, mock_async_omni
             ServerCase(
                 served_model="ByteDance-Seed/BAGEL-7B-MoT",
                 stage_list=["diffusion"],
-                stage_configs=[{"stage_type": "diffusion"}],
+                stage_configs=[_make_stage_config("diffusion")],
                 outputs=[_build_diffusion_image_output_for_chat_endpoint()],
             ),
             "ByteDance-Seed/BAGEL-7B-MoT",
@@ -466,7 +512,7 @@ def api_server(unused_tcp_port_factory, server_case: ServerCase, mock_async_omni
             ServerCase(
                 served_model="Qwen/Qwen-Image-Edit",
                 stage_list=["diffusion"],
-                stage_configs=[{"stage_type": "diffusion"}],
+                stage_configs=[_make_stage_config("diffusion")],
                 outputs=[_build_diffusion_image_output_for_images_endpoint()],
             ),
             "Qwen/Qwen-Image-Edit",
@@ -477,7 +523,7 @@ def api_server(unused_tcp_port_factory, server_case: ServerCase, mock_async_omni
             ServerCase(
                 served_model="ByteDance-Seed/BAGEL-7B-MoT",
                 stage_list=["diffusion"],
-                stage_configs=[{"stage_type": "diffusion"}],
+                stage_configs=[_make_stage_config("diffusion")],
                 outputs=[_build_diffusion_image_output_for_chat_endpoint()],
             ),
             "ByteDance-Seed/BAGEL-7B-MoT",
@@ -542,9 +588,9 @@ async def test_image_generation_node(api_server: str, model: str, image_input: b
                     MagicMock(is_comprehension=False, model_stage="llm"),
                 ],
                 stage_configs=[
-                    {"stage_type": "llm"},
-                    {"stage_type": "llm"},
-                    {"stage_type": "llm"},
+                    _make_stage_config("llm", is_comprehension=True, model_stage="thinker"),
+                    _make_stage_config("llm", is_comprehension=False, model_stage="talker"),
+                    _make_stage_config("llm", is_comprehension=False, model_stage="code2wav"),
                 ],
                 outputs=[_build_audio_chat_output(), _build_text_output("Understanding response")],
             ),
@@ -599,7 +645,7 @@ async def test_understanding_node(api_server: str, sampling_case: SamplingCase):
             ServerCase(
                 served_model="Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
                 stage_list=["llm"],
-                stage_configs=[{"stage_type": "llm"}],
+                stage_configs=[_make_stage_config("llm", model_stage="qwen3_tts")],
                 outputs=[_build_audio_speech_output()],
             ),
             VLLMOmniTTS,
@@ -617,7 +663,7 @@ async def test_understanding_node(api_server: str, sampling_case: SamplingCase):
             ServerCase(
                 served_model="Qwen/Qwen3-TTS-12Hz-1.7B-Base",
                 stage_list=["llm"],
-                stage_configs=[{"stage_type": "llm"}],
+                stage_configs=[_make_stage_config("llm", model_stage="qwen3_tts")],
                 outputs=[_build_audio_speech_output()],
             ),
             VLLMOmniVoiceClone,

tests/conftest.py

@@ -843,7 +843,7 @@ def modify_stage_config(
                     'async_chunk': True,
                     'stage_args': {
                         0: {'engine_args.max_model_len': 5800},
-                        1: {'runtime.max_batch_size': 2}
+                        1: {'engine_args.max_num_seqs': 2}
                     }
                 }
         deletes: Dictionary containing configurations to delete.
@@ -1236,15 +1236,12 @@ def omni_server(request: pytest.FixtureRequest, run_level: str, model_prefix: st
         port = params.port
         stage_config_path = params.stage_config_path
         if run_level == "advanced_model" and stage_config_path is not None:
+            with open(stage_config_path, encoding="utf-8") as f:
+                _cfg = yaml.safe_load(f) or {}
+            _stage_ids = [s["stage_id"] for s in _cfg.get("stage_args", []) if "stage_id" in s]
             stage_config_path = modify_stage_config(
                 stage_config_path,
-                deletes={
-                    "stage_args": {
-                        0: ["engine_args.load_format"],
-                        1: ["engine_args.load_format"],
-                        2: ["engine_args.load_format"],
-                    }
-                },
+                deletes={"stage_args": {sid: ["engine_args.load_format"] for sid in _stage_ids}},
             )
 
         server_args = params.server_args or []
@@ -1733,7 +1730,9 @@ class OmniRunner:
         Returns:
             List of SamplingParams with default decoding for each stage
         """
-        return [st.default_sampling_params for st in self.omni.stage_list]
+        if not hasattr(self.omni, "default_sampling_params_list"):
+            raise AttributeError("Omni.default_sampling_params_list is not available")
+        return list(self.omni.default_sampling_params_list)
 
     def get_omni_inputs(
         self,
@@ -1911,6 +1910,7 @@ class OmniRunner:
     def _cleanup_process(self):
         try:
             keywords = ["enginecore"]
+            matched = []
 
             for proc in psutil.process_iter(["pid", "name", "cmdline", "username"]):
                 try:
@@ -1923,16 +1923,32 @@ class OmniRunner:
 
                     if is_process:
                         print(f"Found vllm process: PID={proc.pid}, cmd={cmdline[:100]}")
-
-                        try:
-                            proc.terminate()
-                            time.sleep(2)
-                        except Exception:
-                            proc.kill()
-
+                        matched.append(proc)
                 except (psutil.NoSuchProcess, psutil.AccessDenied):
                     pass
 
+            for proc in matched:
+                try:
+                    proc.terminate()
+                except (psutil.NoSuchProcess, psutil.AccessDenied):
+                    pass
+
+            _, still_alive = psutil.wait_procs(matched, timeout=5)
+            for proc in still_alive:
+                try:
+                    proc.kill()
+                except (psutil.NoSuchProcess, psutil.AccessDenied):
+                    pass
+
+            if still_alive:
+                _, stubborn = psutil.wait_procs(still_alive, timeout=3)
+                if stubborn:
+                    print(f"Warning: failed to kill residual vllm pids: {[p.pid for p in stubborn]}")
+                else:
+                    print(f"Force-killed residual vllm pids: {[p.pid for p in still_alive]}")
+            elif matched:
+                print(f"Terminated vllm pids: {[p.pid for p in matched]}")
+
         except Exception as e:
             print(f"Error in psutil vllm cleanup: {e}")
 
@@ -1974,9 +1990,9 @@ class OmniRunnerHandler:
             audio_content = None
             for stage_output in outputs:
                 if getattr(stage_output, "final_output_type", None) == "text":
-                    text_content = stage_output.request_output[0].outputs[0].text
+                    text_content = stage_output.request_output.outputs[0].text
                 if getattr(stage_output, "final_output_type", None) == "audio":
-                    audio_content = stage_output.request_output[0].outputs[0].multimodal_output["audio"]
+                    audio_content = stage_output.request_output.outputs[0].multimodal_output["audio"]
 
             result.audio_content = audio_content
             result.text_content = text_content

tests/diffusion/cache/test_cache_backends.py

@@ -16,6 +16,7 @@ from unittest.mock import Mock, patch
 import pytest
 
 from vllm_omni.diffusion.cache.cache_dit_backend import (
+    CUSTOM_DIT_ENABLERS,
     CacheDiTBackend,
 )
 from vllm_omni.diffusion.cache.selector import get_cache_backend
@@ -122,6 +123,52 @@ class TestCacheDiTBackend:
         assert "cache_config" in call_args[1]
         assert backend._last_num_inference_steps == 100
 
+    def test_hunyuan_custom_enabler_registered(self):
+        """Test HunyuanImage3 custom cache-dit enabler is registered."""
+        assert "HunyuanImage3Pipeline" in CUSTOM_DIT_ENABLERS
+
+    @patch("vllm_omni.diffusion.cache.cache_dit_backend.BlockAdapter")
+    @patch("vllm_omni.diffusion.cache.cache_dit_backend.cache_dit")
+    def test_enable_hunyuan_pipeline_uses_model_transformer(self, mock_cache_dit, mock_block_adapter):
+        """Test HunyuanImage3 custom enabler uses pipeline.model for cache enable/refresh."""
+        mock_pipeline = Mock()
+        mock_pipeline.__class__.__name__ = "HunyuanImage3Pipeline"
+        mock_pipeline.model = Mock()
+        mock_pipeline.model.layers = Mock()
+
+        mock_cache_dit.enable_cache = Mock()
+        mock_cache_dit.refresh_context = Mock()
+
+        backend = CacheDiTBackend({"Fn_compute_blocks": 2})
+        backend.enable(mock_pipeline)
+
+        assert backend.enabled is True
+        assert backend._refresh_func is not None
+        mock_block_adapter.assert_called_once()
+        adapter_kwargs = mock_block_adapter.call_args.kwargs
+        assert adapter_kwargs["transformer"] is mock_pipeline.model
+        assert adapter_kwargs["blocks"] is mock_pipeline.model.layers
+        assert adapter_kwargs["forward_pattern"] == adapter_kwargs["forward_pattern"].__class__.Pattern_4
+        assert len(adapter_kwargs["params_modifiers"]) == 1
+        mock_cache_dit.enable_cache.assert_called_once()
+
+        backend.refresh(mock_pipeline, num_inference_steps=12)
+        mock_cache_dit.refresh_context.assert_called_once()
+        call_args = mock_cache_dit.refresh_context.call_args
+        assert call_args[0][0] is mock_pipeline.model
+        assert call_args[1]["num_inference_steps"] == 12
+
+    def test_enable_hunyuan_pipeline_requires_model_layers(self):
+        """Test HunyuanImage3 enabler fails with a formatted pipeline class name."""
+        mock_pipeline = Mock()
+        mock_pipeline.__class__.__name__ = "HunyuanImage3Pipeline"
+        mock_pipeline.model = Mock(spec=[])
+
+        backend = CacheDiTBackend({"Fn_compute_blocks": 2})
+
+        with pytest.raises(ValueError, match="HunyuanImage3Pipeline"):
+            backend.enable(mock_pipeline)
+
 
 class TestTeaCacheBackend:
     """Test TeaCacheBackend implementation."""

tests/diffusion/distributed/test_distributed_vae_executor.py

@@ -70,6 +70,7 @@ def mock_dist():
         patch.object(dist, "get_world_size", return_value=2),
         patch.object(dist, "get_rank", return_value=0),
         patch.object(dist, "is_initialized", return_value=True),
+        patch.object(dist, "all_reduce", return_value=None),
         patch.object(dist, "gather", return_value=None),
         patch.object(dist, "broadcast", return_value=None),
     ):

tests/diffusion/models/t5_encoder/test_t5_encoder_tp.py

@@ -0,0 +1,161 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import pytest
+import torch
+from transformers import T5Config
+from vllm.config import DeviceConfig, VllmConfig, set_current_vllm_config
+
+from vllm_omni.diffusion.models.t5_encoder.t5_encoder import (
+    T5EncoderModel,
+    T5SelfAttention,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_T5_MODULE = "vllm_omni.diffusion.models.t5_encoder.t5_encoder"
+
+SMALL_T5_CONFIG = dict(
+    d_model=64,
+    d_kv=8,
+    d_ff=128,
+    num_heads=8,
+    num_layers=2,
+    vocab_size=256,
+    relative_attention_num_buckets=32,
+    relative_attention_max_distance=128,
+    is_gated_act=True,
+    dense_act_fn="gelu_new",
+    layer_norm_epsilon=1e-6,
+    feed_forward_proj="gated-gelu",
+)
+
+
+@pytest.fixture(scope="module")
+def t5_config() -> T5Config:
+    return T5Config(**SMALL_T5_CONFIG)
+
+
+@pytest.fixture(scope="function", autouse=True)
+def setup_tp_group(monkeypatch, mocker):
+    """Set up TP=2, rank=0, VllmConfig, and mock activation for all tests."""
+    device_config = DeviceConfig(device="cpu")
+
+    # TP world size
+    monkeypatch.setattr("vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size", lambda: 2)
+    monkeypatch.setattr(f"{_T5_MODULE}.get_tensor_model_parallel_world_size", lambda: 2)
+    monkeypatch.setattr(
+        "vllm.model_executor.layers.vocab_parallel_embedding.get_tensor_model_parallel_world_size",
+        lambda: 2,
+    )
+
+    monkeypatch.setattr(f"{_T5_MODULE}.get_tensor_model_parallel_rank", lambda: 0)
+    monkeypatch.setattr(
+        "vllm.model_executor.layers.vocab_parallel_embedding.get_tensor_model_parallel_rank",
+        lambda: 0,
+    )
+
+    # TP group
+    mock_tp_group = mocker.MagicMock()
+    mock_tp_group.world_size = 2
+    mocker.patch("vllm.distributed.parallel_state.get_tp_group", return_value=mock_tp_group)
+
+    monkeypatch.setattr(f"{_T5_MODULE}.get_act_fn", lambda _: torch.nn.GELU())
+
+    with set_current_vllm_config(VllmConfig(device_config=device_config)):
+        yield
+
+
+class TestRelativePositionBiasTPSlicing:
+    """Verify compute_bias slices heads correctly per TP rank."""
+
+    def test_compute_bias_shape(self, t5_config):
+        attn = T5SelfAttention(t5_config, has_relative_attention_bias=True)
+
+        seq_len = 6
+        bias = attn.compute_bias(seq_len, seq_len, device=torch.device("cpu"))
+
+        local_heads = t5_config.num_heads // 2
+        assert bias.shape == (1, local_heads, seq_len, seq_len)
+
+    def test_all_ranks_cover_all_heads(self, t5_config, monkeypatch):
+        seq_len = 4
+
+        biases = []
+        ref_weight = None
+        for rank in range(2):
+            monkeypatch.setattr(f"{_T5_MODULE}.get_tensor_model_parallel_rank", lambda r=rank: r)
+            attn = T5SelfAttention(t5_config, has_relative_attention_bias=True)
+            if rank > 0:
+                attn.relative_attention_bias.weight.data.copy_(ref_weight)
+            else:
+                ref_weight = attn.relative_attention_bias.weight.data.clone()
+            biases.append(attn.compute_bias(seq_len, seq_len, device=torch.device("cpu")))
+
+        full_bias = torch.cat(biases, dim=1)
+        assert full_bias.shape == (1, t5_config.num_heads, seq_len, seq_len)
+
+
+class TestT5EncoderModelWeightLoading:
+    """Test weight loading at the top-level T5EncoderModel."""
+
+    def test_model_instantiation(self, t5_config):
+        model = T5EncoderModel(t5_config, prefix="text_encoder")
+
+        assert model.config is t5_config
+        assert model.encoder is not None
+        assert len(model.encoder.block) == t5_config.num_layers
+
+    def test_embedding_shape(self, t5_config):
+        model = T5EncoderModel(t5_config, prefix="text_encoder")
+
+        assert model.shared.embedding_dim == t5_config.d_model
+
+    def test_embed_input_ids(self, t5_config, monkeypatch):
+        # Verify method and output shape
+        model = T5EncoderModel(t5_config, prefix="text_encoder")
+
+        # Mock all-reduce to be identity (no actual TP communication)
+        monkeypatch.setattr(
+            "vllm.model_executor.layers.vocab_parallel_embedding.tensor_model_parallel_all_reduce",
+            lambda x: x,
+        )
+
+        input_ids = torch.randint(0, t5_config.vocab_size, (2, 8))
+        embeddings = model.embed_input_ids(input_ids)
+
+        assert embeddings.shape == (2, 8, t5_config.d_model)
+
+    def test_qkv_weights_loaded_through_blocks(self):
+        # Verify that HF-style separate Q/K/V weights can be loaded through
+        # the block hierarchy
+        config = T5Config(**{**SMALL_T5_CONFIG, "num_layers": 1})
+        model = T5EncoderModel(config, prefix="text_encoder")
+
+        inner_dim = config.num_heads * config.d_kv
+        prefix = "encoder.block.0.layer.0.SelfAttention."
+
+        loaded = model.load_weights(
+            [
+                (prefix + "q.weight", torch.randn(inner_dim, config.d_model)),
+                (prefix + "k.weight", torch.randn(inner_dim, config.d_model)),
+                (prefix + "v.weight", torch.randn(inner_dim, config.d_model)),
+            ]
+        )
+
+        assert len(loaded) > 0
+        attn = model.encoder.block[0].layer[0].SelfAttention
+        expected_qkv_dim = 3 * (config.num_heads // 2) * config.d_kv
+        assert attn.qkv_proj.weight.shape == (expected_qkv_dim, config.d_model)
+
+
+class TestTPConstraints:
+    """Verify that invalid TP configurations raise clear errors."""
+
+    def test_num_heads_not_divisible_by_tp(self):
+        config = T5Config(**{**SMALL_T5_CONFIG, "num_heads": 7})
+        with pytest.raises(AssertionError, match=r"num_heads.*must be divisible by tp_size"):
+            T5SelfAttention(config)
+
+    def test_num_heads_divisible_by_tp(self, t5_config):
+        attn = T5SelfAttention(t5_config)
+        assert attn.n_heads_per_partition == 4

tests/diffusion/quantization/test_int8_config.py

@@ -0,0 +1,463 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for Int8 quantization config."""
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+from pytest_mock import MockerFixture
+from torch.nn import Module, Parameter
+from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
+
+from vllm_omni.diffusion.quantization import (
+    get_diffusion_quant_config,
+    get_vllm_quant_config_for_layers,
+)
+from vllm_omni.platforms import current_omni_platform
+
+npu_available = pytest.mark.skipif(not current_omni_platform.is_npu(), reason="NPU platform not available.")
+
+cuda_available = pytest.mark.skipif(not current_omni_platform.is_cuda(), reason="GPU platform not available.")
+
+
+def test_int8_config_creation():
+    """Test that Int8 config can be created."""
+    config = get_diffusion_quant_config("int8")
+    assert config is not None
+    assert config.get_name() == "int8"
+
+
+def test_vllm_config_extraction():
+    """Test that vLLM config can be extracted from diffusion config."""
+    diff_config = get_diffusion_quant_config("int8")
+    vllm_config = get_vllm_quant_config_for_layers(diff_config)
+    assert vllm_config is not None
+    assert vllm_config.activation_scheme == "dynamic"
+
+
+def test_none_quantization():
+    """Test that None quantization returns None config."""
+    config = get_diffusion_quant_config(None)
+    assert config is None
+    vllm_config = get_vllm_quant_config_for_layers(config)
+    assert vllm_config is None
+
+
+def test_invalid_quantization():
+    """Test that invalid quantization method raises error."""
+    with pytest.raises(ValueError, match="Unknown quantization method"):
+        get_diffusion_quant_config("invalid_method")
+
+
+def test_int8_config_with_custom_params():
+    """Test Int8 config with custom parameters."""
+    config = get_diffusion_quant_config(
+        "int8",
+        activation_scheme="dynamic",
+        ignored_layers=["proj_out"],
+    )
+    assert config is not None
+    assert config.activation_scheme == "dynamic"
+    assert "proj_out" in config.ignored_layers
+
+
+def test_supported_methods():
+    """Test that supported methods list is correct."""
+    from vllm_omni.diffusion.quantization import SUPPORTED_QUANTIZATION_METHODS
+
+    assert "int8" in SUPPORTED_QUANTIZATION_METHODS
+
+
+def test_quantization_integration():
+    """Test end-to-end quantization flow through OmniDiffusionConfig."""
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+    # Test with quantization string only
+    config = OmniDiffusionConfig(model="test", quantization="int8")
+    assert config.quantization_config is not None
+    assert config.quantization_config.get_name() == "int8"
+
+    # Test with quantization_config dict
+    config2 = OmniDiffusionConfig(
+        model="test",
+        quantization_config={"method": "int8", "activation_scheme": "dynamic"},
+    )
+    assert config2.quantization_config is not None
+    assert config2.quantization_config.get_name() == "int8"
+    assert config2.quantization_config.activation_scheme == "dynamic"
+
+    # Test that vLLM config can be extracted
+    vllm_config = config.quantization_config.get_vllm_quant_config()
+    assert vllm_config is not None
+
+
+def test_quantization_dict_not_mutated():
+    """Test that passing a dict to quantization_config doesn't mutate it."""
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+    original_dict = {"method": "int8", "activation_scheme": "dynamic"}
+    dict_copy = original_dict.copy()
+
+    OmniDiffusionConfig(model="test", quantization_config=original_dict)
+
+    # Original dict should be unchanged
+    assert original_dict == dict_copy
+
+
+def test_quantization_conflicting_methods_warning(caplog):
+    """Test warning when quantization and quantization_config['method'] conflict."""
+    import logging
+
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+    with caplog.at_level(logging.WARNING):
+        config = OmniDiffusionConfig(
+            model="test",
+            quantization="int8",  # This should be overridden
+            quantization_config={"method": "int8", "activation_scheme": "dynamic"},
+        )
+    # No warning when methods match
+    assert config.quantization_config is not None
+
+
+def test_get_quant_method(mocker: MockerFixture):
+    """Test for get_quant_method method for GPU"""
+    from vllm_omni.diffusion.quantization.int8 import Int8OnlineLinearMethod
+
+    diff_config = get_diffusion_quant_config("int8")
+    vllm_config = get_vllm_quant_config_for_layers(diff_config)
+
+    def _fake_init(self, quant_config):
+        pass
+
+    layer = MagicMock(spec=LinearBase)
+    mocker.patch.object(Int8OnlineLinearMethod, "__init__", _fake_init)
+
+    prefix = "test_layer"
+
+    # Mock the platform to be GPU
+    with (
+        patch("vllm_omni.platforms.current_omni_platform.is_cuda", return_value=True),
+        patch("vllm_omni.platforms.current_omni_platform.is_npu", return_value=False),
+    ):
+        method = vllm_config.get_quant_method(layer, prefix)
+        assert isinstance(method, Int8OnlineLinearMethod)
+
+    # Test skipping quantization for a layer
+    vllm_config.ignored_layers = [prefix]
+    method = vllm_config.get_quant_method(layer, prefix)
+    assert isinstance(method, UnquantizedLinearMethod)
+
+
+def test_get_npu_quant_method():
+    """Test for get_quant_method method for NPU"""
+    from vllm_omni.diffusion.quantization.int8 import NPUInt8OnlineLinearMethod
+
+    diff_config = get_diffusion_quant_config("int8")
+    vllm_config = get_vllm_quant_config_for_layers(diff_config)
+
+    layer = MagicMock(spec=LinearBase)
+    prefix = "test_layer"
+
+    # Mock the platform to be NPU
+    with (
+        patch("vllm_omni.platforms.current_omni_platform.is_cuda", return_value=False),
+        patch("vllm_omni.platforms.current_omni_platform.is_npu", return_value=True),
+    ):
+        method = vllm_config.get_quant_method(layer, prefix)
+        assert isinstance(method, NPUInt8OnlineLinearMethod)
+
+    # Test skipping quantization for a layer
+    vllm_config.ignored_layers = [prefix]
+    method = vllm_config.get_quant_method(layer, prefix)
+    assert isinstance(method, UnquantizedLinearMethod)
+
+
+class TestInt8LinearMethod:
+    @pytest.fixture
+    def mock_quant_config(self, mocker):
+        return mocker.Mock()
+
+    @pytest.fixture
+    def mock_kernel(self, mocker):
+        kernel = mocker.Mock()
+        kernel.process_weights_after_loading = mocker.Mock()
+        kernel.apply_weights = mocker.Mock(return_value=torch.randn(1, 10))
+        return kernel
+
+    @pytest.fixture
+    def patch_deps(self, mocker, mock_kernel):
+        # mock init_int8_linear_kernel
+        mocker.patch("vllm_omni.diffusion.quantization.int8.init_int8_linear_kernel", return_value=mock_kernel)
+        return mock_kernel
+
+    def test_init(self, patch_deps, mock_quant_config):
+        # test for Int8LinearMethod init
+        from vllm_omni.diffusion.quantization.int8 import Int8LinearMethod, init_int8_linear_kernel
+
+        method = Int8LinearMethod(mock_quant_config)
+
+        assert method.quant_config == mock_quant_config
+        init_int8_linear_kernel.assert_called_once_with(
+            is_channelwise=False, is_static_input_scheme=False, input_symmetric=True, module_name="Int8LinearMethod"
+        )
+        assert method.int8_linear == patch_deps
+
+    def test_process_weights_after_loading(self, patch_deps, mock_quant_config):
+        from vllm_omni.diffusion.quantization.int8 import Int8LinearMethod
+
+        method = Int8LinearMethod(mock_quant_config)
+        layer = Module()
+
+        method.process_weights_after_loading(layer)
+        patch_deps.process_weights_after_loading.assert_called_once_with(layer)
+
+    def test_apply(self, patch_deps, mock_quant_config):
+        from vllm_omni.diffusion.quantization.int8 import Int8LinearMethod
+
+        method = Int8LinearMethod(mock_quant_config)
+        layer = Module()
+        x = torch.randn(1, 128)
+        bias = torch.randn(128)
+
+        output = method.apply(layer, x, bias)
+
+        patch_deps.apply_weights.assert_called_once_with(layer, x, bias)
+        assert isinstance(output, torch.Tensor)
+
+
+class TestInt8OnlineLinearMethod:
+    @pytest.fixture
+    def mock_quant_config(self, mocker):
+        return mocker.Mock()
+
+    @pytest.fixture
+    def mock_deps(self, mocker):
+        # mock kernel
+        mock_kernel = mocker.Mock()
+        mock_kernel.layer_param_names = ("weight", "weight_scale", "input_scale", "input_zero_point", "azp_adj")
+        mocker.patch("vllm_omni.diffusion.quantization.int8.init_int8_linear_kernel", return_value=mock_kernel)
+        mocker.patch("vllm_omni.diffusion.quantization.int8.replace_parameter")
+
+        # mock scaled_int8_quant return value
+        mock_qweight = torch.ones((128, 64), dtype=torch.int8)
+        mock_scale = torch.randn(128)
+        mock_quant = mocker.patch(
+            "vllm_omni.diffusion.quantization.int8.ops.scaled_int8_quant", return_value=(mock_qweight, mock_scale, None)
+        )
+        return {"kernel": mock_kernel, "quant": mock_quant, "mock_qweight": mock_qweight, "mock_scale": mock_scale}
+
+    def test_process_weights_after_loading(self, mock_deps, mock_quant_config):
+        from vllm_omni.diffusion.quantization.int8 import Int8OnlineLinearMethod
+
+        method = Int8OnlineLinearMethod(mock_quant_config)
+        layer = Module()
+        layer.weight = Parameter(torch.randn(128, 64))
+        method.process_weights_after_loading(layer)
+        mock_deps["quant"].assert_called_once_with(layer.weight, scale=None)
+
+
+@npu_available
+class TestNPUInt8LinearMethod:
+    qweight_mock = torch.randn((128, 64)).to(dtype=torch.int8)
+    scale_mock = torch.randn(128)
+    out_mock = torch.randn((16, 128))
+
+    @pytest.fixture
+    def mock_torch_npu(self, mocker):
+        torch_npu = MagicMock()
+
+        mocker.patch("vllm_omni.diffusion.quantization.int8.torch_npu", return_value=torch_npu)
+        mocker.patch(
+            "vllm_omni.diffusion.quantization.int8.torch_npu.npu_dynamic_quant",
+            return_value=(self.qweight_mock, self.scale_mock),
+        )
+        mocker.patch("vllm_omni.diffusion.quantization.int8.torch_npu.npu_quant_matmul", return_value=self.out_mock)
+        return torch_npu
+
+    @pytest.fixture
+    def mock_quant_config(self, mocker):
+        return mocker.Mock()
+
+    @pytest.fixture
+    def mock_layer(self, mocker):
+        layer = torch.nn.Module()
+        layer.weight = torch.nn.Parameter(self.qweight_mock, requires_grad=False)
+        layer.weight_scale = torch.nn.Parameter(self.scale_mock, requires_grad=False)
+        return layer
+
+    def test_npu_int8_process_weights_after_loading(self, mock_layer, mock_quant_config, mock_torch_npu):
+        from vllm_omni.diffusion.quantization.int8 import NPUInt8LinearMethod
+
+        method = NPUInt8LinearMethod(mock_quant_config)
+        ori_weight_shape = mock_layer.weight.shape
+
+        method.process_weights_after_loading(mock_layer)
+
+        assert mock_layer.weight.shape == ori_weight_shape[::-1]
+        assert mock_layer.weight.is_contiguous()
+
+    def test_npu_int8_apply(self, mock_layer, mock_quant_config, mock_torch_npu):
+        from vllm_omni.diffusion.quantization.int8 import NPUInt8LinearMethod
+
+        method = NPUInt8LinearMethod(mock_quant_config)
+        x = torch.randn(1, 16, 64)
+
+        output = method.apply(mock_layer, x)
+        assert output.shape == (1, 16, 128)
+
+    def test_npu_int8_online_process_weights(self, mock_layer, mock_quant_config, mock_torch_npu):
+        from vllm_omni.diffusion.quantization.int8 import NPUInt8OnlineLinearMethod
+
+        method = NPUInt8OnlineLinearMethod(mock_quant_config)
+        method.process_weights_after_loading(mock_layer)
+
+        assert mock_layer.weight.shape == (64, 128)
+        assert torch.equal(mock_layer.weight_scale, self.scale_mock)
+
+
+@pytest.fixture
+def quant_config():
+    """Shared quant config fixture for smoke tests."""
+    from vllm_omni.diffusion.quantization.int8 import Int8Config
+
+    return Int8Config(
+        is_checkpoint_int8_serialized=False,
+        activation_scheme="dynamic",
+    )
+
+
+@npu_available
+class TestNPUInt8Smoke:
+    """Smoke tests using real torch_npu, only run on NPU."""
+
+    @pytest.fixture
+    def real_layer(self):
+        """Create a real linear layer with fp16 weights on NPU"""
+        layer = torch.nn.Module()
+        layer.weight = torch.nn.Parameter(
+            torch.randn(128, 64, dtype=torch.float16, device="npu"),
+            requires_grad=False,
+        )
+        layer.logical_widths = [128]
+        layer.input_size_per_partition = 64
+        layer.output_size_per_partition = 128
+        layer.orig_dtype = torch.float16
+        return layer
+
+    def test_real_npu_dynamic_quant_shape_contract(self, quant_config, real_layer):
+        """Smoke test: verify npu_dynamic_quant returns correct shapes."""
+        import torch_npu
+
+        # Call real torch_npu.npu_dynamic_quant
+        weight = real_layer.weight
+        qweight, scale = torch_npu.npu_dynamic_quant(weight)
+
+        assert qweight.shape == weight.shape
+        assert qweight.dtype == torch.int8
+        assert scale.shape == (weight.shape[0],)
+
+    def test_real_npu_online_process_weights_after_loading(self, quant_config, real_layer):
+        """Smoke test: full process_weights_after_loading with real torch_npu."""
+        from vllm_omni.diffusion.quantization.int8 import NPUInt8OnlineLinearMethod
+
+        method = NPUInt8OnlineLinearMethod(quant_config)
+
+        method.process_weights_after_loading(real_layer)
+
+        assert real_layer.weight.shape == (64, 128)
+        assert real_layer.weight.dtype == torch.int8
+        assert hasattr(real_layer, "weight_scale")
+        assert real_layer.weight_scale.shape == (128,)
+
+    def test_real_npu_int8_apply_forward(self, quant_config):
+        """Smoke test: forward pass with real npu_quant_matmul."""
+        import torch_npu
+
+        from vllm_omni.diffusion.quantization.int8 import NPUInt8LinearMethod
+
+        method = NPUInt8LinearMethod(quant_config)
+
+        # Create layer with pre-processed weights on NPU
+        layer = torch.nn.Module()
+        weight_fp16 = torch.randn(128, 64, dtype=torch.float16, device="npu")
+        qweight, scale = torch_npu.npu_dynamic_quant(weight_fp16)
+        layer.weight = torch.nn.Parameter(qweight.t().contiguous(), requires_grad=False)
+        layer.weight_scale = torch.nn.Parameter(scale.squeeze(), requires_grad=False)
+
+        # Forward pass on NPU
+        x = torch.randn(2, 16, 64, dtype=torch.float16, device="npu")
+        output = method.apply(layer, x)
+
+        assert output.shape == (2, 16, 128)
+        assert output.dtype == torch.float16
+
+
+@cuda_available
+class TestCudaInt8Smoke:
+    """Smoke tests using real CUDA kernels, only on CUDA"""
+
+    @pytest.fixture
+    def real_layer(self):
+        """Create a real linear layer with fp16 weights on CUDA"""
+        layer = torch.nn.Module()
+        layer.weight = torch.nn.Parameter(
+            torch.randn(128, 64, dtype=torch.float16, device="cuda"),
+            requires_grad=False,
+        )
+        layer.logical_widths = [128]
+        layer.input_size_per_partition = 64
+        layer.output_size_per_partition = 128
+        layer.orig_dtype = torch.float16
+        return layer
+
+    def test_real_cuda_scaled_int8_quant_shape_contract(self, quant_config):
+        """Smoke test: verify scaled_int8_quant returns correct shapes."""
+        from vllm import _custom_ops as ops
+
+        weight = torch.randn(128, 64, dtype=torch.float16, device="cuda")
+        qweight, scale, _ = ops.scaled_int8_quant(weight, scale=None)
+
+        assert qweight.shape == weight.shape
+        assert qweight.dtype == torch.int8
+        assert scale.shape == (weight.shape[0], 1)
+
+    def test_real_cuda_online_process_weights_after_loading(self, quant_config, real_layer):
+        """Smoke test: full process_weights_after_loading with real CUDA ops."""
+        from vllm_omni.diffusion.quantization.int8 import Int8OnlineLinearMethod
+
+        method = Int8OnlineLinearMethod(quant_config)
+
+        method.process_weights_after_loading(real_layer)
+
+        assert real_layer.weight.shape == (64, 128)
+        assert real_layer.weight.dtype == torch.int8
+        assert hasattr(real_layer, "weight_scale")
+
+    def test_real_cuda_int8_apply_forward(self, quant_config):
+        """Smoke test: forward pass with real CUDA int8 kernel."""
+        from vllm import _custom_ops as ops
+
+        from vllm_omni.diffusion.quantization.int8 import Int8LinearMethod
+
+        method = Int8LinearMethod(quant_config)
+
+        # Create layer with pre-processed weights
+        layer = torch.nn.Module()
+        weight_fp16 = torch.randn(128, 64, dtype=torch.float16, device="cuda")
+        qweight, scale, _ = ops.scaled_int8_quant(weight_fp16, scale=None)
+        layer.weight = torch.nn.Parameter(qweight.t(), requires_grad=False)
+        layer.weight_scale = torch.nn.Parameter(scale, requires_grad=False)
+
+        # Set required attributes for kernel
+        layer.input_scale = None
+        layer.input_zero_point = None
+        layer.azp_adj = None
+
+        # Forward pass
+        x = torch.randn(2, 16, 64, dtype=torch.float16, device="cuda")
+        output = method.apply(layer, x)
+
+        assert output.shape == (2, 16, 128)
+        assert output.dtype == torch.float16

tests/diffusion/test_diffusion_step_pipeline.py

@@ -0,0 +1,507 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for step-level diffusion runner and worker execution."""
+
+import os
+from contextlib import contextmanager
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+import vllm_omni.diffusion.worker.diffusion_model_runner as model_runner_module
+from tests.utils import hardware_test
+from vllm_omni.diffusion.data import DiffusionOutput
+from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
+from vllm_omni.diffusion.distributed.comm import RingComm, SeqAllToAll4D
+from vllm_omni.diffusion.distributed.parallel_state import (
+    destroy_distributed_env,
+    get_sp_group,
+    init_distributed_environment,
+    initialize_model_parallel,
+)
+from vllm_omni.diffusion.ipc import (
+    pack_diffusion_output_shm,
+    unpack_diffusion_output_shm,
+)
+from vllm_omni.diffusion.sched.interface import (
+    DiffusionRequestState as SchedulerRequestState,
+)
+from vllm_omni.diffusion.sched.interface import (
+    DiffusionSchedulerOutput,
+)
+from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
+from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
+from vllm_omni.diffusion.worker.utils import RunnerOutput
+from vllm_omni.platforms import current_omni_platform
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
+
+# ---------------------------------------------------------------------------
+# Helpers & fixtures
+# ---------------------------------------------------------------------------
+
+
+@contextmanager
+def _noop_forward_context(*args, **kwargs):
+    del args, kwargs
+    yield
+
+
+def _update_environment_variables(envs_dict: dict[str, str]) -> None:
+    for key, value in envs_dict.items():
+        os.environ[key] = value
+
+
+class _StepPipeline:
+    """Minimal pipeline stub that supports step-wise execution."""
+
+    supports_step_execution = True
+
+    def __init__(self):
+        self.prepare_calls = 0
+        self.denoise_calls = 0
+        self.scheduler_calls = 0
+        self.decode_calls = 0
+
+    def prepare_encode(self, state, **kwargs):
+        del kwargs
+        self.prepare_calls += 1
+        state.timesteps = [torch.tensor(10), torch.tensor(5)]
+        state.latents = torch.tensor([0.0])
+        return state
+
+    def denoise_step(self, state, **kwargs):
+        del state, kwargs
+        self.denoise_calls += 1
+        return torch.tensor([1.0])
+
+    def step_scheduler(self, state, noise_pred, **kwargs):
+        del noise_pred, kwargs
+        self.scheduler_calls += 1
+        state.step_index += 1
+
+    def post_decode(self, state, **kwargs):
+        del kwargs
+        self.decode_calls += 1
+        return DiffusionOutput(output=torch.tensor([state.step_index], dtype=torch.float32))
+
+
+class _IdentityNoiseTransformer(torch.nn.Module):
+    def forward(self, x: torch.Tensor, **kwargs):
+        del kwargs
+        return (x,)
+
+
+class _AdditiveScheduler:
+    def step(self, noise_pred: torch.Tensor, t: torch.Tensor, latents: torch.Tensor, return_dict: bool = False):
+        del t, return_dict
+        return (latents + noise_pred,)
+
+
+class _DistributedStepPipeline(CFGParallelMixin):
+    supports_step_execution = True
+
+    def __init__(self, mode: str, device: torch.device):
+        self.mode = mode
+        self.device = device
+        self._interrupt = False
+        self.scheduler = _AdditiveScheduler()
+        self.transformer = _IdentityNoiseTransformer()
+
+    @property
+    def interrupt(self):
+        return self._interrupt
+
+    def prepare_encode(self, state, **kwargs):
+        del kwargs
+        state.timesteps = [torch.tensor(1.0, device=self.device)]
+        state.latents = torch.ones((1, 1), device=self.device)
+        state.step_index = 0
+        state.scheduler = self.scheduler
+        state.do_true_cfg = self.mode == "cfg"
+        return state
+
+    def denoise_step(self, state, **kwargs):
+        del kwargs
+        if self.mode == "ulysses":
+            sp_group = get_sp_group().ulysses_group
+            seq_world_size = torch.distributed.get_world_size(sp_group)
+            input_tensor = torch.randn(1, 2, 2 * seq_world_size, 2, device=self.device)
+            original = input_tensor.clone()
+            intermediate = SeqAllToAll4D.apply(sp_group, input_tensor, 2, 1, False)
+            output = SeqAllToAll4D.apply(sp_group, intermediate, 1, 2, False)
+            torch.testing.assert_close(output, original, rtol=1e-5, atol=1e-5)
+            return torch.ones_like(state.latents)
+
+        if self.mode == "ring":
+            ring_group = get_sp_group().ring_group
+            rank = torch.distributed.get_rank(ring_group)
+            world_size = torch.distributed.get_world_size(ring_group)
+            comm = RingComm(ring_group)
+            input_tensor = torch.full((1, 2, 2), float(rank + 1), device=self.device)
+            recv_tensor = comm.send_recv(input_tensor)
+            comm.commit()
+            comm.wait()
+            expected = torch.full_like(recv_tensor, float(((rank - 1) % world_size) + 1))
+            torch.testing.assert_close(recv_tensor, expected, rtol=1e-5, atol=1e-5)
+            return torch.ones_like(state.latents)
+
+        positive_kwargs = {"x": state.latents + 1}
+        negative_kwargs = {"x": state.latents - 1}
+        return self.predict_noise_maybe_with_cfg(
+            do_true_cfg=True,
+            true_cfg_scale=1.0,
+            positive_kwargs=positive_kwargs,
+            negative_kwargs=negative_kwargs,
+            cfg_normalize=False,
+        )
+
+    def step_scheduler(self, state, noise_pred, **kwargs):
+        del kwargs
+        if self.mode == "cfg":
+            state.latents = self.scheduler_step_maybe_with_cfg(
+                noise_pred,
+                state.current_timestep,
+                state.latents,
+                do_true_cfg=True,
+                per_request_scheduler=state.scheduler,
+            )
+        else:
+            state.latents = state.latents + noise_pred
+        state.step_index += 1
+
+    def post_decode(self, state, **kwargs):
+        del kwargs
+        return DiffusionOutput(output=state.latents.detach().cpu())
+
+
+def _make_step_request(num_inference_steps: int = 2):
+    return SimpleNamespace(
+        prompts=["a prompt"],
+        request_ids=["req-1"],
+        sampling_params=SimpleNamespace(
+            generator=None,
+            seed=None,
+            generator_device=None,
+            num_inference_steps=num_inference_steps,
+        ),
+    )
+
+
+def _make_runner():
+    runner = object.__new__(DiffusionModelRunner)
+    runner.vllm_config = object()
+    runner.od_config = SimpleNamespace(
+        cache_backend=None,
+        parallel_config=SimpleNamespace(use_hsdp=False),
+    )
+    runner.device = torch.device("cpu")
+    runner.pipeline = _StepPipeline()
+    runner.cache_backend = None
+    runner.offload_backend = None
+    runner.state_cache = {}
+    runner.kv_transfer_manager = SimpleNamespace()
+    return runner
+
+
+def _make_distributed_runner(mode: str, device: torch.device):
+    runner = object.__new__(DiffusionModelRunner)
+    runner.vllm_config = object()
+    runner.od_config = SimpleNamespace(
+        cache_backend=None,
+        parallel_config=SimpleNamespace(use_hsdp=False),
+    )
+    runner.device = device
+    runner.pipeline = _DistributedStepPipeline(mode=mode, device=device)
+    runner.cache_backend = None
+    runner.offload_backend = None
+    runner.state_cache = {}
+    runner.kv_transfer_manager = SimpleNamespace()
+    return runner
+
+
+def _make_scheduler_output(req, sched_req_id="req-1", step_id=0, finished_req_ids=None):
+    return DiffusionSchedulerOutput(
+        step_id=step_id,
+        req_states=[SchedulerRequestState(sched_req_id=sched_req_id, req=req)],
+        finished_req_ids=set() if finished_req_ids is None else set(finished_req_ids),
+        num_running_reqs=1,
+        num_waiting_reqs=0,
+    )
+
+
+def _expected_output_for_mode(mode: str) -> torch.Tensor:
+    if mode == "cfg":
+        return torch.tensor([[3.0]])
+    return torch.tensor([[2.0]])
+
+
+def _distributed_step_worker(local_rank: int, world_size: int, mode: str, master_port: str):
+    device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
+    current_omni_platform.set_device(device)
+    _update_environment_variables(
+        {
+            "RANK": str(local_rank),
+            "LOCAL_RANK": str(local_rank),
+            "WORLD_SIZE": str(world_size),
+            "MASTER_ADDR": "localhost",
+            "MASTER_PORT": master_port,
+        }
+    )
+    model_runner_module.set_forward_context = _noop_forward_context
+
+    try:
+        init_distributed_environment()
+        if mode == "ulysses":
+            initialize_model_parallel(ulysses_degree=world_size)
+        elif mode == "ring":
+            initialize_model_parallel(ring_degree=world_size)
+        elif mode == "cfg":
+            initialize_model_parallel(cfg_parallel_size=world_size)
+        else:
+            raise ValueError(f"Unsupported distributed test mode: {mode}")
+
+        runner = _make_distributed_runner(mode, device)
+        output = DiffusionModelRunner.execute_stepwise(
+            runner,
+            _make_scheduler_output(_make_step_request(num_inference_steps=1), step_id=0),
+        )
+
+        assert output.finished is True
+        assert output.result is not None
+        torch.testing.assert_close(output.result.output, _expected_output_for_mode(mode), rtol=1e-5, atol=1e-5)
+        assert "req-1" not in runner.state_cache
+    finally:
+        destroy_distributed_env()
+
+
+# ---------------------------------------------------------------------------
+# Runner / Worker
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.cpu
+class TestRunner:
+    """DiffusionModelRunner.execute_stepwise"""
+
+    def test_completes_request_and_clears_state(self, monkeypatch):
+        runner = _make_runner()
+        req = _make_step_request()
+        monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+
+        first = DiffusionModelRunner.execute_stepwise(runner, _make_scheduler_output(req, step_id=0))
+        assert first.req_id == "req-1"
+        assert first.step_index == 1
+        assert first.finished is False
+        assert first.result is None
+        assert "req-1" in runner.state_cache
+
+        second = DiffusionModelRunner.execute_stepwise(runner, _make_scheduler_output(req, step_id=1))
+        assert second.req_id == "req-1"
+        assert second.step_index == 2
+        assert second.finished is True
+        assert second.result is not None
+        assert second.result.error is None
+        assert torch.equal(second.result.output, torch.tensor([2.0]))
+        assert "req-1" not in runner.state_cache
+
+        assert runner.pipeline.prepare_calls == 1
+        assert runner.pipeline.denoise_calls == 2
+        assert runner.pipeline.scheduler_calls == 2
+        assert runner.pipeline.decode_calls == 1
+
+    def test_load_model_rejects_unsupported_step_execution(self, monkeypatch):
+        class _RequestOnlyPipeline:
+            pass
+
+        class _FakeLoader:
+            def __init__(self, *args, **kwargs):
+                del args, kwargs
+
+            def load_model(self, **kwargs):
+                del kwargs
+                return _RequestOnlyPipeline()
+
+        class _FakeProfiler:
+            consumed_memory = 0
+
+            def __enter__(self):
+                return self
+
+            def __exit__(self, exc_type, exc, tb):
+                del exc_type, exc, tb
+                return False
+
+        runner = object.__new__(DiffusionModelRunner)
+        runner.vllm_config = object()
+        runner.od_config = SimpleNamespace(
+            enable_cpu_offload=False,
+            enable_layerwise_offload=False,
+            enforce_eager=True,
+            cache_backend=None,
+            cache_config=None,
+            step_execution=True,
+            model_class_name="RequestOnlyPipeline",
+            parallel_config=SimpleNamespace(use_hsdp=False),
+        )
+        runner.device = torch.device("cpu")
+        runner.pipeline = None
+        runner.cache_backend = None
+        runner.offload_backend = None
+        runner.state_cache = {}
+        runner.kv_transfer_manager = SimpleNamespace()
+
+        monkeypatch.setattr(model_runner_module, "DiffusersPipelineLoader", _FakeLoader)
+        monkeypatch.setattr(model_runner_module, "DeviceMemoryProfiler", _FakeProfiler)
+        monkeypatch.setattr(model_runner_module, "get_offload_backend", lambda *args, **kwargs: None)
+        monkeypatch.setattr(model_runner_module, "get_cache_backend", lambda *args, **kwargs: None)
+
+        with pytest.raises(ValueError, match="RequestOnlyPipeline"):
+            DiffusionModelRunner.load_model(runner)
+
+
+@pytest.mark.cpu
+class TestWorker:
+    """DiffusionWorker.execute_stepwise"""
+
+    def test_delegates_to_model_runner(self):
+        worker = object.__new__(DiffusionWorker)
+        expected = RunnerOutput(req_id="req-1", step_index=1, finished=False, result=None)
+        scheduler_output = SimpleNamespace(
+            req_states=[
+                SimpleNamespace(
+                    req=SimpleNamespace(
+                        sampling_params=SimpleNamespace(lora_request=None),
+                    )
+                )
+            ]
+        )
+        worker.lora_manager = None
+        worker.model_runner = SimpleNamespace(
+            execute_stepwise=lambda arg: expected if arg is scheduler_output else None
+        )
+
+        output = DiffusionWorker.execute_stepwise(worker, scheduler_output)
+
+        assert output is expected
+
+    def test_clears_active_lora_before_stepwise_execution(self):
+        worker = object.__new__(DiffusionWorker)
+        scheduler_output = SimpleNamespace(
+            req_states=[
+                SimpleNamespace(
+                    req=SimpleNamespace(
+                        sampling_params=SimpleNamespace(lora_request=None),
+                    )
+                )
+            ]
+        )
+        calls: list[object | None] = []
+
+        class _FakeLoRAManager:
+            def set_active_adapter(self, adapter):
+                calls.append(adapter)
+
+        worker.lora_manager = _FakeLoRAManager()
+        worker.model_runner = SimpleNamespace(execute_stepwise=lambda arg: RunnerOutput(req_id="req-1"))
+
+        DiffusionWorker.execute_stepwise(worker, scheduler_output)
+
+        assert calls == [None]
+
+    def test_rejects_lora_requests_in_step_mode(self):
+        worker = object.__new__(DiffusionWorker)
+        scheduler_output = SimpleNamespace(
+            req_states=[
+                SimpleNamespace(
+                    req=SimpleNamespace(
+                        sampling_params=SimpleNamespace(lora_request=object()),
+                    )
+                )
+            ]
+        )
+        worker.lora_manager = None
+        worker.model_runner = SimpleNamespace(execute_stepwise=lambda arg: RunnerOutput(req_id="req-1"))
+
+        with pytest.raises(ValueError, match="does not support LoRA"):
+            DiffusionWorker.execute_stepwise(worker, scheduler_output)
+
+
+@pytest.mark.cpu
+class TestIPC:
+    def test_pack_unpack_runner_output_shm(self):
+        tensor = torch.zeros(300_000, dtype=torch.float32)
+        output = RunnerOutput(req_id="req-1", finished=True, result=DiffusionOutput(output=tensor))
+
+        packed = pack_diffusion_output_shm(output)
+        assert isinstance(packed.result.output, dict)
+        assert packed.result.output["__tensor_shm__"] is True
+
+        unpacked = unpack_diffusion_output_shm(packed)
+        assert isinstance(unpacked.result.output, torch.Tensor)
+        torch.testing.assert_close(unpacked.result.output, tensor)
+
+
+@pytest.mark.cpu
+class TestSupportedPipelines:
+    """Step-execution protocol checks for supported pipelines."""
+
+    def test_qwen_image_supports_step_execution(self):
+        from vllm_omni.diffusion.models.interface import SupportsStepExecution, supports_step_execution
+        from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image import QwenImagePipeline
+
+        # Avoid loading model weights; protocol membership depends on the class contract.
+        pipeline = object.__new__(QwenImagePipeline)
+
+        assert pipeline.supports_step_execution is True
+        assert supports_step_execution(pipeline) is True
+        assert isinstance(pipeline, SupportsStepExecution) is True
+
+
+@hardware_test(
+    res={"cuda": "L4"},
+    num_cards=2,
+)
+def test_execute_stepwise_with_ulysses_parallel():
+    world_size = 2
+    if current_omni_platform.get_device_count() < world_size:
+        pytest.skip(f"Test requires {world_size} devices")
+
+    torch.multiprocessing.spawn(
+        _distributed_step_worker,
+        args=(world_size, "ulysses", "29540"),
+        nprocs=world_size,
+    )
+
+
+@hardware_test(
+    res={"cuda": "L4"},
+    num_cards=2,
+)
+def test_execute_stepwise_with_ring_parallel():
+    world_size = 2
+    if current_omni_platform.get_device_count() < world_size:
+        pytest.skip(f"Test requires {world_size} devices")
+
+    torch.multiprocessing.spawn(
+        _distributed_step_worker,
+        args=(world_size, "ring", "29541"),
+        nprocs=world_size,
+    )
+
+
+@hardware_test(
+    res={"cuda": "L4"},
+    num_cards=2,
+)
+def test_execute_stepwise_with_cfg_parallel():
+    world_size = 2
+    if current_omni_platform.get_device_count() < world_size:
+        pytest.skip(f"Test requires {world_size} devices")
+
+    torch.multiprocessing.spawn(
+        _distributed_step_worker,
+        args=(world_size, "cfg", "29542"),
+        nprocs=world_size,
+    )

tests/distributed/omni_connectors/test_chunk_transfer_adapter.py

@@ -54,6 +54,8 @@ def build_adapter(monkeypatch, mocker: MockerFixture):
             self._pending_save_reqs = deque()
             self._finished_save_reqs = set()
             self.stop_event = threading.Event()
+            self._recv_cond = threading.Condition()
+            self._save_cond = threading.Condition()
 
         monkeypatch.setattr(OmniTransferAdapterBase, "__init__", _fake_base_init)
         monkeypatch.setattr(

tests/e2e/offline_inference/stage_configs/bagel_mooncake_ci.yaml

@@ -7,9 +7,9 @@ stage_args:
     prompt_expand_func: vllm_omni.model_executor.stage_input_processors.bagel.expand_cfg_prompts
     runtime:
       devices: "0"
-      max_batch_size: 1
     engine_args:
       model_stage: thinker
+      max_num_seqs: 1
       model_arch: BagelForConditionalGeneration
       worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -21,6 +21,7 @@ stage_args:
       enable_prefix_caching: false
       max_num_batched_tokens: 32768
       tensor_parallel_size: 1
+      load_format: dummy
       omni_kv_config:
         need_send_cache: true
         kv_transfer_criteria:
@@ -43,9 +44,9 @@ stage_args:
     cfg_kv_collect_func: vllm_omni.model_executor.stage_input_processors.bagel.collect_cfg_kv_caches
     runtime:
       devices: "0"
-      max_batch_size: 1
     engine_args:
       model_stage: dit
+      max_num_seqs: 1
       gpu_memory_utilization: 0.45
       enforce_eager: true
       trust_remote_code: true
@@ -54,6 +55,7 @@ stage_args:
       enable_prefix_caching: false
       max_num_batched_tokens: 32768
       tensor_parallel_size: 1
+      load_format: dummy
       omni_kv_config:
         need_recv_cache: true
     engine_input_source: [0]