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

Signed-off-by: princepride <wangzhipeng628@gmail.com>
Merge branch 'main' into dev/rebase-v0.18.0
2026-03-21 21:38:59 +08:00 · 2026-03-21 19:13:12 +08:00 · 2026-03-21 19:11:34 +08:00 · 2026-03-21 15:56:29 +08:00 · 2026-03-21 15:53:04 +08:00 · 2026-03-21 13:55:42 +08:00
183 changed files with 11283 additions and 1040 deletions
@@ -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" \
@@ -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 \
@@ -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"
@@ -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"
@@ -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]
@@ -17,7 +17,7 @@ steps:
            - "/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"
@@ -35,7 +35,7 @@ steps:
            - "/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
@@ -164,18 +164,13 @@ steps:
          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
@@ -340,3 +335,42 @@ steps:
                hostPath:
                  path: /mnt/hf-cache
                  type: DirectoryOrCreate
+
+  # - label: "Voxtral TTS E2E Test"
+  #   timeout_in_minutes: 20
+  #   depends_on: upload-merge-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
@@ -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"
@@ -70,7 +73,7 @@ steps:
  - 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"
@@ -161,6 +164,8 @@ steps:
                env:
                  - name: HF_HOME
                    value: /root/.cache/huggingface
+                  - name: VLLM_HTTP_TIMEOUT_KEEP_ALIVE
+                    value: "120"
                  - name: HF_TOKEN
                    valueFrom:
                      secretKeyRef:
@@ -180,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
@@ -16,10 +16,27 @@ steps:
          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
+          always-pull: true
+          propagate-environment: true
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
  - 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:
@@ -121,18 +138,13 @@ steps:
            - "/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
@@ -234,6 +246,45 @@ steps:
          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:
@@ -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
@@ -1,3 +1,4 @@
+import asyncio
 import base64
 import mimetypes
 import os
@@ -206,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"),
+    },
 }
@@ -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()

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

@@ -4,17 +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
      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
@@ -37,18 +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
      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
@@ -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
@@ -1,21 +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
      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
@@ -38,15 +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
      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
@@ -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 []
@@ -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
@@ -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

@@ -32,8 +32,8 @@ stage_args:
    is_prefill_only: true
    runtime:
      devices: "0"
-      max_batch_size: 16
    engine_args:
+      max_num_seqs: 16
      model_stage: thinker
      model_arch: Qwen3OmniMoeForConditionalGeneration
      worker_type: ar
@@ -70,8 +70,8 @@ stage_args:
    is_decode_only: true
    runtime:
      devices: "1"
-      max_batch_size: 64
    engine_args:
+      max_num_seqs: 64
      model_stage: thinker
      model_arch: Qwen3OmniMoeForConditionalGeneration
      worker_type: ar
@@ -128,7 +128,8 @@ After inserting the extra thinker stage, renumber the remaining stages:
  - stage_id: 3
    runtime:
      devices: "2"
-      max_batch_size: 1
+    engine_args:
+      max_num_seqs: 1
    engine_input_source: [2]
    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav
 ```
@@ -36,9 +36,9 @@ stage_args:
    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`

@@ -4,9 +4,9 @@ stage_args:
    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
@@ -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
@@ -967,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:**
@@ -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
@@ -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
@@ -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)
@@ -68,18 +68,18 @@ 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.17.0
+# Pin vLLM version to 0.18.0
 cd /vllm-workspace/vllm
 git pull origin main
 git fetch origin --tags
-git checkout v0.17.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-workspace/vllm-ascend
 git pull origin main
 git fetch origin --tags
-git checkout v0.17.0
+git checkout 1e05c4908f31737bc4eef865a9f351d030a77c9d
 pip install -v -e .

 # Install vLLM-Omni from the latest main branch
@@ -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.

@@ -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
@@ -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` |
@@ -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
@@ -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` | ✅ | ✅ | ✅ | ❌ | ❌ |

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

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

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

@@ -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 |
@@ -34,7 +34,14 @@ def parse_args():
        help="Path to input image for img2img.",
    )

-    # Omni runtime init args
+    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)
    parser.add_argument("--batch-timeout", type=int, default=5)
@@ -70,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:
@@ -172,13 +180,13 @@ def main():
    img_idx = 0
    for req_output in omni_outputs:
        images = getattr(req_output, "images", None)
+
        if not images:
            continue

        for j, img in enumerate(images):
-            save_path = f"output_{img_idx}_{j}.png"
+            save_path = os.path.join(args.output, f"output_{img_idx}_{j}.png")
            img.save(save_path)
-            print(f"[Info] Saved image to {save_path}")
        img_idx += 1

    print(omni_outputs)
@@ -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

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

@@ -131,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",
@@ -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 |
@@ -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)
@@ -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
@@ -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,
+    )
@@ -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}"
@@ -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,8 +1236,6 @@ 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:
-            # Dynamically detect stages from config to avoid KeyError
-            # for models with fewer stages (e.g., BAGEL has 2, Qwen3-Omni has 3)
            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]
@@ -1912,6 +1910,7 @@ class OmniRunner:
    def _cleanup_process(self):
        try:
            keywords = ["enginecore"]
+            matched = []

            for proc in psutil.process_iter(["pid", "name", "cmdline", "username"]):
                try:
@@ -1924,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}")

@@ -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
@@ -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
@@ -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,
+    )
@@ -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
@@ -44,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
@@ -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: OmniBagelForConditionalGeneration
      worker_type: ar
      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -43,9 +43,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
@@ -6,7 +6,6 @@ stage_args:
    runtime:
      process: true            # Run this stage in a separate process
      devices: "0"
-      max_batch_size: 1
    engine_args:
      model_stage: thinker
      model_arch: Qwen2_5OmniForConditionalGeneration
@@ -37,7 +36,6 @@ stage_args:
    runtime:
      process: true
      devices: "1"
-      max_batch_size: 1
    engine_args:
      model_stage: talker
      model_arch: Qwen2_5OmniForConditionalGeneration
@@ -67,9 +65,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
@@ -5,7 +5,7 @@
 End-to-end test for Bagel img2img generation.

 This test validates that the Bagel model generates images from an input image
-and text prompt that match expected reference pixel values within a ±5 tolerance.
+and text prompt that match expected reference pixel values within a ±10 tolerance.

 Equivalent to running:
    python3 examples/offline_inference/bagel/end2end.py \
@@ -25,6 +25,7 @@ from vllm.assets.image import ImageAsset
 from tests.conftest import modify_stage_config
 from tests.utils import hardware_test
 from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.platforms import current_omni_platform

 # Reference pixel data extracted from the known-good output image
 # Generated with seed=52, num_inference_steps=15,
@@ -43,7 +44,21 @@ REFERENCE_PIXELS = [
    {"position": (256, 256), "rgb": (181, 202, 222)},
 ]

-PIXEL_TOLERANCE = 5
+if current_omni_platform.is_rocm():
+    REFERENCE_PIXELS = [
+        {"position": (100, 100), "rgb": (156, 172, 215)},
+        {"position": (400, 50), "rgb": (106, 144, 216)},
+        {"position": (700, 100), "rgb": (118, 158, 231)},
+        {"position": (150, 400), "rgb": (183, 23, 48)},
+        {"position": (512, 336), "rgb": (218, 215, 191)},
+        {"position": (700, 400), "rgb": (194, 14, 42)},
+        {"position": (100, 600), "rgb": (105, 10, 16)},
+        {"position": (400, 600), "rgb": (167, 33, 46)},
+        {"position": (700, 600), "rgb": (102, 86, 92)},
+        {"position": (256, 256), "rgb": (181, 201, 220)},
+    ]
+
+PIXEL_TOLERANCE = 10

 DEFAULT_PROMPT = "<|fim_middle|><|im_start|>Change the grass color to red<|im_end|>"

@@ -191,7 +206,7 @@ def _resolve_stage_config(config_path: str, run_level: str) -> str:
@pytest.mark.core_model
@pytest.mark.advanced_model
@pytest.mark.diffusion
-@hardware_test(res={"cuda": "H100"})
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"})
 def test_bagel_img2img_shared_memory_connector(run_level):
    """Test Bagel img2img with shared memory connector."""
    input_image = _load_input_image()
@@ -31,6 +31,7 @@ from PIL import Image
 from tests.conftest import modify_stage_config
 from tests.utils import hardware_test
 from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.platforms import current_omni_platform

 # Reference pixel data extracted from the known-good output image
 # Each entry contains (x, y) position and expected (R, G, B) values
@@ -49,6 +50,20 @@ REFERENCE_PIXELS = [
    {"position": (256, 256), "rgb": (171, 160, 153)},
 ]

+if current_omni_platform.is_rocm():
+    REFERENCE_PIXELS = [
+        {"position": (100, 100), "rgb": (123, 119, 100)},
+        {"position": (400, 50), "rgb": (162, 161, 142)},
+        {"position": (700, 100), "rgb": (171, 156, 127)},
+        {"position": (150, 400), "rgb": (131, 128, 112)},
+        {"position": (512, 512), "rgb": (134, 61, 59)},
+        {"position": (700, 400), "rgb": (204, 107, 43)},
+        {"position": (100, 700), "rgb": (201, 180, 165)},
+        {"position": (400, 700), "rgb": (140, 108, 87)},
+        {"position": (700, 700), "rgb": (247, 205, 145)},
+        {"position": (256, 256), "rgb": (171, 160, 153)},
+    ]
+
 # Maximum allowed difference per color channel
 PIXEL_TOLERANCE = 5

@@ -181,7 +196,7 @@ def _resolve_stage_config(config_path: str, run_level: str) -> str:
@pytest.mark.core_model
@pytest.mark.advanced_model
@pytest.mark.diffusion
-@hardware_test(res={"cuda": "H100"})
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"})
 def test_bagel_text2img_shared_memory_connector(run_level):
    """Test Bagel text2img with shared memory connector."""
    config_path = str(Path(__file__).parent / "stage_configs" / "bagel_sharedmemory_ci.yaml")
@@ -24,6 +24,7 @@ os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
 # This test is specific to Z-Image LoRA behavior. Keep it focused on a single
 # model to reduce runtime and avoid extra downloads.
 models = ["Tongyi-MAI/Z-Image-Turbo"]
+DIFFUSION_INIT_TIMEOUT_S = 600


@pytest.mark.parametrize("model_name", models)
@@ -76,7 +77,11 @@ def test_diffusion_model(model_name: str, tmp_path: Path):
        )
        return str(adapter_dir)

-    m = Omni(model=model_name)
+    m = Omni(
+        model=model_name,
+        stage_init_timeout=DIFFUSION_INIT_TIMEOUT_S,
+        init_timeout=DIFFUSION_INIT_TIMEOUT_S,
+    )
    try:
        # high resolution may cause OOM on L4
        height = 256
@@ -32,7 +32,7 @@ from tests.conftest import OmniServerParams
 from tests.utils import hardware_test

 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
-os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"

 MODEL = "ByteDance-Seed/BAGEL-7B-MoT"
 STAGE_CONFIGS_PATH = str(
@@ -97,7 +97,7 @@ def test_bagel_text2img_online(omni_server, openai_client) -> None:
@pytest.mark.core_model
@pytest.mark.advanced_model
@pytest.mark.diffusion
-@hardware_test(res={"cuda": "H100"})
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"})
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_bagel_img2img_online(omni_server, openai_client) -> None:
    """Test Bagel img2img via OpenAI-compatible chat completions API."""
@@ -28,6 +28,7 @@ from tests.utils import hardware_test
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"

 MODEL = "Tongyi-MAI/Z-Image-Turbo"
+DIFFUSION_INIT_TIMEOUT_S = 600


 PROMPT = "a photo of a cat sitting on a laptop keyboard"
@@ -37,7 +38,17 @@ SEED = 42

@pytest.fixture(scope="module")
 def omni_server():
-    with OmniServer(MODEL, ["--num-gpus", "1"]) as server:
+    with OmniServer(
+        MODEL,
+        [
+            "--num-gpus",
+            "1",
+            "--stage-init-timeout",
+            str(DIFFUSION_INIT_TIMEOUT_S),
+            "--init-timeout",
+            str(DIFFUSION_INIT_TIMEOUT_S),
+        ],
+    ) as server:
        yield server


@@ -44,13 +44,9 @@ def get_chunk_config():
    return path


-# CI stage config for 2xH100-80G GPUs or AMD GPU MI325
-if current_omni_platform.is_rocm():
-    # ROCm stage config optimized for MI325 GPU
-    stage_configs = [str(Path(__file__).parent.parent / "stage_configs" / "rocm" / "qwen3_omni_ci.yaml")]
-elif current_omni_platform.is_xpu():
+if current_omni_platform.is_xpu():
    stage_configs = [str(Path(__file__).parent.parent / "stage_configs" / "xpu" / "qwen3_omni_ci.yaml")]
-else:
+else:  # MI325 GPU should share the same config as H100
    stage_configs = [get_chunk_config()]

 # Create parameter combinations for model and stage config
@@ -0,0 +1,192 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+E2E Online tests for Voxtral TTS model with text input and audio output.
+
+These tests verify the /v1/audio/speech endpoint works correctly with
+actual model inference, not mocks.
+"""
+
+import os
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
+
+from pathlib import Path
+
+import httpx
+import pytest
+
+from tests.conftest import OmniServerParams
+from tests.utils import hardware_test
+
+MODEL = "mistralai/Voxtral-4B-TTS-2603"
+
+STAGE_CONFIG = str(
+    Path(__file__).parent.parent.parent.parent / "vllm_omni" / "model_executor" / "stage_configs" / "voxtral_tts.yaml"
+)
+EXTRA_ARGS = ["--trust-remote-code", "--enforce-eager", "--disable-log-stats"]
+TEST_PARAMS = [OmniServerParams(model=MODEL, stage_config_path=STAGE_CONFIG, server_args=EXTRA_ARGS)]
+
+
+def make_speech_request(
+    host: str,
+    port: int,
+    text: str,
+    voice: str = "casual_female",
+    language: str = "English",
+    timeout: float = 120.0,
+) -> httpx.Response:
+    """Make a request to the /v1/audio/speech endpoint."""
+    url = f"http://{host}:{port}/v1/audio/speech"
+    payload = {
+        "input": text,
+        "voice": voice,
+        "language": language,
+    }
+
+    with httpx.Client(timeout=timeout) as client:
+        return client.post(url, json=payload)
+
+
+def verify_wav_audio(content: bytes) -> bool:
+    """Verify that content is valid WAV audio data."""
+    # WAV files start with "RIFF" header
+    if len(content) < 44:  # Minimum WAV header size
+        return False
+    return content[:4] == b"RIFF" and content[8:12] == b"WAVE"
+
+
+# Minimum expected audio size for a short sentence (~1 second of 24kHz 16-bit mono WAV)
+MIN_AUDIO_BYTES = 10000
+
+
+@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
+class TestVoxtralTTSFixedVoice:
+    """E2E tests for Voxtral TTS model."""
+
+    @pytest.mark.core_model
+    @pytest.mark.omni
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_speech_english_basic(self, omni_server) -> None:
+        """Test basic English TTS generation."""
+        response = make_speech_request(
+            host=omni_server.host,
+            port=omni_server.port,
+            text="Hello, how are you?",
+            voice="casual_female",
+            language="English",
+        )
+
+        assert response.status_code == 200, f"Request failed: {response.text}"
+        assert response.headers.get("content-type") == "audio/wav"
+        assert verify_wav_audio(response.content), "Response is not valid WAV audio"
+        assert len(response.content) > MIN_AUDIO_BYTES, (
+            f"Audio content too small ({len(response.content)} bytes), expected at least {MIN_AUDIO_BYTES} bytes"
+        )
+
+    @pytest.mark.advanced_model
+    @pytest.mark.omni
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_speech_different_voices(self, omni_server) -> None:
+        """Test TTS with different voice presets."""
+        voices = ["casual_female", "neutral_male"]
+        for voice in voices:
+            response = make_speech_request(
+                host=omni_server.host,
+                port=omni_server.port,
+                text="Testing voice selection.",
+                voice=voice,
+                language="English",
+            )
+
+            assert response.status_code == 200, f"Request failed for voice {voice}: {response.text}"
+            assert verify_wav_audio(response.content), f"Invalid WAV for voice {voice}"
+
+    @pytest.mark.advanced_model
+    @pytest.mark.omni
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_speech_invalid_voice_rejected(self, omni_server) -> None:
+        """Request with a non-existent voice should return an error response."""
+        response = make_speech_request(
+            host=omni_server.host,
+            port=omni_server.port,
+            text="This voice does not exist.",
+            voice="nonexistent_voice_xyz",
+            language="English",
+        )
+
+        # vLLM returns ErrorResponse as JSON with HTTP 200;
+        # verify the body is a JSON error, not valid WAV audio.
+        data = response.json()
+        assert "error" in data or "message" in data, f"Expected error response for invalid voice, got: {data}"
+        assert not verify_wav_audio(response.content), "Should not return valid audio for invalid voice"
+
+    @pytest.mark.advanced_model
+    @pytest.mark.omni
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_speech_binary_response_not_utf8_error(self, omni_server) -> None:
+        """
+        Regression test: Verify binary audio is returned, not UTF-8 error.
+
+        This test ensures the multimodal_output property correctly retrieves
+        audio from completion outputs, preventing the "TTS model did not
+        produce audio output" error.
+        """
+        response = make_speech_request(
+            host=omni_server.host,
+            port=omni_server.port,
+            text="This should return binary audio, not a JSON error.",
+            voice="casual_female",
+            language="English",
+        )
+
+        # Should NOT be a JSON error response
+        assert response.status_code == 200, f"Request failed: {response.text}"
+
+        # Verify it's binary audio, not JSON
+        try:
+            # If this succeeds and starts with {"error", it's a bug
+            text = response.content.decode("utf-8")
+            assert not text.startswith('{"error"'), f"Got error response instead of audio: {text}"
+        except UnicodeDecodeError:
+            # This is expected - binary audio can't be decoded as UTF-8
+            pass
+
+        assert verify_wav_audio(response.content), "Response is not valid WAV audio"
+
+
+@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
+class TestVoxtralTTSAPIEndpoints:
+    """Test API endpoint functionality."""
+
+    @pytest.mark.advanced_model
+    @pytest.mark.omni
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_list_voices_endpoint(self, omni_server) -> None:
+        """Test the /v1/audio/voices endpoint returns available voices."""
+        url = f"http://{omni_server.host}:{omni_server.port}/v1/audio/voices"
+
+        with httpx.Client(timeout=30.0) as client:
+            response = client.get(url)
+
+        assert response.status_code == 200
+        data = response.json()
+        assert "voices" in data
+        assert isinstance(data["voices"], list)
+        assert len(data["voices"]) > 0
+
+    @pytest.mark.advanced_model
+    @pytest.mark.omni
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_models_endpoint(self, omni_server) -> None:
+        """Test the /v1/models endpoint returns loaded model."""
+        url = f"http://{omni_server.host}:{omni_server.port}/v1/models"
+
+        with httpx.Client(timeout=30.0) as client:
+            response = client.get(url)
+
+        assert response.status_code == 200
+        data = response.json()
+        assert "data" in data
+        assert len(data["data"]) > 0
@@ -7,7 +7,6 @@ stage_args:
    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
      model_arch: Qwen2_5OmniForConditionalGeneration
@@ -38,7 +37,6 @@ stage_args:
    runtime:
      process: true
      devices: "1"
-      max_batch_size: 1
    engine_args:
      model_stage: talker
      model_arch: Qwen2_5OmniForConditionalGeneration
@@ -47,7 +45,7 @@ stage_args:
      max_model_len: 16384
      max_num_batched_tokens: 16384
      max_num_seqs: 1
-      gpu_memory_utilization: 0.9
+      gpu_memory_utilization: 0.4
      skip_mm_profiling: true
      enforce_eager: true
      trust_remote_code: true
@@ -68,13 +66,13 @@ stage_args:
    runtime:
      process: true
      devices: "2"            # 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
-      gpu_memory_utilization: 0.9 #increase the gpu memory utilization to enable the test on H800
+      gpu_memory_utilization: 0.5 #increase the gpu memory utilization to enable the test on H800
      enforce_eager: true
      trust_remote_code: true
      enable_prefix_caching: false
@@ -3,7 +3,6 @@ stage_args:
    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
      model_arch: Qwen2_5OmniForConditionalGeneration
@@ -8,9 +8,9 @@ stage_args:
 - stage_id: 0
  runtime:
    devices: "0"
-    max_batch_size: 5
  engine_args:
    model_stage: thinker
+    max_num_seqs: 5
    model_arch: Qwen3OmniMoeForConditionalGeneration
    worker_type: ar
    scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -42,9 +42,9 @@ stage_args:
 - stage_id: 1
  runtime:
    devices: "1"
-    max_batch_size: 5
  engine_args:
    model_stage: talker
+    max_num_seqs: 5
    model_arch: Qwen3OmniMoeForConditionalGeneration
    worker_type: ar
    scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -72,9 +72,9 @@ stage_args:
 - stage_id: 2
  runtime:
    devices: "1"
-    max_batch_size: 1
  engine_args:
    model_stage: code2wav
+    max_num_seqs: 1
    model_arch: Qwen3OmniMoeForConditionalGeneration
    worker_type: generation
    scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
@@ -7,14 +7,13 @@ stage_args:
    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
      model_arch: Qwen2_5OmniForConditionalGeneration
      worker_type: ar
      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      max_model_len: 2400
-      max_num_batched_tokens: 2400
+      max_model_len: 16384
+      max_num_batched_tokens: 16384
      max_num_seqs: 1
      gpu_memory_utilization: 0.8
      skip_mm_profiling: true
@@ -38,14 +37,13 @@ stage_args:
    runtime:
      process: true
      devices: "1"
-      max_batch_size: 1
    engine_args:
      model_stage: talker
      model_arch: Qwen2_5OmniForConditionalGeneration
      worker_type: ar
      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      max_model_len: 2400
-      max_num_batched_tokens: 2400
+      max_model_len: 16384
+      max_num_batched_tokens: 16384
      max_num_seqs: 1
      gpu_memory_utilization: 0.8
      skip_mm_profiling: true
@@ -59,7 +57,7 @@ stage_args:
      temperature: 0.9
      top_p: 0.8
      top_k: 40
-      max_tokens: 128
+      max_tokens: 4096
      seed: 42
      detokenize: True
      repetition_penalty: 1.05
@@ -68,9 +66,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
@@ -79,6 +77,8 @@ stage_args:
      trust_remote_code: true
      enable_prefix_caching: false
      engine_output_type: audio
+      max_num_batched_tokens: 4096
+      max_model_len: 4096
    engine_input_source: [1]
    final_output: true
    final_output_type: audio
@@ -86,7 +86,7 @@ stage_args:
      temperature: 0.0
      top_p: 1.0
      top_k: -1
-      max_tokens: 128
+      max_tokens: 4096
      seed: 42
      detokenize: True
      repetition_penalty: 1.1
@@ -8,9 +8,9 @@ stage_args:
  - stage_id: 0
    runtime:
      devices: "0"
-      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
@@ -39,9 +39,9 @@ stage_args:
  - stage_id: 1
    runtime:
       devices: "1"
-       max_batch_size: 1
    engine_args:
       model_stage: talker
+       max_num_seqs: 1
       model_arch: Qwen3OmniMoeForConditionalGeneration
       worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -70,9 +70,9 @@ stage_args:
  - stage_id: 2
    runtime:
      devices: "1"
-      max_batch_size: 1
    engine_args:
      model_stage: code2wav
+      max_num_seqs: 1
      model_arch: Qwen3OmniMoeForConditionalGeneration
      worker_type: generation
      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
@@ -7,8 +7,8 @@ stage_args:
    runtime:
      process: true # Run this stage in a separate process
      devices: "0" # Visible devices for this stage
-      max_batch_size: 1
    engine_args:
+      max_num_seqs: 1
      model_stage: thinker
      model_arch: Qwen2_5OmniForConditionalGeneration
      worker_type: ar
@@ -39,8 +39,8 @@ stage_args:
    runtime:
      process: true
      devices: "1"
-      max_batch_size: 1
    engine_args:
+      max_num_seqs: 1
      model_stage: talker
      model_arch: Qwen2_5OmniForConditionalGeneration
      worker_type: ar
@@ -71,8 +71,8 @@ stage_args:
    runtime:
      process: true
      devices: "2"
-      max_batch_size: 1
    engine_args:
+      max_num_seqs: 1
      model_stage: code2wav
      model_arch: Qwen2_5OmniForConditionalGeneration
      worker_type: generation
@@ -9,8 +9,8 @@ stage_args:
  stage_type: llm # Use llm stage type for AR stages
  runtime:
    devices: "0,1,2,3"
-    max_batch_size: 1
  engine_args:
+    max_num_seqs: 1
    model_stage: thinker
    model_arch: Qwen3OmniMoeForConditionalGeneration
    worker_type: ar
@@ -44,8 +44,8 @@ stage_args:
  stage_type: llm # Use llm stage type for AR stages
  runtime:
    devices: "4"
-    max_batch_size: 1
  engine_args:
+    max_num_seqs: 1
    model_stage: talker
    model_arch: Qwen3OmniMoeForConditionalGeneration
    worker_type: ar
@@ -78,8 +78,8 @@ stage_args:
  stage_type: llm # Use llm stage type for AR stages
  runtime:
    devices: "5"
-    max_batch_size: 1
  engine_args:
+    max_num_seqs: 1
    model_stage: code2wav
    model_arch: Qwen3OmniMoeForConditionalGeneration
    worker_type: generation
@@ -123,6 +123,7 @@ def make_mock_model(hidden: int = 8):
    cfg.video_token_index = VIDEO_TOKEN_ID
    cfg.audio_token_index = AUDIO_TOKEN_ID
    model.config = cfg
+    model._has_oov_mm_tokens = False

    def fake_lm_embed(ids: torch.Tensor) -> torch.Tensor:
        # Use .clone() so the tensor is contiguous (expand() creates a strided
@@ -137,13 +138,12 @@ def make_mock_model(hidden: int = 8):

    model._embed_text_input_ids = lambda *a, **kw: SupportsMultiModal._embed_text_input_ids(model, *a, **kw)

-    def fake_super_embed(ids, mm_embs=None, *, is_multimodal=None, handle_oov_mm_token=False):
+    def fake_super_embed(ids, mm_embs=None, *, is_multimodal=None):
        return SupportsMultiModal.embed_input_ids(
            model,
            ids,
            mm_embs,
            is_multimodal=is_multimodal,
-            handle_oov_mm_token=handle_oov_mm_token,
        )

    model.embed_input_ids = lambda *a, **kw: Qwen2_5OmniThinkerForConditionalGeneration.embed_input_ids(model, *a, **kw)
@@ -270,3 +270,63 @@ def test_deterministic_across_calls(decoder, wrapper):
        out1 = wrapper.decode(codes)
        out2 = wrapper.decode(codes)
    torch.testing.assert_close(out1, out2, atol=0, rtol=0)
+
+
+# ──────────────────────────────────────────────────────────────────
+# 6. compute_capture_sizes
+# ──────────────────────────────────────────────────────────────────
+
+
+@pytest.mark.parametrize(
+    "kwargs,expected_in,not_expected",
+    [
+        ({}, [2, 4, 8, 16, 32, 64, 128, 256, 325], [512]),
+        (
+            {"codec_chunk_frames": 33, "codec_left_context_frames": 25},
+            [2, 4, 8, 16, 32, 33, 58, 64, 128, 256, 325],
+            [512],
+        ),
+        (
+            {"codec_chunk_frames": 25, "codec_left_context_frames": 25},
+            [2, 4, 8, 16, 25, 32, 50, 64, 128, 256, 325],
+            [512],
+        ),
+    ],
+    ids=["default", "streaming_c33", "streaming_c25"],
+)
+def test_compute_capture_sizes(kwargs, expected_in, not_expected):
+    """compute_capture_sizes produces expected sizes capped by max useful size."""
+    sizes = CUDAGraphDecoderWrapper.compute_capture_sizes(**kwargs)
+    for val in expected_in:
+        assert val in sizes, f"{val} not in {sizes}"
+    for val in not_expected:
+        assert val not in sizes, f"{val} should not be in {sizes}"
+
+
+# ──────────────────────────────────────────────────────────────────
+# 7. SnakeBeta Triton kernel vs eager equivalence
+# ──────────────────────────────────────────────────────────────────
+
+
+@pytest.mark.parametrize(
+    "batch,channels,seq_len",
+    [(2, 64, 1000), (1, 32, 1), (1, 32, 7), (1, 32, 128), (1, 32, 1024), (1, 32, 4096)],
+)
+def test_snakebeta_triton_vs_eager(batch, channels, seq_len):
+    """Fused Triton SnakeBeta kernel must match eager PyTorch output."""
+    from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import (
+        SnakeBeta,
+    )
+
+    if not SnakeBeta._init_triton():
+        pytest.skip("Triton not available")
+
+    torch.manual_seed(42)
+    snake = SnakeBeta(in_features=channels).to(DEVICE).eval()
+    x = torch.randn(batch, channels, seq_len, device=DEVICE)
+
+    with torch.no_grad():
+        eager_out = snake._eager_forward(x)
+        triton_out = snake._triton_forward(x)
+
+    torch.testing.assert_close(triton_out, eager_out, atol=1e-5, rtol=1e-5)
@@ -0,0 +1,161 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Tests for parse_batched_audio_input and apply_ctx_frames_cutting.
+
+These functions parse the [ctx_frames, context_length, ...tokens] wire format
+produced by generator2tokenizer_async_chunk and cut leading context samples
+from decoded audio arrays.
+"""
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.voxtral_tts.voxtral_tts import (
+    apply_ctx_frames_cutting,
+    parse_batched_audio_input,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+NUM_CODEBOOKS = 37
+
+
+# ──────────────────────────────────────────────────────────────────
+# parse_batched_audio_input tests
+# ──────────────────────────────────────────────────────────────────
+
+
+def _make_input_ids(*requests: tuple[int, int, torch.Tensor]) -> torch.Tensor:
+    """Build a flat input_ids tensor from (ctx_frames, context_length, tokens) tuples."""
+    parts = []
+    for ctx_frames, context_length, tokens in requests:
+        parts.append(torch.tensor([ctx_frames, context_length]))
+        parts.append(tokens.flatten())
+    return torch.cat(parts)
+
+
+def test_parse_single_request():
+    """Single request is parsed correctly."""
+    ctx_frames = 3
+    context_length = 5
+    total_frames = ctx_frames + context_length
+    tokens = torch.arange(total_frames * NUM_CODEBOOKS).reshape(total_frames, NUM_CODEBOOKS)
+    input_ids = _make_input_ids((ctx_frames, context_length, tokens))
+
+    all_audio_tokens, all_ctx_frames = parse_batched_audio_input(input_ids, num_codebooks=NUM_CODEBOOKS)
+
+    assert len(all_audio_tokens) == 1
+    assert len(all_ctx_frames) == 1
+    assert all_ctx_frames[0] == ctx_frames
+    assert all_audio_tokens[0].shape == (total_frames, NUM_CODEBOOKS)
+    torch.testing.assert_close(all_audio_tokens[0], tokens)
+
+
+def test_parse_multiple_requests():
+    """Multiple requests in a batch are parsed correctly."""
+    req1 = (2, 4, torch.ones((6, NUM_CODEBOOKS), dtype=torch.long))
+    req2 = (0, 3, torch.full((3, NUM_CODEBOOKS), 7, dtype=torch.long))
+    input_ids = _make_input_ids(req1, req2)
+
+    all_audio_tokens, all_ctx_frames = parse_batched_audio_input(input_ids, num_codebooks=NUM_CODEBOOKS)
+
+    assert len(all_audio_tokens) == 2
+    assert all_ctx_frames == [2, 0]
+    assert all_audio_tokens[0].shape == (6, NUM_CODEBOOKS)
+    assert all_audio_tokens[1].shape == (3, NUM_CODEBOOKS)
+    assert torch.all(all_audio_tokens[0] == 1)
+    assert torch.all(all_audio_tokens[1] == 7)
+
+
+def test_parse_zero_ctx_frames():
+    """ctx_frames=0 means no context, just chunk data."""
+    ctx_frames = 0
+    context_length = 10
+    tokens = torch.randint(0, 100, (context_length, NUM_CODEBOOKS))
+    input_ids = _make_input_ids((ctx_frames, context_length, tokens))
+
+    all_audio_tokens, all_ctx_frames = parse_batched_audio_input(input_ids, num_codebooks=NUM_CODEBOOKS)
+
+    assert all_ctx_frames[0] == 0
+    assert all_audio_tokens[0].shape == (context_length, NUM_CODEBOOKS)
+
+
+def test_parse_fails_on_misaligned_input():
+    """Input not divisible by num_codebooks should raise AssertionError."""
+    # 2 header elements + 5 data elements (not divisible by 37)
+    input_ids = torch.tensor([0, 1] + [0] * 5)
+    with pytest.raises(AssertionError, match="divisible by"):
+        parse_batched_audio_input(input_ids, num_codebooks=NUM_CODEBOOKS)
+
+
+def test_parse_custom_num_codebooks():
+    """Works with a non-default num_codebooks value."""
+    cb = 4
+    ctx_frames = 1
+    context_length = 2
+    total_frames = ctx_frames + context_length
+    tokens = torch.arange(total_frames * cb).reshape(total_frames, cb)
+    input_ids = _make_input_ids((ctx_frames, context_length, tokens))
+
+    all_audio_tokens, all_ctx_frames = parse_batched_audio_input(input_ids, num_codebooks=cb)
+
+    assert all_audio_tokens[0].shape == (total_frames, cb)
+    torch.testing.assert_close(all_audio_tokens[0], tokens)
+
+
+# ──────────────────────────────────────────────────────────────────
+# apply_ctx_frames_cutting tests
+# ──────────────────────────────────────────────────────────────────
+
+
+def test_cut_removes_leading_samples():
+    """Context frames are removed from the front of the audio array."""
+    downsample_factor = 240
+    ctx_frames = 5
+    total_samples = 2400  # 10 frames * 240
+    audio = torch.arange(total_samples, dtype=torch.float32)
+
+    result = apply_ctx_frames_cutting([audio], [ctx_frames], downsample_factor)
+
+    expected_cut = ctx_frames * downsample_factor  # 1200
+    assert len(result) == 1
+    assert result[0].shape[0] == total_samples - expected_cut
+    torch.testing.assert_close(result[0], audio[expected_cut:])
+
+
+def test_cut_zero_ctx_frames_unchanged():
+    """ctx_frames=0 leaves the audio unchanged."""
+    downsample_factor = 240
+    audio = torch.randn(2400)
+
+    result = apply_ctx_frames_cutting([audio], [0], downsample_factor)
+
+    assert result[0].shape == audio.shape
+    torch.testing.assert_close(result[0], audio)
+
+
+def test_cut_multiple_requests():
+    """Each request in the batch gets its own ctx_frames cut."""
+    downsample_factor = 100
+    audio1 = torch.arange(1000, dtype=torch.float32)  # 10 frames
+    audio2 = torch.arange(500, dtype=torch.float32)  # 5 frames
+
+    result = apply_ctx_frames_cutting([audio1, audio2], [3, 0], downsample_factor)
+
+    assert len(result) == 2
+    assert result[0].shape[0] == 1000 - 3 * 100  # 700
+    assert result[1].shape[0] == 500  # unchanged
+    torch.testing.assert_close(result[0], audio1[300:])
+    torch.testing.assert_close(result[1], audio2)
+
+
+def test_cut_all_frames_returns_empty():
+    """Cutting all frames returns an empty tensor."""
+    downsample_factor = 100
+    ctx_frames = 5
+    audio = torch.randn(500)  # exactly 5 frames
+
+    result = apply_ctx_frames_cutting([audio], [ctx_frames], downsample_factor)
+
+    assert result[0].numel() == 0
@@ -0,0 +1,330 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Tests for CUDA Graph acoustic transformer wrapper numerical equivalence.
+
+Verifies that CUDA Graph-accelerated decoding produces results equivalent
+to eager mode, with special attention to padding cases where zero-padding
+may introduce small numerical differences due to the flow matching ODE.
+
+Architecture note: the real FlowMatchingAudioTransformer uses a multi-step
+Euler ODE with CFG.  The synthetic model here uses simple linear layers
+to exercise the wrapper mechanism while keeping the test lightweight.
+"""
+
+import importlib.util
+import os
+
+import pytest
+import torch
+import torch.nn as nn
+
+pytestmark = [
+    pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"),
+    pytest.mark.core_model,
+    pytest.mark.cuda,
+    pytest.mark.L4,
+]
+
+DEVICE = torch.device("cuda:0")
+HIDDEN_DIM = 64
+N_ACOUSTIC_CODEBOOK = 7
+SEMANTIC_CODEBOOK_SIZE = 128
+ACOUSTIC_EMBEDDINGS_LEVELS = 1024
+
+# Load CUDAGraphAcousticTransformerWrapper: try package import first, fall back to direct file load
+try:
+    from vllm_omni.model_executor.models.voxtral_tts.cuda_graph_acoustic_transformer_wrapper import (
+        CUDAGraphAcousticTransformerWrapper,
+    )
+    from vllm_omni.model_executor.models.voxtral_tts.voxtral_tts_audio_generation import (
+        AudioSpecialTokens,
+    )
+except Exception:
+    _WRAPPER_PATH = os.path.join(
+        os.path.dirname(__file__),
+        os.pardir,
+        os.pardir,
+        os.pardir,
+        os.pardir,
+        "vllm_omni",
+        "model_executor",
+        "models",
+        "voxtral_tts",
+        "cuda_graph_acoustic_transformer_wrapper.py",
+    )
+    _spec = importlib.util.spec_from_file_location(
+        "cuda_graph_acoustic_transformer_wrapper", os.path.abspath(_WRAPPER_PATH)
+    )
+    _mod = importlib.util.module_from_spec(_spec)
+    _spec.loader.exec_module(_mod)
+    CUDAGraphAcousticTransformerWrapper = _mod.CUDAGraphAcousticTransformerWrapper
+
+    _AUDIO_GEN_PATH = os.path.join(
+        os.path.dirname(__file__),
+        os.pardir,
+        os.pardir,
+        os.pardir,
+        os.pardir,
+        "vllm_omni",
+        "model_executor",
+        "models",
+        "voxtral_tts",
+        "voxtral_tts_audio_generation.py",
+    )
+    _spec2 = importlib.util.spec_from_file_location("voxtral_tts_audio_generation", os.path.abspath(_AUDIO_GEN_PATH))
+    _mod2 = importlib.util.module_from_spec(_spec2)
+    _spec2.loader.exec_module(_mod2)
+    AudioSpecialTokens = _mod2.AudioSpecialTokens
+
+
+class SyntheticModelArgs:
+    """Mimics MultimodalAudioModelArgs interface."""
+
+    def __init__(self):
+        self.semantic_codebook_size = SEMANTIC_CODEBOOK_SIZE
+        self.n_acoustic_codebook = N_ACOUSTIC_CODEBOOK
+
+
+class SyntheticAcousticTransformer(nn.Module):
+    """A small acoustic transformer mimicking FlowMatchingAudioTransformer's interface.
+
+    Uses simple linear layers so the test stays lightweight while exercising
+    the full wrapper mechanism (semantic logit, flow matching ODE, quantization).
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.model_args = SyntheticModelArgs()
+        self.acoustic_embeddings_levels = ACOUSTIC_EMBEDDINGS_LEVELS
+
+        # semantic_codebook_output: hidden_dim -> padded_codebook_size
+        padded_semantic_size = len(AudioSpecialTokens) + SEMANTIC_CODEBOOK_SIZE
+        # Round up to multiple of 128 like the real model
+        padded_semantic_size = 128 * ((padded_semantic_size + 127) // 128)
+        self.semantic_codebook_output = nn.Linear(HIDDEN_DIM, padded_semantic_size, bias=False)
+
+        # time_embedding: (B, 1) -> (B, dim)
+        self.time_embedding = nn.Linear(1, HIDDEN_DIM, bias=False)
+
+        # _predict_velocity: takes x_t (B, C), llm_output (B, D), t_emb (B, D) -> (B, C)
+        self._velocity_proj = nn.Linear(N_ACOUSTIC_CODEBOOK + HIDDEN_DIM + HIDDEN_DIM, N_ACOUSTIC_CODEBOOK, bias=False)
+
+    def _predict_velocity(self, x_t: torch.Tensor, llm_output: torch.Tensor, t_emb: torch.Tensor) -> torch.Tensor:
+        combined = torch.cat([x_t, llm_output, t_emb], dim=-1)
+        return self._velocity_proj(combined)
+
+
+class SyntheticModel(nn.Module):
+    """Mimics VoxtralTTSAudioGenerationForConditionalGeneration interface."""
+
+    def __init__(self):
+        super().__init__()
+        self.acoustic_transformer = SyntheticAcousticTransformer()
+        end_audio_id = AudioSpecialTokens.id(AudioSpecialTokens.end_audio)
+        empty_audio_id = AudioSpecialTokens.id(AudioSpecialTokens.empty_audio)
+        self.end_audio_id = end_audio_id
+        self.empty_audio_id = empty_audio_id
+
+    def compute_mm_logits(
+        self,
+        hidden_states: torch.Tensor,
+        mm_sampling_tensors=None,
+    ):
+        """Eager fallback path: replicate what the wrapper does."""
+        at = self.acoustic_transformer
+        B = hidden_states.shape[0]
+
+        empty_audio_id = self.empty_audio_id
+        end_audio_id = self.end_audio_id
+        semantic_mask_start = len(AudioSpecialTokens) + at.model_args.semantic_codebook_size
+
+        # Semantic logits
+        semantic_logit = at.semantic_codebook_output(hidden_states).float()
+        semantic_logit[:, empty_audio_id] = -float("inf")
+        semantic_logit[:, semantic_mask_start:] = -float("inf")
+        semantic_code = semantic_logit.argmax(dim=-1, keepdim=True)
+
+        # Flow matching Euler ODE
+        should_decode = semantic_code.squeeze(1) != end_audio_id
+        x = torch.randn(B, at.model_args.n_acoustic_codebook, device=hidden_states.device, dtype=hidden_states.dtype)
+        hidden_zero = torch.zeros_like(hidden_states)
+        timesteps = torch.linspace(0, 1, 16, device=hidden_states.device, dtype=hidden_states.dtype)
+
+        for i in range(len(timesteps) - 1):
+            t = timesteps[i]
+            dt = timesteps[i + 1] - timesteps[i]
+            t_emb = at.time_embedding(t.view(-1, 1).repeat(B, 1)).to(hidden_states.dtype)
+            x_batched = torch.cat([x, x], dim=0)
+            llm_batched = torch.cat([hidden_states, hidden_zero], dim=0)
+            t_emb_batched = t_emb.repeat(2, 1)
+            v_all = at._predict_velocity(x_t=x_batched, llm_output=llm_batched, t_emb=t_emb_batched)
+            v_t, uncond_v_t = v_all[:B], v_all[B:]
+            cfg_alpha = 1.2
+            v_t = cfg_alpha * v_t + (1 - cfg_alpha) * uncond_v_t
+            x = x + v_t * dt
+
+        sampled = torch.clamp(x, -1, 1)
+        scaled_x = ((sampled + 1) / 2) * (at.acoustic_embeddings_levels - 1)
+        output_codes = scaled_x.round().long()
+        output_codes[~should_decode] = empty_audio_id
+        acoustic_codes = output_codes + len(AudioSpecialTokens)
+
+        audio_codes = torch.cat([semantic_code, acoustic_codes], dim=1)
+        fake_eos = torch.where(
+            audio_codes[:, 0] == end_audio_id,
+            torch.tensor(1.0, dtype=hidden_states.dtype, device=hidden_states.device),
+            torch.tensor(0.0, dtype=hidden_states.dtype, device=hidden_states.device),
+        )
+
+        audio_list = list(torch.split(audio_codes.unsqueeze(1), 1, dim=0))
+        mm_tokens = {"audio": audio_list}
+        return fake_eos, mm_tokens
+
+
+@pytest.fixture(scope="module")
+def model():
+    """Create a synthetic model on CUDA with fixed weights in bfloat16."""
+    torch.manual_seed(42)
+    return SyntheticModel().to(device=DEVICE, dtype=torch.bfloat16).eval()
+
+
+@pytest.fixture(scope="module")
+def wrapper(model):
+    """Create a warmed-up CUDAGraphAcousticTransformerWrapper."""
+    w = CUDAGraphAcousticTransformerWrapper(
+        model=model,
+        capture_sizes=[1, 2, 4, 8, 16, 32],
+    )
+    w._warmup_and_capture(
+        device=DEVICE,
+        dtype=torch.bfloat16,
+        hidden_dim=HIDDEN_DIM,
+    )
+    return w
+
+
+def _random_hidden(batch_size, device=DEVICE, dtype=torch.bfloat16):
+    return torch.randn(batch_size, HIDDEN_DIM, device=device, dtype=dtype)
+
+
+def _unpack_audio_codes(result):
+    """Unpack (fake_eos, {"audio": [list of tensors]}) into (fake_eos, audio_codes)."""
+    fake_eos, mm_tokens = result
+    audio_list = mm_tokens["audio"]
+    # Each element is (1, 1, 1+n_acoustic_codebook), concat along batch dim
+    audio_codes = torch.cat(audio_list, dim=0).squeeze(1)  # (B, 1+n_acoustic)
+    return fake_eos, audio_codes
+
+
+# ──────────────────────────────────────────────────────────────────
+# 1. Exact-size inputs (no padding) → correct format and bounds
+# ──────────────────────────────────────────────────────────────────
+
+
+@pytest.mark.parametrize("batch_size", [1, 2, 4, 8, 16, 32])
+def test_exact_size_output_format(model, wrapper, batch_size):
+    """Graph path returns correctly shaped and bounded outputs."""
+    hidden = _random_hidden(batch_size)
+    with torch.no_grad():
+        graph_eos, graph_codes = _unpack_audio_codes(wrapper(hidden))
+    assert graph_eos.shape == (batch_size,)
+    assert graph_codes.shape == (batch_size, 1 + N_ACOUSTIC_CODEBOOK)
+    # fake_eos should be 0.0 or 1.0
+    assert torch.all((graph_eos == 0.0) | (graph_eos == 1.0))
+    # Audio codes should be non-negative
+    assert torch.all(graph_codes >= 0)
+
+
+@pytest.mark.parametrize("batch_size", [1, 2, 4, 8, 16, 32])
+def test_exact_size_deterministic(model, wrapper, batch_size):
+    """Same input + same RNG state produces identical CUDA graph output."""
+    hidden = _random_hidden(batch_size)
+    with torch.no_grad():
+        torch.manual_seed(42)
+        eos1, codes1 = _unpack_audio_codes(wrapper(hidden))
+        torch.manual_seed(42)
+        eos2, codes2 = _unpack_audio_codes(wrapper(hidden))
+    torch.testing.assert_close(eos1, eos2, atol=0, rtol=0)
+    torch.testing.assert_close(codes1, codes2, atol=0, rtol=0)
+
+
+# ──────────────────────────────────────────────────────────────────
+# 2. Padded inputs (zero-padding to nearest capture size)
+# ──────────────────────────────────────────────────────────────────
+
+
+@pytest.mark.parametrize("batch_size", [3, 5, 7, 10, 24])
+def test_padded_output_shape(model, wrapper, batch_size):
+    """Padded decode must return output trimmed to actual batch size."""
+    hidden = _random_hidden(batch_size)
+    with torch.no_grad():
+        graph_eos, graph_codes = _unpack_audio_codes(wrapper(hidden))
+    assert graph_eos.shape == (batch_size,)
+    assert graph_codes.shape == (batch_size, 1 + N_ACOUSTIC_CODEBOOK)
+
+
+@pytest.mark.parametrize("batch_size", [3, 5, 7, 10, 24])
+def test_padded_output_bounded(model, wrapper, batch_size):
+    """Padded output audio codes should be non-negative integers."""
+    hidden = _random_hidden(batch_size)
+    with torch.no_grad():
+        graph_eos, graph_codes = _unpack_audio_codes(wrapper(hidden))
+    # fake_eos should be 0.0 or 1.0
+    assert torch.all((graph_eos == 0.0) | (graph_eos == 1.0))
+    # Audio codes should be non-negative
+    assert torch.all(graph_codes >= 0)
+
+
+# ──────────────────────────────────────────────────────────────────
+# 3. Batch size exceeds all capture sizes so fallback to eager
+# ──────────────────────────────────────────────────────────────────
+
+
+@pytest.mark.parametrize("batch_size", [33, 48, 64])
+def test_fallback_eager_exact_match(model, wrapper, batch_size):
+    """Cudagraph fallback to eager. Two eager runs should produce identical results."""
+    hidden = _random_hidden(batch_size)
+    with torch.no_grad():
+        torch.manual_seed(100)
+        eager_eos, eager_codes = _unpack_audio_codes(model.compute_mm_logits(hidden))
+        torch.manual_seed(100)
+        graph_eos, graph_codes = _unpack_audio_codes(wrapper(hidden))
+    torch.testing.assert_close(graph_eos, eager_eos, atol=0, rtol=0)
+    torch.testing.assert_close(graph_codes, eager_codes, atol=0, rtol=0)
+
+
+# ──────────────────────────────────────────────────────────────────
+# 4. Disabled wrapper so fallback to eager
+# ──────────────────────────────────────────────────────────────────
+
+
+def test_disabled_wrapper_matches_eager(model, wrapper):
+    """Cudagraph fallback to eager. Two eager runs should produce identical results."""
+    hidden = _random_hidden(4)
+    wrapper.enabled = False
+    with torch.no_grad():
+        torch.manual_seed(200)
+        eager_eos, eager_codes = _unpack_audio_codes(model.compute_mm_logits(hidden))
+        torch.manual_seed(200)
+        graph_eos, graph_codes = _unpack_audio_codes(wrapper(hidden))
+    wrapper.enabled = True
+    torch.testing.assert_close(graph_eos, eager_eos, atol=0, rtol=0)
+    torch.testing.assert_close(graph_codes, eager_codes, atol=0, rtol=0)
+
+
+# ──────────────────────────────────────────────────────────────────
+# 5. Determinism - same input + same RNG state on cudagraph runs
+# ──────────────────────────────────────────────────────────────────
+
+
+def test_deterministic_across_calls(model, wrapper):
+    """Same input + same RNG state. Two cudagraph runs should produce identical results."""
+    hidden = _random_hidden(4)
+    with torch.no_grad():
+        torch.manual_seed(300)
+        eos1, codes1 = _unpack_audio_codes(wrapper(hidden))
+        torch.manual_seed(300)
+        eos2, codes2 = _unpack_audio_codes(wrapper(hidden))
+    torch.testing.assert_close(eos1, eos2, atol=0, rtol=0)
+    torch.testing.assert_close(codes1, codes2, atol=0, rtol=0)
@@ -0,0 +1,281 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from collections import defaultdict
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.stage_input_processors.voxtral_tts import (
+    generator2tokenizer,
+    generator2tokenizer_async_chunk,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+# ---- Helpers for generator2tokenizer (non-streaming) ----
+
+
+def _make_stage(engine_outputs):
+    return SimpleNamespace(engine_outputs=engine_outputs)
+
+
+def _make_output(audio_tensors):
+    """Build a single generator output with multimodal_output["audio"] = audio_tensors."""
+    return SimpleNamespace(
+        outputs=[SimpleNamespace(multimodal_output={"audio": audio_tensors})],
+    )
+
+
+# ---- Tests for generator2tokenizer ----
+
+
+class TestGenerator2Tokenizer:
+    def test_single_output_single_chunk(self):
+        """Single output with one audio tensor produces correct prompt_token_ids."""
+        audio = torch.tensor([10, 20, 30, 40])
+        stage = _make_stage([_make_output([audio])])
+
+        result = generator2tokenizer([stage], engine_input_source=[0])
+
+        assert len(result) == 1
+        assert result[0]["prompt_token_ids"] == [10, 20, 30, 40]
+        assert result[0]["multi_modal_data"] is None
+
+    def test_2d_audio_tensors_flattened(self):
+        """2D audio tensors (e.g., multi-codebook frames) are flattened."""
+        audio = torch.tensor([[1, 2, 3], [4, 5, 6]])  # (2, 3)
+        stage = _make_stage([_make_output([audio])])
+
+        result = generator2tokenizer([stage], engine_input_source=[0])
+
+        assert result[0]["prompt_token_ids"] == [1, 2, 3, 4, 5, 6]
+
+    def test_empty_engine_input_source_raises(self):
+        """Empty engine_input_source raises ValueError."""
+        stage = _make_stage([_make_output([torch.tensor([1])])])
+        with pytest.raises(ValueError, match="cannot be empty"):
+            generator2tokenizer([stage], engine_input_source=[])
+
+    def test_invalid_stage_id_raises(self):
+        """Stage id beyond stage_list length raises IndexError."""
+        stage = _make_stage([_make_output([torch.tensor([1])])])
+        with pytest.raises(IndexError, match="Invalid stage_id"):
+            generator2tokenizer([stage], engine_input_source=[5])
+
+    def test_no_outputs_yet_raises(self):
+        """Stage with engine_outputs=None raises RuntimeError."""
+        stage = _make_stage(engine_outputs=None)
+        with pytest.raises(RuntimeError, match="no outputs yet"):
+            generator2tokenizer([stage], engine_input_source=[0])
+
+
+# ---- Helpers for generator2tokenizer_async_chunk (streaming) ----
+
+
+def _req(external_req_id: str, *, finished: bool):
+    return SimpleNamespace(
+        external_req_id=external_req_id,
+        is_finished=lambda: finished,
+    )
+
+
+def _make_transfer_manager(codec_chunk_frames=25, codec_left_context_frames=25, codec_chunk_frames_at_begin=5):
+    return SimpleNamespace(
+        code_prompt_token_ids=defaultdict(list),
+        connector=SimpleNamespace(
+            config={
+                "extra": {
+                    "codec_chunk_frames": codec_chunk_frames,
+                    "codec_left_context_frames": codec_left_context_frames,
+                    "codec_chunk_frames_at_begin": codec_chunk_frames_at_begin,
+                }
+            }
+        ),
+    )
+
+
+def test_empty_chunk_when_not_finished():
+    """Returns None when no audio data and not finished."""
+    transfer_manager = _make_transfer_manager()
+    request = _req("rid-empty", finished=False)
+
+    payload = generator2tokenizer_async_chunk(
+        transfer_manager=transfer_manager,
+        pooling_output={"audio": torch.zeros((0,))},
+        request=request,
+    )
+
+    assert payload is None
+
+
+def test_flush_tail_when_finished():
+    """Emits remaining frames on finish with finished=True."""
+    transfer_manager = _make_transfer_manager()
+    request_id = "rid-tail"
+    # Pre-populate with 24 frames of 8-element codes
+    transfer_manager.code_prompt_token_ids[request_id] = [[1, 2, 3, 4, 5, 6, 7, 8] for _ in range(24)]
+    request = _req(request_id, finished=True)
+
+    payload = generator2tokenizer_async_chunk(
+        transfer_manager=transfer_manager,
+        pooling_output=None,  # e.g. EOS step with no audio
+        request=request,
+    )
+
+    assert payload is not None
+    assert payload["finished"].item() is True
+    codes = payload["code_predictor_codes"]
+    # Format: [ctx_frames, context_length, ...flat_codes]
+    assert len(codes) >= 2  # At least ctx_frames + context_length header
+    ctx_frames = codes[0]
+    context_length = codes[1]
+    flat_codes = codes[2:]
+    # Total frames in window = ctx_frames + context_length
+    total_window_frames = ctx_frames + context_length
+    assert len(flat_codes) == total_window_frames * 8
+
+
+def test_eof_marker_when_finished_with_no_frames():
+    """Emits EOF marker with empty codes when finished with no accumulated frames."""
+    transfer_manager = _make_transfer_manager()
+    request = _req("rid-eof", finished=True)
+
+    payload = generator2tokenizer_async_chunk(
+        transfer_manager=transfer_manager,
+        pooling_output=None,
+        request=request,
+    )
+
+    assert payload == {
+        "code_predictor_codes": [],
+        "finished": torch.tensor(True, dtype=torch.bool),
+    }
+
+
+def test_normal_chunk_emission():
+    """Emits chunks at codec_chunk_frames boundary."""
+    transfer_manager = _make_transfer_manager(
+        codec_chunk_frames=25,
+        codec_left_context_frames=25,
+        codec_chunk_frames_at_begin=5,
+    )
+    request_id = "rid-chunk"
+    request = _req(request_id, finished=False)
+
+    # Feed 25 frames one by one
+    for i in range(25):
+        pooling_output = {"audio": torch.tensor([float(i)] * 4)}
+        payload = generator2tokenizer_async_chunk(
+            transfer_manager=transfer_manager,
+            pooling_output=pooling_output,
+            request=request,
+        )
+
+    # A chunk should be emitted
+    assert payload is not None
+    codes = payload["code_predictor_codes"]
+    ctx_frames = codes[0]
+    context_length = codes[1]
+    assert ctx_frames == 20  # 25 - 5(chunk_size_at_begin)
+    assert context_length == 5  # 25 - 20(chunk_size_at_begin * 4)
+
+
+def test_small_initial_chunks():
+    """Uses codec_chunk_frames_at_begin for initial frames when length <= chunk_size."""
+    transfer_manager = _make_transfer_manager(
+        codec_chunk_frames=25,
+        codec_left_context_frames=25,
+        codec_chunk_frames_at_begin=5,
+    )
+    request_id = "rid-begin"
+    request = _req(request_id, finished=False)
+
+    # Feed 5 frames (should trigger emission because codec_chunk_frames_at_begin=5)
+    for i in range(5):
+        pooling_output = {"audio": torch.tensor([float(i + 1)] * 3)}
+        payload = generator2tokenizer_async_chunk(
+            transfer_manager=transfer_manager,
+            pooling_output=pooling_output,
+            request=request,
+        )
+
+    assert payload is not None
+    codes = payload["code_predictor_codes"]
+    ctx_frames = codes[0]
+    context_length = codes[1]
+    assert ctx_frames == 0
+    assert context_length == 5  # chunk_size_at_begin
+
+
+def test_no_emission_between_boundaries():
+    """No chunk emitted when not at chunk boundary and not finished."""
+    transfer_manager = _make_transfer_manager(
+        codec_chunk_frames=25,
+        codec_left_context_frames=25,
+        codec_chunk_frames_at_begin=5,
+    )
+    request_id = "rid-mid"
+    request = _req(request_id, finished=False)
+
+    # Feed 3 frames (3 % 5 != 0, should not emit)
+    for i in range(3):
+        pooling_output = {"audio": torch.tensor([float(i)] * 4)}
+        payload = generator2tokenizer_async_chunk(
+            transfer_manager=transfer_manager,
+            pooling_output=pooling_output,
+            request=request,
+        )
+
+    assert payload is None
+
+
+def test_context_handling_format():
+    """Verifies [ctx_frames, context_length, ...codes] format."""
+    transfer_manager = _make_transfer_manager(
+        codec_chunk_frames=10,
+        codec_left_context_frames=5,
+        codec_chunk_frames_at_begin=5,
+    )
+    request_id = "rid-ctx"
+    request = _req(request_id, finished=False)
+
+    # Feed 5 frames to trigger a chunk (chunk_size_at_begin=5)
+    for i in range(5):
+        pooling_output = {"audio": torch.tensor([float(i + 10)] * 2)}  # codebook_dim=2
+        payload = generator2tokenizer_async_chunk(
+            transfer_manager=transfer_manager,
+            pooling_output=pooling_output,
+            request=request,
+        )
+
+    assert payload is not None
+    codes = payload["code_predictor_codes"]
+    # First two elements are ctx_frames and context_length
+    ctx_frames = codes[0]
+    context_length = codes[1]
+    flat_codes = codes[2:]
+    # The window has ctx_frames + context_length frames
+    assert isinstance(ctx_frames, int)
+    assert isinstance(context_length, int)
+    assert ctx_frames >= 0
+    assert context_length > 0
+    # flat_codes = total_window_frames * codebook_dim
+    total_window_frames = ctx_frames + context_length
+    assert len(flat_codes) == total_window_frames * 2  # codebook_dim=2
+
+
+def test_none_pooling_output_not_finished_returns_none():
+    """None pooling_output when not finished returns None."""
+    transfer_manager = _make_transfer_manager()
+    request = _req("rid-none", finished=False)
+
+    payload = generator2tokenizer_async_chunk(
+        transfer_manager=transfer_manager,
+        pooling_output=None,
+        request=request,
+    )
+
+    assert payload is None
@@ -10,9 +10,9 @@ stage_args:
    stage_type: llm  # Use llm stage type for AR stages
    runtime:
      devices: "0"
-      max_batch_size: 64
    engine_args:
      model_stage: thinker
+      max_num_seqs: 64
      model_arch: Qwen3OmniMoeForConditionalGeneration
      worker_type: ar
      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -41,9 +41,9 @@ stage_args:
    stage_type: llm  # Use llm stage type for AR stages
    runtime:
      devices: "1"
-      max_batch_size: 64
    engine_args:
      model_stage: talker
+      max_num_seqs: 64
      model_arch: Qwen3OmniMoeForConditionalGeneration
      worker_type: ar
      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -72,9 +72,9 @@ stage_args:
    stage_type: llm  # Use llm stage type for AR stages
    runtime:
      devices: "1"
-      max_batch_size: 64
    engine_args:
      model_stage: code2wav
+      max_num_seqs: 64
      model_arch: Qwen3OmniMoeForConditionalGeneration
      worker_type: generation
      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
@@ -10,8 +10,8 @@ stage_args:
    is_comprehension: true
    runtime:
      devices: "0"
-      max_batch_size: 4
    engine_args:
+      max_num_seqs: 4
      model_stage: qwen3_tts
      model_arch: Qwen3TTSTalkerForConditionalGeneration
      worker_type: ar
@@ -41,8 +41,8 @@ stage_args:
    stage_type: llm
    runtime:
      devices: "0"
-      max_batch_size: 4
    engine_args:
+      max_num_seqs: 4
      model_stage: code2wav
      model_arch: Qwen3TTSCode2Wav
      worker_type: generation
@@ -12,11 +12,13 @@
                "endpoint": "/v1/chat/completions",
                "num_prompts": [
                    10,
-                    40
+                    40,
+                    100
                ],
                "max_concurrency": [
                    1,
-                    4
+                    4,
+                    10
                ],
                "random_input_len": 100,
                "random_output_len": 100,
@@ -34,11 +36,13 @@
                "endpoint": "/v1/chat/completions",
                "num_prompts": [
                    10,
-                    40
+                    40,
+                    100
                ],
                "max_concurrency": [
                    1,
-                    4
+                    4,
+                    10
                ],
                "random_input_len": 100,
                "random_output_len": 100,
@@ -96,11 +100,13 @@
                "endpoint": "/v1/chat/completions",
                "num_prompts": [
                    10,
-                    40
+                    40,
+                    100
                ],
                "max_concurrency": [
                    1,
-                    4
+                    4,
+                    10
                ],
                "random_input_len": 100,
                "random_output_len": 100,
@@ -118,11 +124,13 @@
                "endpoint": "/v1/chat/completions",
                "num_prompts": [
                    10,
-                    40
+                    40,
+                    100
                ],
                "max_concurrency": [
                    1,
-                    4
+                    4,
+                    10
                ],
                "random_input_len": 100,
                "random_output_len": 100,
@@ -83,6 +83,8 @@ class TestStageConfig:

    def test_to_omegaconf_with_runtime_overrides(self):
        """Test that runtime overrides are applied to OmegaConf output."""
+        import warnings
+
        config = StageConfig(
            stage_id=0,
            model_stage="thinker",
@@ -93,12 +95,41 @@ class TestStageConfig:
                "max_batch_size": 64,
            },
        )
-        omega_config = config.to_omegaconf()
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore", FutureWarning)
+            omega_config = config.to_omegaconf()

        assert omega_config.engine_args.gpu_memory_utilization == 0.9
        assert omega_config.engine_args.tensor_parallel_size == 2
        assert omega_config.runtime.devices == "0,1"
-        assert omega_config.runtime.max_batch_size == 64
+        # max_batch_size is migrated to engine_args.max_num_seqs
+        assert omega_config.engine_args.max_num_seqs == 64
+
+    def test_to_omegaconf_max_batch_size_deprecation(self):
+        """Test that runtime.max_batch_size emits a FutureWarning."""
+        import warnings
+
+        config = StageConfig(
+            stage_id=0,
+            model_stage="thinker",
+            runtime_overrides={"max_batch_size": 8},
+        )
+        with warnings.catch_warnings(record=True) as w:
+            warnings.simplefilter("always")
+            config.to_omegaconf()
+            deprecation_warnings = [x for x in w if issubclass(x.category, FutureWarning)]
+            assert len(deprecation_warnings) == 1
+            assert "max_batch_size" in str(deprecation_warnings[0].message)
+
+    def test_to_omegaconf_max_num_seqs_in_engine_args(self):
+        """Test that max_num_seqs in yaml_engine_args takes precedence."""
+        config = StageConfig(
+            stage_id=0,
+            model_stage="thinker",
+            yaml_engine_args={"max_num_seqs": 32},
+        )
+        omega_config = config.to_omegaconf()
+        assert omega_config.engine_args.max_num_seqs == 32


 class TestModelPipeline:
@@ -425,6 +456,29 @@ stages:
        errors = pipeline.validate_pipeline()
        assert errors == [], f"Unexpected validation errors: {errors}"

+    def test_parse_yaml_migrates_runtime_max_batch_size(self, tmp_path):
+        """Test that runtime.max_batch_size is migrated to engine_args.max_num_seqs."""
+        yaml_content = """\
+model_type: legacy_batch
+stages:
+  - stage_id: 0
+    model_stage: entry
+    stage_type: llm
+    input_sources: []
+    runtime:
+      devices: "0"
+      max_batch_size: 16
+    engine_args:
+      model_arch: SomeModel
+"""
+        yaml_file = tmp_path / "legacy_batch.yaml"
+        yaml_file.write_text(yaml_content)
+
+        pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "legacy_batch")
+        s0 = pipeline.stages[0]
+        assert "max_batch_size" not in s0.yaml_runtime
+        assert s0.yaml_engine_args.get("max_num_seqs") == 16
+
    def test_parse_diffusion_stage_type(self, tmp_path):
        """Test parsing a diffusion stage type from YAML."""
        yaml_content = """\
@@ -0,0 +1,185 @@
+"""Tests for CUDAGraphWrapper.__getattr__ performance optimization.
+
+This module tests that the patched CUDAGraphWrapper avoids expensive __repr__
+calls when hasattr() is used for non-existent attributes. The original vLLM
+implementation includes {self.runnable} in the AttributeError message, which
+triggers model tree traversal and can take ~6ms on large models.
+"""
+
+import time
+
+import pytest
+import torch
+import torch.nn as nn
+
+from vllm_omni.worker.gpu_model_runner import CUDAGraphWrapper
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class SlowReprModel(nn.Module):
+    """A mock model with artificially slow __repr__ to detect unwanted calls."""
+
+    def __init__(self, repr_delay_ms: float = 10.0):
+        super().__init__()
+        self.linear = nn.Linear(16, 16)
+        self.repr_delay_ms = repr_delay_ms
+        self.repr_call_count = 0
+
+    def forward(self, x):
+        return self.linear(x)
+
+    def __repr__(self):
+        self.repr_call_count += 1
+        # Simulate expensive repr by sleeping
+        time.sleep(self.repr_delay_ms / 1000.0)
+        return f"SlowReprModel(delay={self.repr_delay_ms}ms)"
+
+
+class MockCUDAGraphWrapper:
+    """A minimal mock that mimics CUDAGraphWrapper structure for CPU testing."""
+
+    def __init__(self, runnable):
+        # Store in __dict__ directly to avoid triggering __getattr__
+        object.__setattr__(self, "runnable", runnable)
+
+    def __getattr__(self, key: str):
+        # This is the optimized implementation we're testing
+        runnable = object.__getattribute__(self, "runnable")
+        if hasattr(runnable, key):
+            return getattr(runnable, key)
+        # Key optimization: DO NOT include {self.runnable} in error message
+        # as it triggers expensive __repr__ on large models
+        raise AttributeError(f"Attribute {key} not exists in the runnable of cudagraph wrapper")
+
+
+def test_hasattr_nonexistent_does_not_trigger_repr():
+    """Verify that hasattr for non-existent attributes doesn't call __repr__."""
+    model = SlowReprModel(repr_delay_ms=100.0)  # Very slow repr
+    wrapper = MockCUDAGraphWrapper(model)
+
+    # Reset counter
+    model.repr_call_count = 0
+
+    # Call hasattr for non-existent attribute multiple times
+    for _ in range(10):
+        result = hasattr(wrapper, "nonexistent_attribute_xyz")
+        assert result is False
+
+    # __repr__ should never have been called
+    assert model.repr_call_count == 0, (
+        f"__repr__ was called {model.repr_call_count} times when checking "
+        "for non-existent attributes. This indicates the AttributeError "
+        "message contains {self.runnable} which triggers expensive repr."
+    )
+
+
+def test_hasattr_nonexistent_is_fast():
+    """Verify that hasattr for non-existent attributes is fast (<1ms per call)."""
+    model = SlowReprModel(repr_delay_ms=100.0)
+    wrapper = MockCUDAGraphWrapper(model)
+
+    num_iterations = 100
+    start = time.perf_counter()
+    for _ in range(num_iterations):
+        hasattr(wrapper, "nonexistent_attribute_xyz")
+    elapsed_ms = (time.perf_counter() - start) * 1000
+
+    avg_ms = elapsed_ms / num_iterations
+    # If __repr__ were being called, each would take ~100ms
+    # We expect <1ms per call with the fix
+    assert avg_ms < 1.0, (
+        f"hasattr for non-existent attribute took {avg_ms:.2f}ms on average. "
+        "Expected <1ms. This suggests __repr__ is being triggered."
+    )
+
+
+def test_hasattr_existing_attribute_works():
+    """Verify that hasattr for existing attributes returns True and works correctly."""
+    model = SlowReprModel()
+    wrapper = MockCUDAGraphWrapper(model)
+
+    # 'forward' exists on nn.Module
+    assert hasattr(wrapper, "forward") is True
+
+    # 'linear' exists on our model
+    assert hasattr(wrapper, "linear") is True
+
+    # Can actually access the attribute
+    linear = wrapper.linear
+    assert isinstance(linear, nn.Linear)
+
+
+def test_getattr_existing_attribute_returns_value():
+    """Verify that getattr for existing attributes returns the correct value."""
+    model = SlowReprModel()
+    wrapper = MockCUDAGraphWrapper(model)
+
+    # Access forward method
+    forward_method = wrapper.forward
+    assert callable(forward_method)
+
+    # Access linear layer
+    linear = wrapper.linear
+    assert isinstance(linear, nn.Linear)
+    assert linear.in_features == 16
+    assert linear.out_features == 16
+
+
+def test_getattr_nonexistent_raises_attribute_error():
+    """Verify that getattr for non-existent attributes raises AttributeError."""
+    model = SlowReprModel()
+    wrapper = MockCUDAGraphWrapper(model)
+
+    with pytest.raises(AttributeError) as exc_info:
+        _ = wrapper.nonexistent_attribute
+
+    # Verify error message format (should NOT contain model repr)
+    error_msg = str(exc_info.value)
+    assert "nonexistent_attribute" in error_msg
+    assert "cudagraph wrapper" in error_msg
+    # Should NOT contain the slow repr output
+    assert "SlowReprModel(delay=" not in error_msg
+
+
+def test_attribute_error_message_does_not_contain_runnable_repr():
+    """Explicitly verify the error message doesn't trigger runnable repr."""
+    model = SlowReprModel(repr_delay_ms=100.0)
+    wrapper = MockCUDAGraphWrapper(model)
+    model.repr_call_count = 0
+
+    try:
+        _ = wrapper.nonexistent_attr
+    except AttributeError:
+        pass
+
+    # __repr__ should not have been called during error construction
+    assert model.repr_call_count == 0, (
+        "AttributeError message construction triggered __repr__. The error message should not include {self.runnable}."
+    )
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+def test_real_cudagraph_wrapper_hasattr_performance():
+    """Test the actual CUDAGraphWrapper from vllm_omni (requires CUDA)."""
+    from vllm.config import CUDAGraphMode
+
+    model = SlowReprModel(repr_delay_ms=50.0).cuda()
+    model.repr_call_count = 0
+
+    # Create actual CUDAGraphWrapper
+    try:
+        wrapper = CUDAGraphWrapper(model, runtime_mode=CUDAGraphMode.NONE)
+    except Exception:
+        pytest.skip("Could not create CUDAGraphWrapper")
+
+    # Test hasattr performance
+    num_iterations = 50
+    start = time.perf_counter()
+    for _ in range(num_iterations):
+        hasattr(wrapper, "nonexistent_xyz")
+    elapsed_ms = (time.perf_counter() - start) * 1000
+
+    avg_ms = elapsed_ms / num_iterations
+    assert avg_ms < 1.0, f"Real CUDAGraphWrapper hasattr took {avg_ms:.2f}ms avg. Expected <1ms with the optimization."
+    assert model.repr_call_count == 0, f"__repr__ called {model.repr_call_count} times"
@@ -114,91 +114,127 @@ async def async_request_openai_chat_omni_completions(

    output = MixRequestFuncOutput()
    output.prompt_len = request_func_input.prompt_len
-
-    generated_text = ""
-    generated_audio = None
-    ttft = 0.0
-    st = time.perf_counter()
-    output.start_time = st
-    most_recent_timestamp = st
-    timestamp = st
-    audio_generate_time = 0.0
-    try:
-        async with session.post(url=api_url, json=payload, headers=headers) as response:
-            if response.status == 200:
-                handler = StreamedResponseHandler()
-                async for chunk_bytes in response.content.iter_any():
-                    chunk_bytes = chunk_bytes.strip()
-                    if not chunk_bytes:
-                        continue
-
-                    messages = handler.add_chunk(chunk_bytes)
-                    for message in messages:
-                        if type(message) is bytes:
-                            message = message.decode("utf-8")
-                        # NOTE: SSE comments (often used as pings) start with
-                        # a colon. These are not JSON data payload and should
-                        # be skipped.
-                        if message.startswith(":"):
+    max_retries = 3
+    retry_delay = 0.1
+    for attempt in range(max_retries + 1):
+        # Reset per-attempt state so that retries do not mix partial
+        # outputs or metrics from previous attempts.
+        generated_text = ""
+        generated_audio = None
+        ttft = 0.0
+        st = time.perf_counter()
+        output.start_time = st
+        most_recent_timestamp = st
+        timestamp = st
+        audio_generate_time = 0.0
+        output.itl = []
+        output.generated_text = ""
+        output.ttft = 0.0
+        output.audio_ttfp = 0.0
+        output.audio_duration = 0.0
+        output.audio_frames = 0
+        output.audio_rtf = 0.0
+        output.text_latency = 0.0
+        output.output_tokens = 0
+        output.error = ""
+        output.success = False
+        try:
+            async with session.post(url=api_url, json=payload, headers=headers) as response:
+                if response.status == 200:
+                    handler = StreamedResponseHandler()
+                    async for chunk_bytes in response.content.iter_any():
+                        chunk_bytes = chunk_bytes.strip()
+                        if not chunk_bytes:
                            continue

-                        chunk = message.removeprefix("data: ")
-                        if chunk != "[DONE]":
-                            timestamp = time.perf_counter()
-                            data = json.loads(chunk)
-                            if choices := data.get("choices"):
-                                modality = data.get("modality")
-                                content = choices[0]["delta"].get("content")
-                                if modality == "text":
-                                    # First token
-                                    if ttft == 0.0:
-                                        ttft = timestamp - st
-                                        output.ttft = ttft
-                                    else:
-                                        output.itl.append(timestamp - most_recent_timestamp)
-                                    generated_text += content or ""
-                                    most_recent_timestamp = timestamp
-                                    output.text_latency = timestamp - st
-                                elif modality == "audio":
-                                    if output.audio_ttfp == 0.0:
-                                        output.audio_ttfp = timestamp - st
-                                    audio_generate_time = timestamp - st
-                                    if content != "":
-                                        audio_bytes = base64.b64decode(content)
-                                        seg = AudioSegment.from_file(io.BytesIO(audio_bytes))
-                                        if seg is not None:
-                                            if generated_audio is None:
-                                                generated_audio = seg
-                                            else:
-                                                generated_audio = generated_audio + seg
+                        messages = handler.add_chunk(chunk_bytes)
+                        for message in messages:
+                            if type(message) is bytes:
+                                message = message.decode("utf-8")
+                            # NOTE: SSE comments (often used as pings) start with
+                            # a colon. These are not JSON data payload and should
+                            # be skipped.
+                            if message.startswith(":"):
+                                continue

-                            if metrics := data.get("metrics"):
-                                output.output_tokens = metrics.get("num_tokens_out", 0)
+                            chunk = message.removeprefix("data: ")
+                            if chunk != "[DONE]":
+                                timestamp = time.perf_counter()
+                                data = json.loads(chunk)
+                                if choices := data.get("choices"):
+                                    modality = data.get("modality")
+                                    content = choices[0]["delta"].get("content")
+                                    if modality == "text":
+                                        # First token
+                                        if ttft == 0.0:
+                                            ttft = timestamp - st
+                                            output.ttft = ttft
+                                        else:
+                                            output.itl.append(timestamp - most_recent_timestamp)
+                                        generated_text += content or ""
+                                        most_recent_timestamp = timestamp
+                                        output.text_latency = timestamp - st
+                                    elif modality == "audio":
+                                        if output.audio_ttfp == 0.0:
+                                            output.audio_ttfp = timestamp - st
+                                        audio_generate_time = timestamp - st
+                                        if content != "":
+                                            audio_bytes = base64.b64decode(content)
+                                            seg = AudioSegment.from_file(io.BytesIO(audio_bytes))
+                                            if seg is not None:
+                                                if generated_audio is None:
+                                                    generated_audio = seg
+                                                else:
+                                                    generated_audio = generated_audio + seg

-                output.latency = timestamp - st
-                output.generated_text = generated_text
-                if generated_audio is not None:
-                    output.audio_duration = len(generated_audio) / 1000.0
-                    frame_width = generated_audio.frame_width
-                    if frame_width > 0:
-                        output.audio_frames = len(generated_audio.raw_data) // frame_width
-                    else:
-                        output.audio_frames = 0
-                        logger.warning("Audio frame width is zero")
-                    audio_duration = output.audio_duration
-                    if audio_duration > 0:
-                        output.audio_rtf = audio_generate_time / output.audio_duration
-                    else:
-                        output.audio_rtf = 0
-                        logger.warning("Audio duration is zero")
-                output.success = True
-            else:
-                output.error = response.reason or ""
-                output.success = False
-    except Exception:
-        output.success = False
-        output.error = traceback.format_exc()
-        logger.error(f"ERROR: send request failed, reason is: {output.error}")
+                                if metrics := data.get("metrics"):
+                                    output.output_tokens = metrics.get("num_tokens_out", 0)
+
+                    output.latency = timestamp - st
+                    output.generated_text = generated_text
+                    if generated_audio is not None:
+                        output.audio_duration = len(generated_audio) / 1000.0
+                        frame_width = generated_audio.frame_width
+                        if frame_width > 0:
+                            output.audio_frames = len(generated_audio.raw_data) // frame_width
+                        else:
+                            output.audio_frames = 0
+                            logger.warning("Audio frame width is zero")
+                        audio_duration = output.audio_duration
+                        if audio_duration > 0:
+                            output.audio_rtf = audio_generate_time / output.audio_duration
+                        else:
+                            output.audio_rtf = 0
+                            logger.warning("Audio duration is zero")
+                    output.success = True
+                else:
+                    output.error = response.reason or ""
+                    output.success = False
+            break
+        except aiohttp.ClientError as e:
+            # transient transport error: may retry
+            output.success = False
+            output.error = traceback.format_exc()
+            if attempt < max_retries:
+                logger.warning(
+                    "ClientError in omni benchmark request (will retry): attempt=%d/%d delay=%.2fs: %s",
+                    attempt + 1,
+                    max_retries + 1,
+                    retry_delay,
+                    str(e),
+                )
+                await asyncio.sleep(retry_delay)
+                continue
+            logger.error(
+                "ClientError in omni benchmark request (giving up):\n%s",
+                output.error,
+            )
+            break
+        except Exception:
+            output.success = False
+            output.error = traceback.format_exc()
+            logger.error(f"ERROR: send request failed, reason is: {output.error}")
+            break

    if pbar:
        pbar.update(1)
@@ -11,6 +11,7 @@ Runtime parameters (gpu_memory_utilization, tp_size, etc.) come from CLI.
 from __future__ import annotations

 import re
+import warnings
 from dataclasses import asdict, dataclass, field
 from enum import Enum
 from pathlib import Path
@@ -74,7 +75,7 @@ class StageConfig:

    # Per-stage engine args from pipeline YAML (defaults)
    yaml_engine_args: dict[str, Any] = field(default_factory=dict)
-    # Per-stage runtime config from pipeline YAML (devices, max_batch_size)
+    # Per-stage runtime config from pipeline YAML (devices, etc.)
    yaml_runtime: dict[str, Any] = field(default_factory=dict)
    # Pass-through fields from pipeline YAML (default_sampling_params,
    # output_connectors, input_connectors, tts_args, etc.)
@@ -109,13 +110,23 @@ class StageConfig:
        # Build runtime config from YAML defaults + CLI overrides
        runtime: dict[str, Any] = dict(self.yaml_runtime)
        runtime.setdefault("process", True)
-        runtime.setdefault("max_batch_size", self.runtime_overrides.get("max_batch_size", 1))
        if "devices" in self.runtime_overrides:
            runtime["devices"] = self.runtime_overrides["devices"]

-        # Inject max_num_seqs from max_batch_size (legacy compat)
-        max_batch_size = int(runtime.get("max_batch_size", 1) or 1)
-        engine_args.setdefault("max_num_seqs", max_batch_size)
+        # Legacy compat: migrate runtime.max_batch_size → engine_args.max_num_seqs
+        legacy_mbs = runtime.pop("max_batch_size", None)
+        cli_mbs = self.runtime_overrides.get("max_batch_size")
+        if legacy_mbs is not None or cli_mbs is not None:
+            warnings.warn(
+                "runtime.max_batch_size is deprecated and will be removed in a "
+                "future release. Use engine_args.max_num_seqs instead.",
+                FutureWarning,
+                stacklevel=2,
+            )
+            effective_mbs = int(cli_mbs or legacy_mbs or 1)
+            engine_args.setdefault("max_num_seqs", effective_mbs)
+
+        engine_args.setdefault("max_num_seqs", 1)

        # Build full config dict
        config_dict: dict[str, Any] = {
@@ -222,6 +233,7 @@ class StageConfigFactory:
        "qwen2_5_omni": "qwen2_5_omni",
        "bagel": "bagel",
        "qwen3_tts": "qwen3_tts",
+        "voxtral_tts": "voxtral_tts",
        "mimo_audio": "mimo_audio",
        "glm-image": "glm_image",
        "cosyvoice3": "cosyvoice3",
@@ -321,13 +333,14 @@ class StageConfigFactory:
        if "dtype" in engine_args:
            engine_args["dtype"] = str(engine_args["dtype"])

+        engine_args.setdefault("max_num_seqs", 1)
+
        config_dict: dict[str, Any] = {
            "stage_id": 0,
            "stage_type": StageType.DIFFUSION.value,
            "runtime": {
                "process": True,
                "devices": devices,
-                "max_batch_size": 1,
            },
            "engine_args": create_config(engine_args),
            "final_output": True,
@@ -425,6 +438,16 @@ class StageConfigFactory:
            raw_rt = stage_data.get("runtime", None)
            yaml_runtime = to_dict(raw_rt) if raw_rt is not None else {}

+            # Migrate legacy runtime.max_batch_size → engine_args.max_num_seqs
+            if "max_batch_size" in yaml_runtime:
+                mbs = yaml_runtime.pop("max_batch_size")
+                yaml_engine_args.setdefault("max_num_seqs", int(mbs))
+                logger.debug(
+                    "Stage %s: migrated runtime.max_batch_size=%s to engine_args.max_num_seqs",
+                    stage_data.get("stage_id", "?"),
+                    mbs,
+                )
+
            # Topology-level fields that also live inside engine_args in legacy
            # YAMLs (worker_type, scheduler_cls, etc.) — read from both places.
            worker_type = stage_data.get("worker_type", None) or yaml_engine_args.pop("worker_type", None)
@@ -364,6 +364,7 @@ class OmniARScheduler(VLLMScheduler):
        if stopped_preempted_reqs:
            # This is a rare case and unlikely to impact performance.
            self.waiting.remove_requests(stopped_preempted_reqs)
+            self.skipped_waiting.remove_requests(stopped_preempted_reqs)

        # [Main] Handle failed KV load requests
        if failed_kv_load_req_ids and not self.recompute_kv_load_failures:
@@ -247,6 +247,15 @@ class OmniGenerationScheduler(VLLMScheduler):
        )

        total_num_scheduled_tokens = sum(num_scheduled_tokens.values())
+
+        # Record the request ids scheduled in this step (v0.14.0 behavior).
+        self.prev_step_scheduled_req_ids.clear()
+        self.prev_step_scheduled_req_ids.update(num_scheduled_tokens.keys())
+
+        new_block_ids_to_zero = (
+            (self.kv_cache_manager.take_new_block_ids() or None) if self.needs_kv_cache_zeroing else None
+        )
+
        scheduler_output = SchedulerOutput(
            scheduled_new_reqs=new_reqs_data,
            scheduled_cached_reqs=cached_reqs_data,
@@ -258,12 +267,9 @@ class OmniGenerationScheduler(VLLMScheduler):
            finished_req_ids=self.finished_req_ids,
            free_encoder_mm_hashes=self.encoder_cache_manager.get_freed_mm_hashes(),
            preempted_req_ids=set(),
+            new_block_ids_to_zero=new_block_ids_to_zero,
        )

-        # Record the request ids scheduled in this step (v0.14.0 behavior).
-        self.prev_step_scheduled_req_ids.clear()
-        self.prev_step_scheduled_req_ids.update(num_scheduled_tokens.keys())
-
        # KVTransfer: package metadata
        if self.connector is not None:
            meta = self.connector.build_connector_meta(scheduler_output)
@@ -496,6 +502,7 @@ class OmniGenerationScheduler(VLLMScheduler):
        if stopped_preempted_reqs:
            # This is a rare case and unlikely to impact performance.
            self.waiting.remove_requests(stopped_preempted_reqs)
+            self.skipped_waiting.remove_requests(stopped_preempted_reqs)

        # Handle failed KV load requests
        if failed_kv_load_req_ids and not self.recompute_kv_load_failures:
@@ -464,6 +464,9 @@ class OmniDiffusionConfig:
    # Diffusion pipeline Profiling config
    enable_diffusion_pipeline_profiler: bool = False

+    # Step mode settings
+    step_execution: bool = False
+
    @property
    def is_moe(self) -> bool:
        num_experts = self.tf_model_config.get("num_experts", None)
@@ -182,7 +182,13 @@ class CFGParallelMixin(metaclass=ABCMeta):
        """
        raise NotImplementedError("Subclasses must implement diffuse")

-    def scheduler_step(self, noise_pred: torch.Tensor, t: torch.Tensor, latents: torch.Tensor) -> torch.Tensor:
+    def scheduler_step(
+        self,
+        noise_pred: torch.Tensor,
+        t: torch.Tensor,
+        latents: torch.Tensor,
+        per_request_scheduler: Any | None = None,
+    ) -> torch.Tensor:
        """
        Step the scheduler.

@@ -190,14 +196,30 @@ class CFGParallelMixin(metaclass=ABCMeta):
            noise_pred: Predicted noise
            t: Current timestep
            latents: Current latents
+            per_request_scheduler: Optional request-scoped scheduler that
+                overrides ``self.scheduler`` for this call. This is
+                primarily used by step-wise execution, where each request
+                may keep scheduler state in its own runner-managed state
+                object. Request-level execution should usually leave this
+                as ``None`` and continue using ``self.scheduler``.

        Returns:
            Updated latents after scheduler step
        """
-        return self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+        sched = per_request_scheduler if per_request_scheduler is not None else getattr(self, "scheduler", None)
+        if sched is None:
+            raise ValueError("No scheduler is available. Set self.scheduler or pass per_request_scheduler.")
+        if not callable(getattr(sched, "step", None)):
+            raise TypeError("per_request_scheduler must provide a callable step(...) method.")
+        return sched.step(noise_pred, t, latents, return_dict=False)[0]

    def scheduler_step_maybe_with_cfg(
-        self, noise_pred: torch.Tensor, t: torch.Tensor, latents: torch.Tensor, do_true_cfg: bool
+        self,
+        noise_pred: torch.Tensor,
+        t: torch.Tensor,
+        latents: torch.Tensor,
+        do_true_cfg: bool,
+        per_request_scheduler: Any | None = None,
    ) -> torch.Tensor:
        """
        Step the scheduler with (maybe) automatic CFG parallel synchronization.
@@ -210,6 +232,11 @@ class CFGParallelMixin(metaclass=ABCMeta):
            t: Current timestep
            latents: Current latents
            do_true_cfg: Whether CFG is enabled
+            per_request_scheduler: Optional request-scoped scheduler that
+                overrides ``self.scheduler`` for this call. This is mainly
+                needed by step-wise execution, where scheduler state may be
+                stored per request. Request-level execution should normally
+                leave this as ``None``.

        Returns:
            Updated latents (synchronized across all CFG ranks)
@@ -223,13 +250,23 @@ class CFGParallelMixin(metaclass=ABCMeta):

            # Only rank 0 computes the scheduler step
            if cfg_rank == 0:
-                latents = self.scheduler_step(noise_pred, t, latents)
+                latents = self.scheduler_step(
+                    noise_pred,
+                    t,
+                    latents,
+                    per_request_scheduler=per_request_scheduler,
+                )

            # Broadcast the updated latents to all ranks
            latents = latents.contiguous()
            cfg_group.broadcast(latents, src=0)
        else:
            # No CFG parallel: directly compute scheduler step
-            latents = self.scheduler_step(noise_pred, t, latents)
+            latents = self.scheduler_step(
+                noise_pred,
+                t,
+                latents,
+                per_request_scheduler=per_request_scheduler,
+            )

        return latents
@@ -65,12 +65,7 @@ def _tensor_from_shm(handle: dict[str, Any]) -> torch.Tensor:
    return tensor


-def pack_diffusion_output_shm(output: DiffusionOutput) -> DiffusionOutput:
-    """Replace large tensors in *output* with shared-memory handles.
-
-    The DiffusionOutput is modified **in-place** so that the (now lightweight)
-    object can be serialised cheaply through a MessageQueue.
-    """
+def _pack_diffusion_fields(output: DiffusionOutput) -> DiffusionOutput:
    if output.output is not None and isinstance(output.output, torch.Tensor):
        if output.output.nelement() * output.output.element_size() > _SHM_TENSOR_THRESHOLD:
            output.output = _tensor_to_shm(output.output)
@@ -80,10 +75,35 @@ def pack_diffusion_output_shm(output: DiffusionOutput) -> DiffusionOutput:
    return output


-def unpack_diffusion_output_shm(output: DiffusionOutput) -> DiffusionOutput:
-    """Reconstruct tensors from shared-memory handles produced by ``pack_diffusion_output_shm``."""
+def pack_diffusion_output_shm(output: object) -> object:
+    """Replace large tensors in diffusion worker outputs with SHM handles.
+
+    Supports either a bare ``DiffusionOutput`` or a wrapper object carrying one
+    in ``.result`` (for example ``RunnerOutput``).
+    """
+    if isinstance(output, DiffusionOutput):
+        return _pack_diffusion_fields(output)
+
+    result = getattr(output, "result", None)
+    if isinstance(result, DiffusionOutput):
+        output.result = _pack_diffusion_fields(result)
+    return output
+
+
+def _unpack_diffusion_fields(output: DiffusionOutput) -> DiffusionOutput:
    if isinstance(output.output, dict) and output.output.get("__tensor_shm__"):
        output.output = _tensor_from_shm(output.output)
    if isinstance(output.trajectory_latents, dict) and output.trajectory_latents.get("__tensor_shm__"):
        output.trajectory_latents = _tensor_from_shm(output.trajectory_latents)
    return output
+
+
+def unpack_diffusion_output_shm(output: object) -> object:
+    """Reconstruct tensors from SHM handles in diffusion worker outputs."""
+    if isinstance(output, DiffusionOutput):
+        return _unpack_diffusion_fields(output)
+
+    result = getattr(output, "result", None)
+    if isinstance(result, DiffusionOutput):
+        output.result = _unpack_diffusion_fields(result)
+    return output
@@ -18,7 +18,7 @@ from diffusers.schedulers.scheduling_flow_match_euler_discrete import (
 )
 from diffusers.utils.torch_utils import randn_tensor
 from torch import nn
-from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast
+from transformers import AutoConfig, CLIPTextModel, CLIPTokenizer, T5TokenizerFast
 from vllm.model_executor.models.utils import AutoWeightsLoader

 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
@@ -27,6 +27,7 @@ from vllm_omni.diffusion.distributed.parallel_state import get_classifier_free_g
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.models.flux import FluxTransformer2DModel
+from vllm_omni.diffusion.models.t5_encoder import T5EncoderModel
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.quantization import get_vllm_quant_config_for_layers
 from vllm_omni.diffusion.request import OmniDiffusionRequest
@@ -148,7 +149,14 @@ class FluxPipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfilerMixin):
                revision=None,
                prefix="transformer.",
                fall_back_to_pt=True,
-            )
+            ),
+            DiffusersPipelineLoader.ComponentSource(
+                model_or_path=od_config.model,
+                subfolder="text_encoder_2",
+                revision=None,
+                prefix="text_encoder_2.",
+                fall_back_to_pt=True,
+            ),
        ]

        self.device = get_local_device()
@@ -162,9 +170,8 @@ class FluxPipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfilerMixin):
        self.text_encoder = CLIPTextModel.from_pretrained(
            model, subfolder="text_encoder", local_files_only=local_files_only
        ).to(self.device)
-        self.text_encoder_2 = T5EncoderModel.from_pretrained(
-            model, subfolder="text_encoder_2", local_files_only=local_files_only
-        ).to(self.device)
+        t5_config = AutoConfig.from_pretrained(model, subfolder="text_encoder_2", local_files_only=local_files_only)
+        self.text_encoder_2 = T5EncoderModel(t5_config, prefix="text_encoder_2").to(self.device)
        self.vae = AutoencoderKL.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to(
            self.device
        )
@@ -291,7 +298,7 @@ class FluxPipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfilerMixin):
                f" {max_sequence_length} tokens: {removed_text}"
            )

-        prompt_embeds = self.text_encoder_2(text_input_ids.to(self.device), output_hidden_states=False)[0]
+        prompt_embeds = self.text_encoder_2(text_input_ids.to(self.device))[0]

        dtype = self.text_encoder_2.dtype
        prompt_embeds = prompt_embeds.to(dtype=dtype, device=self.device)
@@ -897,13 +897,15 @@ class HeliosPipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPip
        _, ph, pw = patch_size
        block_size = ph * pw

+        device = generator.device if generator is not None else self.device
+
        cov = torch.eye(block_size) * (1 + gamma) - torch.ones(block_size, block_size) * gamma
        cov += torch.eye(block_size) * 1e-8
        cov = cov.float()  # Upcast to fp32 for numerical stability — cholesky is unreliable in fp16/bf16.

-        L = torch.linalg.cholesky(cov)
+        L = torch.linalg.cholesky(cov).to(device)
        block_number = batch_size * channel * num_frames * (height // ph) * (width // pw)
-        z = torch.randn(block_number, block_size, generator=generator)
+        z = torch.randn(block_number, block_size, generator=generator, device=device)
        noise = z @ L.T

        noise = noise.view(batch_size, channel, num_frames, height // ph, width // pw, ph, pw)
@@ -1,11 +1,21 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from __future__ import annotations
+
 from typing import (
+    TYPE_CHECKING,
+    Any,
    ClassVar,
    Protocol,
    runtime_checkable,
 )

+if TYPE_CHECKING:
+    import torch
+
+    from vllm_omni.diffusion.data import DiffusionOutput
+    from vllm_omni.diffusion.worker.utils import DiffusionRequestState
+

@runtime_checkable
 class SupportImageInput(Protocol):
@@ -21,3 +31,34 @@ class SupportAudioInput(Protocol):
@runtime_checkable
 class SupportAudioOutput(Protocol):
    support_audio_output: ClassVar[bool] = True
+
+
+@runtime_checkable
+class SupportsStepExecution(Protocol):
+    """State-driven step-level execution protocol for diffusion pipelines.
+
+    Pipelines should split request-level ``forward()`` into:
+    ``prepare_encode()`` (one-time request setup), ``denoise_step()``
+    (one denoise forward), ``step_scheduler()`` (one scheduler update),
+    and ``post_decode()`` (final decode).
+    """
+
+    supports_step_execution: ClassVar[bool] = True
+
+    def prepare_encode(self, state: DiffusionRequestState, **kwargs: Any) -> DiffusionRequestState:
+        """Prepare request-level inputs and return initialized state."""
+
+    def denoise_step(self, state: DiffusionRequestState, **kwargs: Any) -> torch.Tensor | None:
+        """Run one denoise step."""
+
+    def step_scheduler(self, state: DiffusionRequestState, noise_pred: torch.Tensor, **kwargs: Any) -> None:
+        """Run one scheduler step."""
+
+    def post_decode(self, state: DiffusionRequestState, **kwargs: Any) -> DiffusionOutput:
+        """Decode output after denoise loop."""
+
+
+def supports_step_execution(pipeline: object) -> bool:
+    """Return whether `pipeline` implements :class:`SupportsStepExecution`."""
+
+    return isinstance(pipeline, SupportsStepExecution)
@@ -685,6 +685,7 @@ class OvisImagePipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfilerMi

        negative_text_ids = None
        if do_classifier_free_guidance:
+            negative_prompt = negative_prompt if negative_prompt is not None else ""
            negative_prompt_embeds, negative_text_ids = self.encode_prompt(
                prompt=negative_prompt,
                prompt_embeds=negative_prompt_embeds,
@@ -1,13 +1,14 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

+import copy
 import inspect
 import json
 import logging
 import math
 import os
 from collections.abc import Iterable
-from typing import Any
+from typing import TYPE_CHECKING, Any, ClassVar

 import numpy as np
 import torch
@@ -35,6 +36,10 @@ from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPi
 from vllm_omni.diffusion.quantization import get_vllm_quant_config_for_layers
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs
+
+if TYPE_CHECKING:
+    from vllm_omni.diffusion.worker.utils import DiffusionRequestState
+
 from vllm_omni.model_executor.model_loader.weight_utils import (
    download_weights_from_hf_specific,
 )
@@ -238,6 +243,8 @@ def apply_rotary_emb_qwen(


 class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineProfilerMixin):
+    supports_step_execution: ClassVar[bool] = True
+
    def __init__(
        self,
        *,
@@ -539,6 +546,372 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP
    def interrupt(self):
        return self._interrupt

+    def _extract_prompts(self, prompts):
+        """Extract prompt and negative_prompt from OmniPromptType list."""
+        prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in prompts] or None
+        if all(isinstance(p, str) or p.get("negative_prompt") is None for p in prompts):
+            negative_prompt = None
+        elif prompts:
+            negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in prompts]
+        else:
+            negative_prompt = None
+        return prompt, negative_prompt
+
+    def _prepare_generation_context(
+        self,
+        *,
+        prompt,
+        negative_prompt,
+        height,
+        width,
+        num_inference_steps,
+        sigmas,
+        guidance_scale,
+        num_images_per_prompt,
+        generator,
+        true_cfg_scale,
+        max_sequence_length,
+        prompt_embeds=None,
+        prompt_embeds_mask=None,
+        negative_prompt_embeds=None,
+        negative_prompt_embeds_mask=None,
+        latents=None,
+        attention_kwargs=None,
+        callback_on_step_end_tensor_inputs=None,
+    ):
+        """Shared preparation logic for forward() and prepare_encode().
+
+        Validates inputs, encodes prompts, prepares latents, computes timesteps,
+        and returns all intermediate values as a dict.
+        """
+        self.check_inputs(
+            prompt,
+            height,
+            width,
+            negative_prompt,
+            prompt_embeds,
+            negative_prompt_embeds,
+            prompt_embeds_mask,
+            negative_prompt_embeds_mask,
+            callback_on_step_end_tensor_inputs,
+            max_sequence_length,
+        )
+
+        self._guidance_scale = guidance_scale
+        self._attention_kwargs = attention_kwargs or {}
+        self._current_timestep = None
+        self._interrupt = False
+
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        elif prompt_embeds is not None:
+            batch_size = prompt_embeds.shape[0]
+        else:
+            batch_size = 1
+
+        has_neg_prompt = negative_prompt is not None or (
+            negative_prompt_embeds is not None and negative_prompt_embeds_mask is not None
+        )
+        do_true_cfg = true_cfg_scale > 1 and has_neg_prompt
+        self.check_cfg_parallel_validity(true_cfg_scale, has_neg_prompt)
+
+        prompt_embeds, prompt_embeds_mask = self.encode_prompt(
+            prompt=prompt,
+            prompt_embeds=prompt_embeds,
+            prompt_embeds_mask=prompt_embeds_mask,
+            num_images_per_prompt=num_images_per_prompt,
+            max_sequence_length=max_sequence_length,
+        )
+        if do_true_cfg:
+            negative_prompt_embeds, negative_prompt_embeds_mask = self.encode_prompt(
+                prompt=negative_prompt,
+                prompt_embeds=negative_prompt_embeds,
+                prompt_embeds_mask=negative_prompt_embeds_mask,
+                num_images_per_prompt=num_images_per_prompt,
+                max_sequence_length=max_sequence_length,
+            )
+        else:
+            negative_prompt_embeds = None
+            negative_prompt_embeds_mask = None
+
+        num_channels_latents = self.transformer.in_channels // 4
+        latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            prompt_embeds.dtype,
+            self.device,
+            generator,
+            latents,
+        )
+
+        img_shapes = [[(1, height // self.vae_scale_factor // 2, width // self.vae_scale_factor // 2)]] * batch_size
+
+        timesteps, num_inference_steps = self.prepare_timesteps(
+            num_inference_steps,
+            sigmas,
+            latents.shape[1],
+        )
+        self._num_timesteps = len(timesteps)
+
+        if self.transformer.guidance_embeds:
+            guidance = torch.full([1], guidance_scale, dtype=torch.float32)
+            guidance = guidance.expand(latents.shape[0])
+        else:
+            guidance = None
+
+        txt_seq_lens = prompt_embeds_mask.sum(dim=1).tolist() if prompt_embeds_mask is not None else None
+        negative_txt_seq_lens = (
+            negative_prompt_embeds_mask.sum(dim=1).tolist() if negative_prompt_embeds_mask is not None else None
+        )
+
+        return {
+            "prompt_embeds": prompt_embeds,
+            "prompt_embeds_mask": prompt_embeds_mask,
+            "negative_prompt_embeds": negative_prompt_embeds,
+            "negative_prompt_embeds_mask": negative_prompt_embeds_mask,
+            "latents": latents,
+            "img_shapes": img_shapes,
+            "timesteps": timesteps,
+            "do_true_cfg": do_true_cfg,
+            "guidance": guidance,
+            "txt_seq_lens": txt_seq_lens,
+            "negative_txt_seq_lens": negative_txt_seq_lens,
+        }
+
+    def prepare_encode(
+        self,
+        state: "DiffusionRequestState",
+        **kwargs: Any,
+    ) -> "DiffusionRequestState":
+        """Populate *state* with encoded prompts, latents, timesteps, and CFG config."""
+        sampling = state.sampling
+        prompt, negative_prompt = self._extract_prompts(state.prompts or [])
+
+        ctx = self._prepare_generation_context(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            height=sampling.height or self.default_sample_size * self.vae_scale_factor,
+            width=sampling.width or self.default_sample_size * self.vae_scale_factor,
+            num_inference_steps=sampling.num_inference_steps or 50,
+            sigmas=sampling.sigmas,
+            guidance_scale=sampling.guidance_scale if sampling.guidance_scale_provided else 1.0,
+            num_images_per_prompt=sampling.num_outputs_per_prompt if sampling.num_outputs_per_prompt > 0 else 1,
+            generator=sampling.generator,
+            true_cfg_scale=sampling.true_cfg_scale or 4.0,
+            max_sequence_length=sampling.max_sequence_length or 512,
+            attention_kwargs=kwargs.get("attention_kwargs"),
+        )
+
+        # prepare_timesteps() has already materialized request-specific timestep
+        # state on self.scheduler, so deepcopy preserves dynamic-shifting state
+        # without replaying set_timesteps() on the per-request scheduler.
+        # Per-request scheduler (must not share state with self.scheduler)
+        req_scheduler = copy.deepcopy(self.scheduler)
+        req_scheduler.set_begin_index(0)
+
+        # Populate state from generation context
+        state.prompt_embeds = ctx["prompt_embeds"]
+        state.prompt_embeds_mask = ctx["prompt_embeds_mask"]
+        state.negative_prompt_embeds = ctx["negative_prompt_embeds"]
+        state.negative_prompt_embeds_mask = ctx["negative_prompt_embeds_mask"]
+        state.latents = ctx["latents"]
+        state.timesteps = ctx["timesteps"]
+        state.step_index = 0
+        state.scheduler = req_scheduler
+        state.do_true_cfg = ctx["do_true_cfg"]
+        state.guidance = ctx["guidance"]
+        state.img_shapes = ctx["img_shapes"]
+        state.txt_seq_lens = ctx["txt_seq_lens"]
+        state.negative_txt_seq_lens = ctx["negative_txt_seq_lens"]
+        # QwenImage always normalizes CFG output (matching forward())
+        state.sampling.cfg_normalize = True
+
+        return state
+
+    def _build_denoise_kwargs(
+        self,
+        latents: torch.Tensor,
+        timestep: torch.Tensor,
+        guidance: torch.Tensor | None,
+        prompt_embeds: torch.Tensor,
+        prompt_embeds_mask: torch.Tensor,
+        img_shapes: list,
+        txt_seq_lens: list[int] | None,
+        do_true_cfg: bool,
+        negative_prompt_embeds: torch.Tensor | None,
+        negative_prompt_embeds_mask: torch.Tensor | None,
+        negative_txt_seq_lens: list[int] | None,
+        image_latents: torch.Tensor | None = None,
+        extra_transformer_kwargs: dict[str, Any] | None = None,
+    ) -> tuple[dict[str, Any], dict[str, Any] | None, int | None]:
+        """Build positive/negative kwargs and output_slice for one denoise step.
+
+        Returns:
+            (positive_kwargs, negative_kwargs, output_slice)
+        """
+        extra_transformer_kwargs = extra_transformer_kwargs or {}
+
+        # Broadcast timestep to match batch size
+        t_for_model = timestep.expand(latents.shape[0]).to(
+            device=latents.device,
+            dtype=latents.dtype,
+        )
+
+        # Concatenate image latents if available (editing pipelines)
+        latent_model_input = latents
+        if image_latents is not None:
+            latent_model_input = torch.cat([latents, image_latents], dim=1)
+
+        positive_kwargs = {
+            "hidden_states": latent_model_input,
+            "timestep": t_for_model / 1000,
+            "guidance": guidance,
+            "encoder_hidden_states_mask": prompt_embeds_mask,
+            "encoder_hidden_states": prompt_embeds,
+            "img_shapes": img_shapes,
+            "txt_seq_lens": txt_seq_lens,
+            **extra_transformer_kwargs,
+        }
+        if do_true_cfg:
+            negative_kwargs = {
+                "hidden_states": latent_model_input,
+                "timestep": t_for_model / 1000,
+                "guidance": guidance,
+                "encoder_hidden_states_mask": negative_prompt_embeds_mask,
+                "encoder_hidden_states": negative_prompt_embeds,
+                "img_shapes": img_shapes,
+                "txt_seq_lens": negative_txt_seq_lens,
+                **extra_transformer_kwargs,
+            }
+        else:
+            negative_kwargs = None
+
+        output_slice = latents.size(1) if image_latents is not None else None
+        return positive_kwargs, negative_kwargs, output_slice
+
+    def _decode_latents(
+        self,
+        latents: torch.Tensor,
+        height: int,
+        width: int,
+        output_type: str = "pil",
+    ) -> DiffusionOutput:
+        """Unpack, normalize, and VAE-decode latents into a DiffusionOutput."""
+        if output_type == "latent":
+            return DiffusionOutput(
+                output=latents,
+                stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
+            )
+
+        latents = self._unpack_latents(latents, height, width, self.vae_scale_factor)
+        latents = latents.to(self.vae.dtype)
+        latents_mean = (
+            torch.tensor(self.vae.config.latents_mean)
+            .view(1, self.vae.config.z_dim, 1, 1, 1)
+            .to(latents.device, latents.dtype)
+        )
+        latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to(
+            latents.device, latents.dtype
+        )
+        latents = latents / latents_std + latents_mean
+        image = self.vae.decode(latents, return_dict=False)[0][:, :, 0]
+        return DiffusionOutput(
+            output=image,
+            stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
+        )
+
+    def denoise_step(
+        self,
+        state: "DiffusionRequestState",
+        **kwargs: Any,
+    ) -> torch.Tensor | None:
+        """One denoise step: read from *state*, delegate to CFGParallelMixin.
+
+        Reuses ``predict_noise_maybe_with_cfg`` so that CFG-parallel,
+        sequential-CFG, and no-CFG paths are handled identically to
+        ``diffuse()``.
+        """
+        if self.interrupt:
+            return None
+
+        t = state.current_timestep
+        self._current_timestep = t
+        self.transformer.do_true_cfg = state.do_true_cfg
+
+        # Normalize timestep to [batch_size] tensor
+        if not torch.is_tensor(t):
+            t = torch.tensor(t, device=state.latents.device, dtype=state.latents.dtype)
+
+        positive_kwargs, negative_kwargs, output_slice = self._build_denoise_kwargs(
+            latents=state.latents,
+            timestep=t,
+            guidance=state.guidance,
+            prompt_embeds=state.prompt_embeds,
+            prompt_embeds_mask=state.prompt_embeds_mask,
+            img_shapes=state.img_shapes,
+            txt_seq_lens=state.txt_seq_lens,
+            do_true_cfg=state.do_true_cfg,
+            negative_prompt_embeds=state.negative_prompt_embeds,
+            negative_prompt_embeds_mask=state.negative_prompt_embeds_mask,
+            negative_txt_seq_lens=state.negative_txt_seq_lens,
+            image_latents=state.sampling.image_latent,
+            extra_transformer_kwargs={
+                "attention_kwargs": self.attention_kwargs,
+                "return_dict": False,
+            },
+        )
+
+        true_cfg_scale = state.sampling.true_cfg_scale or 4.0
+        cfg_normalize = state.sampling.cfg_normalize
+
+        return self.predict_noise_maybe_with_cfg(
+            state.do_true_cfg,
+            true_cfg_scale,
+            positive_kwargs,
+            negative_kwargs,
+            cfg_normalize,
+            output_slice,
+        )
+
+    def step_scheduler(
+        self,
+        state: "DiffusionRequestState",
+        noise_pred: torch.Tensor,
+        **kwargs: Any,
+    ) -> None:
+        """One scheduler step: update ``state.latents`` and advance ``step_index``."""
+        if self.interrupt:
+            return
+
+        t = state.current_timestep
+        state.latents = self.scheduler_step_maybe_with_cfg(
+            noise_pred,
+            t,
+            state.latents,
+            state.do_true_cfg,
+            per_request_scheduler=state.scheduler,
+        )
+
+        state.step_index += 1
+
+    def post_decode(
+        self,
+        state: "DiffusionRequestState",
+        **kwargs: Any,
+    ) -> DiffusionOutput:
+        """Decode final latents from *state*."""
+        self._current_timestep = None
+
+        height = state.sampling.height or self.default_sample_size * self.vae_scale_factor
+        width = state.sampling.width or self.default_sample_size * self.vae_scale_factor
+        output_type = kwargs.get("output_type", "pil")
+
+        return self._decode_latents(state.latents, height, width, output_type)
+
    def forward(
        self,
        req: OmniDiffusionRequest,
@@ -562,13 +935,8 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP
        callback_on_step_end_tensor_inputs: list[str] = ["latents"],
        max_sequence_length: int = 512,
    ) -> DiffusionOutput:
-        # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "")
-        # TODO: May be some data formatting operations on the API side. Hack for now.
-        prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt
-        if all(isinstance(p, str) or p.get("negative_prompt") is None for p in req.prompts):
-            negative_prompt = None
-        elif req.prompts:
-            negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in req.prompts]
+        extracted_prompt, negative_prompt = self._extract_prompts(req.prompts)
+        prompt = extracted_prompt or prompt

        height = req.sampling_params.height or self.default_sample_size * self.vae_scale_factor
        width = req.sampling_params.width or self.default_sample_size * self.vae_scale_factor
@@ -584,113 +952,40 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP
            if req.sampling_params.num_outputs_per_prompt > 0
            else num_images_per_prompt
        )
-        # 1. check inputs
-        # 2. encode prompts
-        # 3. prepare latents and timesteps
-        # 4. diffusion process
-        # 5. decode latents
-        # 6. post-process outputs
-        self.check_inputs(
-            prompt,
-            height,
-            width,
-            negative_prompt,
-            prompt_embeds,
-            negative_prompt_embeds,
-            prompt_embeds_mask,
-            negative_prompt_embeds_mask,
-            callback_on_step_end_tensor_inputs,
-            max_sequence_length,
-        )

-        self._guidance_scale = guidance_scale
-        self._attention_kwargs = attention_kwargs
-        self._current_timestep = None
-        self._interrupt = False
-
-        if prompt is not None and isinstance(prompt, str):
-            batch_size = 1
-        elif prompt is not None and isinstance(prompt, list):
-            batch_size = len(prompt)
-        else:
-            batch_size = prompt_embeds.shape[0]
-
-        has_neg_prompt = negative_prompt is not None or (
-            negative_prompt_embeds is not None and negative_prompt_embeds_mask is not None
-        )
-
-        do_true_cfg = true_cfg_scale > 1 and has_neg_prompt
-        self.check_cfg_parallel_validity(true_cfg_scale, has_neg_prompt)
-
-        prompt_embeds, prompt_embeds_mask = self.encode_prompt(
+        ctx = self._prepare_generation_context(
            prompt=prompt,
+            negative_prompt=negative_prompt,
+            height=height,
+            width=width,
+            num_inference_steps=num_inference_steps,
+            sigmas=sigmas,
+            guidance_scale=guidance_scale,
+            num_images_per_prompt=num_images_per_prompt,
+            generator=generator,
+            true_cfg_scale=true_cfg_scale,
+            max_sequence_length=max_sequence_length,
            prompt_embeds=prompt_embeds,
            prompt_embeds_mask=prompt_embeds_mask,
-            num_images_per_prompt=num_images_per_prompt,
-            max_sequence_length=max_sequence_length,
+            negative_prompt_embeds=negative_prompt_embeds,
+            negative_prompt_embeds_mask=negative_prompt_embeds_mask,
+            latents=latents,
+            attention_kwargs=attention_kwargs,
+            callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
        )
-        if do_true_cfg:
-            negative_prompt_embeds, negative_prompt_embeds_mask = self.encode_prompt(
-                prompt=negative_prompt,
-                prompt_embeds=negative_prompt_embeds,
-                prompt_embeds_mask=negative_prompt_embeds_mask,
-                num_images_per_prompt=num_images_per_prompt,
-                max_sequence_length=max_sequence_length,
-            )
-
-        num_channels_latents = self.transformer.in_channels // 4
-        latents = self.prepare_latents(
-            batch_size * num_images_per_prompt,
-            num_channels_latents,
-            height,
-            width,
-            prompt_embeds.dtype,
-            self.device,
-            generator,
-            latents,
-        )
-        img_shapes = [
-            [
-                (
-                    1,
-                    height // self.vae_scale_factor // 2,
-                    width // self.vae_scale_factor // 2,
-                )
-            ]
-        ] * batch_size
-
-        timesteps, num_inference_steps = self.prepare_timesteps(num_inference_steps, sigmas, latents.shape[1])
-        # num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
-        self._num_timesteps = len(timesteps)
-
-        # handle guidance
-        if self.transformer.guidance_embeds:
-            guidance = torch.full([1], guidance_scale, dtype=torch.float32)
-            guidance = guidance.expand(latents.shape[0])
-        else:
-            guidance = None
-
-        if self.attention_kwargs is None:
-            self._attention_kwargs = {}
-
-        txt_seq_lens = prompt_embeds_mask.sum(dim=1).tolist() if prompt_embeds_mask is not None else None
-        negative_txt_seq_lens = (
-            negative_prompt_embeds_mask.sum(dim=1).tolist() if negative_prompt_embeds_mask is not None else None
-        )
-        # print inputp params

        latents = self.diffuse(
-            prompt_embeds,
-            prompt_embeds_mask,
-            negative_prompt_embeds,
-            negative_prompt_embeds_mask,
-            latents,
-            img_shapes,
-            txt_seq_lens,
-            negative_txt_seq_lens,
-            timesteps,
-            do_true_cfg,
-            guidance,
+            ctx["prompt_embeds"],
+            ctx["prompt_embeds_mask"],
+            ctx["negative_prompt_embeds"],
+            ctx["negative_prompt_embeds_mask"],
+            ctx["latents"],
+            ctx["img_shapes"],
+            ctx["txt_seq_lens"],
+            ctx["negative_txt_seq_lens"],
+            ctx["timesteps"],
+            ctx["do_true_cfg"],
+            ctx["guidance"],
            true_cfg_scale,
            image_latents=None,
            cfg_normalize=True,
@@ -701,26 +996,7 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP
        )

        self._current_timestep = None
-        if output_type == "latent":
-            image = latents
-        else:
-            latents = self._unpack_latents(latents, height, width, self.vae_scale_factor)
-            latents = latents.to(self.vae.dtype)
-            latents_mean = (
-                torch.tensor(self.vae.config.latents_mean)
-                .view(1, self.vae.config.z_dim, 1, 1, 1)
-                .to(latents.device, latents.dtype)
-            )
-            latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to(
-                latents.device, latents.dtype
-            )
-            latents = latents / latents_std + latents_mean
-            image = self.vae.decode(latents, return_dict=False)[0][:, :, 0]
-            # processed_image = self.image_processor.postprocess(image, output_type=output_type)
-
-        return DiffusionOutput(
-            output=image, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None
-        )
+        return self._decode_latents(latents, height, width, output_type)

    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
        loader = AutoWeightsLoader(self)
@@ -0,0 +1,7 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tensor-parallel T5 encoder model."""
+
+from vllm_omni.diffusion.models.t5_encoder.t5_encoder import T5EncoderModel
+
+__all__ = ["T5EncoderModel"]
@@ -0,0 +1,389 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from __future__ import annotations
+
+import math
+from collections.abc import Iterable
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import T5Config
+from vllm.distributed import get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size
+from vllm.model_executor.layers.activation import get_act_fn
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.linear import (
+    ColumnParallelLinear,
+    MergedColumnParallelLinear,
+    QKVParallelLinear,
+    RowParallelLinear,
+)
+from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+
+
+class T5SelfAttention(nn.Module):
+    def __init__(
+        self,
+        config: T5Config,
+        has_relative_attention_bias: bool = False,
+        prefix: str = "",
+    ):
+        super().__init__()
+        self.d_model = config.d_model
+        self.d_kv = config.d_kv
+        self.n_heads = config.num_heads
+        self.inner_dim = self.n_heads * self.d_kv
+        self.has_relative_attention_bias = has_relative_attention_bias
+        self.relative_attention_num_buckets = config.relative_attention_num_buckets
+        self.relative_attention_max_distance = config.relative_attention_max_distance
+
+        tp_size = get_tensor_model_parallel_world_size()
+        assert self.n_heads % tp_size == 0, f"num_heads ({self.n_heads}) must be divisible by tp_size ({tp_size})"
+        self.n_heads_per_partition = self.n_heads // tp_size
+
+        # Fused Q/K/V projection, sharded across heads.
+        # HF has separate q, k, v — we fuse into qkv_proj and use
+        # stacked_params_mapping in load_weights() to load each shard.
+        self.qkv_proj = QKVParallelLinear(
+            hidden_size=self.d_model,
+            head_size=self.d_kv,
+            total_num_heads=self.n_heads,
+            total_num_kv_heads=self.n_heads,  # T5 uses MHA
+            bias=False,
+        )
+
+        # Output projection: all-reduce back to full d_model
+        # Named ``o`` to match HF's ``SelfAttention.o``
+        self.o = RowParallelLinear(
+            self.inner_dim,
+            self.d_model,
+            bias=False,
+            input_is_parallel=True,
+            return_bias=False,
+        )
+
+        if has_relative_attention_bias:
+            # Store full embedding; slice heads per rank in forward
+            self.relative_attention_bias = nn.Embedding(self.relative_attention_num_buckets, self.n_heads)
+
+    @staticmethod
+    def _relative_position_bucket(
+        relative_position: torch.Tensor,
+        bidirectional: bool = True,
+        num_buckets: int = 32,
+        max_distance: int = 128,
+    ) -> torch.Tensor:
+        relative_buckets = 0
+        if bidirectional:
+            num_buckets //= 2
+            relative_buckets += (relative_position > 0).to(torch.long) * num_buckets
+            relative_position = torch.abs(relative_position)
+        else:
+            relative_position = -torch.min(relative_position, torch.zeros_like(relative_position))
+        max_exact = num_buckets // 2
+        is_small = relative_position < max_exact
+        relative_position_if_large = max_exact + (
+            torch.log(relative_position.float() / max_exact)
+            / math.log(max_distance / max_exact)
+            * (num_buckets - max_exact)
+        ).to(torch.long)
+        relative_position_if_large = torch.min(
+            relative_position_if_large,
+            torch.full_like(relative_position_if_large, num_buckets - 1),
+        )
+        relative_buckets += torch.where(is_small, relative_position, relative_position_if_large)
+        return relative_buckets
+
+    def compute_bias(self, query_length: int, key_length: int, device: torch.device) -> torch.Tensor:
+        """Compute relative position bias, returning only the local head shard."""
+        context_position = torch.arange(query_length, dtype=torch.long, device=device)[:, None]
+        memory_position = torch.arange(key_length, dtype=torch.long, device=device)[None, :]
+        relative_position = memory_position - context_position
+        relative_position_bucket = self._relative_position_bucket(
+            relative_position,
+            bidirectional=True,
+            num_buckets=self.relative_attention_num_buckets,
+            max_distance=self.relative_attention_max_distance,
+        )
+        # values: (query_length, key_length, num_heads)
+        values = self.relative_attention_bias(relative_position_bucket)
+        # Slice to local heads for this TP rank
+        tp_rank = get_tensor_model_parallel_rank()
+        head_start = tp_rank * self.n_heads_per_partition
+        head_end = head_start + self.n_heads_per_partition
+        values = values[:, :, head_start:head_end]
+        # (1, local_heads, query_length, key_length)
+        values = values.permute(2, 0, 1).unsqueeze(0)
+        return values
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        mask: torch.Tensor | None = None,
+        position_bias: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        batch_size, seq_length = hidden_states.shape[:2]
+
+        # Fused QKV projection
+        qkv, _ = self.qkv_proj(hidden_states)
+        q_size = self.n_heads_per_partition * self.d_kv
+        kv_size = self.n_heads_per_partition * self.d_kv
+        query_states, key_states, value_states = qkv.split([q_size, kv_size, kv_size], dim=-1)
+
+        # Reshape: (batch, seq, local_heads, d_kv) -> (batch, local_heads, seq, d_kv)
+        query_states = query_states.view(batch_size, seq_length, self.n_heads_per_partition, self.d_kv).transpose(1, 2)
+        key_states = key_states.view(batch_size, seq_length, self.n_heads_per_partition, self.d_kv).transpose(1, 2)
+        value_states = value_states.view(batch_size, seq_length, self.n_heads_per_partition, self.d_kv).transpose(1, 2)
+
+        # Attention scores: (batch, local_heads, seq, seq)
+        scores = torch.matmul(query_states, key_states.transpose(3, 2))
+
+        if position_bias is None:
+            if self.has_relative_attention_bias:
+                position_bias = self.compute_bias(seq_length, seq_length, device=scores.device)
+            else:
+                position_bias = torch.zeros(
+                    (1, self.n_heads_per_partition, seq_length, seq_length),
+                    device=scores.device,
+                    dtype=scores.dtype,
+                )
+            if mask is not None:
+                position_bias = position_bias + mask
+
+        scores += position_bias
+
+        attn_weights = F.softmax(scores.float(), dim=-1).type_as(scores)
+        attn_output = torch.matmul(attn_weights, value_states)
+
+        # (batch, local_heads, seq, d_kv) -> (batch, seq, local_heads * d_kv)
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.view(batch_size, seq_length, -1)
+
+        attn_output = self.o(attn_output)
+
+        return attn_output, position_bias
+
+
+class T5DenseGatedActDense(nn.Module):
+    def __init__(self, config: T5Config, prefix: str = ""):
+        super().__init__()
+        self.wi = MergedColumnParallelLinear(
+            config.d_model,
+            [config.d_ff, config.d_ff],
+            bias=False,
+            gather_output=False,
+        )
+        self.wo = RowParallelLinear(
+            config.d_ff,
+            config.d_model,
+            bias=False,
+            input_is_parallel=True,
+            return_bias=False,
+        )
+        self.act = get_act_fn(config.dense_act_fn)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        # MergedColumnParallelLinear outputs concatenated [gate, up]
+        gate_up, _ = self.wi(hidden_states)
+        gate, up = gate_up.chunk(2, dim=-1)
+        hidden_states = self.act(gate) * up
+        hidden_states = self.wo(hidden_states)
+        return hidden_states
+
+
+class T5DenseActDense(nn.Module):
+    def __init__(self, config: T5Config, prefix: str = ""):
+        super().__init__()
+        self.wi = ColumnParallelLinear(
+            config.d_model,
+            config.d_ff,
+            bias=False,
+            gather_output=False,
+            return_bias=False,
+        )
+        self.wo = RowParallelLinear(
+            config.d_ff,
+            config.d_model,
+            bias=False,
+            input_is_parallel=True,
+            return_bias=False,
+        )
+        self.act = get_act_fn(config.dense_act_fn)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.wi(hidden_states)
+        hidden_states = self.act(hidden_states)
+        hidden_states = self.wo(hidden_states)
+        return hidden_states
+
+
+class T5LayerSelfAttention(nn.Module):
+    def __init__(self, config: T5Config, has_relative_attention_bias: bool = False, prefix: str = ""):
+        super().__init__()
+        self.SelfAttention = T5SelfAttention(config, has_relative_attention_bias, prefix=f"{prefix}.SelfAttention")
+        self.layer_norm = RMSNorm(config.d_model, eps=config.layer_norm_epsilon)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        mask: torch.Tensor | None = None,
+        position_bias: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        normed = self.layer_norm(hidden_states)
+        attn_output, position_bias = self.SelfAttention(normed, mask=mask, position_bias=position_bias)
+        hidden_states = hidden_states + attn_output
+
+        # Clamp for fp16 stability (matching HF T5)
+        if hidden_states.dtype == torch.float16:
+            clamp_value = torch.finfo(hidden_states.dtype).max - 1000
+            hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value)
+
+        return hidden_states, position_bias
+
+
+class T5LayerFF(nn.Module):
+    def __init__(self, config: T5Config, prefix: str = ""):
+        super().__init__()
+        if config.is_gated_act:
+            self.DenseReluDense = T5DenseGatedActDense(config, prefix=f"{prefix}.DenseReluDense")
+        else:
+            self.DenseReluDense = T5DenseActDense(config, prefix=f"{prefix}.DenseReluDense")
+        self.layer_norm = RMSNorm(config.d_model, eps=config.layer_norm_epsilon)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        normed = self.layer_norm(hidden_states)
+        ff_output = self.DenseReluDense(normed)
+        hidden_states = hidden_states + ff_output
+
+        if hidden_states.dtype == torch.float16:
+            clamp_value = torch.finfo(hidden_states.dtype).max - 1000
+            hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value)
+
+        return hidden_states
+
+
+class T5Block(nn.Module):
+    def __init__(self, config: T5Config, has_relative_attention_bias: bool = False, prefix: str = ""):
+        super().__init__()
+        self.layer = nn.ModuleList(
+            [
+                T5LayerSelfAttention(config, has_relative_attention_bias, prefix=f"{prefix}.layer.0"),
+                T5LayerFF(config, prefix=f"{prefix}.layer.1"),
+            ]
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        mask: torch.Tensor | None = None,
+        position_bias: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        hidden_states, position_bias = self.layer[0](hidden_states, mask=mask, position_bias=position_bias)
+        hidden_states = self.layer[1](hidden_states)
+        return hidden_states, position_bias
+
+
+class T5Stack(nn.Module):
+    def __init__(self, config: T5Config, shared: nn.Embedding, prefix: str = ""):
+        super().__init__()
+        self.embed_tokens = shared
+        self.block = nn.ModuleList(
+            [
+                T5Block(config, has_relative_attention_bias=(i == 0), prefix=f"{prefix}.block.{i}")
+                for i in range(config.num_layers)
+            ]
+        )
+        self.final_layer_norm = RMSNorm(config.d_model, eps=config.layer_norm_epsilon)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        hidden_states = self.embed_tokens(input_ids)
+
+        # Build attention mask: (batch, seq) -> (batch, 1, 1, seq)
+        if attention_mask is not None:
+            extended_mask = attention_mask[:, None, None, :].to(dtype=hidden_states.dtype)
+            extended_mask = (1.0 - extended_mask) * torch.finfo(hidden_states.dtype).min
+        else:
+            extended_mask = None
+
+        position_bias = None
+        for block in self.block:
+            hidden_states, position_bias = block(
+                hidden_states,
+                mask=extended_mask,
+                position_bias=position_bias,
+            )
+
+        hidden_states = self.final_layer_norm(hidden_states)
+        return hidden_states
+
+
+class T5EncoderModel(nn.Module):
+    """T5 encoder model applying upstream vLLM layers"""
+
+    def __init__(self, config: T5Config, prefix: str = ""):
+        super().__init__()
+        self.config = config
+        self.shared = VocabParallelEmbedding(config.vocab_size, config.d_model)
+        self.encoder = T5Stack(config, self.shared, prefix=f"{prefix}.encoder")
+
+    @property
+    def dtype(self) -> torch.dtype:
+        return next(self.parameters()).dtype
+
+    @property
+    def device(self) -> torch.device:
+        return next(self.parameters()).device
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.shared(input_ids)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, ...]:
+        hidden_states = self.encoder(input_ids, attention_mask=attention_mask)
+        return (hidden_states,)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        stacked_params_mapping = [
+            ("qkv_proj", "q", "q"),
+            ("qkv_proj", "k", "k"),
+            ("qkv_proj", "v", "v"),
+            ("wi", "wi_0", 0),
+            ("wi", "wi_1", 1),
+        ]
+
+        params_dict = dict(self.named_parameters())
+        loaded_params: set[str] = set()
+
+        for name, loaded_weight in weights:
+            original_name = name
+            lookup_name = name
+
+            for param_name, weight_name, shard_id in stacked_params_mapping:
+                if f".{weight_name}." not in name:
+                    continue
+                lookup_name = name.replace(weight_name, param_name)
+                if lookup_name not in params_dict:
+                    continue
+                param = params_dict[lookup_name]
+                weight_loader = param.weight_loader
+                weight_loader(param, loaded_weight, shard_id)
+                break
+            else:
+                if name not in params_dict:
+                    continue
+                param = params_dict[name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight)
+
+            loaded_params.add(original_name)
+            loaded_params.add(lookup_name)
+
+        return loaded_params
@@ -305,7 +305,7 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin):

    @property
    def do_classifier_free_guidance(self):
-        return self._guidance_scale > 1
+        return self._guidance_scale > 0

    @property
    def joint_attention_kwargs(self):
@@ -365,7 +365,7 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin):
                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
                `guidance_scale` is defined as `w` of equation 2. of [Imagen
                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
-                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
+                0`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
                usually at the expense of lower image quality.
            cfg_normalization (`bool`, *optional*, defaults to False):
                Whether to apply configuration normalization.
@@ -374,7 +374,7 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin):
            negative_prompt (`str` or `list[str]`, *optional*):
                The prompt or prompts not to guide the image generation. If not defined, one has to pass
                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
-                less than `1`).
+                less than or equal to `0`).
            num_images_per_prompt (`int`, *optional*, defaults to 1):
                The number of images to generate per prompt.
            generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
@@ -279,6 +279,7 @@ class ZImageAttention(nn.Module):
            total_num_kv_heads=num_kv_heads,
            bias=False,
            quant_config=quant_config,
+            prefix="to_qkv",
        )

        assert qk_norm is True
@@ -297,6 +298,7 @@ class ZImageAttention(nn.Module):
                    input_is_parallel=True,
                    return_bias=False,
                    quant_config=quant_config,
+                    prefix="to_out",
                )
            ]
        )
@@ -361,6 +363,7 @@ class FeedForward(nn.Module):
        dim: int,
        hidden_dim: int,
        quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
    ):
        super().__init__()
        self.w13 = MergedColumnParallelLinear(
@@ -369,6 +372,7 @@ class FeedForward(nn.Module):
            bias=False,
            return_bias=False,
            quant_config=quant_config,
+            prefix=prefix,
        )
        self.act = SiluAndMul()
        self.w2 = RowParallelLinear(
@@ -378,6 +382,7 @@ class FeedForward(nn.Module):
            input_is_parallel=True,
            return_bias=False,
            quant_config=quant_config,
+            prefix=prefix,
        )

    def forward(self, x):
@@ -409,9 +414,7 @@ class ZImageTransformerBlock(nn.Module):
        )

        self.feed_forward = FeedForward(
-            dim=dim,
-            hidden_dim=int(dim / 3 * 8),
-            quant_config=quant_config,
+            dim=dim, hidden_dim=int(dim / 3 * 8), quant_config=quant_config, prefix="feed_forward"
        )
        self.layer_id = layer_id

@@ -29,6 +29,7 @@ from vllm.logger import init_logger
 from .base import DiffusionQuantizationConfig
 from .fp8 import DiffusionFp8Config
 from .gguf import DiffusionGgufConfig
+from .int8 import DiffusionInt8Config

 if TYPE_CHECKING:
    from vllm.model_executor.layers.quantization.base_config import (
@@ -41,6 +42,7 @@ logger = init_logger(__name__)
 # To add a new method, create a new config class and register it here
 _QUANT_CONFIG_REGISTRY: dict[str, type[DiffusionQuantizationConfig]] = {
    "fp8": DiffusionFp8Config,
+    "int8": DiffusionInt8Config,
    "gguf": DiffusionGgufConfig,
 }

@@ -110,6 +112,7 @@ def get_vllm_quant_config_for_layers(
 __all__ = [
    "DiffusionQuantizationConfig",
    "DiffusionFp8Config",
+    "DiffusionInt8Config",
    "DiffusionGgufConfig",
    "get_diffusion_quant_config",
    "get_vllm_quant_config_for_layers",