Update CUDA Docker base image to vLLM v0.21.0

Align Dockerfile.cuda with the v0.21.0 rebase on main and other
GPU Dockerfiles/docs. Related to #3858.

Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>

.buildkite/pipeline-intel.yaml

@@ -10,7 +10,7 @@ steps:
           DOCKER_BUILDKIT: "1"
           # Buildkite will automatically replace this with the actual commit hash
           VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}"
-          VLLM_VERSION: "v0.22.0"
+          VLLM_VERSION: "v0.21.0"
         priority: 100
         timeout_in_minutes: 60
         soft_fail: false

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

@@ -52,7 +52,6 @@ time timeout -k 30 30m docker run \
     pip install tblib==3.1.0
     cd /workspace/vllm-omni
     pytest -v -s -m "core_model and xpu and B60"
-    pytest -v -s tests/diffusion/quantization/test_mxfp8_config.py
     export VLLM_XPU_USE_SAMPLER_KERNEL=0    # NOTE: Remove this after vLLM v0.21.1 is merged. Fixes Qwen2-5 omni-expansion tests.
     pytest -v -s -m "advanced_model and xpu and B60"
 '

.buildkite/test-amd-merge.yml

@@ -1,28 +1,15 @@
 steps:
 
-- label: "Simple · Diffusion & Model Executor Test"
+- label: "Simple Unit Test"
   agent_pool: mi325_1
   depends_on: amd-build
-  timeout_in_minutes: 50
   mirror_hardwares: [amdproduction]
   grade: Blocking
   commands:
     - export VLLM_ROCM_USE_AITER=0
     # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
     # that is not supported on CPU
-    - "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
-
-- label: "Simple · Other Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  timeout_in_minutes: 50
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - export VLLM_ROCM_USE_AITER=0
-    # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
-    # that is not supported on CPU
-    - "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
+    - "timeout 60m pytest -v -s -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
 
 - label: "Diffusion Images API LoRA E2E"
   agent_pool: mi325_1
@@ -62,6 +49,16 @@ steps:
 #     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
 #     - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_expansion.py -m "advanced_model and diffusion and L4" --run-level advanced_model
 
+- label: "Diffusion Cache Backend Test"
+  agent_pool: mi325_1
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - 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 (Need 4 GPUs)"
   agent_pool: mi325_4
   depends_on: amd-build
@@ -126,7 +123,7 @@ steps:
   grade: Blocking
   commands:
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py -m 'advanced_model and cuda' --run-level 'advanced_model'
+    - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py tests/e2e/online_serving/test_mimo_audio.py -m "advanced_model" --run-level "advanced_model"
 
 - label: "Qwen3-TTS CustomVoice E2E Test"
   agent_pool: mi325_1

.buildkite/test-amd-ready.yaml

@@ -1,6 +1,6 @@
 steps:
 
-- label: "Simple · Diffusion & Model Executor Test"
+- label: "Simple Unit Test"
   agent_pool: mi325_1
   depends_on: amd-build
   mirror_hardwares: [amdproduction]
@@ -9,18 +9,7 @@ steps:
     - export VLLM_ROCM_USE_AITER=0
     # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
     # that is not supported on CPU
-    - "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
-
-- label: "Simple · Other Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - export VLLM_ROCM_USE_AITER=0
-    # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
-    # that is not supported on CPU
-    - "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
+    - "timeout 60m pytest -v -s -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
 
 - label: "Voxtral TTS CUDA Unit Test"
   agent_pool: mi325_1
@@ -30,13 +19,13 @@ steps:
   commands:
     - timeout 10m pytest -s -v tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py
 
-- label: "Diffusion · Batch Test"
+- label: "Diffusion Batching Test"
   agent_pool: mi325_1
   depends_on: amd-build
   mirror_hardwares: [amdproduction]
   grade: Blocking
   commands:
-    - timeout 20m pytest -sv tests/diffusion/batching -m 'core_model and cuda' --run-level "core_model"
+    - timeout 20m pytest -s -v tests/diffusion/batching/test_diffusion_batching.py -m "core_model and diffusion" --run-level "core_model"
 
 - label: "Custom Pipeline Test"
   agent_pool: mi325_1

.buildkite/test-merge.yml

@@ -107,7 +107,7 @@ steps:
     depends_on: upload-merge-pipeline
     steps:
       - label: "Entrypoint Test with H100"
-        timeout_in_minutes: 60
+        timeout_in_minutes: 35
         commands:
           - pytest -s -v tests/entrypoints/ -m "advanced_model and cuda and H100" --run-level "advanced_model"
         agents:
@@ -389,54 +389,6 @@ steps:
               volumes:
                 - "/fsx/hf_cache:/fsx/hf_cache"
 
-      - label: "TTS · VoxCPM2 Online Test"
-        timeout_in_minutes: 30
-        commands:
-          - |
-            timeout 30m bash -c "
-              export VLLM_LOGGING_LEVEL=DEBUG
-              export VLLM_WORKER_MULTIPROC_METHOD=spawn
-              pytest -s -v tests/e2e/online_serving/test_voxcpm2_tts.py -m 'advanced_model and cuda' --run-level 'advanced_model'
-            "
-        agents:
-          queue: "gpu_1_queue"
-        plugins:
-          - docker#v5.2.0:
-              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-              always-pull: true
-              propagate-environment: true
-              shm-size: "8gb"
-              environment:
-                - "HF_HOME=/fsx/hf_cache"
-                - "HF_TOKEN"
-              volumes:
-                - "/fsx/hf_cache:/fsx/hf_cache"
-
-      - label: "Omni · Higgs-Audio V2 Test"
-        timeout_in_minutes: 30
-        commands:
-          - |
-            timeout 30m bash -c "
-              export VLLM_LOGGING_LEVEL=DEBUG
-              export VLLM_WORKER_MULTIPROC_METHOD=spawn
-              export VLLM_USE_DEEP_GEMM=0
-              export VLLM_MOE_USE_DEEP_GEMM=0
-              pytest -s -v tests/e2e/online_serving/test_higgs_audio_v2.py tests/e2e/offline_inference/test_higgs_audio_v2.py -m 'advanced_model and cuda' --run-level 'advanced_model'
-            "
-        agents:
-          queue: "gpu_1_queue"
-        plugins:
-          - docker#v5.2.0:
-              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-              always-pull: true
-              propagate-environment: true
-              shm-size: "8gb"
-              environment:
-                - "HF_HOME=/fsx/hf_cache"
-                - "HF_TOKEN"
-              volumes:
-                - "/fsx/hf_cache:/fsx/hf_cache"
-
       - label: "Diffusion · Qwen Image Edit Test"
         timeout_in_minutes: 20
         commands:

.buildkite/test-nightly.yml

@@ -305,51 +305,6 @@ steps:
                       path: /mnt/hf-cache
                       type: DirectoryOrCreate
 
-      - label: ":full_moon: Omni · Perf Test · Multi-Replica"
-        key: nightly-omni-performance-multi-replicas
-        timeout_in_minutes: 180
-        commands:
-          - export BENCHMARK_DIR=tests/dfx/perf/results
-          - |
-            set +e
-            pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen3_omni_multi_replicas.json
-            EXIT=$$?
-            buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
-            exit $$EXIT
-        agents:
-          queue: "mithril-h100-pool"
-        plugins:
-          - kubernetes:
-              podSpec:
-                containers:
-                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-                    resources:
-                      limits:
-                        nvidia.com/gpu: 3
-                    volumeMounts:
-                      - name: devshm
-                        mountPath: /dev/shm
-                      - name: hf-cache
-                        mountPath: /root/.cache/huggingface
-                    env:
-                      - name: HF_HOME
-                        value: /root/.cache/huggingface
-                      - name: HF_TOKEN
-                        valueFrom:
-                          secretKeyRef:
-                            name: hf-token-secret
-                            key: token
-                nodeSelector:
-                  node.kubernetes.io/instance-type: gpu-h100-sxm
-                volumes:
-                  - name: devshm
-                    emptyDir:
-                      medium: Memory
-                  - name: hf-cache
-                    hostPath:
-                      path: /mnt/hf-cache
-                      type: DirectoryOrCreate
-
 
       - label: ":full_moon: Omni · Multi-Replica Startup Test with 4x H100"
         timeout_in_minutes: 45
@@ -398,7 +353,7 @@ steps:
         timeout_in_minutes: 90
         commands:
           - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
-          - pytest -s -v tests/e2e/ -m "full_model and L4 and tts" --run-level "full_model" --ignore=tests/e2e/accuracy --ignore=tests/e2e/online_serving/test_voxcpm2_tts_expansion.py
+          - pytest -s -v tests/e2e/ -m "full_model and L4 and tts" --run-level "full_model" --ignore=tests/e2e/accuracy
         agents:
           queue: "gpu_1_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
         plugins:
@@ -413,30 +368,6 @@ steps:
               volumes:
                 - "/fsx/hf_cache:/fsx/hf_cache"
 
-      - label: ":full_moon: TTS · Accuracy Test · VoxCPM2"
-        timeout_in_minutes: 30
-        commands:
-          - |
-            timeout 30m bash -c "
-              export VLLM_LOGGING_LEVEL=DEBUG
-              export VLLM_WORKER_MULTIPROC_METHOD=spawn
-              export VLLM_ALLOW_LONG_MAX_MODEL_LEN=1
-              pytest -s -v tests/e2e/online_serving/test_voxcpm2_tts_expansion.py -m 'full_model and cuda and tts' --run-level 'full_model'
-            "
-        agents:
-          queue: "gpu_1_queue"
-        plugins:
-          - docker#v5.2.0:
-              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-              always-pull: true
-              propagate-environment: true
-              shm-size: "8gb"
-              environment:
-                - "HF_HOME=/fsx/hf_cache"
-                - "HF_TOKEN"
-              volumes:
-                - "/fsx/hf_cache:/fsx/hf_cache"
-
       - label: ":full_moon: TTS · Perf Test"
         key: nightly-tts-performance
         timeout_in_minutes: 180
@@ -483,52 +414,6 @@ steps:
                       path: /mnt/hf-cache
                       type: DirectoryOrCreate
 
-      - label: ":full_moon: TTS · VoxCPM2 · Perf Test"
-        key: nightly-tts-performance-voxcpm2
-        timeout_in_minutes: 180
-        commands:
-          - export BENCHMARK_DIR=tests/dfx/perf/results
-          - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
-          - |
-            set +e
-            pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_voxcpm2.json
-            EXIT=$$?
-            buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
-            exit $$EXIT
-        agents:
-          queue: "mithril-h100-pool"
-        plugins:
-          - kubernetes:
-              podSpec:
-                containers:
-                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-                    resources:
-                      limits:
-                        nvidia.com/gpu: 1
-                    volumeMounts:
-                      - name: devshm
-                        mountPath: /dev/shm
-                      - name: hf-cache
-                        mountPath: /root/.cache/huggingface
-                    env:
-                      - name: HF_HOME
-                        value: /root/.cache/huggingface
-                      - name: HF_TOKEN
-                        valueFrom:
-                          secretKeyRef:
-                            name: hf-token-secret
-                            key: token
-                nodeSelector:
-                  node.kubernetes.io/instance-type: gpu-h100-sxm
-                volumes:
-                  - name: devshm
-                    emptyDir:
-                      medium: Memory
-                  - name: hf-cache
-                    hostPath:
-                      path: /mnt/hf-cache
-                      type: DirectoryOrCreate
-
   # Diffusion X2I suite: x2i / x2a / x2t and related non-video paths; x2v is only in "Diffusion X2V Model Test" below.
   - group: ":card_index_dividers: Diffusion X2I(&A&T) Model Test"
     key: nightly-diffusion-x2iat-group
@@ -640,7 +525,6 @@ steps:
             tests/e2e/online_serving/test_longcat_image_expansion.py
             tests/e2e/online_serving/test_longcat_image_edit_expansion.py
             tests/e2e/online_serving/test_flux_2_dev_expansion.py
-            tests/e2e/offline_inference/test_glm_image_autoround_w4a16_expansion.py
             -m "full_model and diffusion and H100 and distributed_cuda"
             --run-level "full_model"
         agents:
@@ -694,49 +578,7 @@ steps:
                   - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
                     resources:
                       limits:
-                        nvidia.com/gpu: 2
-                    volumeMounts:
-                      - name: devshm
-                        mountPath: /dev/shm
-                      - name: hf-cache
-                        mountPath: /root/.cache/huggingface
-                    env:
-                      - name: HF_HOME
-                        value: /root/.cache/huggingface
-                      - name: HF_TOKEN
-                        valueFrom:
-                          secretKeyRef:
-                            name: hf-token-secret
-                            key: token
-                nodeSelector:
-                  node.kubernetes.io/instance-type: gpu-h100-sxm
-                volumes:
-                  - name: devshm
-                    emptyDir:
-                      medium: Memory
-                  - name: hf-cache
-                    hostPath:
-                      path: /mnt/hf-cache
-                      type: DirectoryOrCreate
-
-      - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100 · Multi-GPU DreamZero"
-        timeout_in_minutes: 120
-        commands:
-          - >-
-            pytest -sv
-            tests/e2e/online_serving/test_dreamzero_expansion.py
-            -m "full_model and diffusion and H100 and distributed_cuda"
-            --run-level "full_model"
-        agents:
-          queue: "mithril-h100-pool"
-        plugins:
-          - kubernetes:
-              podSpec:
-                containers:
-                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-                    resources:
-                      limits:
-                        nvidia.com/gpu: 2
+                        nvidia.com/gpu: 3
                     volumeMounts:
                       - name: devshm
                         mountPath: /dev/shm
@@ -955,8 +797,8 @@ steps:
                       path: /mnt/hf-cache
                       type: DirectoryOrCreate
 
-      - label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image-Edit-2511"
-        key: nightly-diffusion-x2iat-performance-qwen-image-edit-2511
+      - label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image-Edit-2509"
+        key: nightly-diffusion-x2iat-performance-qwen-image-edit-2509
         timeout_in_minutes: 180
         commands:
           - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
@@ -964,7 +806,7 @@ steps:
           - export CACHE_DIT_VERSION=1.3.0
           - |
             set +e
-            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json
+            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_2509_vllm_omni.json
             EXIT=$$?
             buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
             buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
@@ -1051,20 +893,26 @@ steps:
                       path: /mnt/hf-cache
                       type: DirectoryOrCreate
 
-      - label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · BAGEL"
-        key: nightly-diffusion-x2iat-performance-bagel
-        timeout_in_minutes: 180
+      - label: ":full_moon: Diffusion X2I(&A&T) · HunyuanImage3 · DiT Perf Test"
+        key: nightly-hunyuan-image3-performance
+        soft_fail: true
+        timeout_in_minutes: 120
+        if: build.env("RUN_HUNYUAN_IMAGE3_PERF") == "1"
         commands:
           - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
           - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
           - export CACHE_DIT_VERSION=1.3.0
           - |
             set +e
-            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_bagel_vllm_omni.json
-            EXIT=$$?
+            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_hunyuan_image_tp4_fp8.json
+            EXIT1=$$?
+            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_hunyuan_image_tp2_fp8_sp2.json
+            EXIT2=$$?
+            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_hunyuan_image_tp2_fp8_cfgp2.json
+            EXIT3=$$?
             buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
             buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
-            exit $$EXIT
+            exit $$((EXIT1 | EXIT2 | EXIT3))
         agents:
           queue: "mithril-h100-pool"
         plugins:
@@ -1074,7 +922,7 @@ steps:
                   - 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
+                        nvidia.com/gpu: 4
                     volumeMounts:
                       - name: devshm
                         mountPath: /dev/shm
@@ -1150,8 +998,6 @@ steps:
         timeout_in_minutes: 180
         commands:
           - pytest -s -v tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py -m full_model --run-level full_model
-          - pytest -s -v tests/e2e/accuracy/hunyuanvideo15_t2v/test_hunyuanvideo15_t2v_video_similarity.py -m full_model --run-level full_model
-          - pytest -s -v tests/e2e/accuracy/hunyuanvideo15_i2v/test_hunyuanvideo15_i2v_video_similarity.py -m full_model --run-level full_model
           - pytest -s -v tests/e2e/accuracy/test_ltx2_3_video_similarity.py -m full_model --run-level full_model
           - pytest -s -v tests/e2e/accuracy/test_diffusers_backend_similarity.py -k '2v_matches_diffusers' -m full_model --run-level full_model
         agents:
@@ -1342,13 +1188,11 @@ steps:
       - nightly-omni-performance-no-async-chunk
       - nightly-omni-performance-async-chunk
       - nightly-omni-performance-vllm-text
-      - nightly-omni-performance-multi-replicas
       - nightly-tts-performance
       - nightly-diffusion-x2iat-performance-qwen-image
       - nightly-diffusion-x2iat-performance-qwen-image-edit
-      - nightly-diffusion-x2iat-performance-qwen-image-edit-2511
+      - nightly-diffusion-x2iat-performance-qwen-image-edit-2509
       - nightly-diffusion-x2iat-performance-qwen-image-layered
-      - nightly-diffusion-x2iat-performance-bagel
       - nightly-diffusion-x2v-performance
       - nightly-testcase-statistics
     if: build.env("NIGHTLY") == "1" && build.env("EMAIL_DISTRIBUTION") == "1"
@@ -1360,12 +1204,10 @@ steps:
       - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-no-async-chunk
       - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-async-chunk
       - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-vllm-text
-      - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-multi-replicas
       - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image
       - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-edit
-      - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-edit-2511
+      - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-edit-2509
       - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-layered
-      - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-bagel
       - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2v-performance
       - buildkite-agent artifact download "tests/dfx/perf/results/*.html" . --step nightly-testcase-statistics
       - python tools/nightly/generate_nightly_perf_excel.py

.buildkite/test-ready.yml

@@ -183,7 +183,7 @@ steps:
     depends_on: upload-ready-pipeline
     commands:
       - |
-        timeout 20m bash -c "
+        timeout 15m bash -c "
                   pytest -sv tests/distributed/ -m 'core_model and cuda' --run-level 'core_model'
         "
     agents:
@@ -296,29 +296,6 @@ steps:
               volumes:
                 - "/fsx/hf_cache:/fsx/hf_cache"
 
-      - label: "TTS · VoxCPM2 Online Test"
-        timeout_in_minutes: 30
-        commands:
-          - |
-            timeout 30m bash -c "
-              export VLLM_LOGGING_LEVEL=DEBUG
-              export VLLM_WORKER_MULTIPROC_METHOD=spawn
-              pytest -s -v tests/e2e/online_serving/test_voxcpm2_tts.py -m 'core_model and tts' --run-level 'core_model'
-            "
-        agents:
-          queue: "gpu_1_queue"
-        plugins:
-          - docker#v5.2.0:
-              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-              always-pull: true
-              propagate-environment: true
-              shm-size: "8gb"
-              environment:
-                - "HF_HOME=/fsx/hf_cache"
-                - "HF_TOKEN"
-              volumes:
-                - "/fsx/hf_cache:/fsx/hf_cache"
-
       - label: "TTS · Qwen3-TTS CustomVoice Test"
         commands:
           - |

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

@@ -53,9 +53,6 @@
 {% endif %}
             {{ indented_cmd | safe }}
         priority: 100
-        {% if step.timeout_in_minutes %}
-        timeout_in_minutes: {{ step.timeout_in_minutes }}
-        {% endif %}
         {% if step.grade and step.grade == "Blocking" %}
         soft_fail: false
         {% else %}

.claude/skills/diffusion-perf-opt/SKILL.md

@@ -21,9 +21,6 @@ Before proposing changes, ask for the optimization scene if it is not already kn
 - Current enabled strategies: USP/SP, CFG parallel, HSDP/FSDP, VAE patch parallel, torch.compile, profiler options.
 - Optimization target: latency, throughput, memory, cost, or quality-preserving speed.
 - Precision/quality tolerance: bf16/fp8/quantization/sparsity/approximate attention allowed or not.
-- Quality validation method: exact-output comparison when feasible, or metrics
-  such as SSIM, PSNR, LPIPS, cosine similarity, MAE, MSE, temporal flicker
-  checks, baseline self-run variance, and reviewable output artifacts.
 
 ## Workflow
 
@@ -82,19 +79,7 @@ Before proposing changes, ask for the optimization scene if it is not already kn
    - Record latency, stage timings, memory, output quality, and logs.
    - Report comparison tables in ms. Include at least end-to-end client time, server `inference_time`, server stage generation time if available, `vae.encode`, `diffuse`, `vae.decode`, and peak memory.
 
-6. **Enforce the quality and precision gate for every optimization.**
-   - Do not mark an optimization ready until quality is checked on the same
-     workload: model, prompt/media, seed, shape, steps, scheduler, dtype,
-     backend, parallelism, and output encoding.
-   - For math-preserving changes, require exact/near-exact agreement or show
-     that differences are within baseline self-run variance.
-   - For precision, quantization, approximate attention, sparse/custom kernels,
-     or backend changes, include reviewable artifacts and metrics such as SSIM,
-     PSNR, LPIPS, cosine similarity, MAE, MSE, and temporal flicker checks.
-   - If quality tolerance is not stated, default to quality-preserving behavior.
-     Failed or inconclusive quality validation blocks a ready-to-merge claim.
-
-7. **Collect diagnostic trace only after narrowing hypotheses.**
+6. **Collect diagnostic trace only after narrowing hypotheses.**
    - Use torch profiler for a small number of requests.
    - Run two separate diagnostic traces instead of mixing concerns:
      - **Operator/shape trace**: enable `torch_profiler_record_shapes=True` and keep stack collection disabled. Use this to rank CUDA kernels, NCCL collectives, attention/MLP/norm/RoPE work, and shape-specific hot operators.
@@ -106,16 +91,15 @@ Before proposing changes, ask for the optimization scene if it is not already kn
    - Diagnostic reports must be written to disk and preserve: server command, profiler config, warmup/request/polling commands, trace artifact paths, rank analyzed, analyzer output or summary, and the decision about whether additional ranks are necessary.
    - Analyze both rank-level balance and device-level free bubbles when additional ranks are opened.
 
-8. **Analyze host, communication, and operators.**
+7. **Analyze host, communication, and operators.**
    - Find GPU idle/free intervals and map each large gap to the enclosing CPU/Python code.
    - Separate real GPU idle from profiler overhead such as CUPTI `Command Buffer Full`.
    - Compare NCCL kernel time to user annotations; annotations can overcount nested intervals.
    - Rank operator work by total CUDA time and by repeated small-kernel launch count.
 
-9. **Produce an optimization plan.**
+8. **Produce an optimization plan.**
    - Classify candidates as P0/P1/P2.
    - For each candidate, state necessity, expected benefit, implementation path, validation plan, and quality risk.
-   - Include a concrete quality/precision validation gate for each candidate.
    - Do not implement high-risk operator rewrites before proving the operator is a bottleneck for the target shapes.
    - End the plan with a user-facing candidate selection table. The assistant
      should not automatically choose a risky optimization just because it is
@@ -140,13 +124,7 @@ Before proposing changes, ask for the optimization scene if it is not already kn
 - **P1:** meaningful code changes with contained risk. Examples: cross-attention KV caching, VAE gather/broadcast reduction, AdaLayerNorm/RMSNorm/RoPE fusion after trace evidence.
 - **P2:** high implementation or quality risk. Examples: FA to LA replacement, custom Triton/CUDA fused kernels, FP8/quantization, sparsity/Rainfusion-style acceleration.
 
-Every implemented optimization needs A/B validation and a passing
-quality/precision gate, including math-preserving changes such as padding trim,
-layout cleanup, cache reuse, or host/runtime cleanup. Use objective metrics
-when available, keep reviewable artifacts, and compare against baseline
-self-run variance for precision, quantization, approximate kernels, or backend
-changes. If quality validation fails or is inconclusive, do not present the
-optimization as ready to merge.
+Every implemented optimization needs A/B validation and quality regression. A/B means same workload and hardware before/after. Quality regression means checking generated image/video stability, artifacts, temporal flicker, and seed behavior when precision or approximate kernels change.
 
 ## Optimization Layers
 

.gitignore

@@ -270,6 +270,3 @@ vllm_omni/_version.py
 .worktrees/
 # CI overlay yamls materialized from tests/utils.py:_CI_OVERLAYS at test time
 tests/.ci_generated/
-
-# Local Lance precision-matching dump scripts (not part of the repo)
-lance_compare/

.readthedocs.yml

@@ -10,11 +10,6 @@ build:
   jobs:
     post_checkout:
       - git fetch --unshallow || true
-    pre_install:
-      # Pre-install CPU torch so the subsequent `pip install .[docs]` does not
-      # pull ~2 GB of CUDA 13 libs (cuBLAS/cuDNN/NCCL/...) that exceed RTD's
-      # 1 GB memory cap and time out the mkdocs build.
-      - python -m pip install --index-url https://download.pytorch.org/whl/cpu torch
 
 mkdocs:
   configuration: mkdocs.yml

benchmarks/README.md

@@ -1,60 +1,44 @@
-# Benchmarks
+# Benchmarks Overview and Architecture
 
-This directory contains benchmark suites for evaluating different model families and infrastructure components in vLLM-Omni. Each subfolder targets a different benchmark family with its own scripts, configs, and metrics. See the per-subfolder READMEs for detailed usage.
+This document explains the benchmark architecture across all benchmark assets in this repo. It describes what we measure, and where to find or plug in new scenarios. Per-task details remain in subfolder READMEs (e.g., `benchmarks/<model>/README.md`).
 
-## Benchmark families
+## Scope and goals
+- Establish repeatable latency/throughput measurements for multimodal LLM pipelines.
+- Provide both HF Transformers (offline) and vLLM-Omni (multi-stage/pipeline) baselines.
+- Make it easy to plug in new datasets and models with minimal changes to the runner scripts.
 
-### [TTS](tts/README.md) — Text-to-Speech
+## Dataset and inputs
+- Default example: SeedTTS top-100 prompts (`benchmarks/build_dataset/top100.txt`) via `benchmarks/build_dataset/`.
+- Extensible: drop in new prompt files or modality-aligned payloads; keep the expected format for the consuming scripts (e.g., one prompt per line).
+- If you add a new dataset, document it under `benchmarks/<model>/README.md` and point scripts to your data path.
 
-Model-agnostic serving benchmarks for TTS models, including Qwen3-TTS and VoxCPM2.
+## Directory layout
+- `benchmarks/build_dataset/` — dataset prep utilities (e.g., SeedTTS top100).
+- `benchmarks/<model>/vllm_omni/` — vLLM-Omni pipeline benchmarks, logs, outputs.
+- `benchmarks/accuracy/` — accuracy benchmark integrations that adapt external
+  benchmark suites to vLLM-Omni serving and evaluation flows.
+- Add new tasks under `benchmarks/<model>/...` with the same pattern: `transformers/`, `vllm_omni/`, task-specific README, and (optionally) dataset prep notes.
 
-- **Layout**: `tts/bench_tts.py` (serving benchmark driver), `tts/model_configs.yaml` (model registry), `tts/plot_results.py` (result plotting)
-- **Dataset**: Seed-TTS full or text-only datasets, plus bundled smoke/design prompts under `build_dataset/`
-- **Key metrics**: TTFP (time to first audio packet), E2E latency, RTF (real-time factor), throughput (audio seconds / wall-clock second)
+## Reference workflows
+- **HF Transformers (offline, single process)**  
+  Script (example): `benchmarks/<model>/transformers/eval_qwen3_moe_omni_transformers.sh`  
+  Outputs: `benchmark_results/perf_stats.json`, `benchmark_results/results.json`, `benchmark_results/audio/` (if audio is produced).
 
-### [Diffusion](diffusion/README.md) — Image and Video Generation
+- **vLLM-Omni end-to-end pipeline**  
+  Script (example): `benchmarks/<model>/vllm_omni/eval_qwen3_moe_omni.sh`  
+  Outputs: `vllm_omni/logs/*.stats.jsonl` (per-stage/overall latency & TPS), `vllm_omni/logs/stage*.log`, `vllm_omni/outputs/` (text/audio artifacts).
 
-Online-serving benchmark for diffusion image/video models, sending requests to the configured vLLM serving endpoint (`/v1/chat/completions`, `/v1/images/generations`, or `/v1/videos`, depending on backend/task).
+- **Adding a new task/model**  
+  1) Create `benchmarks/<model>/transformers/` and/or `benchmarks/<model>/vllm_omni/` with scripts referencing your model and dataset.  
+  2) Add a task README describing dataset, configs, and expected outputs.  
+  3) Keep the output/log structure similar for easy comparison (perf_stats/results/audio or text outputs; stats.jsonl/logs for pipeline).
 
-- **Tasks**: text-to-image, text-to-video, image-to-image, image-to-video, text+image-to-image, text+image-to-video
-- **Datasets**: `vbench`, `trace`, `random`
-- **Key metrics**: request throughput, latency percentiles, optional SLO attainment
+## Metrics to watch
+- **Throughput**: `overall_tps`, `*_tps_avg` per stage.
+- **Latency distribution**: look for long tails in `*.stats.jsonl`.
+- **Quality/completeness**: missing outputs or errors in stage logs indicate pipeline failures or misconfigurations.
 
-### [GLM-Image](glm_image/README.md) — Text-to-Image and Image-to-Image
-
-Benchmarks for GLM-Image performance across HuggingFace baseline, vLLM-Omni offline inference, and vLLM-Omni online serving.
-
-- **Layout**: `glm_image/huggingface/` (HF baseline), `glm_image/vllm-omni/` (offline inference), `glm_image/benchmark_glm_image.py` (online serving)
-- **Tasks**: text-to-image and image-to-image
-- **Key metrics**: request/image throughput, latency percentiles, optional per-stage pipeline timings
-
-### [Distributed](distributed/omni_connectors/README.md) — RDMA Connector Testing
-
-RDMA environment setup and transfer tests for `MooncakeTransferEngineConnector`, including pytest-based single-node checks and manual cross-node benchmarks.
-
-- **Transfer modes**: `copy`, `zerocopy`, `gpu` (GPUDirect)
-- **Supports**: single-node pytest suites and manual multi-node/cross-node transfer testing
-
-### [Accuracy](accuracy/README.md) — Image Generation and Editing Quality
-
-Accuracy benchmarks for image generation/editing models, adapting external suites to vLLM-Omni serving and local judge-evaluation flows.
-
-- **Layout**: `accuracy/text_to_image/` (GEBench), `accuracy/image_to_image/` (GEdit-Bench)
-- **Method**: generation and judge scoring both run through local `vllm-omni serve` endpoints
-
-### Common serving metrics framework
-
-`vllm_omni/benchmarks/` extends `vllm bench serve --omni` with Omni-specific datasets, backends, and multimodal metrics. Key metrics include:
-
-- **Text output**: TTFT (time to first token), TPOT (time per output token), ITL (inter-token latency)
-- **Audio output**: TTFP (time to first audio packet), E2E latency, RTF (real-time factor)
-- **Throughput**: request throughput, output token throughput, total token throughput, audio throughput
-
-See `vllm_omni/benchmarks/serve.py` for the `vllm bench serve --omni` runner wrapper and `vllm_omni/benchmarks/metrics/` for Omni metric definitions.
-
-## Adding a new benchmark
-
-1. Create a subfolder under `benchmarks/<name>/` with scripts, configs if needed, and a `README.md`.
-2. If comparing against another runtime, use clear backend subfolders where applicable, such as `huggingface/` and `vllm-omni/`, or follow the shared TTS serving benchmark pattern in `tts/`.
-3. Add reusable dataset or prompt-building utilities to `build_dataset/` if applicable.
-4. Update this README with a link to the new benchmark family.
+## Troubleshooting
+- Verify GPU/driver/FlashAttention2 requirements for your chosen model/config.
+- Ensure network access for dataset/model downloads (Google Drive, Hugging Face, etc.).
+- If outputs are missing or slow, inspect per-stage logs and `*.stats.jsonl` for errors, stragglers, or contention.

benchmarks/build_dataset/sound_effect_smoke/en/meta.lst

@@ -1,10 +0,0 @@
-se001|ocean waves crashing on a rocky beach|10.0
-se002|heavy rain on a metal roof at night|8.0
-se003|a dog barking in the distance|6.0
-se004|a crackling campfire in a quiet forest|10.0
-se005|footsteps on dry leaves|6.0
-se006|thunder rolling across a wide valley|8.0
-se007|a busy coffee shop with chatter and espresso machines|10.0
-se008|wind howling through tall pine trees|8.0
-se009|a small stream flowing over pebbles|10.0
-se010|a creaking wooden door slowly opening|6.0

benchmarks/build_dataset/ttsd_smoke/en/meta.lst

@@ -1,10 +0,0 @@
-ttsd001|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] Hello and welcome to the show today. [S2] Thanks for having me, it's a pleasure to be here. [S1] Let's dive right in.
-ttsd002|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] Did you finish the report? [S2] Almost — I'm just polishing the conclusion. [S1] Great, send it over when you're done.
-ttsd003|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] Have you tried that new restaurant downtown? [S2] Yes, the pasta is amazing. [S1] I'll book a table for Friday.
-ttsd004|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] What did you think of yesterday's meeting? [S2] Productive, but a bit long. [S1] Agreed, we should set an hour cap next time.
-ttsd005|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] The weather is beautiful today. [S2] Perfect for a hike. [S1] Want to drive out to the trails this afternoon?
-ttsd006|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] How did the demo go? [S2] Better than expected — the client signed off on the prototype. [S1] Fantastic news!
-ttsd007|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] Coffee or tea this morning? [S2] Coffee please, with a splash of milk. [S1] Coming right up.
-ttsd008|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] Did you read the email about the deadline change? [S2] I did, we now have until next Friday. [S1] That should give us enough buffer.
-ttsd009|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] Are you joining the team dinner tonight? [S2] I'll be there around seven. [S1] See you then.
-ttsd010|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] Have we backed up the database? [S2] Last backup ran two hours ago. [S1] Good, I'll trigger a fresh snapshot before the migration.

benchmarks/diffusion/bench_attention_backends.py

@@ -1,328 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Standalone diagnostic for the diffusion attention backends added in #3079.
-
-Exercises the same synthetic attention shape through each backend and each
-torch SDPA sub-kernel, with and without an attn_mask, so we can pinpoint why
-CUDNN_ATTN underperforms on some SKUs. Two typical causes:
-
-  (a) cuDNN lacks a tuned kernel for this (SM, head_dim, seq) combination
-      and silently falls back to MATH.
-  (b) attn_mask is non-None and the cuDNN SDPA dispatch rejects it,
-      walking CUDNN -> FLASH -> MATH.
-
-Run:
-    python benchmarks/diffusion/bench_attention_backends.py --preset hv15
-    python benchmarks/diffusion/bench_attention_backends.py --preset wan22
-    python benchmarks/diffusion/bench_attention_backends.py \
-        --batch 1 --heads 24 --seq 14336 --head-dim 128
-
-Optional extras:
-    pip install --pre flash-attn-4        # FA4 is currently pre-release only
-    pip install -U flashinfer              # latest FlashInfer (0.6.9)
-
-The table at the end is the data we want on the PR — surface the row where a
-backend is >1.5x the SDPA baseline, that's the one to gate off the auto-route.
-"""
-
-from __future__ import annotations
-
-import argparse
-import time
-
-import torch
-from torch.nn.attention import SDPBackend, sdpa_kernel
-
-# Presets mirror the DiT hot-path attention call (single-stream block) at the
-# shapes we run in the PR's validation configs.
-_PRESETS = {
-    # HunyuanVideo-1.5 480p/33f: (30 * 52 * 9) latent + ~256 text tokens.
-    "hv15": {"batch": 1, "heads": 24, "seq": 14336, "head_dim": 128},
-    # Wan 2.2 480p/33f rough estimate; adjust when we measure real shapes.
-    "wan22": {"batch": 1, "heads": 40, "seq": 16384, "head_dim": 128},
-    # Smaller image-gen shape for quick smoke tests.
-    "flux": {"batch": 1, "heads": 24, "seq": 4096, "head_dim": 128},
-}
-
-_SDPA_BACKENDS = [
-    ("CUDNN_ATTENTION", [SDPBackend.CUDNN_ATTENTION]),
-    ("FLASH_ATTENTION", [SDPBackend.FLASH_ATTENTION]),
-    ("EFFICIENT_ATTENTION", [SDPBackend.EFFICIENT_ATTENTION]),
-    ("MATH", [SDPBackend.MATH]),
-    # The PR's CUDNN_ATTN impl uses this priority chain; keep it here to show
-    # what torch actually picks when multiple are allowed.
-    ("CUDNN_ATTN_CHAIN", [SDPBackend.CUDNN_ATTENTION, SDPBackend.FLASH_ATTENTION, SDPBackend.MATH]),
-]
-
-
-def _env_report() -> None:
-    if not torch.cuda.is_available():
-        raise SystemExit("CUDA not available; this diagnostic needs a GPU.")
-    cc = torch.cuda.get_device_capability()
-    print("=" * 72)
-    print(f"GPU           : {torch.cuda.get_device_name()}")
-    print(f"SM capability : sm_{cc[0]}{cc[1]}")
-    print(f"torch         : {torch.__version__}")
-    print(f"cuDNN version : {torch.backends.cudnn.version()}")
-    print("flashinfer    : ", end="")
-    try:
-        import flashinfer
-
-        print(getattr(flashinfer, "__version__", "present"))
-    except Exception as e:
-        print(f"not installed ({type(e).__name__})")
-    print("=" * 72)
-
-
-def _make_qkv(batch: int, heads: int, seq: int, head_dim: int, device: str, dtype: torch.dtype):
-    # Layout: (B, S, H, D) — what the PR's backends accept as input.
-    q = torch.randn(batch, seq, heads, head_dim, device=device, dtype=dtype)
-    k = torch.randn(batch, seq, heads, head_dim, device=device, dtype=dtype)
-    v = torch.randn(batch, seq, heads, head_dim, device=device, dtype=dtype)
-    return q, k, v
-
-
-def _make_mask(batch: int, seq: int, device: str, dtype: torch.dtype, pad_tokens: int = 128) -> torch.Tensor:
-    # Mirrors a text-encoder padding mask: last `pad_tokens` positions masked.
-    mask = torch.zeros(batch, 1, seq, seq, device=device, dtype=dtype)
-    if pad_tokens > 0:
-        mask[..., -pad_tokens:] = float("-inf")
-    return mask
-
-
-def _time_call(fn, *args, warmup: int = 3, iters: int = 10) -> tuple[float, str]:
-    """Return (median ms, error string). Skips on any exception so that a
-    backend that rejects our args (wrong dtype, missing JIT module, unsupported
-    kwarg value) doesn't abort the whole sweep."""
-    try:
-        for _ in range(warmup):
-            fn(*args)
-        torch.accelerator.synchronize()
-        times = []
-        for _ in range(iters):
-            torch.accelerator.synchronize()
-            t0 = time.perf_counter()
-            fn(*args)
-            torch.accelerator.synchronize()
-            times.append((time.perf_counter() - t0) * 1000.0)
-        times.sort()
-        return times[len(times) // 2], ""
-    except Exception as e:  # noqa: BLE001 — probe script, keep the whole table going
-        msg = str(e).split("\n", 1)[0][:60]
-        return float("nan"), f"{type(e).__name__}: {msg}" if msg else type(e).__name__
-
-
-def _run_sdpa_variants(q, k, v, attn_mask, scale: float) -> list[tuple[str, float, str]]:
-    # q/k/v come in (B, S, H, D); F.sdpa wants (B, H, S, D).
-    q_t = q.permute(0, 2, 1, 3)
-    k_t = k.permute(0, 2, 1, 3)
-    v_t = v.permute(0, 2, 1, 3)
-    rows: list[tuple[str, float, str]] = []
-    for name, backends in _SDPA_BACKENDS:
-
-        def _call(mask=attn_mask):
-            with sdpa_kernel(backends):
-                return torch.nn.functional.scaled_dot_product_attention(
-                    q_t, k_t, v_t, attn_mask=mask, dropout_p=0.0, is_causal=False, scale=scale
-                )
-
-        ms, err = _time_call(_call)
-        rows.append((name, ms, err))
-    return rows
-
-
-def _run_flashinfer(q, k, v, scale: float, backend: str | None = None, attn_mask=None) -> tuple[float, str]:
-    """Call FlashInfer's dense single-prefill.
-
-    ``backend`` hints at cutlass/fa3/trtllm if the installed version exposes
-    that kwarg. ``attn_mask`` is the same additive 4D mask we pass to SDPA —
-    we convert it to the 2D boolean ``custom_mask`` form FlashInfer accepts.
-    Per FlashInfer docs, ``custom_mask`` only applies when ``causal=False``.
-    """
-    try:
-        import inspect
-
-        from flashinfer.prefill import single_prefill_with_kv_cache
-    except Exception as e:
-        return float("nan"), f"import-{type(e).__name__}"
-
-    kwargs: dict = {"sm_scale": scale, "causal": False, "return_lse": False}
-    if backend is not None:
-        sig = inspect.signature(single_prefill_with_kv_cache)
-        if "backend" not in sig.parameters:
-            return float("nan"), "no-backend-kwarg"
-        kwargs["backend"] = backend
-
-    if attn_mask is not None:
-        # Collapse (B, 1, S, S) additive float mask to (S, S) boolean.
-        # FlashInfer expects True = keep, False = masked out.
-        mask_2d = attn_mask[0, 0]
-        kwargs["custom_mask"] = mask_2d != float("-inf")
-
-    def _call():
-        out = single_prefill_with_kv_cache(q[0], k[0], v[0], **kwargs)
-        return out.unsqueeze(0)
-
-    return _time_call(_call)
-
-
-def _run_fa4(q, k, v, scale: float) -> tuple[float, str]:
-    """Call FlashAttention-4 directly (``pip install flash-attn-4``). FA4
-    ships a Blackwell-native kernel via CuTe-DSL; on sm_120 it should beat
-    cuDNN by ~20%. API lives under ``flash_attn.cute``; the older
-    ``flash_attn.flash_attn_func`` path is FA2/FA3 only."""
-    try:
-        from flash_attn.cute import flash_attn_func
-    except Exception as e:
-        return float("nan"), f"import-{type(e).__name__}"
-
-    def _call():
-        # FA4 accepts (B, S, H, D) directly — same layout the PR's backends use.
-        return flash_attn_func(q, k, v, softmax_scale=scale, causal=False)
-
-    return _time_call(_call)
-
-
-def _run_flashinfer_cudnn_batch(q, k, v, scale: float) -> tuple[float, str]:
-    """Call FlashInfer's direct cuDNN wrapper. Bypasses PyTorch SDPA dispatch
-    (which has a few hundred ns of overhead per call) and talks to cuDNN FMHA
-    straight. Useful as a ceiling for 'pure cuDNN, no SDPA' on Blackwell."""
-    try:
-        from flashinfer.prefill import cudnn_batch_prefill_with_kv_cache
-    except Exception as e:
-        return float("nan"), f"import-{type(e).__name__}"
-
-    b, s, h, d = q.shape
-    qo_indptr = torch.tensor([0, s], dtype=torch.int32, device=q.device)
-    kv_indptr = torch.tensor([0, s], dtype=torch.int32, device=q.device)
-
-    def _call():
-        return cudnn_batch_prefill_with_kv_cache(
-            q.reshape(b * s, h, d),
-            k.reshape(b * s, h, d),
-            v.reshape(b * s, h, d),
-            qo_indptr=qo_indptr,
-            kv_indptr=kv_indptr,
-            max_qo_len=s,
-            max_kv_len=s,
-            sm_scale=scale,
-            causal=False,
-        )
-
-    return _time_call(_call)
-
-
-def _print_table(title: str, rows: list[tuple[str, float, str]], baseline_name: str | None = None) -> None:
-    print(f"\n{title}")
-    print("-" * 88)
-    header = f"{'backend':<24} {'median (ms)':>14}    {'vs baseline':>12}    status"
-    print(header)
-    print("-" * 88)
-    baseline_ms = None
-    if baseline_name is not None:
-        for name, ms, _ in rows:
-            if name == baseline_name and ms == ms:  # not NaN
-                baseline_ms = ms
-                break
-    for name, ms, err in rows:
-        ms_str = f"{ms:>14.3f}" if ms == ms else f"{'n/a':>14}"
-        if baseline_ms is not None and ms == ms:
-            ratio = baseline_ms / ms
-            ratio_str = f"{ratio:>11.2f}x"
-        else:
-            ratio_str = f"{'—':>12}"
-        status = "ok" if not err else f"FAILED ({err})"
-        print(f"{name:<24} {ms_str}    {ratio_str}    {status}")
-
-
-def _pick_winner(rows: list[tuple[str, float, str]]) -> tuple[str, float] | None:
-    """Return (backend, ms) of the fastest non-failing row, or None."""
-    ok_rows = [(n, m) for n, m, err in rows if not err and m == m]
-    if not ok_rows:
-        return None
-    return min(ok_rows, key=lambda x: x[1])
-
-
-def _bench_one_shape(shape: dict, dtype: torch.dtype, device: str) -> tuple[list, list]:
-    torch.manual_seed(0)
-    q, k, v = _make_qkv(**shape, device=device, dtype=dtype)
-    scale = 1.0 / (shape["head_dim"] ** 0.5)
-
-    rows_nomask = _run_sdpa_variants(q, k, v, attn_mask=None, scale=scale)
-    rows_nomask.append(("FLASHINFER (default)", *_run_flashinfer(q, k, v, scale)))
-    # `trtllm-gen` has no sm_120 cubins (NVIDIA/TensorRT-LLM#11799) — skipped.
-    for fi_backend in ("fa2", "fa3", "cutlass", "auto"):
-        rows_nomask.append((f"FLASHINFER ({fi_backend})", *_run_flashinfer(q, k, v, scale, backend=fi_backend)))
-    rows_nomask.append(("FLASHINFER (cudnn-batch)", *_run_flashinfer_cudnn_batch(q, k, v, scale)))
-    rows_nomask.append(("FA4 (direct)", *_run_fa4(q, k, v, scale)))
-
-    mask = _make_mask(shape["batch"], shape["seq"], device=device, dtype=dtype, pad_tokens=256)
-    rows_mask = _run_sdpa_variants(q, k, v, attn_mask=mask, scale=scale)
-    rows_mask.append(("FLASHINFER (dense)", *_run_flashinfer(q, k, v, scale, attn_mask=mask)))
-    rows_mask.append(("FA4 (direct)", float("nan"), "mask-not-supported"))
-
-    return rows_nomask, rows_mask
-
-
-def main() -> None:
-    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
-    parser.add_argument("--preset", choices=list(_PRESETS.keys()), default="hv15")
-    parser.add_argument("--sweep", action="store_true", help="Run all presets and print a ranking")
-    parser.add_argument("--batch", type=int, default=None)
-    parser.add_argument("--heads", type=int, default=None)
-    parser.add_argument("--seq", type=int, default=None)
-    parser.add_argument("--head-dim", type=int, default=None)
-    parser.add_argument("--dtype", choices=["bf16", "fp16"], default="bf16")
-    parser.add_argument("--device", default="cuda:0")
-    args = parser.parse_args()
-
-    dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float16
-    _env_report()
-    print(f"dtype         : {args.dtype}")
-
-    presets = list(_PRESETS.keys()) if args.sweep else [args.preset]
-    summary: dict[str, tuple[str, float] | None] = {}
-
-    for preset_name in presets:
-        shape = dict(_PRESETS[preset_name])
-        if not args.sweep:
-            for k in ("batch", "heads", "seq", "head_dim"):
-                v = getattr(args, k if k != "head_dim" else "head_dim")
-                if v is not None:
-                    shape[k] = v
-
-        print("\n" + "=" * 88)
-        print(
-            f"Preset: {preset_name}  |  batch={shape['batch']} heads={shape['heads']} "
-            f"seq={shape['seq']} head_dim={shape['head_dim']}"
-        )
-        print("=" * 88)
-
-        rows_nomask, rows_mask = _bench_one_shape(shape, dtype, args.device)
-        _print_table(f"[{preset_name}] No attention mask", rows_nomask, baseline_name="CUDNN_ATTENTION")
-        _print_table(
-            f"[{preset_name}] With attention mask (pad 256 tokens)", rows_mask, baseline_name="CUDNN_ATTENTION"
-        )
-
-        summary[preset_name] = _pick_winner(rows_nomask)
-
-    if args.sweep or len(presets) > 1:
-        print("\n" + "=" * 88)
-        print("Winners per preset (no-mask path)")
-        print("=" * 88)
-        for preset_name, winner in summary.items():
-            if winner is None:
-                print(f"  {preset_name:<10} — no successful backend")
-            else:
-                name, ms = winner
-                print(f"  {preset_name:<10} {name:<24} {ms:>8.3f} ms")
-
-    print("\nNotes:")
-    print("  * Ratios are relative to CUDNN_ATTENTION. >1.0x means faster than cuDNN.")
-    print("  * Mask-path winner inherits CUDNN_ATTN's fallback in the PR's backends.")
-    print("  * `trtllm-gen` and FA4 4.0.0b10 are known-broken on sm_120 as of Apr 2026.")
-    print("  * For e2e timings, run text_to_video.py / text_to_image.py while varying DIFFUSION_ATTENTION_BACKEND.")
-
-
-if __name__ == "__main__":
-    main()

benchmarks/diffusion/bench_sage_attn3_matrix.sh

@@ -1,108 +0,0 @@
-#!/usr/bin/env bash
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-#
-# 2x2 timing matrix (backend x compile mode) for SAGE_ATTN_3 vs TORCH_SDPA,
-# on HunyuanVideo-1.5 480p T2V and Wan 2.2 TI2V-5B. Run from the repo root:
-#
-#     bash benchmarks/diffusion/bench_sage_attn3_matrix.sh
-#
-# Requires: sageattn3 installed (for SAGE_ATTN_3 rows) and a Blackwell GPU.
-
-set -u
-
-STEPS=${STEPS:-50}
-SEED=${SEED:-42}
-PROMPT=${PROMPT:-"A dog running across a field of golden wheat."}
-OUT_DIR=${OUT_DIR:-/tmp/bench_sage3}
-mkdir -p "$OUT_DIR"
-
-declare -A TOTAL PER_STEP TEXT_ENC TRANSFORMER VAE
-
-run() {
-    local name="$1"; shift
-    local model="$1"; shift
-    local h="$1"; shift
-    local w="$1"; shift
-    local f="$1"; shift
-    local g="$1"; shift
-    local log="$OUT_DIR/${name}.log"
-    echo "[$(date +%T)] Running $name ..."
-    env "$@" python examples/offline_inference/text_to_video/text_to_video.py \
-        --model "$model" \
-        --prompt "$PROMPT" \
-        --height "$h" --width "$w" --num-frames "$f" \
-        --num-inference-steps "$STEPS" --seed "$SEED" --guidance-scale "$g" \
-        --enable-diffusion-pipeline-profiler \
-        --output "$OUT_DIR/${name}.mp4" > "$log" 2>&1 || {
-            echo "  FAILED — see $log"
-            TOTAL[$name]="FAIL"; PER_STEP[$name]="FAIL"
-            TEXT_ENC[$name]="-"; TRANSFORMER[$name]="-"; VAE[$name]="-"
-            return
-        }
-
-    local wait_ms
-    wait_ms=$(grep "add_req_and_wait=" "$log" | tail -1 | sed -nE 's/.*add_req_and_wait=([0-9.]+) ms.*/\1/p')
-    local total_s
-    total_s=$(grep -oE "Total generation time: [0-9.]+ seconds" "$log" | tail -1 | sed -nE 's/.*: ([0-9.]+) seconds.*/\1/p')
-    local text_enc_s
-    text_enc_s=$(grep -oE "text_encoder.forward took [0-9.]+s" "$log" | awk '{sub(/s$/,"",$NF); sum+=$NF} END {printf "%.3f", sum+0}')
-    local vae_s
-    vae_s=$(grep -oE "vae.decode took [0-9.]+s" "$log" | awk '{sub(/s$/,"",$NF); sum+=$NF} END {printf "%.3f", sum+0}')
-    local pipeline_s
-    pipeline_s=$(grep -oE "Pipeline\.forward took [0-9.]+s" "$log" | tail -1 | sed -nE 's/.*took ([0-9.]+)s.*/\1/p')
-    local transformer_s
-    transformer_s=$(awk "BEGIN {printf \"%.3f\", ${pipeline_s:-0} - ${vae_s:-0} - ${text_enc_s:-0}}")
-
-    TOTAL[$name]="${total_s:-0}"
-    PER_STEP[$name]=$(awk "BEGIN {printf \"%.3f\", ${wait_ms:-0} / 1000 / $STEPS}")
-    TEXT_ENC[$name]="$text_enc_s"
-    TRANSFORMER[$name]="$transformer_s"
-    VAE[$name]="$vae_s"
-
-    echo "  total=${TOTAL[$name]}s  per-step=${PER_STEP[$name]}s/it  transformer=${TRANSFORMER[$name]}s  vae=${VAE[$name]}s"
-}
-
-# ========== HunyuanVideo-1.5 480p T2V (480x832x33, CFG=6.0) ==========
-HV_MODEL=hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v
-run hv15_sdpa_eager     "$HV_MODEL" 480 832 33 6.0 TORCH_COMPILE_DISABLE=1
-run hv15_sdpa_compiled  "$HV_MODEL" 480 832 33 6.0
-run hv15_sage3_eager    "$HV_MODEL" 480 832 33 6.0 TORCH_COMPILE_DISABLE=1 DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN_3
-run hv15_sage3_compiled "$HV_MODEL" 480 832 33 6.0                         DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN_3
-
-# ========== Wan 2.2 TI2V-5B (704x1280x33, CFG=5.0) ==========
-WAN_MODEL=Wan-AI/Wan2.2-TI2V-5B-Diffusers
-run wan22_sdpa_eager     "$WAN_MODEL" 704 1280 33 5.0 TORCH_COMPILE_DISABLE=1
-run wan22_sdpa_compiled  "$WAN_MODEL" 704 1280 33 5.0
-run wan22_sage3_eager    "$WAN_MODEL" 704 1280 33 5.0 TORCH_COMPILE_DISABLE=1 DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN_3
-run wan22_sage3_compiled "$WAN_MODEL" 704 1280 33 5.0                         DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN_3
-
-print_table() {
-    local title="$1"; shift
-    echo
-    echo "$title (${STEPS} steps, seed=${SEED})"
-    echo
-    printf "| %-23s | %-9s | %-15s | %-12s | %-15s | %-14s |\n" \
-        "Config" "Total (s)" "Per-step (s/it)" "Text Enc (s)" "Transformer (s)" "VAE Decode (s)"
-    printf "| %-23s | %-9s | %-15s | %-12s | %-15s | %-14s |\n" \
-        "-----------------------" "---------" "---------------" "------------" "---------------" "--------------"
-    for row in "$@"; do
-        case "$row" in
-            *_sdpa_eager)     label="SDPA + Eager" ;;
-            *_sdpa_compiled)  label="SDPA + Compiled" ;;
-            *_sage3_eager)    label="SAGE_ATTN_3 + Eager" ;;
-            *_sage3_compiled) label="SAGE_ATTN_3 + Compiled" ;;
-        esac
-        printf "| %-23s | %-9s | %-15s | %-12s | %-15s | %-14s |\n" \
-            "$label" "${TOTAL[$row]}" "${PER_STEP[$row]}" "${TEXT_ENC[$row]}" "${TRANSFORMER[$row]}" "${VAE[$row]}"
-    done
-}
-
-print_table "HunyuanVideo-1.5 480p T2V, 832x480x33" \
-    hv15_sdpa_eager hv15_sdpa_compiled hv15_sage3_eager hv15_sage3_compiled
-
-print_table "Wan 2.2 TI2V-5B, 1280x704x33" \
-    wan22_sdpa_eager wan22_sdpa_compiled wan22_sage3_eager wan22_sage3_compiled
-
-echo
-echo "Logs: $OUT_DIR/*.log"

benchmarks/diffusion/diffusion_benchmark_serving.py

@@ -47,13 +47,6 @@ Usage:
         --endpoint /v1/images/edits --dataset random --task ti2i --num-prompts 10 \
         --bot-task think
 
-    python benchmarks/diffusion/diffusion_benchmark_serving.py \
-      --endpoint /v1/images/edits \
-      --dataset custom \
-      --dataset-path custom_requests.jsonl \
-      --task ti2i \
-      --bot-task think
-
     i2i (chat-based models such as Qwen-Image-Edit):
     python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
         --endpoint /v1/chat/completions --dataset vbench --task i2i --num-prompts 10
@@ -573,135 +566,6 @@ class TraceDataset(BaseDataset):
         return [self[i] for i in range(len(self))]
 
 
-class CustomDataset(BaseDataset):
-    """
-    Custom dataset that loads requests from a JSONL file.
-
-    Each line in the JSONL file should be a JSON object with the following fields:
-    - prompt (required): The text prompt for the request
-    - width (optional): Image/video width
-    - height (optional): Image height
-    - num_inference_steps (optional): Number of diffusion steps
-    - seed (optional): Random seed
-    - image_paths (optional): List of input image paths for i2i/i2v/ti2i tasks
-    - image_urls (optional): List of input image URLs (alternative to image_paths)
-
-    Example JSONL for ti2i:
-    {"prompt": "Add sunset lighting", "width": 1024, "height": 1024, "image_urls": ["https://example.com/image.jpg"]}
-    """
-
-    def __init__(self, args, api_url: str, model: str):
-        super().__init__(args, api_url, model)
-        self.dataset_path = args.dataset_path
-        if not self.dataset_path:
-            raise ValueError("--dataset-path must be provided when using 'custom' dataset")
-        self.load_data()
-
-    def load_data(self) -> None:
-        """Load data from JSONL file."""
-        import pandas as pd
-
-        if not self.dataset_path.endswith(".jsonl"):
-            raise ValueError("Custom dataset must be a JSONL file")
-
-        try:
-            df = pd.read_json(path_or_buf=self.dataset_path, lines=True)
-        except Exception as e:
-            raise ValueError(f"Failed to load JSONL file {self.dataset_path}: {e}")
-
-        if "prompt" not in df.columns:
-            raise ValueError("JSONL file must contain a 'prompt' column")
-
-        self.data = df.to_dict("records")
-        print(f"Loaded {len(self.data)} requests from {self.dataset_path}")
-
-    def _resolve_image_paths(self, item: dict) -> list[str] | None:
-        # Handle local file paths
-        if "image_paths" in item and item["image_paths"]:
-            paths = item["image_paths"]
-            if isinstance(paths, str):
-                paths = [paths]
-
-            # Verify all paths exist
-            valid_paths = []
-            for path in paths:
-                if os.path.exists(path):
-                    valid_paths.append(path)
-                else:
-                    raise ValueError(f"Image file not found: {path}")
-
-            return valid_paths if valid_paths else None
-
-        # Handle URLs - download to temp files
-        if "image_urls" in item and item["image_urls"]:
-            urls = item["image_urls"]
-            if isinstance(urls, str):
-                urls = [urls]
-
-            downloaded_paths = []
-            for url in urls:
-                try:
-                    response = requests.get(url, timeout=30)
-                    response.raise_for_status()
-
-                    # Create temp file with appropriate extension
-                    suffix = os.path.splitext(url)[1] or ".png"
-                    with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
-                        tmp.write(response.content)
-                        downloaded_paths.append(tmp.name)
-                except Exception as e:
-                    raise ValueError(f"Failed to download image from {url}: {e}")
-
-            return downloaded_paths if downloaded_paths else None
-
-        return None
-
-    def __len__(self) -> int:
-        return len(self.data)
-
-    def __getitem__(self, idx: int) -> RequestFuncInput:
-        item = self.data[idx]
-
-        # Extract fields with fallback to args defaults
-        prompt = item["prompt"]
-        width = item.get("width", self.args.width)
-        height = item.get("height", self.args.height)
-        num_inference_steps = item.get("num_inference_steps", self.args.num_inference_steps)
-        seed = item.get("seed", self.args.seed)
-        reserved_keys = {"prompt", "width", "height", "num_inference_steps", "seed", "image_paths", "image_urls"}
-        extra_body = {k: v for k, v in item.items() if k not in reserved_keys}
-
-        # Handle image paths/URLs
-        image_paths = self._resolve_image_paths(item)
-
-        return RequestFuncInput(
-            prompt=prompt,
-            api_url=self.api_url,
-            model=self.model,
-            seed=seed,
-            image_paths=image_paths,
-            extra_body=extra_body,
-            width=width,
-            height=height,
-            num_inference_steps=num_inference_steps,
-        )
-
-    def get_requests(self) -> list[RequestFuncInput]:
-        """Get all requests, cycling through data if num_prompts > dataset size."""
-        num_requests = getattr(self.args, "num_prompts", len(self.data))
-
-        if num_requests <= len(self.data):
-            return [self[i] for i in range(num_requests)]
-
-        # Cycle through the dataset to reach num_prompts
-        requests = []
-        for i in range(num_requests):
-            idx = i % len(self.data)
-            requests.append(self[idx])
-
-        return requests
-
-
 class RandomDataset(BaseDataset):
     def __init__(self, args, api_url: str, model: str, enable_negative_prompt: bool = False):
         super().__init__(args, api_url, model)
@@ -1086,8 +950,6 @@ async def benchmark(args):
         dataset = TraceDataset(args, api_url, args.model)
     elif args.dataset == "random":
         dataset = RandomDataset(args, api_url, args.model, args.enable_negative_prompt)
-    elif args.dataset == "custom":
-        dataset = CustomDataset(args, api_url, args.model)
     else:
         raise ValueError(f"Unknown dataset: {args.dataset}")
 
@@ -1235,7 +1097,7 @@ if __name__ == "__main__":
         "--dataset",
         type=str,
         default="vbench",
-        choices=["vbench", "trace", "random", "custom"],
+        choices=["vbench", "trace", "random"],
         help="Dataset to use.",
     )
     parser.add_argument(

benchmarks/tts/bench_tts.py

@@ -53,8 +53,6 @@ _TASK_TO_DATASET: dict[str, str] = {
     "voice_clone": "seed-tts",
     "default_voice": "seed-tts-text",
     "voice_design": "seed-tts-design",
-    "dialogue": "ttsd",
-    "sound_effect": "sound-effect",
 }
 
 # Default design dataset path (bundled with the repo)
@@ -83,7 +81,6 @@ def build_bench_args(
     output_dir: str | None,
     result_filename: str | None,
     extra_cli_args: list[str],
-    output_len: int | None = None,
 ) -> list[str]:
     """Build the ``vllm bench serve --omni`` command for one (task, concurrency) run."""
     dataset_name = _TASK_TO_DATASET[task]
@@ -133,18 +130,12 @@ def build_bench_args(
     if task_extra_body:
         cmd += ["--extra-body", json.dumps(task_extra_body, separators=(",", ":"))]
 
-    if model_cfg.get("trust_remote_code"):
-        cmd.append("--trust-remote-code")
-
     if concurrency is not None:
         cmd += ["--max-concurrency", str(concurrency), "--request-rate", "inf"]
 
     if wer_eval:
         cmd.append("--seed-tts-wer-eval")
 
-    if output_len is not None:
-        cmd += ["--hf-output-len", str(output_len)]
-
     if output_dir or result_filename:
         out_dir = output_dir or "."
         os.makedirs(out_dir, exist_ok=True)
@@ -244,15 +235,6 @@ def main() -> None:
     parser.add_argument("--host", default="localhost")
     parser.add_argument("--port", type=int, default=8000)
     parser.add_argument("--model-configs", default=str(_DEFAULT_MODEL_CONFIGS), help="Path to model_configs.yaml")
-    parser.add_argument(
-        "--output-len",
-        type=int,
-        default=None,
-        help="Override per-request max_new_tokens forwarded to the server "
-        "(seed-tts datasets default to 2048; the MOSS-TTS full talker takes "
-        "this as 'max_new_frames' so smaller values cap audio length per "
-        "request and keep bench runtime tractable).",
-    )
     parser.add_argument("extra", nargs=argparse.REMAINDER, help="Extra args passed directly to vllm bench serve")
     args = parser.parse_args()
 
@@ -307,7 +289,6 @@ def main() -> None:
                 output_dir=args.output_dir,
                 result_filename=result_filename,
                 extra_cli_args=args.extra or [],
-                output_len=args.output_len,
             )
             result = run_one_benchmark(cmd)
             if result is not None:

benchmarks/tts/model_configs.yaml

@@ -14,7 +14,6 @@ models:
     supported_tasks: [default_voice, voice_design]
     backend: openai-audio-speech
     endpoint: /v1/audio/speech
-    trust_remote_code: true
     task_extra_body:
       default_voice:
         voice: Vivian
@@ -28,76 +27,13 @@ models:
     supported_tasks: [voice_clone]
     backend: openai-audio-speech
     endpoint: /v1/audio/speech
-    trust_remote_code: true
     task_extra_body:
       voice_clone:
         task_type: Base
 
   openbmb/VoxCPM2:
-    supported_tasks: [voice_clone, default_voice]
-    trust_remote_code: true
-    backend: openai-audio-speech
-    endpoint: /v1/audio/speech
-    trust_remote_code: true
-    task_extra_body:
-      voice_clone: {}
-      default_voice:
-        voice: default
-
-  # MOSS-TTS family. After the serving_speech.py variant-detection refactor
-  # all five variants now go through OpenAI /v1/audio/speech with the
-  # variant-specific request fields (ambient_sound, ref_audio_2, instructions).
-  OpenMOSS-Team/MOSS-TTS:
     supported_tasks: [voice_clone]
     backend: openai-audio-speech
     endpoint: /v1/audio/speech
     task_extra_body:
       voice_clone: {}
-
-  # MOSS-TTS-v1.5: continued-training upgrade of MOSS-TTS 1.0 — same
-  # MossTTSDelay-8B architecture and API, so identical wiring. Adds 31-language
-  # synthesis via the `language` field and inline `[pause Xs]` markers.
-  OpenMOSS-Team/MOSS-TTS-v1.5:
-    supported_tasks: [voice_clone]
-    backend: openai-audio-speech
-    endpoint: /v1/audio/speech
-    task_extra_body:
-      voice_clone: {}
-
-  OpenMOSS-Team/MOSS-TTS-Realtime:
-    supported_tasks: [voice_clone]
-    backend: openai-audio-speech
-    endpoint: /v1/audio/speech
-    task_extra_body:
-      voice_clone: {}
-
-  # MOSS-TTSD synthesizes [S1]/[S2] dialogue. The seed-tts voice_clone dataset
-  # has single-speaker rows; bench piggybacks on it (both speakers share the
-  # same ref_audio). For real dialogue evaluation use the ttsd dataset.
-  OpenMOSS-Team/MOSS-TTSD-v1.0:
-    supported_tasks: [voice_clone, dialogue]
-    backend: openai-audio-speech
-    endpoint: /v1/audio/speech
-    task_extra_body:
-      voice_clone: {}
-      dialogue: {}
-
-  # MOSS-SoundEffect ignores text/ref_audio; expects ambient_sound +
-  # duration_seconds. The sound_effect dataset module emits those.
-  OpenMOSS-Team/MOSS-SoundEffect:
-    supported_tasks: [sound_effect]
-    backend: openai-audio-speech
-    endpoint: /v1/audio/speech
-    task_extra_body:
-      sound_effect:
-        duration_seconds: 10
-
-  # MOSS-VoiceGenerator: instructions-driven (no ref_audio). seed-tts-design
-  # dataset already emits {instructions, target_text}, so voice_design works
-  # out of the box once the serving validator accepts instructions.
-  OpenMOSS-Team/MOSS-VoiceGenerator:
-    supported_tasks: [voice_design]
-    backend: openai-audio-speech
-    endpoint: /v1/audio/speech
-    task_extra_body:
-      voice_design: {}

docker/Dockerfile.ci

@@ -1,8 +1,5 @@
-# Target vllm version: v0.21.0 (releases/v0.21.0 branch).
-# The pinned VLLM_PRECOMPILED_WHEEL_COMMIT and VLLM_BASE_TAG below must
-# be updated on every rebase to match the new upstream vllm commit/tag.
 ARG VLLM_BASE_IMAGE=vllm/vllm-openai
-ARG VLLM_BASE_TAG=v0.22.0
+ARG VLLM_BASE_TAG=v0.21.0
 FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG}
 ARG APP_DIR=/workspace/vllm-omni
 WORKDIR ${APP_DIR}

docker/Dockerfile.cuda

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

docker/Dockerfile.rocm

@@ -1,4 +1,4 @@
-ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.22.0
+ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.21.0
 FROM ${BASE_IMAGE} AS base
 
 # Declare a variable to know if we want to use the nightly build or the stable build.

docker/Dockerfile.xpu

@@ -88,7 +88,7 @@ ENV UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
 ENV UV_INDEX_STRATEGY="unsafe-best-match"
 ENV UV_LINK_MODE="copy"
 
-ARG VLLM_VERSION=v0.22.0
+ARG VLLM_VERSION=v0.21.0
 RUN git clone -b ${VLLM_VERSION} https://github.com/vllm-project/vllm
 WORKDIR /workspace/vllm
 

docs/design/index.md

@@ -13,10 +13,6 @@ This section contains design documents and architecture specifications for vLLM-
 - [Adding Step Execution Support for Diffusion Pipelines](feature/diffusion_step_execution.md)
 - [Continuous Batching for Step-Wise Diffusion](feature/diffusion_continuous_batching.md)
 
-## Infrastructure Design Documents
-
-- [Prometheus Metrics](metrics.md)
-
 ## Module Design Documents
 
 - [AR Module](module/ar_module.md)

docs/design/metrics.md

@@ -1,208 +0,0 @@
-# Prometheus Metrics Design
-
-This document describes how vLLM-Omni exposes Prometheus metrics for multi-stage pipelines, the constraints that shaped the design, and how the pipeline-level metrics coexist with upstream vLLM per-engine metrics. This iteration covers the text and audio surface; diffusion / image / video families land in a follow-up.
-
-## Objectives
-
-- Expose pipeline-level request and latency metrics that span the full multi-stage execution (orchestrator scope).
-- Preserve all upstream vLLM per-engine metrics (`vllm:*`) for stages backed by an AR LLM engine, and reshape their `engine` label into `stage` + `replica` so multi-replica deployments gain per-replica visibility automatically.
-- Expose per-modality SLO metrics that the upstream `vllm:*` families do not capture — audio TTFP / RTF / duration / frames / streaming continuity / silent-loss.
-- Expose per-replica-edge cross-stage transfer metrics so the slack between E2E latency and the sum of per-stage `gen_time` (queueing, serialization, network) becomes attributable.
-- Keep the metrics collection overhead low enough that it does not regress TTFA or throughput.
-
-## Background
-
-### Upstream vLLM Metrics
-
-Upstream vLLM defines ~37 Prometheus metric families under the `vllm:` prefix. These are registered by `PrometheusStatLogger` and cover engine-level state: KV cache usage, running/waiting request counts, token throughput, TTFT, inter-token latency, e2e latency, and so on. They are served via the `/metrics` HTTP endpoint provided by `prometheus_fastapi_instrumentator` and the default `prometheus_client` WSGI handler.
-
-vLLM's `unregister_vllm_metrics()` function strips every `prometheus_client` collector whose `_name` attribute contains the substring `"vllm"`. This runs during engine initialization to clean up stale collectors from prior instantiations within the same process.
-
-### The Problem
-
-vLLM-Omni runs multiple engine instances (stages × replicas) within a single process, coordinated by an Orchestrator. The pipeline needs its own metrics — aggregate request counts, end-to-end latency across all stages, per-modality SLO signals, and cross-stage transfer attribution — that do not exist in upstream vLLM. All pipeline-level metrics use the `vllm:omni_` prefix to distinguish them from upstream per-engine metrics. At import time (see `vllm_omni/patch.py`) we replace `unregister_vllm_metrics()` with a scoped version that still strips upstream `vllm:*` collectors before each new `PrometheusStatLogger` registers (so multi-engine processes don't crash on duplicate registration), but preserves anything prefixed `vllm:omni_` / `vllm_omni`.
-
-Upstream per-engine metrics retain the `vllm:` prefix but are now registered by `OmniPrometheusStatLogger`, a thin subclass of upstream's `PrometheusStatLogger` that reshapes the single `engine` label into a `stage` + `replica` pair (see "OmniPrometheusStatLogger wrap" below).
-
-## Architecture
-
-### Component Overview
-
-```text
-                       +------------------------+
-                       |  API Server (FastAPI)  |
-                       |   GET /metrics         |
-                       +-----------+------------+
-                                   |
-                  prometheus_client default registry
-                                   |
-        +--------+--------+--------+--------+--------+
-        |                                            |
-   vllm:omni_*                                    vllm:*
-   collectors                                  collectors
-        |                                            |
-   +----+--------+   +-----------+   +----------+   +-----------+
-   | OmniPromet- |   | OmniMod-  |   | OmniTra- |   | OmniProm- |
-   | heusMetrics |   | alityMet- |   | nsferMe- |   | etheusSt- |
-   |             |   | rics      |   | trics    |   | atLogger  |
-   +----+--------+   +-----+-----+   +----+-----+   +----+------+
-        |                  |              |              |
-     OmniBase           OmniBase     Orchestrator    Orchestrator
-   (request life-     (finalize +   (record_trans-  (per-(stage,
-    cycle, success/    streaming     fer_tx/rx        replica)
-    fail counter)      hooks via     hooks via        scheduler/
-                       observe_*     emit hook in     iteration
-                       APIs)         OrchestratorAg-  stats)
-                                     gregator)
-```
-
-### Data Flow
-
-There are four independent paths for metric collection.
-
-**Path 1: Pipeline-level metrics (`vllm:omni_*`)**
-
-`OmniPrometheusMetrics` registers the Gauge / Counter / Histogram collectors at import time. It is instantiated once per entrypoint, labeled with the model name. The entrypoint calls its methods as requests progress:
-
-- `set_running(n)` / `set_waiting(n)` — updated after each request completes. The running count comes from `OmniRequestCounter`, a simple counter incremented/decremented by the Orchestrator as it tracks requests. Waiting is derived as `total - running`.
-- `request_succeeded(e2e_seconds, finished_reason="stop")` — recorded when a request finishes at the final stage. `finished_reason` is extracted from `engine_outputs.outputs[0].finish_reason` (vLLM `CompletionOutput` convention) and increments `vllm:omni_requests_success_total{finished_reason}`.
-- `request_failed()` — recorded by the cleanup path when a request exits without natural completion. Internally maps to `finished_reason="abort"` so a single Counter family covers both natural and aborted completion.
-
-**Path 2: Audio modality metrics (`vllm:omni_audio_*`)**
-
-`OmniModalityMetrics` registers seven audio families with `{model_name, stage, replica}` (plus an extra `threshold_ms` / `reason` label on the two extra-cardinality Counters). Three observation sites:
-
-- `observe_modality_at_finalize(...)` — called from `omni_base._process_single_result` inside the existing `e2e_done` finalize guard. For `output_type == "audio"` it emits `audio_frames_total`, `audio_duration_s`, `audio_rtf` (or `audio_skipped_requests_total{reason="no_audio_data"}` when no audio was produced). Sample rate is resolved from `engine_outputs.multimodal_output` via `definitions.resolve_audio_sample_rate(...)` (fallback chain mirrors `serving_chat.py`'s audio response path).
-- `observe_audio_first_packet(...)` — called from the OpenAI SSE audio branch in `serving_chat.py` on the first audio packet for a request. The once-per-request guard is held by `ClientRequestState.first_audio_ts`. The `request_arrival_ts` anchor is stored in `ClientRequestState` by `async_omni.generate()`, computed at request entry.
-- `observe_audio_streaming_finalize(...)` — called from `serving_chat.py` after the streaming chunk loop exhausts. It runs the per-chunk player simulation from `vllm_omni/benchmarks/audio_continuity.py` to compute the worst-case underrun and emits `audio_underrun_s` plus (when the request stayed below the threshold) `audio_continuity_ok_total{threshold_ms}`. Per-chunk PCM byte counts and arrival timestamps are recorded by the same audio branch that updates `first_audio_ts`.
-
-**Path 3: Cross-stage transfer metrics (`vllm:omni_transfer_*`)**
-
-`OmniTransferMetrics` registers four Histogram families with `{model_name, from_stage, from_replica, to_stage, to_replica}` labels. Each observation corresponds to one physical transfer hop (one chunk between adjacent stages), not the per-request accumulated total — so the histograms track per-transfer distribution.
-
-The hook lives in `OrchestratorAggregator.record_transfer_tx` and `record_transfer_rx`. After the existing `TransferEdgeStats` accumulation, the aggregator calls `_emit_transfer_tx` / `_emit_transfer_rx`. Those:
-
-1. Resolve `from_replica` / `to_replica` via a `replica_resolver` callback supplied by `async_omni.py`. The resolver delegates to `stage_pool.get_bound_replica_id(request_id)` — i.e. the orchestrator's existing sticky-routing binding is the source of truth.
-2. Convert the underlying `_ms` accumulators to seconds and call the `_s`-suffixed observe methods on `OmniTransferMetrics`.
-
-Defensive fail-safe: if `transfer_emitter` or `replica_resolver` is missing, or the resolver returns `None` for either side, the emit is skipped silently (the underlying `TransferEdgeStats` accumulation is unaffected).
-
-**Path 4: Per-engine metrics (`vllm:*`, stage/replica wrap)**
-
-The Orchestrator instantiates `OmniPrometheusStatLogger` (a thin subclass of upstream `vllm.v1.metrics.loggers.PrometheusStatLogger`) and feeds it scheduler stats and iteration stats after processing each batch of engine outputs. This populates the standard ~37 vLLM metric families (TTFT, ITL, TPOT, KV cache usage, etc.) using the same upstream code path — but with the `engine` label reshaped into `stage` + `replica` so multi-replica deployments produce distinct series per replica. See the next section for the wrap mechanics.
-
-### Shared State Between Threads
-
-The Orchestrator runs in a background thread. The API server (OmniBase) runs in the asyncio event loop thread. `OmniRequestCounter` bridges them — a plain Python object with an `int` field. The Orchestrator increments/decrements it; the entrypoint reads it for gauge updates. No lock is needed because the counter is advisory (a stale read by one Prometheus scrape interval is acceptable). It is created by `AsyncOmniEngine.__init__()` and passed to the Orchestrator at construction time.
-
-### Metric Registration and Lifecycle
-
-All `vllm:omni_*` collectors are registered once when their owning class (`OmniPrometheusMetrics` / `OmniModalityMetrics` / `OmniTransferMetrics`) is imported. Per-`(stage, replica)` labels are bound lazily on first observation to avoid registering label sets for combinations that never produce data.
-
-The `prometheus_client` default registry holds all collectors. FastAPI's `/metrics` endpoint serves the default registry, so `vllm:omni_*` and the wrapped `vllm:*` metrics appear in the same scrape response alongside `http_*` and `process_*` metrics from the instrumentator and the Python client runtime.
-
-## OmniPrometheusStatLogger Wrap
-
-Upstream `PrometheusStatLogger.__init__` hard-codes `labelnames = ["model_name", "engine"]` as a local variable, references it across ~37 metric-family construction sites, and uses the `engine` label value in five different `.labels()` call shapes (kwarg with int engine, kwarg with str engine, positional with str engine in the middle, plus a `metrics_info["engine"] = str(...)` dict pattern). To reshape `engine` into `stage` + `replica` without forking the entire upstream `__init__`, the wrap uses three coordinated mechanisms:
-
-1. **Class-level metric class slot overrides.** `OmniPrometheusStatLogger` overrides `_gauge_cls`, `_counter_cls`, `_histogram_cls` (which upstream calls via `self._gauge_cls(...)` etc.) with `_RelabelGauge` / `_RelabelCounter` / `_RelabelHistogram` wrapper classes. These intercept the `labelnames` kwarg at metric family creation time and replace `engine` with `("stage", "replica")`.
-2. **Property descriptor for `per_engine_labelvalues`.** Upstream builds `self.per_engine_labelvalues = {idx: [model_name, str(idx)]}` inside `__init__` and then captures it into a local variable for `create_metric_per_engine` calls. By making `per_engine_labelvalues` a Python property on the subclass, the setter intercepts upstream's assignment and rewrites each 2-tuple into a 3-tuple `[model_name, stage, replica]` using the `stage_replica_map` supplied at construction time. The captured local then sees the rewritten dict.
-3. **Override of `.labels()` on the wrapper classes.** For the five call sites that pass `engine` directly (kwarg or positional, int or str), `_RelabelMixin.labels()` translates the engine value back to `(stage, replica)` via a process-level `_ENGINE_INDEX_MAP` populated by `OmniPrometheusStatLogger.__init__`. This handles `gauge_engine_sleep_state.labels(engine=idx, ...)`, `counter_request_success_base.labels(model_name, str(idx), str(reason))`, `info_gauge.labels(**metrics_info)`, etc.
-
-The three sub-helpers that upstream `PrometheusStatLogger.__init__` constructs (`spec_decoding_prom` / `kv_connector_prom` / `perf_metrics_prom`) use their own `_counter_cls` / `_gauge_cls` / `_histogram_cls` slots and would otherwise build families with the raw 2-element labelnames. `_OmniPerfMetricsProm` / `_OmniSpecDecodingProm` / `_OmniKVConnectorProm` subclass each helper to route the same relabel mixin through their internal family construction.
-
-The `Orchestrator` constructs `stage_replica_map` from the static `stage_pools` configuration at startup:
-
-```python
-stage_replica_map = {
-    flat_idx: (str(stage_id), str(replica_id))
-    for flat_idx, (stage_id, replica_id) in enumerate(
-        (s, r)
-        for s, pool in enumerate(stage_pools)
-        for r in range(pool.num_replicas)
-    )
-}
-```
-
-A reverse map `(stage_id, replica_id) -> flat_idx` is maintained on the Orchestrator so the per-replica `record(engine_idx=...)` call site can look up the right flat index.
-
-> Dynamic add/remove of replicas at runtime is intentionally out of scope — the upstream `PrometheusStatLogger` materializes per-engine_idx child metrics at init time, and supporting hot-add would require non-trivial intervention into upstream's per-family child dictionaries.
-
-## Gating: `--log-stats`
-
-All metrics — both the 15 `vllm:omni_*` families and the ~65 upstream `vllm:*` wrap families — are gated by the user's `--log-stats` CLI flag (default off). The flag is plumbed from `OmniBase.__init__(log_stats=...)` through `AsyncOmniEngine` to the stage-spawn helpers and to the three Prometheus metric classes (`OmniPrometheusMetrics` / `OmniModalityMetrics` / `OmniTransferMetrics`), and is forwarded to `Orchestrator._init_metrics_state(...)`.
-
-Behavior with `--log-stats=off` (default):
-
-- The three `Omni*Metrics` classes register their module-level Gauge / Counter / Histogram families at import time (prometheus_client requires up-front registration), but each `observe / inc / set` method early-returns. The per-label child series for `vllm:omni_*` stay materialized but never have data written to them.
-- `OmniPrometheusStatLogger` is not constructed in `_init_metrics_state`, so the ~65 upstream `vllm:*` wrap families are not registered in the default registry at all.
-- The engine core's `Scheduler.make_stats()` also short-circuits inside upstream (`if not self.log_stats: return None`), so no `SchedulerStats` is produced per step — the per-iteration cost is bounded by the existing upstream gate.
-
-Behavior with `--log-stats=on`: all metric paths fire normally; the orchestrator's per-replica recording is bounded only by `OmniSchedulerMixin.make_stats()`'s per-scheduler 1 Hz throttle (see next section).
-
-The overhead with the flag on is small enough that an A/B benchmark on Qwen3-Omni-30B single replica (30 sequential audio requests) showed a mean latency delta of +0.6% (Welch's t = 0.318, n=30, not statistically significant at α=0.05); the `/metrics` line count drops from 1358 to 124 lines when the flag is off.
-
-## Throttling: `make_stats()` Override
-
-Upstream vLLM's `Scheduler.make_stats()` runs on every AR generation step, returning a SchedulerStats object for the orchestrator. Under vLLM's architecture, this is fine. But since vLLM-Omni requires that the object be serialized and transferred over ZMQ, receiving a SchedulerStats object on every step can introduce unacceptable overhead to the system.
-
-`OmniSchedulerMixin.make_stats()` (in `vllm_omni/core/sched/omni_scheduler_mixin.py`) throttles stats emission to at most once per second **per scheduler** — i.e. per `(stage, replica)` since each replica owns its own scheduler instance. Between intervals it returns `None`, which the engine core skips serializing. This keeps gauges fresh enough for Prometheus scrapes (typically 15-30s intervals) while eliminating the per-step overhead.
-
-The orchestrator side does not add its own throttle on top: the per-replica recording loop gates only on `raw_outputs.scheduler_stats is not None` (i.e. this replica's scheduler passed its own 1 Hz gate). A previous global `_last_stats_ts` on the orchestrator was removed because it starved every replica other than the first to emit within each second.
-
-## Metric Definitions
-
-### Pipeline (4)
-
-| Metric | Type | Labels | Description |
-|--------|------|--------|-------------|
-| `vllm:omni_num_requests_running` | Gauge | `model_name` | Requests currently executing across all stages |
-| `vllm:omni_num_requests_waiting` | Gauge | `model_name` | Requests queued but not yet scheduled |
-| `vllm:omni_requests_success_total` | Counter | `model_name`, `finished_reason` | Total requests by completion reason ({stop, length, abort, ...}); aborts cover client-disconnect / cancellation paths in addition to upstream `FinishReason.ABORT` |
-| `vllm:omni_e2e_request_latency_s` | Histogram | `model_name` | Pipeline-global end-to-end latency in seconds |
-
-### Audio (7)
-
-Labels: `{model_name, stage, replica}` plus the listed extra label.
-
-| Metric | Type | Extra label | Description |
-|--------|------|-------------|-------------|
-| `vllm:omni_audio_ttfp_s` | Histogram | — | Time from request arrival to first audio packet/frame |
-| `vllm:omni_audio_duration_s` | Histogram | — | Audio content duration (`audio_frames / sample_rate`) |
-| `vllm:omni_audio_rtf` | Histogram | — | Real-time factor `stage_gen_time_s / audio_duration_s` (SLO `< 1`); uses `RTF_BUCKETS` |
-| `vllm:omni_audio_frames_total` | Counter | — | Cumulative audio frames generated |
-| `vllm:omni_audio_underrun_s` | Histogram | — | Per-request worst-case player deficit; `> 0` indicates listener heard silent gaps |
-| `vllm:omni_audio_continuity_ok_total` | Counter | `threshold_ms` | Incremented when the request's worst underrun stayed below `threshold_ms` |
-| `vllm:omni_audio_skipped_requests_total` | Counter | `reason` | Silent-loss counter — code2wav rejected malformed codec input and returned `200 OK` with empty audio |
-
-### Cross-stage transfer (4)
-
-Labels: `{model_name, from_stage, from_replica, to_stage, to_replica}`.
-
-| Metric | Type | Description |
-|--------|------|-------------|
-| `vllm:omni_transfer_size_bytes` | Histogram | Per-transfer payload size in bytes |
-| `vllm:omni_transfer_tx_s` | Histogram | Sender-side time (serialize + submit to connector) |
-| `vllm:omni_transfer_rx_s` | Histogram | Receiver-side time (recv + deserialize) |
-| `vllm:omni_transfer_in_flight_s` | Histogram | Network in-flight time (TX done → RX recv start) |
-
-### LLM stage-level (wrapped `vllm:*`)
-
-After the wrap, every upstream `vllm:*` family — TTFT, ITL, TPOT, e2e latency, KV cache usage, scheduler running/waiting, request success counts, etc. — carries `{model_name, stage, replica}` labels. For the full upstream catalog see [the vLLM docs](https://github.com/vllm-project/vllm/blob/main/docs/usage/metrics.md); note that metrics depending on features unsupported in vLLM-Omni (e.g. speculative decoding, LoRA) will not be available.
-
-## Naming Convention
-
-- All time-bearing metrics use the `_s` suffix (values in seconds). Two bucket families are used:
-  - `SECONDS_BUCKETS` (0.05 s – 300 s) for e2e / generation / TTFP style values.
-  - `SECONDS_FAST_BUCKETS` (0.001 s – 60 s) for fine-grained cross-stage transfer and audio-underrun values that need millisecond-level resolution.
-- Counters use the `_total` suffix (auto-appended by `prometheus_client`).
-- Sizes use the `_bytes` suffix.
-- All omni-specific families are prefixed `vllm:omni_`. The upstream `unregister_vllm_metrics()` function is monkey-patched to a scoped version that still strips upstream `vllm:*` collectors (so multi-engine init within one process does not crash on duplicate registration) but preserves anything prefixed `vllm:omni_` / `vllm_omni`.
-
-## Logging vs. Prometheus
-
-`OrchestratorAggregator` (in `vllm_omni/metrics/stats.py`) is the logging-oriented metrics path. It collects detailed per-request, per-stage, and per-transfer statistics and prints formatted tables to the `INFO` log. This is designed for development and debugging — individual request traces, transfer bandwidth, inter-stage timing.
-
-`OmniPrometheusMetrics` / `OmniModalityMetrics` / `OmniTransferMetrics` form the Prometheus-oriented path. They record aggregate counters, gauges, and histograms suitable for time-series monitoring and alerting. Both paths share the same source data (`StageRequestStats`, `TransferEdgeStats`) — `OrchestratorAggregator.record_transfer_tx/rx` in particular calls both the existing accumulator code and the Prometheus emit hook in the same method body. The two consumption models can run simultaneously without coupling.
-
-The separation follows upstream vLLM's pattern of `LoggingStatLogger` vs. `PrometheusStatLogger` — same underlying data, different consumption models.

docs/design/module/dit_module.md

@@ -169,13 +169,15 @@ class DiffusionRequestStatus(enum.IntEnum):
 
 @dataclass
 class DiffusionRequestState:
-    request_id: str
+    sched_req_id: str
     req: OmniDiffusionRequest
     status: DiffusionRequestStatus = DiffusionRequestStatus.WAITING
 ```
 
 **Design Features**:
 
+- **Scheduler-owned ID**: Each `OmniDiffusionRequest` is tracked by an internal `sched_req_id`, separated from public `request_id` values.
+
 - **Explicit lifecycle**: Requests move through waiting, running, optional preemption, and terminal states.
 
 - **Centralized error handling**: Completion, abort, and error states are all normalized in the scheduler layer.
@@ -186,6 +188,7 @@ class DiffusionRequestState:
 class _BaseScheduler(SchedulerInterface):
     def __init__(self) -> None:
         self._request_states = {}
+        self._request_id_to_sched_req_id = {}
         self._waiting = deque()
         self._running = []
         self._finished_req_ids = set()
@@ -196,8 +199,7 @@ class _BaseScheduler(SchedulerInterface):
 
 - **Common state storage**: Shared request maps and waiting/running sets live in the base class.
 
-- **Shared cleanup logic**: Duplicate-request checks, finish handling, and state removal are centralized instead of
-  duplicated in each policy.
+- **Shared cleanup logic**: Request-id registration, finish handling, and state removal are centralized instead of duplicated in each policy.
 
 - **Current constraint boundary**: `_BaseScheduler` derives `max_num_running_reqs` from `max_num_seqs`, but request-mode diffusion is still clamped back to `1` by the engine. The step-wise path can keep this above `1` for compatible-request batching.
 
@@ -217,12 +219,12 @@ class RequestScheduler(_BaseScheduler):
 
 - **Single-request admission**: `RequestScheduler` still admits one active request at a time because request-mode execution completes a whole request per dispatch.
 
-- **Executor result feedback**: `update_from_output()` converts executor output into `FINISHED_COMPLETED` or `FINISHED_ERROR` and returns finished request ids.
+- **Executor result feedback**: `update_from_output()` converts executor output into `FINISHED_COMPLETED` or `FINISHED_ERROR` and returns finished scheduler ids.
 
 #### 2.5 Engine-Driven Execution Loop
 
 ```python
-request_id = scheduler.add_request(request)
+sched_req_id = scheduler.add_request(request)
 while True:
     sched_output = scheduler.schedule()
     output = executor.add_req(req)

docs/design/qwen3_omni_tts_performance_optimization.md

@@ -411,8 +411,6 @@ Notes:
 - `runtime.max_batch_size` controls stage-level batching.
 - Thinker/Talker commonly use `enforce_eager: false` for CUDA Graph paths.
 - Code2Wav often remains eager (`enforce_eager: true`) depending on runtime behavior.
-- Qwen3-Omni defaults `VLLM_USE_FLASHINFER_MOE_FP16=0`. The Triton has been more stable & faster
- than the FlashInfer CUTLASS unquantized MoE backend on recent vLLM rebases.
 
 #### 2) Enable async chunk
 

docs/getting_started/installation/gpu.md

@@ -34,7 +34,7 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr
 
 ### Pre-built wheels
 
-Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, 0.21.0, and 0.22.0. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
+Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, and 0.21.0. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
 
 === "NVIDIA CUDA"
 

docs/getting_started/installation/gpu/cuda.inc.md

@@ -20,7 +20,7 @@ Therefore, it is recommended to install vLLM and vLLM-Omni with a **fresh new**
 
 vLLM-Omni is built based on vLLM. Please install it with command below.
 ```bash
-uv pip install vllm==0.22.0 --torch-backend=auto
+uv pip install vllm==0.21.0 --torch-backend=auto
 ```
 
 #### Installation of vLLM-Omni
@@ -39,13 +39,13 @@ uv pip install 'vllm-omni[demo]'
 # --8<-- [start:build-wheel-from-source]
 
 #### Installation of vLLM
-If you do not need to modify source code of vLLM, you can directly install the stable 0.22.0 release version of the library
+If you do not need to modify source code of vLLM, you can directly install the stable 0.21.0 release version of the library
 
 ```bash
-uv pip install vllm==0.22.0 --torch-backend=auto
+uv pip install vllm==0.21.0 --torch-backend=auto
 ```
 
-The 0.22.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
+The 0.21.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
 
 #### Installation of vLLM-Omni
 Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -66,12 +66,12 @@ If you want to check, modify or debug with source code of vLLM, install the libr
 ```bash
 git clone https://github.com/vllm-project/vllm.git
 cd vllm
-git checkout v0.22.0
+git checkout v0.21.0
 ```
 Set up environment variables to get pre-built wheels. If there are internet problems, just download the whl file manually. And set `VLLM_PRECOMPILED_WHEEL_LOCATION` as your local absolute path of whl file.
 ```bash
-#For CUDA 13.0 (the default for v0.22.0; the wheel filename has no `+cu130` suffix)
-export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.22.0/vllm-0.22.0-cp38-abi3-manylinux_2_35_x86_64.whl
+#For CUDA 13.0 (the default for v0.21.0; the wheel filename has no `+cu130` suffix)
+export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.21.0/vllm-0.21.0-cp38-abi3-manylinux_2_35_x86_64.whl
 ```
 Install vllm with command below (If you have no existing PyTorch).
 ```bash
@@ -102,7 +102,7 @@ docker run --runtime nvidia --gpus 2 \
     --env "HF_TOKEN=$HF_TOKEN" \
     -p 8091:8091 \
     --ipc=host \
-    vllm/vllm-omni:v0.22.0 \
+    vllm/vllm-omni:v0.21.0 \
     vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 ```
 

docs/getting_started/installation/gpu/rocm.inc.md

@@ -13,7 +13,7 @@ vLLM-Omni current recommends the steps in under setup through Docker Images.
 
 vLLM-Omni is built based on vLLM. Please install it with command below.
 ```bash
-uv pip install vllm==0.22.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.22.0/rocm722
+uv pip install vllm==0.21.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.21.0/rocm722
 ```
 
 #### Installation of vLLM-Omni
@@ -34,13 +34,13 @@ uv pip install onnxruntime-rocm
 # --8<-- [start:build-wheel-from-source]
 
 #### Installation of vLLM
-If you do not need to modify source code of vLLM, you can directly install the stable 0.22.0 release version of the library
+If you do not need to modify source code of vLLM, you can directly install the stable 0.21.0 release version of the library
 
 ```bash
-uv pip install vllm==0.22.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.22.0/rocm722
+uv pip install vllm==0.21.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.21.0/rocm722
 ```
 
-The pre-built 0.22.0 vLLM wheel targets ROCm 7.2.2. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
+The pre-built 0.21.0 vLLM wheel targets ROCm 7.2.2. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
 
 #### Installation of vLLM-Omni
 Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -58,7 +58,7 @@ If you want to check, modify or debug with source code of vLLM, install the libr
 ```bash
 git clone https://github.com/vllm-project/vllm.git
 cd vllm
-git checkout v0.22.0
+git checkout v0.21.0
 python3 -m pip install -r requirements/rocm.txt
 python3 setup.py develop
 ```
@@ -130,7 +130,7 @@ docker run --rm \
   -v ~/.cache/huggingface:/root/.cache/huggingface \
   --env "HF_TOKEN=$HF_TOKEN" \
   -p 8091:8091 \
-  vllm/vllm-omni-rocm:v0.22.0 \
+  vllm/vllm-omni-rocm:v0.21.0 \
   --model Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 ```
 
@@ -149,7 +149,7 @@ docker run --rm -it \
   -v ~/.cache/huggingface:/root/.cache/huggingface \
   --env "HF_TOKEN=$HF_TOKEN" \
   --entrypoint bash \
-  vllm/vllm-omni-rocm:v0.22.0
+  vllm/vllm-omni-rocm:v0.21.0
 ```
 
 # --8<-- [end:pre-built-images]

docs/getting_started/quickstart.md

@@ -19,10 +19,10 @@ uv venv --python 3.12 --seed
 source .venv/bin/activate
 
 # On CUDA
-uv pip install vllm==0.22.0 --torch-backend=auto
+uv pip install vllm==0.21.0 --torch-backend=auto
 
 # On ROCm
-uv pip install vllm==0.22.0+rocm721 --extra-index-url https://wheels.vllm.ai/rocm/0.22.0/rocm721
+uv pip install vllm==0.21.0+rocm721 --extra-index-url https://wheels.vllm.ai/rocm/0.21.0/rocm721
 
 git clone https://github.com/vllm-project/vllm-omni.git
 cd vllm-omni
@@ -35,7 +35,7 @@ For additional installation methods — please see the [installation guide](inst
 !!! note
     It is important to install the same major & minor version of vLLM and vLLM Omni, otherwise things may not work as expected. If the versions are misaligned, you will see a warning when you import vLLM Omni.
 
-    If you are seeing strange behavior with the `vllm` command not handling the `--omni` flag correctly, you most likely have a version mismatch with vLLM < `0.22.0` and vLLM Omni `0.22.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue.
+    If you are seeing strange behavior with the `vllm` command not handling the `--omni` flag correctly, you most likely have a version mismatch with vLLM < `0.21.0` and vLLM Omni `0.21.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue.
 
 ## Offline Inference
 

docs/mkdocs/hooks/generate_argparse.py

@@ -119,10 +119,9 @@ def extract_omni_serve_subparser_init():
                     exec_globals = {
                         "_FlexibleArgumentParser": _FlexibleArgumentParser,
                         "FlexibleArgumentParser": _FlexibleArgumentParser,
-                        "TrackingArgumentParser": _FlexibleArgumentParser,  # only needs add_argument, not tracking
-                        "TrackingNamespace": None,  # type annotation only
                         "make_arg_parser": lambda parser: parser,  # no-op for doc
                         "_ensure_vllm_platform": lambda: None,  # no-op for doc
+                        "nullify_stage_engine_defaults": lambda parser: None,  # no-op for doc
                         "VLLM_SUBCMD_PARSER_EPILOG": "",
                         "logger": logger,
                         "DummySubparsers": DummySubparsers,

docs/mkdocs/hooks/generate_examples.py

@@ -16,7 +16,6 @@ ROOT_DIR_RELATIVE = "../../../../.."
 EXAMPLE_DIR = ROOT_DIR / "examples"
 EXAMPLE_DOC_DIR = ROOT_DIR / "docs/user_guide/examples"
 NAV_FILE = ROOT_DIR / "docs/.nav.yml"
-MAX_INLINE_MATERIAL_SIZE = 8 * 1024
 
 
 def fix_case(text: str) -> str:
@@ -192,21 +191,18 @@ class Example:
 
         return re.sub(link_pattern, replace_link, content)
 
-    def github_url(self, path: Path) -> str:
-        url = "https://github.com/vllm-project/vllm-omni/"
-        url += "tree/main" if path.is_dir() else "blob/main"
-        return f"{url}/{path.relative_to(ROOT_DIR).as_posix()}"
-
     def generate(self) -> str:
         content = f"# {self.title}\n\n"
-        content += f"Source <{self.github_url(self.path)}>.\n\n"
+        url = "https://github.com/vllm-project/vllm-omni/"
+        url += "tree/main" if self.path.is_dir() else "blob/main"
+        content += f"Source <{url}/{self.path.relative_to(ROOT_DIR)}>.\n\n"
 
         # Use long code fence to avoid issues with
         # included files containing code fences too
         code_fence = "``````"
 
         if self.is_code:
-            main_file_rel = self.main_file.relative_to(ROOT_DIR).as_posix()
+            main_file_rel = self.main_file.relative_to(ROOT_DIR)
             content += f'{code_fence}{self.main_file.suffix[1:]}\n--8<-- "{main_file_rel}"\n{code_fence}\n'
         else:
             with open(self.main_file, encoding="utf-8") as f:
@@ -220,15 +216,10 @@ class Example:
 
         content += "## Example materials\n\n"
         for file in sorted(self.other_files):
-            content += f'??? abstract "{file.relative_to(self.path).as_posix()}"\n'
-            if file.stat().st_size > MAX_INLINE_MATERIAL_SIZE:
-                content += (
-                    f"    Large file omitted from the rendered docs. View it on GitHub: <{self.github_url(file)}>.\n\n"
-                )
-                continue
+            content += f'??? abstract "{file.relative_to(self.path)}"\n'
             if file.suffix != ".md":
                 content += f"    {code_fence}{file.suffix[1:]}\n"
-            content += f'    --8<-- "{file.relative_to(ROOT_DIR).as_posix()}"\n'
+            content += f'    --8<-- "{file.relative_to(ROOT_DIR)}"\n'
             if file.suffix != ".md":
                 content += f"    {code_fence}\n"
 
@@ -300,7 +291,7 @@ def update_nav_file(examples: list[Example]):
         for example in category_examples:
             doc_path = EXAMPLE_DOC_DIR / example.category / f"{example.path.stem}.md"
             rel_path = doc_path.relative_to(ROOT_DIR / "docs")
-            category_items.append({example.title: rel_path.as_posix()})
+            category_items.append({example.title: str(rel_path)})
 
         if category_items:
             # Format category name (e.g., "offline_inference" -> "Offline Inference")

docs/mkdocs/hooks/search_index.py

@@ -1,29 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Reduce browser search-index memory for generated API reference pages."""
-
-from material.plugins.search.plugin import SearchIndex
-
-_add_entry_from_context = SearchIndex.add_entry_from_context
-
-
-def add_entry_from_context(self: SearchIndex, page) -> None:
-    search = page.meta.get("search") or {}
-    if search.get("exclude"):
-        return
-
-    if str(page.url).startswith("api/"):
-        entry = {
-            "location": page.url,
-            "title": str(page.meta.get("title", page.title)),
-            "text": "",
-        }
-        if "boost" in search:
-            entry["boost"] = search["boost"]
-        self.entries.append(entry)
-        return
-
-    _add_entry_from_context(self, page)
-
-
-SearchIndex.add_entry_from_context = add_entry_from_context

docs/mkdocs/javascript/mathjax.js

@@ -1,44 +1,20 @@
-// Enables MathJax rendering only on pages that contain math.
+// Enables MathJax rendering
 window.MathJax = {
-  tex: {
-    inlineMath: [["\\(", "\\)"]],
-    displayMath: [["\\[", "\\]"]],
-    processEscapes: true,
-    processEnvironments: true,
-  },
-  options: {
-    ignoreHtmlClass: ".*|",
-    processHtmlClass: "arithmatex",
-  },
-};
+    tex: {
+      inlineMath: [["\\(", "\\)"]],
+      displayMath: [["\\[", "\\]"]],
+      processEscapes: true,
+      processEnvironments: true
+    },
+    options: {
+      ignoreHtmlClass: ".*|",
+      processHtmlClass: "arithmatex"
+    }
+  };
 
-let mathJaxLoading;
-
-function loadMathJax() {
-  if (window.MathJax.typesetPromise) {
-    return Promise.resolve();
-  }
-  if (!mathJaxLoading) {
-    mathJaxLoading = new Promise((resolve, reject) => {
-      const script = document.createElement("script");
-      script.src = "https://unpkg.com/mathjax@3.2.2/es5/tex-mml-chtml.js";
-      script.async = true;
-      script.onload = resolve;
-      script.onerror = reject;
-      document.head.appendChild(script);
-    });
-  }
-  return mathJaxLoading;
-}
-
-document$.subscribe(() => {
-  if (!document.querySelector(".arithmatex")) {
-    return;
-  }
-  loadMathJax().then(() => {
-    MathJax.startup.output.clearCache();
-    MathJax.typesetClear();
-    MathJax.texReset();
-    MathJax.typesetPromise();
-  });
-});
+  document$.subscribe(() => {
+    MathJax.startup.output.clearCache()
+    MathJax.typesetClear()
+    MathJax.texReset()
+    MathJax.typesetPromise()
+  })

docs/mkdocs/javascript/memory.js

@@ -1,2 +0,0 @@
-// Opt out of bfcache so large generated API pages are released after navigation.
-window.addEventListener("unload", () => {});

docs/mkdocs/javascript/mermaid.js

@@ -1,42 +1,21 @@
-// Load Mermaid only on pages that contain diagrams.
-let mermaidLoading;
+// Initialize Mermaid for diagram rendering
+mermaid.initialize({
+  startOnLoad: false,
+  theme: 'default',
+  securityLevel: 'loose',
+  flowchart: {
+    useMaxWidth: true,
+    htmlLabels: true
+  }
+});
 
-function loadMermaid() {
-  if (window.mermaid) {
-    return Promise.resolve();
+// Render Mermaid diagrams when page content is ready
+document$.subscribe(() => {
+  const mermaidElements = document.querySelectorAll('.mermaid');
+  if (mermaidElements.length > 0) {
+    mermaid.run({
+      querySelector: '.mermaid',
+      nodes: mermaidElements
+    });
   }
-  if (!mermaidLoading) {
-    mermaidLoading = new Promise((resolve, reject) => {
-      const script = document.createElement("script");
-      script.src = "https://unpkg.com/mermaid@10/dist/mermaid.min.js";
-      script.async = true;
-      script.onload = () => {
-        mermaid.initialize({
-          startOnLoad: false,
-          theme: "default",
-          securityLevel: "loose",
-          flowchart: {
-            useMaxWidth: true,
-            htmlLabels: true,
-          },
-        });
-        resolve();
-      };
-      script.onerror = reject;
-      document.head.appendChild(script);
-    });
-  }
-  return mermaidLoading;
-}
-
-document$.subscribe(() => {
-  if (!document.querySelector(".mermaid")) {
-    return;
-  }
-  loadMermaid().then(() => {
-    mermaid.run({
-      querySelector: ".mermaid",
-      nodes: document.querySelectorAll(".mermaid"),
-    });
-  });
 });

docs/models/supported_models.md

@@ -19,7 +19,6 @@ th {
 |--------------|--------|-------------------|------------|---------|-----|-----------|
 | `Qwen3OmniMoeForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `MingFlashOmniForConditionalGeneration` + `MingImagePipeline` | Ming-flash-omni-2.0 (omni-speech + imagegen<sup>1</sup>) | `Jonathan1909/Ming-flash-omni-2.0` | ✅︎ |   |   |   |
 | `BagelForConditionalGeneration` | BAGEL (DiT-only) | `ByteDance-Seed/BAGEL-7B-MoT` | ✅︎ | ✅︎ | | ✅︎ |
 | `InternVLAA1Pipeline` | InternVLA-A1 | `InternRobotics/InternVLA-A1-3B` | ✅︎ | ✅︎ | | |
 | `HunyuanImage3ForCausalMM` | HunyuanImage3.0 (DiT-only) | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
@@ -33,7 +32,6 @@ th {
 | `ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, T2V with sound | `nvidia/Cosmos3-Nano` | ✅︎ | | | |
 | `WanSpeechToVideoPipeline` | Wan2.2-S2V | `Wan-AI/Wan2.2-S2V-14B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `LTX2Pipeline` | LTX-2-T2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | |
@@ -42,14 +40,13 @@ th {
 | `LTX2ImageToVideoTwoStagesPipeline` | LTX-2-I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |
 | `LTX23Pipeline` | LTX-2.3-T2V | `dg845/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
 | `LTX23ImageToVideoPipeline` | LTX-2.3-I2V | `dg845/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
-| `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | | | |
 | `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | |
 | `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | |
 | `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | ✅︎ | ✅︎ | | ✅︎ |
 | `LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` | ✅︎ | ✅︎ | | ✅︎ |
-| `CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✅︎ | ✅︎ | | ✅︎ |
+| `CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✅︎ | ✅︎ | | |
 | `MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` | ✅︎ | ✅︎ | | |
 | `Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `FluxKontextPipeline` | FLUX.1-Kontext-dev | `black-forest-labs/FLUX.1-Kontext-dev` | ✅︎ | ✅︎ | | |

docs/serving/image_edit_api.md

@@ -13,7 +13,7 @@ For example...
 ```bash
 # Qwen-Image
 vllm serve Qwen/Qwen-Image-Edit-2511 --omni --port 8000
-```
+
 
 ### Generate Images
 
@@ -53,7 +53,7 @@ result = client.images.edit(
     output_format='jpeg',
     # url格式
     extra_body={
-        "url": [input_image_url,input_image_url],
+        "url": [input_image_url1,input_image_url],
         "num_inference_steps": 50,
         "guidance_scale": 1,
         "seed": 777,
@@ -92,8 +92,7 @@ Content-Type: multipart/form-data
 | `user` | string | null | User identifier for tracking |
 | `output_format` | string | "png" | The format in which the generated images are returned. Must be one of "png", "jpg", "jpeg", "webp". |
 | `output_compression` | integer | 100 | The compression level (0-100%) for the generated images. |
-| `background` | string or null | "auto" | Allows to set transparency for the background of the generated image(s). |
-| `stream` | boolean | false | Return Server-Sent Events instead of a single JSON response. Currently supported only by multi-stage image edit pipelines that produce AR text before the final image, such as HunyuanImage3 IT2I AR+DiT. |
+| `background` | string or null | "auto" | Allows to set transparency for the background of the generated image(s).
 
 #### vllm-omni Extension Parameters
 
@@ -110,8 +109,6 @@ Content-Type: multipart/form-data
 
 ### Response Format
 
-When `stream=false` or omitted, the endpoint returns the standard image edit response:
-
 ```json
 {
   "created": 1701234567,
@@ -123,56 +120,10 @@ When `stream=false` or omitted, the endpoint returns the standard image edit res
     }
   ],
   "output_format": null,
-  "size": null
+  "size": null,
 }
 ```
 
-### Streaming Response Format
-
-Set `stream=true` when you want to receive the HunyuanImage3 IT2I AR
-recaption text before the final edited image is ready. Streaming is only
-available for multi-stage image edit pipelines where stage 0 produces AR text
-and a later diffusion stage produces the image. Single-stage diffusion image
-edit pipelines reject `stream=true` with HTTP 400.
-
-The response uses Server-Sent Events with `Content-Type: text/event-stream`.
-Each event is sent as one `data:` line. The stream order is:
-
-1. One or more AR text delta chunks.
-2. One final image chunk.
-3. `data: [DONE]`.
-
-AR text chunks expose only the generated text delta and the AR completion
-`index`. They do not expose token ids, ratio tokens, KV cache data, or internal
-prompt token ids.
-
-Example AR delta event:
-
-```text
-data: {"object":"image.edit.chunk","type":"ar_delta","delta":"A close-up product photo...","index":0,"created":1701234567,"model":"tencent/HunyuanImage-3.0-Instruct"}
-```
-
-Example final image event:
-
-```text
-data: {"object":"image.edit.chunk","type":"image","data":[{"b64_json":"<base64-encoded PNG>","url":null,"revised_prompt":null}],"output_format":"png","size":"1024x1024","created":1701234567,"model":"tencent/HunyuanImage-3.0-Instruct"}
-```
-
-Terminal event:
-
-```text
-data: [DONE]
-```
-
-If the engine fails after the SSE response has started, the stream emits an
-error chunk followed by `[DONE]`:
-
-```text
-data: {"object":"error","created":1701234567,"model":"tencent/HunyuanImage-3.0-Instruct","error":{"message":"<error message>","type":"server_error","code":500}}
-
-data: [DONE]
-```
-
 ## Examples
 
 ### Multiple Images input
@@ -183,28 +134,12 @@ curl -s -D >(grep -i x-request-id >&2) \
   -X POST "http://localhost:8000/v1/images/edits" \
   -F "model=xxx" \
   -F "image=@xx.png" \
-  -F "image=@xx.png" \
+  -F "image=@xx.png"
   -F "prompt='this bear is wearing sportwear. holding a basketball, and bending one leg.'" \
   -F "size=1024x1024" \
   -F "output_format=png"
 ```
 
-### Streaming HunyuanImage3 IT2I AR Text
-
-```bash
-curl -N -X POST "http://localhost:8000/v1/images/edits" \
-  -F "model=tencent/HunyuanImage-3.0-Instruct" \
-  -F "image=@./input.png" \
-  -F "prompt=Turn the product into a clean studio advertisement." \
-  -F "size=1024x1024" \
-  -F "output_format=png" \
-  -F "stream=true"
-```
-
-Use the `ar_delta` events for progressive display of the AR-generated
-recaption. Decode the `b64_json` field from the final `image` event to get the
-edited image.
-
 
 ## Parameter Handling
 

docs/serving/speech_api.md

@@ -381,44 +381,6 @@ are cached in-process with a shared LRU so repeated requests with the same
 all TTS model types; deleting a voice invalidates every model-type slot at
 once.
 
-### Precomputed Custom Voices
-
-Qwen3-TTS Base and VoxCPM2 can load offline-precomputed voices at startup.
-Generate a directory containing `custom_voice_manifest.json` plus one
-`.safetensors` file per voice, then set the pipeline-wide deploy config field:
-
-```yaml
-custom_voice_dir: /path/to/custom_voices
-```
-
-Qwen3-TTS profiles are created with:
-
-```bash
-python examples/online_serving/text_to_speech/qwen3_tts/precompute_custom_voice.py \
-  --model Qwen/Qwen3-TTS-12Hz-1.7B-Base \
-  --voice-name alice \
-  --ref-audio /path/to/reference.wav \
-  --ref-text "Original transcript of the reference audio" \
-  --mode icl \
-  --output-dir /path/to/custom_voices
-```
-
-VoxCPM2 profiles are created with:
-
-```bash
-python examples/online_serving/text_to_speech/voxcpm2/precompute_custom_voice.py \
-  --model openbmb/VoxCPM2 \
-  --voice-name alice \
-  --ref-audio /path/to/reference.wav \
-  --mode ref_continuation \
-  --prompt-text "Original transcript of the reference audio" \
-  --output-dir /path/to/custom_voices
-```
-
-Only profiles whose safetensors payload can be loaded and validated are exposed
-by `GET /v1/audio/voices`. Valid precomputed voices can be used in
-`POST /v1/audio/speech` by passing `voice="alice"` without `ref_audio`.
-
 **Configuration (environment variables):**
 
 | Variable | Default | Description |

docs/usage/metrics.md

@@ -1,98 +0,0 @@
-# Production Metrics
-
-vLLM-Omni exposes Prometheus metrics via the `/metrics` endpoint on the OpenAI-compatible API server. This page covers the text and audio surface; diffusion / image / video metrics are tracked in a follow-up PR.
-
-```bash
-vllm-omni serve Qwen/Qwen3-Omni-30B-A3B-Instruct --port 8000 --log-stats
-curl http://localhost:8000/metrics
-```
-
-**`--log-stats` is required to populate metric data.** Without the flag, the endpoint still returns `200 OK` but the upstream `vllm:*` wrap is not registered at all, and the 15 `vllm:omni_*` families are registered as placeholders with no data written to them. This keeps the runtime cost essentially zero for deployments that don't need monitoring. With the flag, all ~80 families populate.
-
-## Metric Namespaces
-
-| Prefix | Source | Present when |
-|--------|--------|--------------|
-| `vllm:omni_` | vLLM-Omni orchestrator / audio modality / cross-stage transfer | Pipeline-dependent |
-| `vllm:` | Upstream vLLM engine, wrapped by `OmniPrometheusStatLogger` to expose `{stage, replica}` | Pipeline includes an LLM (AR) stage |
-| `http_` / `process_` | Uvicorn / Python runtime | Always |
-
-## Pipeline-Level Metrics (`vllm:omni_`)
-
-Defined in `vllm_omni/metrics/prometheus.py`. Track request lifecycle across the full multi-stage pipeline.
-
-### Request counts
-
-| Metric | Type | Labels | Description |
-|--------|------|--------|-------------|
-| `vllm:omni_num_requests_running` | Gauge | `model_name` | Pipeline-global in-flight requests (dispatched to engine, not yet finalized) |
-| `vllm:omni_num_requests_waiting` | Gauge | `model_name` | Requests waiting in the Orchestrator queue |
-| `vllm:omni_requests_success_total` | Counter | `model_name`, `finished_reason` | Total requests by completion reason. `finished_reason` ∈ {`stop`, `length`, `abort`, ...} mirroring upstream `vllm:request_success_total`; aborts cover client disconnect / cancellation paths in addition to upstream `FinishReason.ABORT` |
-
-### Latency
-
-| Metric | Type | Labels | Description |
-|--------|------|--------|-------------|
-| `vllm:omni_e2e_request_latency_s` | Histogram | `model_name` | Pipeline-global end-to-end request latency in seconds |
-
-## Audio Modality Metrics (`vllm:omni_`)
-
-Emitted at request finalize, except for `audio_ttfp_s` (streaming-hook at the first audio packet) and `audio_underrun_s` / `audio_continuity_ok_total` (streaming finalize, after the chunk stream is exhausted). All carry `{model_name, stage, replica}` plus the listed extra label.
-
-| Metric | Type | Extra label | Description |
-|--------|------|-------------|-------------|
-| `vllm:omni_audio_ttfp_s` | Histogram | — | Time from request arrival to first audio packet/frame |
-| `vllm:omni_audio_duration_s` | Histogram | — | Audio content duration (`audio_frames / sample_rate`) |
-| `vllm:omni_audio_rtf` | Histogram | — | Real-time factor (`stage_gen_time_s / audio_duration_s`); streaming TTS SLO red line `< 1`; uses `RTF_BUCKETS` |
-| `vllm:omni_audio_frames_total` | Counter | — | Cumulative audio frame count; throughput via `rate()` |
-| `vllm:omni_audio_underrun_s` | Histogram | — | Per-request worst-case player deficit; `> 0` indicates listener heard silent gaps |
-| `vllm:omni_audio_continuity_ok_total` | Counter | `threshold_ms` | Incremented when the request's worst underrun stayed below `threshold_ms` |
-| `vllm:omni_audio_skipped_requests_total` | Counter | `reason` | Silent-loss counter — code2wav rejected malformed codec input and returned `200 OK` with empty audio |
-
-The continuity math comes from `vllm_omni/benchmarks/audio_continuity.py::compute_continuity_stats` so the server-side observation aligns with the bench-side definition.
-
-## Cross-Stage Transfer Metrics (`vllm:omni_`)
-
-Per-physical-transfer histograms tracking the data hop between adjacent stages. Labels `{model_name, from_stage, from_replica, to_stage, to_replica}` let dashboards attribute latency to specific replica edges. `from_replica` / `to_replica` are resolved from the orchestrator's sticky-routing binding (`stage_pool.get_bound_replica_id(request_id)`), so no extra plumbing through `TransferEdgeStats` is needed.
-
-| Metric | Type | Description |
-|--------|------|-------------|
-| `vllm:omni_transfer_size_bytes` | Histogram | Per-transfer payload size in bytes |
-| `vllm:omni_transfer_tx_s` | Histogram | Sender-side time (serialize + submit to connector) |
-| `vllm:omni_transfer_rx_s` | Histogram | Receiver-side time (recv + deserialize) |
-| `vllm:omni_transfer_in_flight_s` | Histogram | Network in-flight time (TX done → RX recv start) |
-
-## vLLM Engine Metrics (`vllm:`)
-
-When the pipeline includes an LLM stage, the upstream vLLM engine exposes its full set of ~37 metric families under the `vllm:` prefix.
-
-vLLM-Omni wraps the upstream `vllm.v1.metrics.loggers.PrometheusStatLogger` with `OmniPrometheusStatLogger` so that the original `engine` single label is reshaped into `stage` + `replica`. Every `vllm:*` family — TTFT, ITL, TPOT, e2e latency, KV cache usage, scheduler running/waiting, request success counts, etc. — therefore gains per-`(stage, replica)` visibility automatically. No omni-side duplicate is needed for the text path.
-
-```text
-# Before wrap:
-vllm:num_requests_running{model_name="...", engine="1"}              3.0
-
-# After wrap:
-vllm:num_requests_running{model_name="...", stage="1", replica="0"}  2.0
-vllm:num_requests_running{model_name="...", stage="1", replica="1"}  1.0
-```
-
-For the full list of upstream metrics, see [the vLLM docs](https://github.com/vllm-project/vllm/blob/main/docs/usage/metrics.md).
-
-## Metric Availability by Pipeline Type
-
-| Metric group | Multi-stage LLM (Qwen3-Omni) |
-|---|---|
-| `vllm:omni_` request tracking + latency | With `--log-stats` |
-| `vllm:omni_` audio modality | With `--log-stats`, if pipeline has a talker stage |
-| `vllm:omni_` transfer | With `--log-stats`, if pipeline has ≥ 2 stages |
-| `vllm:` engine metrics (per `(stage, replica)`) | With `--log-stats` |
-| `vllm:` MFU metrics | With `--log-stats --enable-mfu-metrics` |
-
-## Naming Convention
-
-- All time-bearing metrics use the `_s` suffix (values in seconds). Buckets are `SECONDS_BUCKETS` for e2e / generation-style values and `SECONDS_FAST_BUCKETS` (1 ms → 60 s) for the fine-grained transfer and audio-underrun values.
-- Counters use the `_total` suffix (auto-appended by `prometheus_client`).
-- Sizes use the `_bytes` suffix.
-- All omni-specific families are prefixed `vllm:omni_`. The upstream `unregister_vllm_metrics()` function is monkey-patched (see `vllm_omni/patch.py`) to a scoped version that still strips upstream `vllm:*` collectors so multi-engine init within one process does not crash on duplicate registration, but preserves anything prefixed `vllm:omni_` / `vllm_omni`.
-- Text and audio first-output use distinct families (`vllm:time_to_first_token_seconds` reused from upstream for text; `vllm:omni_audio_ttfp_s` for audio) rather than a single metric with a `modality` label.

docs/user_guide/diffusion/attention_backends.md

@@ -1,26 +1,28 @@
 # Diffusion Attention Backends
 
-This document describes the diffusion attention backends available in vLLM-Omni, how to select them globally and per-role, the per-platform defaults, and how to use SageAttention.
+This document describes the diffusion attention backends available in vLLM-Omni, how to select them globally and per-role, and how to use SageAttention.
 
 ## Overview
 
 Diffusion attention backend selection is resolved in `vllm_omni.diffusion.attention.selector`. It looks up the backend from a structured `AttentionConfig` carried on `OmniDiffusionConfig` and falls back to the platform default when nothing is configured.
 
-This backend is used by diffusion attention layers such as the DiT attention in video and image generation models. It does **not** affect autoregressive (LLM) attention paths — those go through vLLM's own attention backend selector.
+This backend is used by diffusion attention layers such as the DiT attention in video and image generation models.
 
-The full set of backends and their platform defaults is in the **Backend Options** and **Platform Defaults** sections below. If no attention backend is configured, vLLM-Omni asks the current platform to choose the default.
+On CUDA, the practical choices today are:
+
+- `FLASH_ATTN`: FlashAttention backend. This is the default on supported CUDA systems when FlashAttention is installed.
+- `TORCH_SDPA`: PyTorch `scaled_dot_product_attention`.
+- `SAGE_ATTN`: SageAttention backend, if `sageattention` is installed.
+
+If no attention backend is configured, vLLM-Omni asks the current platform to choose the default. On CUDA, that normally means `FLASH_ATTN` when available, otherwise `TORCH_SDPA`.
 
 ## Backend Options
 
 | Value | Notes |
 |---|---|
-| `FLASH_ATTN` | Wraps FlashAttention 2. Default on Hopper / Ada / Ampere when `flash-attn` is installed. |
-| `CUDNN_ATTN` | Pins `sdpa_kernel([CUDNN_ATTENTION])`. Default on Blackwell (sm_10x / sm_12x) with cuDNN ≥ 9.5. Wins on mask-heavy DiTs (HunyuanVideo-1.5: 2× e2e vs SDPA). |
-| `FLASHINFER_ATTN` | Calls FlashInfer's dense `single_prefill_with_kv_cache` directly with `custom_mask` for non-causal masked attention. Used as Blackwell fallback when cuDNN is unavailable. Requires `flashinfer`. |
-| `TORCH_SDPA` | PyTorch `scaled_dot_product_attention` with the default backend dispatcher. Most conservative; always available. |
-| `SAGE_ATTN` | SageAttention 2.2 — INT8-quantized attention with FP16 accumulation. Lossy but typically visually indistinguishable on diffusion outputs. Requires `sageattention`. |
-| `SAGE_ATTN_3` | Requires `sageattn3` from `SageAttention/sageattention3_blackwell`. CUDA only, intended for Blackwell GPUs, with GQA/MQA requests falling back to PyTorch SDPA. |
-
+| `FLASH_ATTN` | Default on CUDA when FlashAttention is available. Good default for most diffusion workloads. |
+| `TORCH_SDPA` | Most conservative fallback. Useful for debugging or compatibility. |
+| `SAGE_ATTN` | Requires `sageattention`. Can improve performance on some workloads, but output quality must be validated model-by-model. |
 
 ## Configuration
 
@@ -93,80 +95,6 @@ config = OmniDiffusionConfig(
 
 A plain dict is also accepted and normalized to `AttentionConfig`.
 
-## Platform Defaults
-
-### Blackwell (sm_100 / sm_103 / sm_120 / sm_121)
-
-Auto-route preference, in order:
-
-1. `CUDNN_ATTN` — when cuDNN ≥ 9.5 is available (ships in PyTorch 2.5+ wheels)
-2. `FLASHINFER_ATTN` — when `flashinfer` is installed but cuDNN < 9.5
-3. `FLASH_ATTN` — when `flash-attn` is installed with the Blackwell CUTE kernel
-4. `TORCH_SDPA` — last resort
-
-The startup log line `Defaulting to diffusion attention backend CUDNN_ATTN (Blackwell sm_120, cuDNN 91002)` confirms the route.
-
-**Why CUDNN_ATTN by default**: on mask-heavy diffusion models (HunyuanVideo-1.5, Qwen-Image), cuDNN's pinned FMHA kernel sidesteps a PyTorch SDPA dispatch quirk where the unpinned dispatcher picks `EFFICIENT_ATTENTION` (~25 ms) for masked calls instead of cuDNN (~11 ms). The pin gives 2× e2e on HV-1.5 with no regression on lighter models.
-
-### Hopper (sm_90) / Ada (sm_89) / Ampere (sm_80–sm_86)
-
-Auto-route preference:
-
-1. `FLASH_ATTN` — when `flash-attn` is installed
-2. `TORCH_SDPA` — fallback
-
-`CUDNN_ATTN` and `FLASHINFER_ATTN` are still selectable via env var on these GPUs but are not in the auto-route — FlashAttention 2 is the well-tuned path on pre-Blackwell hardware.
-
-## End-to-End Benchmark (BF16, sm_120 RTX Pro 6000 Blackwell)
-
-Same prompt and seed across runs. `Total generation time` from `text_to_video.py` / `text_to_image.py`.
-
-| Model | Shape | TORCH_SDPA | CUDNN_ATTN | FLASHINFER_ATTN |
-|---|---|---|---|---|
-| HunyuanVideo-1.5 (T2V) | 480p / 33f / 50 steps | 147.05 s | **73.02 s** | 127.84 s |
-| Wan 2.2 14B (T2V) | 480p / 33f / 40 steps | 117.75 s | 117.17 s | **115.07 s** |
-| Qwen-Image (T2I) | 1024² / 50 steps | 17.41 s | **15.14 s** | 16.02 s |
-| FLUX.2-dev (T2I) | 1024² / 50 steps, TP=2 | 53.62 s | **53.30 s** | 54.94 s |
-
-Pattern: mask-heavy DiTs (HV-1.5, Qwen-Image) favor `CUDNN_ATTN`; lighter-mask DiTs and TP-saturated configs (Wan 2.2, FLUX.2 TP=2) tie within noise.
-
-## Known Limitations
-
-### LTX-2.0: `CUDNN_ATTN` crashes under torch.compile
-
-LTX-2's audio attention has a symbolic head_dim under torch.compile tracing. cuDNN's SDPA backend selector rejects symbolic dims and Dynamo aborts compilation. Tracked in [#3121](https://github.com/vllm-project/vllm-omni/issues/3121).
-
-**Workaround**: explicitly select `FLASHINFER_ATTN` or `TORCH_SDPA` for LTX-2.0:
-
-```bash
-DIFFUSION_ATTENTION_BACKEND=FLASHINFER_ATTN python examples/offline_inference/text_to_video/text_to_video.py \
-    --model Lightricks/LTX-2 ...
-```
-
-### FA4 not yet integrated
-
-FlashAttention-4 (released March 2026) targets Blackwell natively and reportedly beats cuDNN by ~20% on B200. As of this writing the `flash-attn-4 4.0.0b10` wheel crashes with `AttributeError: 'NoneType' object has no attribute '_trait'` during JIT on sm_120. Not yet wired into vLLM-Omni; revisit when stable lands.
-
-## Choosing a Backend Manually
-
-### When to override the default
-
-- **Quality validation**: compare a new backend against `TORCH_SDPA` as the reference, since SDPA's default dispatcher is the most extensively tested.
-- **Lossy speedup hunting**: try `SAGE_ATTN` (INT8 quantized) on diffusion outputs — typically indistinguishable visually but always validate.
-- **Workaround for known issues**: see Known Limitations above.
-
-### Verifying which backend is in use
-
-The startup log prints one of:
-
-```
-Using diffusion attention backend 'CUDNN_ATTN'           # explicit override
-Defaulting to diffusion attention backend CUDNN_ATTN ... # auto-route
-Defaulting to diffusion attention backend SDPA           # nothing else available
-```
-
-If you don't see one of these, the model didn't reach diffusion stage init — check earlier logs for failures.
-
 ## SageAttention Installation
 
 vLLM-Omni expects SageAttention to be installed into the same Python environment as vLLM-Omni.
@@ -187,56 +115,11 @@ Quick check:
 python -c "import sageattention; print(sageattention.__file__)"
 ```
 
-## SageAttention3 Installation
+## Usage
 
-vLLM-Omni expects SageAttention3 to be installed into the same Python environment as vLLM-Omni.
+### Enable SageAttention
 
-Build from source:
-
-```bash
-git clone https://github.com/thu-ml/SageAttention.git
-cd SageAttention/sageattention3_blackwell
-python setup.py install
-```
-
-Quick check:
-
-```bash
-python -c "import sageattn3; print(sageattn3.__file__)"
-```
-
-Notes:
-
-- `SAGE_ATTN_3` is only selected on CUDA when `sageattn3` is importable and the GPU is Blackwell-class.
-- SageAttention3's Blackwell kernel assumes `Hq == Hkv`. In vLLM-Omni, GQA/MQA diffusion requests fall back to PyTorch SDPA for correctness.
-
-## Usage Examples
-
-### Default (auto-route)
-
-```bash
-python examples/offline_inference/text_to_video/text_to_video.py \
-    --model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \
-    --prompt "A dog running across a field of golden wheat." \
-    --height 480 --width 832 --num-frames 33 \
-    --num-inference-steps 50 --seed 42 --guidance-scale 6.0 \
-    --output hv15.mp4
-```
-
-On Blackwell this picks `CUDNN_ATTN` automatically. Check the log for the `Defaulting to ...` line.
-
-### Explicit backend selection
-
-```bash
-DIFFUSION_ATTENTION_BACKEND=FLASHINFER_ATTN python examples/offline_inference/text_to_video/text_to_video.py \
-    --model Lightricks/LTX-2 \
-    --prompt "A dog running across a field of golden wheat." \
-    --height 480 --width 832 --num-frames 33 \
-    --num-inference-steps 40 --seed 42 --guidance-scale 4.0 \
-    --output ltx2.mp4
-```
-
-### SageAttention (lossy)
+Example: HunyuanVideo-1.5 text-to-video
 
 ```bash
 DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN python examples/offline_inference/text_to_video/text_to_video.py \
@@ -245,7 +128,7 @@ DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN python examples/offline_inference/text_to_
     --height 480 --width 832 --num-frames 33 \
     --num-inference-steps 30 --seed 42 --guidance-scale 6.0 \
     --tensor-parallel-size 2 \
-    --output hv15_sage.mp4
+    --output ../tmp/hv15_modelopt_sage.mp4
 ```
 
 Example: Wan2.2 TI2V 5B
@@ -260,20 +143,6 @@ DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN python examples/offline_inference/text_to_
     --output outputs/wan22_sage.mp4
 ```
 
-### Enable SageAttention3
-
-Example:
-
-```bash
-DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN_3 python examples/offline_inference/text_to_video/text_to_video.py \
-    --model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \
-    --prompt "A dog running across a field of golden wheat." \
-    --height 480 --width 832 --num-frames 33 \
-    --num-inference-steps 30 --seed 42 --guidance-scale 6.0 \
-    --tensor-parallel-size 2 \
-    --output outputs/hv15_sage3.mp4
-```
-
 ### Mixed backends across roles
 
 Use `FLASH_ATTN` for self-attention and `TORCH_SDPA` for cross-attention:
@@ -304,13 +173,13 @@ python examples/offline_inference/text_to_video/text_to_video.py \
 
 ## Validation Guidance
 
-Don't assume a faster attention backend is numerically interchangeable with `TORCH_SDPA`.
+Do not assume that a faster attention backend is numerically interchangeable with `FLASH_ATTN`.
 
 Always compare:
 
 - End-to-end runtime
-- Diffusion-stage runtime (`add_req_and_wait` line in DiffusionEngine.step breakdown)
-- Output quality against a known-good baseline (CLIP similarity, frame-level diff, or visual review)
+- DiT / diffusion stage runtime
+- Output quality against a known-good baseline
 
 At minimum, keep the same:
 
@@ -318,19 +187,6 @@ At minimum, keep the same:
 - prompt
 - seed
 - resolution
-- frame count / step count
-- parallel config (TP / CFG-parallel / Ulysses degrees)
-
-## Reproducing the Benchmark Table
-
-The end-to-end numbers above were collected by running `text_to_video.py` /
-`text_to_image.py` with the same prompt and seed while varying
-`DIFFUSION_ATTENTION_BACKEND`. For a quick kernel-level comparison of the
-backends without loading a model:
-
-```bash
-python benchmarks/diffusion/bench_attention_backends.py --preset hv15
-```
-
-It runs all three BF16 backends on representative DiT attention shapes and
-prints a ranking table at the end.
+- frame count
+- inference steps
+- parallel config

docs/user_guide/diffusion_features.md

@@ -113,7 +113,7 @@ The following tables show which models support each feature:
 | **FLUX.1-schnell**       |     ❌     |     ✅      |           ❌           |       ✅        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ✅        |        ❌         |
 | **FLUX.2-klein**         |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ✅        |        ❌         |
 | **FLUX.1-Kontext-dev**   |     ❌     |     ❌      |           ❌           |       ❌        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ❌        |        ❌         |
-| **FLUX.2-dev**           |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ❌        |        ❌         |
+| **FLUX.2-dev**           |     ✅     |     ✅      |           ❌           |       ✅        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **GLM-Image**            |     ❌     |     ❌      |           ❌           |       ✅        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **Hidream-I1-Full**        |     ❌     |     ❌      |           ❌           |       ❌        |         ✅         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **HunyuanImage3**        |     ❌     |     ✅      |           ❌           |       ❌        |         ✅         |          ❌          |   ❌    |             ❌             |          ❌           |       ✅        |        ❌         |
@@ -129,16 +129,14 @@ The following tables show which models support each feature:
 | **Qwen-Image-Edit**      |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
 | **Qwen-Image-Edit-2509** |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |        ✅ (decode)         |          ✅           |       ❌        |        ❌         |
 | **Qwen-Image-Layered**   |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
-| **SenseNova-U1**         |     ❌     |     ✅      |           ❌           |       ❌        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
+| **SenseNova-U1**         |     ❌     |     ❌      |           ❌           |       ❌        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **Stable-Diffusion3.5**  |     ❌     |     ✅      |           ❌           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
 | **Z-Image**              |     ✅     |     ✅      |           ✅           |       ❓        |   ✅ (TP=2 only)   |          ❌          |   ✅    |             ❌             |      ✅ (decode)      |       ✅        |        ❌         |
 | **ERNIE-Image**          |     ❌     |     ✅      |           ✅           |       ❓        |         ✅         |          ❌          |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |
-| **Cosmos3**              |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ✅        |        ❌         |
 
 > Notes:
 > 1. Nextstep_1(T2I) does not support cache acceleration methods such as TeaCache or Cache-DiT.
 > 2. `Tongyi-MAI/Z-Image-Turbo` and `SII-GAIR/daVinci-MagiHuman-Base-1080p` are distilled models with minimal NFEs; CFG-Parallel is not necessary.
-> 3. Cosmos3 T2I uses `Cosmos3OmniDiffusersPipeline` with `modalities=["image"]`. Model-level CPU offload is not supported; use layerwise offload.
 
 ### VideoGen
 
@@ -147,12 +145,10 @@ The following tables show which models support each feature:
 | **Wan2.2**                   |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ✅         |   ✅    |             ✅             |  ✅ (encode/decode)   |       ❌        |        ❌         |
 | **Wan2.1-VACE**              |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
 | **LTX-2**                    |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |
-| **LTX-2.3**                  |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
+| **LTX-2.3**                  |     ❌     |     ✅      |           ✅           |       ❌        |         ✅         |         ❌         |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **Helios**                   |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |
-| **HunyuanVideo-1.5 T2V I2V** |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |  ✅ (encode/decode)   |       ✅        |        ❌         |
+| **HunyuanVideo-1.5 T2V I2V** |     ❌     |     ✅      |           ❌           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |      ✅ (decode)      |       ✅        |        ❌         |
 | **DreamID-Omni**             |     ❌     |     ❌      |           ❌           |       ✅        |         ❌         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |
-| **Cosmos3**                  |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |  ✅ (encode/decode)   |       ✅        |        ❌         |
-
 
 **Frame Interpolation Support**
 

docs/user_guide/examples/offline_inference/text_to_speech.md

@@ -23,7 +23,6 @@ list of supported architectures across all modalities, see
 | OmniVoice | `k2-fsa/OmniVoice` | 2 (gen + dec) | ✓ | — | voice design (`--instruct`), language (`--lang`) | 24 kHz |
 | Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | 2 (talker + code2wav) | ✓ (Base) | ✓ | 3 task variants (`--query-type`) | 24 kHz |
 | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | varies | ✓ | ✓ | voice presets (`--voice`) | 24 kHz |
-| higgs-audio v2 | `bosonai/higgs-audio-v2-generation-3B-base` (+ `k2-fsa/OmniVoice/audio_tokenizer/` codec) | 2 (talker + code2wav, DualFFN) | ✓ (`--ref-audio` + `--ref-text`) | — | — | 24 kHz |
 
 ## Common Quick Start
 
@@ -334,46 +333,6 @@ Available voice presets are listed on the HF model card (`mistralai/Voxtral-4B-T
 - `--concurrency M` requires `--streaming` and must evenly divide `--num-prompts`.
 - Run `--help` for the full argument surface.
 
----
-
-## higgs-audio v2
-
-2-stage TTS at 24 kHz: a vLLM-native Llama-3.2-3B talker with a DualFFN audio expert (Stage 0) feeding a HiggsAudio codec decoder (Stage 1) over the shared-memory connector. Stage 1 builds on the HiggsAudio decoder kernel at `vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_decoder.py`.
-
-### Prerequisites
-
-Voice clone uses HF's `HiggsAudioV2TokenizerModel`, instantiated from `k2-fsa/OmniVoice/audio_tokenizer/` — the boson-ai standalone tokenizer Hub repo's `model.safetensors` is actually the 3B talker LM, so we point HF at k2's repackaged codec weights instead. Only the `audio_tokenizer/` subdir (~806 MB) is downloaded.
-
-```bash
-pip install -U "transformers>=5.3.0"
-```
-
-### Quick start (plain TTS)
-
-```bash
-python examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py \
-    --texts "Hello world." "The quick brown fox jumps over the lazy dog." \
-    --output-dir results/wavs
-```
-
-### Voice cloning
-
-Pass both `--ref-audio` and `--ref-text` together:
-
-```bash
-python examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py \
-    --texts "Hello, this is a cloned voice." \
-    --ref-audio /path/to/reference.wav \
-    --ref-text  "Exact transcript spoken in reference.wav." \
-    --output-dir results/wavs
-```
-
-### Notes
-
-- Deploy config auto-loads from `vllm_omni/deploy/higgs_audio_v2.yaml`.
-- `--ref-text` must be the real transcript of `--ref-audio`; mismatched text degrades cloned-voice quality.
-- Out of scope (rejected with 4xx by the request validator): multi-speaker `[SPEAKERn]` dialogue, `profile:` text-only speaker descriptions, the `ref_audio_in_system_message` system-block variant, chunked long-form generation, and per-request `voice` / `instructions` / `task_type` / `language` / `speed != 1.0` / `x_vector_only_mode` / `speaker_embedding`.
-
 ## Example materials
 
 ??? abstract "cosyvoice3/end2end.py"
@@ -384,9 +343,6 @@ python examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py \
     ``````py
     --8<-- "examples/offline_inference/text_to_speech/fish_speech/end2end.py"
     ``````
-??? abstract "higgs_audio_v2/end2end.py"
-    ``````py
-    --8<-- "examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py"
 ??? abstract "glm_tts/end2end.py"
     ``````py
     --8<-- "examples/offline_inference/text_to_speech/glm_tts/end2end.py"

docs/user_guide/examples/online_serving/qwen3_omni.md

@@ -17,9 +17,6 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 
 The default deployment configuration situated at `vllm_omni/deploy/qwen3_omni_moe.yaml` is resolved and loaded
 automatically via the model registry, obviating the necessity for the `--deploy-config` flag in standard deployment topologies.
-The bundled Qwen3-Omni setup defaults `VLLM_USE_FLASHINFER_MOE_FP16=0`. This keeps the Thinker & Talker on vLLM's
-Triton unquantized MoE path and avoids the performance regression observed with the FlashInfer CUTLASS unquantized MoE
-backend.
 Asynchronous chunk streaming is **enabled by default** within the bundled configuration.
 
 To explicitly utilize a custom deployment YAML, specify the configuration path:
@@ -75,12 +72,6 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
     --stage-overrides '{"1": {"gpu_memory_utilization": 0.5}}'
 ```
 
-To experiment with the FlashInfer FP16 MoE path, set `VLLM_USE_FLASHINFER_MOE_FP16=1` before launching the server:
-```bash
-VLLM_USE_FLASHINFER_MOE_FP16=1 \
-vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
-```
-
 For the stage-based CLI, you usually do **not** need `--stage-overrides` for
 that kind of change. Since each command launches one stage, just pass the knob
 directly on that stage command:

docs/user_guide/examples/online_serving/text_to_speech.md

@@ -19,7 +19,6 @@ For the full list of supported architectures across all modalities, see
 |---|---|---|---|---|---|
 | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | — |
 | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | ✓ |
-| higgs-audio v2 | `bosonai/higgs-audio-v2-generation-3B-base` | ✓ (`ref_audio`+`ref_text`) | ✓ (codec_streaming) | — | — |
 | GLM-TTS | `zai-org/GLM-TTS` | ✓ (`ref_audio`+`ref_text`, required) | ✓ (PCM stream) | — | ✓ |
 | OmniVoice | `k2-fsa/OmniVoice` | (offline only) | — | — | — |
 | Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | ✓ (Base) | ✓ (PCM + WebSocket) | ✓ (presets + `/v1/audio/voices` upload) | ✓ (standard + FastRTC) |
@@ -194,45 +193,6 @@ python fish_speech/gradio_demo.py --api-base http://localhost:8091  # if server
 
 ---
 
-## higgs-audio v2
-
-2-stage TTS at 24 kHz: a vLLM-native Llama-3.2-3B talker with a DualFFN audio expert (Stage 0) feeding a HiggsAudio codec decoder (Stage 1) that streams chunks back to the client.
-
-### Prerequisites
-
-Voice clone uses HF's `HiggsAudioV2TokenizerModel` loaded from `k2-fsa/OmniVoice/audio_tokenizer/` (~806 MB subdir; the boson-ai standalone tokenizer Hub repo's `model.safetensors` is the 3B talker LM, not the codec):
-
-```bash
-pip install -U "transformers>=5.3.0"
-```
-
-### Launch
-```bash
-GPUS=6,7 PORT=8094 bash examples/online_serving/text_to_speech/higgs_audio_v2/run_server.sh
-```
-Deploy config auto-loads from `vllm_omni/deploy/higgs_audio_v2.yaml`.
-
-### Sending requests
-```bash
-# Plain TTS
-python higgs_audio_v2/batch_speech_client.py \
-    --base-url http://localhost:8094 \
-    --output-dir /tmp/higgs_out \
-    --prompts "Hello world." "The quick brown fox jumps over the lazy dog."
-
-# Voice cloning — pass a reference clip and its transcript together
-python higgs_audio_v2/batch_speech_client.py \
-    --base-url http://localhost:8094 \
-    --output-dir /tmp/higgs_clone \
-    --ref-audio /path/to/reference.wav \
-    --ref-text  "Exact transcript spoken in reference.wav." \
-    --prompts "Hello, this is a cloned voice."
-```
-
-### Notes
-- Output: 24 kHz mono.
-- `--ref-text` must be the real transcript of `--ref-audio`; mismatched text degrades cloned-voice quality.
-- Out of scope (rejected with explicit 4xx): multi-speaker `[SPEAKERn]` tags inside `input`, `profile:` text-only speaker descriptions, the `ref_audio_in_system_message` system-block variant, chunked long-form generation, and per-request `voice` / `instructions` / `task_type` / `language` / `speed != 1.0` / `x_vector_only_mode` / `speaker_embedding`.
 ## GLM-TTS
 
 2-stage TTS (AR + DiT flow-matching) at 24 kHz. Every request requires `ref_audio` + `ref_text`.
@@ -329,16 +289,6 @@ Single-GPU serves now default to the uniproc executor (lower IPC overhead, the B
 
 To opt out of chunked streaming, pass `--no-async-chunk` — the pipeline auto-dispatches to the end-to-end codec processor.
 
-### Tuning stage 1 `max_num_seqs` per task type
-The bundled `qwen3_tts.yaml` ships stage 1 (Code2Wav) at `max_num_seqs: 10`, tuned for Base voice cloning: stage-1 lifetimes are long (~3 s/req), so admitting up to 10 concurrent codec sequences lets requests progress in parallel in the scheduler — ~2× TTFA p95 at c=4 / c=8 (1× H100, 1.7B-Base, seed-tts) at an 8–12 % audio-throughput cost.
-
-CustomVoice / VoiceDesign have much shorter stage-1 lifetimes (~50–200 ms) and are TTFA-optimal at `max_num_seqs: 1`. Override the default when serving those task types:
-
-```bash
-vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base --omni \
-    --stage-overrides '{"1": {"max_num_seqs": 1}}'
-```
-
 ### Sending requests
 ```bash
 # CustomVoice with a predefined speaker
@@ -380,31 +330,6 @@ curl -X POST http://localhost:8091/v1/audio/voices \
 ```
 Uploaded voices are then usable as `voice="custom_voice_1"` on subsequent requests.
 
-### Precomputed custom voices
-For reused Base voice-cloning speakers, precompute the reference artifacts once and load them at server startup:
-```bash
-python qwen3_tts/precompute_custom_voice.py \
-    --model Qwen/Qwen3-TTS-12Hz-1.7B-Base \
-    --voice-name alice \
-    --ref-audio /path/to/reference.wav \
-    --ref-text "Original transcript of the reference audio" \
-    --mode icl \
-    --output-dir /path/to/custom_voices
-```
-`--mode icl` stores both `speaker_embedding` and `ref_code`; `--mode xvec` stores only the speaker embedding. Add the output directory to a deploy config:
-```yaml
-custom_voice_dir: /path/to/custom_voices
-```
-Then start the server with that config and call the Speech API with only the voice name:
-```bash
-vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base --omni --deploy-config /path/to/qwen3_tts_custom_voice.yaml
-
-curl -X POST http://localhost:8091/v1/audio/speech \
-    -H "Content-Type: application/json" \
-    -d '{"input":"Hello from a precomputed voice.","voice":"alice","task_type":"Base"}' \
-    --output alice.wav
-```
-
 ### Streaming PCM
 ```bash
 curl -X POST http://localhost:8091/v1/audio/speech \
@@ -444,10 +369,8 @@ python qwen3_tts/gradio_fastrtc_demo.py --api-base http://localhost:8000
 `qwen3_tts/batch_speech_client.py` issues many concurrent requests for throughput measurement.
 
 ### Notes
-- Base voice cloning has uniproc-vs-mp tradeoffs depending on per-request reference audio cost; see the executor-backend section above.
-- With async chunking, Qwen3-TTS Base voice cloning sends the full reference context in the first Code2Wav packet, then caches that prefix on the Code2Wav stage for follow-up chunks in the same request.
+- Base voice cloning has per-request reference-audio cost; the uniproc default keeps IPC overhead off the critical path. See the executor-backend section above for background.
 - `vllm_omni/deploy/qwen3_tts.yaml` is the default deploy config (loaded by HF `model_type`); per-stage runtime overrides are available via `--stage-N-<field> <value>`.
-- Under vocoder-bound overload (single-stream `rtf_p99 ≥ 1` at the target concurrency), set `active_stream_window: 2` at the top of the deploy yaml to cap simultaneously active Stage 1 streams. Off by default; trades TTFP for streaming continuity. See [#3592](https://github.com/vllm-project/vllm-omni/pull/3592) for the mechanism and tradeoff numbers.
 
 ---
 
@@ -475,23 +398,6 @@ python voxcpm2/openai_speech_client.py \
 ```
 The `ref_audio` field accepts local file paths (auto-base64), HTTP URLs, or `data:audio/wav;base64,...` data URIs.
 
-### Precomputed custom voices
-For repeated VoxCPM2 speakers, precompute the prompt cache and load it through `custom_voice_dir`:
-```bash
-python voxcpm2/precompute_custom_voice.py \
-    --model openbmb/VoxCPM2 \
-    --voice-name alice \
-    --ref-audio /path/to/reference.wav \
-    --mode ref_continuation \
-    --prompt-text "Original transcript of the reference audio" \
-    --output-dir /path/to/custom_voices
-```
-Add the output directory to the deploy config:
-```yaml
-custom_voice_dir: /path/to/custom_voices
-```
-After startup, `/v1/audio/voices` lists `alice`, and `/v1/audio/speech` can use `voice="alice"` without sending `ref_audio`.
-
 ### Gradio demo (gapless streaming via AudioWorklet)
 ```bash
 python voxcpm2/gradio_demo.py
@@ -553,13 +459,6 @@ The demo handles voice-preset selection and reference-audio upload. `voxtral_tts
     ``````py
     --8<-- "examples/online_serving/text_to_speech/fish_speech/speech_client.py"
     ``````
-??? abstract "higgs_audio_v2/batch_speech_client.py"
-    ``````py
-    --8<-- "examples/online_serving/text_to_speech/higgs_audio_v2/batch_speech_client.py"
-    ``````
-??? abstract "higgs_audio_v2/run_server.sh"
-    ``````sh
-    --8<-- "examples/online_serving/text_to_speech/higgs_audio_v2/run_server.sh"
 ??? abstract "glm_tts/gradio_demo.py"
     ``````py
     --8<-- "examples/online_serving/text_to_speech/glm_tts/gradio_demo.py"

docs/user_guide/quantization/autoround.md

@@ -52,8 +52,8 @@ vLLM-Omni's runtime module names.
 
 | Model | Scope | Status | Notes |
 |-------|-------|--------|-------|
-| GLM-Image | Diffusion transformer  | ✅ | `Intel/GLM-Image-int4-AutoRound` |
 | BAGEL | Checkpoint-defined diffusion or transformer stage | Not validated | Requires a compatible AutoRound checkpoint |
+| GLM-Image | Checkpoint-defined diffusion or transformer stage | Not validated | Requires a compatible AutoRound checkpoint |
 
 ## Configuration
 

docs/user_guide/quantization/gguf.md

@@ -74,7 +74,7 @@ Online:
 vllm serve Qwen/Qwen-Image \
   --omni \
   --port 8000 \
-  --diffusion-quantization-config '{"method":"gguf","gguf_model":"QuantStack/Qwen-Image-GGUF/Qwen_Image-Q4_K_M.gguf"}'
+  --quantization-config '{"method":"gguf","gguf_model":"QuantStack/Qwen-Image-GGUF/Qwen_Image-Q4_K_M.gguf"}'
 ```
 
 ## Parameters

docs/user_guide/quantization/mxfp8.md

@@ -7,28 +7,26 @@ using the OCP MX format: groups of 32 K-dimension elements share a single
 `float8_e8m0fnu` exponent scale. This gives better accuracy than channel-wise
 FP8 while keeping the same 8-bit weight footprint.
 
-This method supports three modes:
+This method supports two modes:
 
 | Mode | Description |
 |------|-------------|
 | **Online** | BF16 weights are quantized to MXFP8 at load time — no pre-processing needed |
-| **Offline (Native)** | msModelSlim-exported MXFP8 weights converted to diffusers format via `merge_mxfp8_checkpoint.py` — weights and scales are loaded directly from the preprocessed checkpoint |
-| **Offline (AutoRound)** | AutoRound MXFP8 checkpoints with `data_type="mx_fp"` — auto-detected from `config.json` |
+| **Offline** | msModelSlim-exported MXFP8 weights converted to diffusers format via `merge_mxfp8_checkpoint.py` — weights and scales are loaded directly from the preprocessed checkpoint |
 
 ## Hardware Support
 
-| Device | Online | Offline (Native) | Offline (AutoRound) |
-|--------|--------|------------------|---------------------|
-| NVIDIA Blackwell GPU (SM 100+) | ⭕ | ⭕ | ⭕ |
-| NVIDIA Ada/Hopper GPU (SM 89+) | ⭕ | ⭕ | ⭕ |
-| NVIDIA Ampere GPU (SM 80+) | ⭕ | ⭕ | ⭕ |
-| AMD ROCm | ⭕ | ⭕ | ⭕ |
-| Intel XPU | ✅ | ❌ | ✅ |
-| Ascend NPU (Atlas 950 A5) | ✅ | ✅ | ⭕ |
+| Device | Support |
+|--------|---------|
+| NVIDIA Blackwell GPU (SM 100+) | ⭕ |
+| NVIDIA Ada/Hopper GPU (SM 89+) | ⭕ |
+| NVIDIA Ampere GPU (SM 80+) | ⭕ |
+| AMD ROCm | ⭕ |
+| Intel XPU | ⭕ |
+| Ascend NPU (Atlas 950 A5) | ✅ |
 
-Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this guide.
-
-**Note**: Intel XPU only supports AutoRound MXFP8 for offline mode. Use AutoRound quantized checkpoints or online mode for XPU.
+Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
+guide.
 
 ## Model Type Support
 
@@ -84,9 +82,9 @@ python text_to_video.py --model <your-model> --quantization mxfp8
 vllm serve <your-model> --omni --quantization mxfp8
 ```
 
-### Offline Mode (Native)
+### Offline Mode
 
-Native offline mode loads a pre-quantized checkpoint from msModelSlim. A preprocessing
+Offline mode loads a pre-quantized checkpoint from msModelSlim. A preprocessing
 step converts the raw quantized output to the diffusers format expected by
 vLLM-Omni and injects the quantization config into `transformer/config.json` so
 that vLLM-Omni auto-detects the offline path without a `--quantization` flag.
@@ -161,41 +159,10 @@ omni = Omni(model="/path/to/Wan2.2-TI2V-5B-MXFP8")
 ```
 
 !!! note
-    No `--quantization` flag is needed for native offline mode. The preprocessing
+    No `--quantization` flag is needed for offline mode. The preprocessing
     script injects `quantization_config` into each `transformer/config.json`,
     which vLLM-Omni reads automatically to activate the offline MXFP8 method.
 
-### Offline Mode (AutoRound)
-
-AutoRound MXFP8 checkpoints declare `quant_method="auto-round"` with `data_type="mx_fp"` in their `config.json`. These are automatically detected and use the `IncMxfp8OfflineLinearMethod` backend.
-
-To use an AutoRound MXFP8 checkpoint:
-
-```bash
-python text_to_video.py --model <autoround-mxfp8-model>
-
-# Online serving
-vllm serve <autoround-mxfp8-model> --omni
-```
-
-Python API:
-
-```python
-omni = Omni(model="<autoround-mxfp8-model>")
-```
-
-!!! note
-    AutoRound MXFP8 checkpoints are auto-detected from `config.json` and do not require a `--quantization` flag. The config must include:
-    ```json
-    {
-      "quantization_config": {
-        "quant_method": "auto-round",
-        "data_type": "mx_fp",
-        ...
-      }
-    }
-    ```
-
 ## Parameters
 
 | Parameter | Type | Default | Description |

docs/user_guide/quantization/online.md

@@ -21,7 +21,7 @@ by msModelSlim and the merge tools.
 | NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ | ⭕ | ⭕ |
 | NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ | ⭕ | ⭕ |
 | AMD ROCm | ⭕ | ⭕ | ⭕ | ⭕ |
-| Intel XPU | ⭕ | ⭕ | ✅ | ⭕ |
+| Intel XPU | ⭕ | ⭕ | ⭕ | ⭕ |
 | Ascend NPU | ❌ | ✅ | ✅ | ✅ |
 
 Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
@@ -36,7 +36,7 @@ MXFP4 are documented for the Ascend NPU path.
 |--------|-------|----------------|--------|
 | FP8 W8A8 | [FP8](fp8.md) | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image family and other DiT models |
 | Int8 W8A8 | [Int8](int8.md) | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image and Z-Image |
-| MXFP8 W8A8 | [MXFP8](mxfp8.md) | Wan2.2-T2V-A14B, Wan2.2-I2V-A14B, Wan2.2-TI2V-5B | Validated on Ascend NPU and Intel XPU |
+| MXFP8 W8A8 | [MXFP8](mxfp8.md) | Wan2.2-T2V-A14B, Wan2.2-I2V-A14B, Wan2.2-TI2V-5B | Ascend NPU only |
 | MXFP4 W4A4 | [MXFP4](mxfp4.md) | Wan2.2-T2V-A14B, Wan2.2-I2V-A14B | Ascend NPU only; TI2V-5B is not supported |
 
 ### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)

examples/offline_inference/bagel/README.md

@@ -195,101 +195,6 @@ stages:
     devices: "0,1"
 ```
 
-### VAE Patch Parallelism
-
-[VAE Patch Parallelism](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/parallelism/vae_patch_parallel.html) splits Bagel VAE **decode/encode** tiles across multiple GPUs on the **DiT stage**, reducing **per-GPU peak memory during VAE decode**. Use it when high-resolution `text2img` or `img2img` hits VAE OOM or large decode spikes.
-
-**Bagel-specific notes:**
-
-- Implemented in `BagelPipeline` via `DistributedAutoEncoder` (DiT stage only).
-- **Single-stage** is the simplest path: one DiT process with TP + VAE patch parallel.
-- **Two-stage**: enable on **stage 1 (DiT)** only; stage 0 (Thinker) keeps encoder-only `VAEEncoder` and does not use VAE patch parallel.
-- You need a DiT `world_size` ≥ `vae_patch_parallel_size` (typically `tensor_parallel_size=2` on that stage). VAE PP reuses the DiT process group; it is not a standalone second-GPU VAE worker.
-
-**Single-stage via deploy YAML** (recommended for `end2end.py`):
-
-```yaml
-pipeline: bagel_single_stage
-async_chunk: false
-
-stages:
-  - stage_id: 0
-    max_num_batched_tokens: 32768
-    max_num_seqs: 1
-    enforce_eager: true
-    trust_remote_code: true
-    enable_prefix_caching: false
-    devices: "0,1"
-    vae_use_tiling: true
-    parallel_config:
-      tensor_parallel_size: 2
-      vae_patch_parallel_size: 2
-    default_sampling_params:
-      seed: 52
-```
-
-```bash
-cd examples/offline_inference/bagel
-
-CUDA_VISIBLE_DEVICES=0,1 python end2end.py \
-    --model /path/to/BAGEL-7B-MoT \
-    --deploy-config /path/to/bagel_single_stage_vae_pp.yaml \
-    --modality text2img \
-    --prompts "A cute cat" \
-    --steps 10 \
-    --output ./out_vae_pp
-```
-
-**Single-stage via `Omni` kwargs** (same flags as online serving):
-
-```python
-from vllm_omni.entrypoints.omni import Omni
-
-omni = Omni(
-    model="ByteDance-Seed/BAGEL-7B-MoT",
-    deploy_config="vllm_omni/deploy/bagel_single_stage.yaml",
-    tensor_parallel_size=2,
-    vae_patch_parallel_size=2,
-    vae_use_tiling=True,
-)
-# Then call omni.generate(...) as in end2end.py
-```
-
-**Two-stage (VAE PP on DiT only):**
-
-```yaml
-stages:
-  - stage_id: 0
-    devices: "0"
-    # AR Thinker — no vae_patch_parallel here
-
-  - stage_id: 1
-    devices: "0,1"
-    vae_use_tiling: true
-    parallel_config:
-      tensor_parallel_size: 2
-      vae_patch_parallel_size: 2
-```
-
-```bash
-python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
-    --deploy-config /path/to/bagel_vae_pp.yaml \
-    --modality text2img \
-    --prompts "A cute cat"
-```
-
-**Startup log checks:**
-
-```text
-INFO ... vae_patch_parallel_size=2 requires vae_use_tiling; automatically enabling it.
-```
-
-| Setting | Role |
-| :------ | :--- |
-| `parallel_config.tensor_parallel_size` | DiT world size / TP (must be ≥ `vae_patch_parallel_size`) |
-| `parallel_config.vae_patch_parallel_size` | Number of ranks for distributed VAE tiles (`1` = off) |
-| `vae_use_tiling` | Enable spatial tiling (auto-enabled when `vae_patch_parallel_size > 1`) |
-
 #### Hybrid Sharded Data Parallel (HSDP)
 
 For larger Bagel deployments on multiple GPUs, you can enable HSDP (Hybrid Sharded Data Parallel) by modifying the stage configuration (for example, [`bagel.yaml`](../../../vllm_omni/deploy/bagel.yaml)). HSDP shards transformer weights across GPUs to reduce per-GPU memory usage.

examples/offline_inference/bagel/end2end.py

@@ -121,6 +121,9 @@ def parse_args():
         help="Temperature for text generation sampling (default: 0.3).",
     )
 
+    from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
+
+    nullify_stage_engine_defaults(parser)
     args = parser.parse_args()
     return args
 
@@ -151,7 +154,6 @@ def main():
     from vllm_omni.entrypoints.omni import Omni
 
     omni_kwargs = vars(args).copy()
-    omni_kwargs.pop("model", None)
     deploy_config = args.deploy_config
     if args.think and deploy_config is None:
         deploy_config = "vllm_omni/deploy/bagel_think.yaml"
@@ -162,7 +164,9 @@ def main():
     if args.quantization:
         omni_kwargs["quantization_config"] = args.quantization
 
-    omni = Omni(model=model_name, **omni_kwargs)
+    # Override CLI --model with the derived model_name.
+    omni_kwargs["model"] = model_name
+    omni = Omni(**omni_kwargs)
 
     formatted_prompts = []
     for p in prompts:

examples/offline_inference/dreamzero/export_prediction_video.py

@@ -1,304 +0,0 @@
-#!/usr/bin/env python3
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from __future__ import annotations
-
-import argparse
-import uuid
-from pathlib import Path
-
-import cv2
-import numpy as np
-import torch
-from PIL import Image
-
-from vllm_omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-from vllm_omni.outputs import OmniRequestOutput
-
-WORKER_EXTENSION = "vllm_omni.diffusion.models.dreamzero.video_export_worker.DreamZeroVideoExportWorkerExtension"
-ASSET_REPO_ID = "YangshenDeng/vllm-omni-dreamzero-assets"
-DEFAULT_SESSION_PREFIX = "dreamzero-export"
-RELATIVE_OFFSETS = [-23, -16, -8, 0]
-ACTION_HORIZON = 24
-CAMERA_FILES = {
-    "observation/exterior_image_0_left": "exterior_image_1_left.mp4",
-    "observation/exterior_image_1_left": "exterior_image_2_left.mp4",
-    "observation/wrist_image_left": "wrist_image_left.mp4",
-}
-
-
-def _parse_args() -> argparse.Namespace:
-    repo_root = Path(__file__).resolve().parents[3]
-    parser = argparse.ArgumentParser(description="Export DreamZero prediction video from downloaded example inputs.")
-    parser.add_argument("--model", default="GEAR-Dreams/DreamZero-DROID")
-    parser.add_argument("--deploy-config", type=Path, required=True)
-    parser.add_argument(
-        "--video-dir",
-        type=Path,
-        default=repo_root / "outputs" / "dreamzero" / "assets",
-        help="Directory containing the three camera MP4 files.",
-    )
-    parser.add_argument(
-        "--output-dir", type=Path, default=repo_root / "outputs" / "dreamzero" / "generated_predictions"
-    )
-    parser.add_argument("--output-stem", default="dreamzero_prediction")
-    parser.add_argument(
-        "--prompt",
-        default="Move the pan forward and use the brush in the middle of the plates to brush the inside of the pan",
-    )
-    parser.add_argument("--session-id", default=None)
-    parser.add_argument("--save-input-video", action="store_true")
-    parser.add_argument("--save-gif", action="store_true")
-    parser.add_argument("--save-actions", action="store_true")
-    parser.add_argument("--fps", type=int, default=5)
-    return parser.parse_args()
-
-
-def _load_all_frames(video_path: Path) -> np.ndarray:
-    cap = cv2.VideoCapture(str(video_path))
-    frames = []
-    while True:
-        ok, frame = cap.read()
-        if not ok:
-            break
-        frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
-    cap.release()
-    if not frames:
-        raise RuntimeError(f"No frames loaded from {video_path}")
-    return np.stack(frames, axis=0)
-
-
-def _load_camera_frames(video_dir: Path) -> dict[str, np.ndarray]:
-    camera_frames: dict[str, np.ndarray] = {}
-    for camera_key, file_name in CAMERA_FILES.items():
-        video_path = video_dir / file_name
-        if not video_path.exists():
-            raise FileNotFoundError(
-                f"Missing DreamZero example asset: {video_path}. "
-                "Download the example videos with: "
-                f"`hf download {ASSET_REPO_ID} --repo-type dataset --local-dir {video_dir}`"
-            )
-        camera_frames[camera_key] = _load_all_frames(video_path)
-    return camera_frames
-
-
-def _build_frame_schedule(total_frames: int, num_chunks: int) -> list[list[int]]:
-    chunks: list[list[int]] = []
-    current_frame = 23
-    for _ in range(num_chunks):
-        indices = [max(current_frame + offset, 0) for offset in RELATIVE_OFFSETS]
-        if indices[-1] >= total_frames:
-            break
-        chunks.append(indices)
-        current_frame += ACTION_HORIZON
-    return chunks
-
-
-def _make_obs_from_video(
-    camera_frames: dict[str, np.ndarray],
-    frame_indices: list[int],
-    *,
-    prompt: str,
-    session_id: str,
-) -> dict:
-    obs: dict = {}
-    for camera_key, all_frames in camera_frames.items():
-        selected = all_frames[frame_indices]
-        obs[camera_key] = selected[0] if len(frame_indices) == 1 else selected
-
-    obs["observation/joint_position"] = np.zeros(7, dtype=np.float32)
-    obs["observation/cartesian_position"] = np.zeros(6, dtype=np.float32)
-    obs["observation/gripper_position"] = np.zeros(1, dtype=np.float32)
-    obs["prompt"] = prompt
-    obs["session_id"] = session_id
-    return obs
-
-
-def _build_observations(video_dir: Path, prompt: str, session_id: str) -> tuple[dict[str, np.ndarray], list[dict]]:
-    camera_frames = _load_camera_frames(video_dir)
-    total_frames = min(frames.shape[0] for frames in camera_frames.values())
-    chunks = _build_frame_schedule(total_frames, 1)
-    observations = [
-        _make_obs_from_video(camera_frames, [0], prompt=prompt, session_id=session_id),
-    ]
-    if chunks:
-        observations.append(
-            _make_obs_from_video(
-                camera_frames,
-                chunks[0],
-                prompt=prompt,
-                session_id=session_id,
-            )
-        )
-    if len(observations) < 2:
-        raise RuntimeError("Need at least two DreamZero example observations to export a prediction video.")
-    return camera_frames, observations[:2]
-
-
-def _extract_latents(output: OmniRequestOutput) -> torch.Tensor:
-    if not isinstance(output, OmniRequestOutput):
-        raise TypeError(f"Expected OmniRequestOutput, got {type(output)!r}")
-    if not output.images:
-        raise RuntimeError("DreamZero output does not contain video latents in `images`.")
-
-    latents = output.images[0]
-    if not isinstance(latents, torch.Tensor):
-        raise TypeError(f"Expected tensor latents, got {type(latents)!r}")
-
-    latents = latents.detach().cpu()
-    if latents.dim() == 4:
-        latents = latents.unsqueeze(0)
-    if latents.dim() != 5:
-        raise ValueError(f"Unexpected latent shape: {tuple(latents.shape)}")
-
-    if latents.shape[1] < latents.shape[2]:
-        latents = latents.transpose(1, 2).contiguous()
-    return latents
-
-
-def _write_mp4(path: Path, frames: np.ndarray, fps: int) -> None:
-    path.parent.mkdir(parents=True, exist_ok=True)
-    height, width = frames.shape[1:3]
-    writer = cv2.VideoWriter(
-        str(path),
-        cv2.VideoWriter_fourcc(*"mp4v"),
-        float(fps),
-        (width, height),
-    )
-    if not writer.isOpened():
-        raise RuntimeError(f"Failed to open video writer for {path}")
-    try:
-        for frame in frames:
-            writer.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
-    finally:
-        writer.release()
-
-
-def _write_gif(path: Path, frames: np.ndarray, fps: int) -> None:
-    path.parent.mkdir(parents=True, exist_ok=True)
-    images = [Image.fromarray(frame) for frame in frames]
-    duration_ms = max(int(round(1000 / max(fps, 1))), 1)
-    images[0].save(
-        path,
-        save_all=True,
-        append_images=images[1:],
-        duration=duration_ms,
-        loop=0,
-    )
-
-
-def _stitch_input_frames(camera_frames: dict[str, np.ndarray]) -> np.ndarray:
-    total_frames = min(frames.shape[0] for frames in camera_frames.values())
-    stitched = []
-    for frame_index in range(total_frames):
-        left = camera_frames["observation/exterior_image_0_left"][frame_index]
-        right = camera_frames["observation/exterior_image_1_left"][frame_index]
-        wrist = camera_frames["observation/wrist_image_left"][frame_index]
-        pad = np.zeros((left.shape[0], left.shape[1], 3), dtype=np.uint8)
-        canvas = np.concatenate([left, right], axis=1)
-        bottom = np.concatenate([wrist, pad], axis=1)
-        stitched.append(np.concatenate([canvas, bottom], axis=0))
-    return np.stack(stitched, axis=0)
-
-
-def _run_generation(
-    model: str, deploy_config_path: Path, observations: list[dict]
-) -> tuple[Omni, list[OmniRequestOutput]]:
-    omni = Omni(
-        model=model,
-        deploy_config=str(deploy_config_path),
-        enforce_eager=True,
-        worker_extension_cls=WORKER_EXTENSION,
-    )
-
-    outputs: list[OmniRequestOutput] = []
-    for index, obs in enumerate(observations):
-        sampling_params = OmniDiffusionSamplingParams(
-            extra_args={
-                "reset": index == 0,
-                "session_id": obs["session_id"],
-                "robot_obs": obs,
-            }
-        )
-        result = omni.generate(obs["prompt"], sampling_params_list=[sampling_params])
-        if not result:
-            raise RuntimeError(f"No output returned for DreamZero request {index}")
-        outputs.append(result[0])
-    return omni, outputs
-
-
-def _decode_with_worker(omni: Omni, full_latents: torch.Tensor) -> np.ndarray:
-    stage_client = omni.engine.stage_clients[0]
-    engine = getattr(stage_client, "_engine", None)
-    if engine is None:
-        raise RuntimeError("DreamZero export requires inline diffusion stage access.")
-
-    decoded = engine.executor.collective_rpc(
-        "decode_video_latents_to_uint8",
-        args=(full_latents,),
-        unique_reply_rank=0,
-        exec_all_ranks=True,
-    )
-    if isinstance(decoded, torch.Tensor):
-        decoded = decoded.numpy()
-    if not isinstance(decoded, np.ndarray):
-        raise TypeError(f"Unexpected decoded output type: {type(decoded)!r}")
-    return decoded
-
-
-def main() -> None:
-    args = _parse_args()
-    session_id = args.session_id or f"{DEFAULT_SESSION_PREFIX}-{uuid.uuid4()}"
-
-    camera_frames, observations = _build_observations(
-        video_dir=args.video_dir,
-        prompt=args.prompt,
-        session_id=session_id,
-    )
-
-    args.output_dir.mkdir(parents=True, exist_ok=True)
-
-    if args.save_input_video:
-        input_frames = _stitch_input_frames(camera_frames)
-        _write_mp4(args.output_dir / f"{args.output_stem}_input.mp4", input_frames, fps=15)
-        if args.save_gif:
-            _write_gif(args.output_dir / f"{args.output_stem}_input.gif", input_frames[::3], fps=5)
-
-    omni = None
-    try:
-        omni, outputs = _run_generation(
-            model=args.model,
-            deploy_config_path=args.deploy_config,
-            observations=observations,
-        )
-        latent_steps = [_extract_latents(output) for output in outputs]
-        full_latents = torch.cat(latent_steps, dim=2)
-        frames = _decode_with_worker(omni, full_latents)
-    finally:
-        if omni is not None:
-            omni.close()
-
-    mp4_path = args.output_dir / f"{args.output_stem}.mp4"
-
-    _write_mp4(mp4_path, frames, fps=args.fps)
-    print(f"SAVED_MP4={mp4_path}")
-
-    if args.save_gif:
-        gif_path = args.output_dir / f"{args.output_stem}.gif"
-        _write_gif(gif_path, frames, fps=args.fps)
-        print(f"SAVED_GIF={gif_path}")
-
-    if args.save_actions:
-        npz_path = args.output_dir / f"{args.output_stem}_actions.npz"
-        np.savez(
-            npz_path,
-            step0=np.asarray(outputs[0].multimodal_output.get("actions")),
-            step1=np.asarray(outputs[1].multimodal_output.get("actions")),
-        )
-        print(f"SAVED_ACTIONS={npz_path}")
-
-
-if __name__ == "__main__":
-    main()

examples/offline_inference/dynin_omni/README.md

@@ -42,7 +42,7 @@ pip install \
 ## 3. Hardware and VRAM Requirements
 
 This example uses a 3-stage pipeline on one GPU by default
-([`dynin_omni.yaml`](../../../vllm_omni/deploy/dynin_omni.yaml)):
+([`dynin_omni.yaml`](../../../vllm_omni/model_executor/stage_configs/dynin_omni.yaml)):
 
 - Stage-0 (`token2text`): `gpu_memory_utilization: 0.5`
 - Stage-1 (`token2image`): `gpu_memory_utilization: 0.1`

examples/offline_inference/dynin_omni/end2end.py

@@ -18,6 +18,8 @@ import numpy as np
 import torch
 from PIL import Image
 
+from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
+
 TASK_CHOICES = ("t2t", "t2i", "t2s", "i2i", "i2t", "s2t", "v2t")
 
 TASK_DEFAULT_RUNTIME = {
@@ -878,8 +880,8 @@ def parse_args(repo_root: Path) -> argparse.Namespace:
     parser.add_argument(
         "--stage-config-path",
         type=str,
-        default=str(repo_root / "vllm_omni/deploy/dynin_omni.yaml"),
-        help="Path to deploy config yaml.",
+        default=str(repo_root / "vllm_omni/model_executor/stage_configs/dynin_omni.yaml"),
+        help="Path to stage config yaml.",
     )
     parser.add_argument(
         "--dynin-config-path",
@@ -970,6 +972,8 @@ def parse_args(repo_root: Path) -> argparse.Namespace:
     parser.add_argument("--vq-model-audio-local-files-only", action=argparse.BooleanOptionalAction, default=None)
 
     parser.add_argument("--disable-hf-xet", action=argparse.BooleanOptionalAction, default=True)
+
+    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 

examples/offline_inference/hunyuan_image3/end2end.py

@@ -17,7 +17,7 @@ from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniPromptType
 
 _REPO_ROOT = Path(__file__).resolve().parents[3]
-_DEFAULT_DEPLOY_CONFIG = str(_REPO_ROOT / "vllm_omni" / "deploy" / "hunyuan_image_3_moe.yaml")
+_DEFAULT_DEPLOY_CONFIG = str(_REPO_ROOT / "vllm_omni" / "deploy" / "hunyuan_image3.yaml")
 _DEFAULT_AR_DEPLOY_CONFIG = str(_REPO_ROOT / "vllm_omni" / "deploy" / "hunyuan_image3_ar.yaml")
 
 _MODALITY_TASK_MAP: dict[str, tuple[str, str | None]] = {
@@ -61,8 +61,8 @@ def parse_args():
     parser.add_argument("--steps", type=int, default=50, help="Number of inference steps.")
     parser.add_argument("--guidance-scale", type=float, default=5.0, help="Classifier-free guidance scale.")
     parser.add_argument("--seed", type=int, default=42, help="Random seed.")
-    parser.add_argument("--height", type=int, default=None, help="Output image height.")
-    parser.add_argument("--width", type=int, default=None, help="Output image width.")
+    parser.add_argument("--height", type=int, default=1024, help="Output image height.")
+    parser.add_argument("--width", type=int, default=1024, help="Output image width.")
     parser.add_argument("--vae-use-tiling", action="store_true", help="Enable VAE tiling.")
     parser.add_argument(
         "--bot-task",
@@ -107,6 +107,9 @@ def parse_args():
         ),
     )
 
+    from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
+
+    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 
@@ -226,21 +229,7 @@ def main():
 
     from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
-    if (args.height is None) != (args.width is None):
-        raise ValueError("--height and --width must both be specified or both omitted.")
-    user_specified_size = args.height is not None and args.width is not None
-    if args.modality in ("img2text", "text2text"):
-        ar_image_size = "auto"
-    elif user_specified_size:
-        ar_image_size = f"{args.width}x{args.height}"
-    else:
-        ar_image_size = None
-    ar_stop_token_ids = resolve_stop_token_ids(
-        task=task, bot_task=bot_task, tokenizer=tokenizer, image_size=ar_image_size
-    )
-    print(
-        f"[AR Config] task={task}, bot_task={bot_task}, image_size={ar_image_size}, stop_token_ids={ar_stop_token_ids}"
-    )
+    ar_stop_token_ids = resolve_stop_token_ids(task=task, bot_task=bot_task, tokenizer=tokenizer)
     for sp in params_list:
         if isinstance(sp, OmniDiffusionSamplingParams):
             sp.num_inference_steps = args.steps

examples/offline_inference/internvla_a1/end2end.py

@@ -127,12 +127,7 @@ def run_one_path(
             ),
             args.device,
         )
-        pred = run_pipeline_forward(
-            pipeline,
-            batch_inputs,
-            noise,
-            request_id=f"internvla-a1-sample-{index}",
-        )
+        pred = run_pipeline_forward(pipeline, batch_inputs, noise)
         pred = pred[:, :, : dataset.physical_action_dim].to(torch.float32).cpu()
         results.append(
             {
@@ -155,7 +150,6 @@ def run_pipeline_forward(
     pipeline,
     batch_inputs: dict[str, torch.Tensor],
     noise: torch.Tensor,
-    request_id: str,
 ) -> torch.Tensor:
     output = pipeline.forward(
         OmniDiffusionRequest(
@@ -167,7 +161,6 @@ def run_pipeline_forward(
                     "decode_image": False,
                 }
             ),
-            request_id=request_id,
         )
     )
     if output.error:
@@ -219,38 +212,23 @@ def benchmark_forward(
 
     _synchronize(args.device)
     cold_start_begin = time.perf_counter()
-    pred = run_pipeline_forward(
-        pipeline,
-        batch_inputs,
-        noise,
-        request_id=f"internvla-a1-benchmark-{index}-cold",
-    )
+    pred = run_pipeline_forward(pipeline, batch_inputs, noise)
     _synchronize(args.device)
     cold_start_ms = (time.perf_counter() - cold_start_begin) * 1000.0
 
     warmup_ms: list[float] = []
-    for iter_idx in range(args.warmup_iters):
+    for _ in range(args.warmup_iters):
         _synchronize(args.device)
         begin = time.perf_counter()
-        _ = run_pipeline_forward(
-            pipeline,
-            batch_inputs,
-            noise,
-            request_id=f"internvla-a1-benchmark-{index}-warmup-{iter_idx}",
-        )
+        _ = run_pipeline_forward(pipeline, batch_inputs, noise)
         _synchronize(args.device)
         warmup_ms.append((time.perf_counter() - begin) * 1000.0)
 
     benchmark_ms: list[float] = []
-    for iter_idx in range(args.benchmark_iters):
+    for _ in range(args.benchmark_iters):
         _synchronize(args.device)
         begin = time.perf_counter()
-        _ = run_pipeline_forward(
-            pipeline,
-            batch_inputs,
-            noise,
-            request_id=f"internvla-a1-benchmark-{index}-iter-{iter_idx}",
-        )
+        _ = run_pipeline_forward(pipeline, batch_inputs, noise)
         _synchronize(args.device)
         benchmark_ms.append((time.perf_counter() - begin) * 1000.0)
 
@@ -317,12 +295,7 @@ def main() -> None:
             dataset=dataset,
             train_meta=train_meta,
             collate_samples=collate_open_loop_samples,
-            run_sample_actions=lambda policy, batch_inputs, noise: run_pipeline_forward(
-                policy,
-                batch_inputs,
-                noise,
-                request_id="internvla-a1-open-loop",
-            ),
+            run_sample_actions=lambda policy, batch_inputs, noise: run_pipeline_forward(policy, batch_inputs, noise),
             num_episodes=args.num_episodes,
             seed=args.seed,
             device=args.device,

examples/offline_inference/lance/README.md

@@ -1,52 +0,0 @@
-# Lance (ByteDance) — offline inference
-
-[Lance](https://huggingface.co/bytedance-research/Lance) is a 3B unified
-autoregressive + diffusion multimodal model on a Qwen2.5-VL backbone. It is
-**BAGEL-lineage** (ByteDance Mixture-of-Transformers): the released `Lance_3B`
-checkpoint uses the same `*_moe_gen` MoT weight layout as BAGEL, so vLLM-Omni
-implements it by reusing the BAGEL transformer core and specializing only the
-ViT (Qwen2.5-VL vision), the VAE (Wan2.2) and the checkpoint layout.
-
-This example covers all six Lance modalities from the upstream HF model card:
-`t2i`, `t2v`, `image_edit`, `video_edit`, `x2t_image` (image understanding) and
-`x2t_video` (video understanding).
-
-## Hardware
-
-Single NVIDIA GPU with 16 GB+ VRAM in BF16 (we test on B300 / A100). CUDA ≥ 12.4.
-
-## Run
-
-```bash
-# Text-to-image
-python examples/offline_inference/lance/end2end.py \
-    --model bytedance-research/Lance \
-    --prompts "a corgi astronaut on the moon, cinematic" \
-    --steps 30 --cfg-text-scale 4.0 --timestep-shift 3.5 \
-    --height 1024 --width 1024 \
-    --output ./out
-
-# Text-to-video (uses the Lance_3B_Video subfolder; see ``--modality``
-# choices for all six task variants)
-python examples/offline_inference/lance/end2end.py \
-    --model bytedance-research/Lance/Lance_3B_Video --modality text2video \
-    --num-frames 25 --video-height 480 --video-width 768 \
-    --prompts "a cat playing piano, cinematic" \
-    --steps 30 --fps 8 --output ./out
-```
-
-`video_edit` requires `--model bytedance-research/Lance/Lance_3B_Video` so
-the 3-D `latent_pos_embed` table is loaded; the other paths can point at
-the top-level `bytedance-research/Lance` repo and resolve the right
-sub-checkpoint automatically.
-
-The HF repo bundles everything (`Lance_3B/`, `Lance_3B_Video/`,
-`Qwen2.5-VL-ViT/`, `Wan2.2_VAE.pth`); no separate downloads are required.
-
-## Defaults
-
-Matches upstream `inference_lance.sh`: 30 denoising steps, timestep-shift 3.5,
-text CFG 4.0, seed 42, 1024×1024 (override with `--height` / `--width`).  For
-the understanding paths (`img2text` / `video2text`), sampling is enabled by
-default at `--text-temperature 0.8` because Lance's greedy decoder emits an
-immediate EOS for many prompts.

examples/offline_inference/lance/end2end.py

@@ -1,339 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Lance-3B offline inference (single-stage).
-
-Lance is BAGEL-lineage; this example covers the single-stage paths wired in
-this PR:
-
-    text2img  — image generation from text
-    img2text  — image understanding (caption / VQA)
-
-Examples:
-
-    python examples/offline_inference/lance/end2end.py \
-        --model bytedance-research/Lance \
-        --prompts "a corgi astronaut on the moon, cinematic" \
-        --output ./out
-
-    python examples/offline_inference/lance/end2end.py \
-        --model bytedance-research/Lance --modality img2text \
-        --image-path /path/to/photo.jpg \
-        --prompts "Describe this image in detail."
-"""
-
-import argparse
-import os
-
-
-def parse_args():
-    p = argparse.ArgumentParser()
-    p.add_argument("--model", default="bytedance-research/Lance", help="HF repo or local path.")
-    p.add_argument("--prompts", nargs="+", default=None, help="Input text prompts.")
-    p.add_argument("--txt-prompts", type=str, default=None, help="File with one prompt per line.")
-    p.add_argument(
-        "--modality",
-        default="text2img",
-        choices=["text2img", "img2text", "text2video", "video2text", "img2img", "video2video", "image2video"],
-        help="Lance single-stage modality.",
-    )
-    p.add_argument("--image-path", type=str, default=None, help="Input image for img2text.")
-    p.add_argument("--video-path", type=str, default=None, help="Input video file for video2text.")
-    p.add_argument("--num-frames", type=int, default=25, help="Number of RGB frames for text2video (max 121).")
-    p.add_argument("--video-height", type=int, default=480, help="Video frame height.")
-    p.add_argument("--video-width", type=int, default=768, help="Video frame width.")
-    p.add_argument("--height", type=int, default=None, help="Image height (t2i). Default = max_hw (1024).")
-    p.add_argument("--width", type=int, default=None, help="Image width (t2i). Default = max_hw (1024).")
-    p.add_argument(
-        "--fps", type=int, default=12, help="Output video FPS when saving MP4 (matches upstream Lance's save_fps=12)."
-    )
-    p.add_argument("--output", type=str, default=".", help="Output directory.")
-    p.add_argument("--steps", type=int, default=30, help="Denoising steps (Lance default 30).")
-    p.add_argument("--cfg-text-scale", type=float, default=4.0, help="Text CFG scale (Lance default 4.0).")
-    p.add_argument("--timestep-shift", type=float, default=3.5, help="Flow-match timestep shift (Lance default 3.5).")
-    p.add_argument("--negative-prompt", type=str, default=None)
-    p.add_argument("--seed", type=int, default=42)
-    p.add_argument("--max-text-tokens", type=int, default=512, help="Max tokens for img2text generation.")
-    p.add_argument(
-        "--do-sample",
-        action="store_true",
-        default=True,
-        help="Sample (vs greedy) for x2t generation; greedy frequently emits an immediate EOS for Lance.",
-    )
-    p.add_argument("--no-sample", dest="do_sample", action="store_false", help="Disable sampling (greedy decoding).")
-    p.add_argument(
-        "--text-temperature",
-        type=float,
-        default=0.8,
-        help="Sampling temperature for x2t generation. Lance frequently emits an immediate EOS below ~0.7; 0.8 is a good default.",
-    )
-    p.add_argument(
-        "--system-prompt",
-        type=str,
-        default=None,
-        help=(
-            "Override the task system prompt (used by x2t_image / x2t_video for "
-            'per-example QA instructions like "Look at the image carefully and '
-            'answer the question."; without it x2t falls back to a '
-            "caption-style default and the model describes instead of answering)."
-        ),
-    )
-    # Lance is single-stage diffusion — no deploy YAML needed.  Required
-    # engine knobs (``pipeline``, ``enforce_eager``, ``trust_remote_code``,
-    # ``max_num_seqs=1`` …) are passed as flat kwargs to ``Omni`` below
-    # and ``create_default_diffusion`` materializes the stage config.
-
-    from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
-
-    nullify_stage_engine_defaults(p)
-    return p.parse_args()
-
-
-from vllm_omni.diffusion.models.lance.prompts import (  # noqa: E402
-    VIDEO_PAD,
-    VISION_END,
-    VISION_START,
-    render_lance_prompt,
-)
-
-# Visual block matching upstream Lance: <|vision_start|><|video_pad|><|vision_end|>
-# (Lance's renderer uses ``<|video_pad|>`` even for image inputs by default.)
-_VISION_BLOCK = f"{VISION_START}{VIDEO_PAD}{VISION_END}"
-
-
-def _format_text2img_prompts(prompts):
-    return [{"prompt": render_lance_prompt("t2i", p), "modalities": ["image"]} for p in prompts]
-
-
-def _format_text2video_prompts(prompts, num_frames, video_h, video_w):
-    extra = {"num_frames": num_frames, "video_height": video_h, "video_width": video_w}
-    return [{"prompt": render_lance_prompt("t2v", p), "modalities": ["video"], "extra_args": extra} for p in prompts]
-
-
-def _format_image2video_prompts(prompts, image_path, num_frames, video_h, video_w):
-    """Image-to-Video (no first-frame pin): image + text → long video.
-
-    Passes the input image via ``multi_modal_data.first_frame``.  The pipeline
-    treats it as a 1-frame reference (VAE+ViT prefill) and generates a fresh
-    multi-frame video at the requested ``num_frames`` × ``video_h`` ×
-    ``video_w`` shape.
-    """
-    import os as _os
-
-    if not image_path or not _os.path.exists(image_path):
-        raise ValueError(f"image2video requires --image-path pointing to an existing file, got: {image_path}")
-    from PIL import Image as _Image
-
-    img = _Image.open(image_path).convert("RGB")
-    extra = {"num_frames": num_frames, "video_height": video_h, "video_width": video_w}
-    return [
-        {
-            "prompt": render_lance_prompt("i2v", p, vision_token=_VISION_BLOCK),
-            "multi_modal_data": {"first_frame": img},
-            "modalities": ["video"],
-            "extra_args": extra,
-        }
-        for p in prompts
-    ]
-
-
-def _format_img2img_prompts(prompts, image_path):
-    import os as _os
-
-    if not image_path or not _os.path.exists(image_path):
-        raise ValueError(f"img2img requires --image-path pointing to an existing file, got: {image_path}")
-    from PIL import Image as _Image
-
-    img = _Image.open(image_path).convert("RGB")
-    return [
-        {
-            "prompt": render_lance_prompt("image_edit", p, vision_token=_VISION_BLOCK),
-            "multi_modal_data": {"img2img": img},
-            "modalities": ["image"],
-        }
-        for p in prompts
-    ]
-
-
-def _format_video2video_prompts(prompts, video_path, num_frames, video_h, video_w):
-    import os as _os
-
-    if not video_path or not _os.path.exists(video_path):
-        raise ValueError(f"video2video requires --video-path pointing to an existing file, got: {video_path}")
-
-    extra = {"num_frames": num_frames, "video_height": video_h, "video_width": video_w}
-    return [
-        {
-            "prompt": render_lance_prompt("video_edit", p, vision_token=_VISION_BLOCK),
-            "multi_modal_data": {"video": video_path},
-            "modalities": ["video"],
-            "extra_args": extra,
-        }
-        for p in prompts
-    ]
-
-
-def _format_video2text_prompts(prompts, video_path, system_prompt=None):
-    import os as _os
-
-    if not video_path or not _os.path.exists(video_path):
-        raise ValueError(f"video2text requires --video-path pointing to an existing file, got: {video_path}")
-    import imageio.v3 as iio
-
-    video = iio.imread(video_path)
-    return [
-        {
-            "prompt": render_lance_prompt("x2t_video", p, vision_token=_VISION_BLOCK, system_prompt=system_prompt),
-            "multi_modal_data": {"video": video},
-            "modalities": ["text"],
-        }
-        for p in prompts
-    ]
-
-
-def _format_img2text_prompts(prompts, image_path, system_prompt=None):
-    from PIL import Image
-
-    if not image_path or not os.path.exists(image_path):
-        raise ValueError(f"img2text requires --image-path pointing to an existing file, got: {image_path}")
-    img = Image.open(image_path).convert("RGB")
-    return [
-        {
-            "prompt": render_lance_prompt("x2t_image", p, vision_token=_VISION_BLOCK, system_prompt=system_prompt),
-            "multi_modal_data": {"image": img},
-            "modalities": ["text"],
-        }
-        for p in prompts
-    ]
-
-
-def main():
-    args = parse_args()
-    os.makedirs(args.output, exist_ok=True)
-
-    # Lance aligns more naturally with upstream when using FLASH_ATTN
-    # (upstream uses flash_attn_varlen_func; SDPA accumulates ~5x more
-    # numerical drift through the 36-layer Qwen2 stack on B300).  Default
-    # to flash_attn when the user hasn't explicitly set a backend.
-    os.environ.setdefault("DIFFUSION_ATTENTION_BACKEND", "FLASH_ATTN")
-
-    if args.txt_prompts:
-        with open(args.txt_prompts, encoding="utf-8") as f:
-            prompts = [ln.strip() for ln in f if ln.strip()]
-    elif args.modality == "img2text":
-        prompts = args.prompts or ["Describe this image in detail."]
-    else:
-        prompts = args.prompts or ["A cute cat, studio lighting, 4k"]
-
-    from vllm_omni.entrypoints.omni import Omni
-
-    # Video modalities require the Lance_3B_Video checkpoint (3D
-    # latent_pos_embed shaped (126976, 2048)).  Without it, vllm-omni silently
-    # loads Lance_3B (4096, 2048) — the image-only table — and any t_lat >= 1
-    # immediately indexes out of bounds.  Rewrite the model path so
-    # LancePipeline._select_video_variant picks the video subfolder.
-    resolved_model = args.model
-    if args.modality in ("text2video", "video2video", "video2text", "image2video"):
-        if not str(resolved_model).rstrip("/").endswith("Lance_3B_Video"):
-            candidate = os.path.join(resolved_model, "Lance_3B_Video")
-            if os.path.isdir(candidate):
-                resolved_model = candidate
-
-    omni_kwargs = vars(args).copy()
-    omni_kwargs["model"] = resolved_model
-    # Lance single-stage defaults (formerly in vllm_omni/deploy/lance.yaml):
-    omni_kwargs.setdefault("pipeline", "lance")
-    omni_kwargs.setdefault("max_num_batched_tokens", 32768)
-    omni_kwargs.setdefault("max_num_seqs", 1)
-    omni_kwargs.setdefault("enforce_eager", True)
-    omni_kwargs.setdefault("trust_remote_code", True)
-    omni_kwargs.setdefault("enable_prefix_caching", False)
-    omni_kwargs.setdefault("async_chunk", False)
-    omni = Omni(**omni_kwargs)
-
-    if args.modality == "img2text":
-        formatted = _format_img2text_prompts(prompts, args.image_path, system_prompt=args.system_prompt)
-    elif args.modality == "text2video":
-        formatted = _format_text2video_prompts(prompts, args.num_frames, args.video_height, args.video_width)
-    elif args.modality == "video2text":
-        formatted = _format_video2text_prompts(prompts, args.video_path, system_prompt=args.system_prompt)
-    elif args.modality == "img2img":
-        formatted = _format_img2img_prompts(prompts, args.image_path)
-    elif args.modality == "video2video":
-        formatted = _format_video2video_prompts(
-            prompts, args.video_path, args.num_frames, args.video_height, args.video_width
-        )
-    elif args.modality == "image2video":
-        formatted = _format_image2video_prompts(
-            prompts, args.image_path, args.num_frames, args.video_height, args.video_width
-        )
-    else:
-        formatted = _format_text2img_prompts(prompts)
-
-    params_list = omni.default_sampling_params_list
-    diffusion_params = params_list[0]  # single-stage: one param set
-    diffusion_params.num_inference_steps = args.steps  # type: ignore
-    if args.seed is not None:
-        diffusion_params.seed = args.seed  # type: ignore
-    if args.height is not None:
-        diffusion_params.height = args.height  # type: ignore
-    if args.width is not None:
-        diffusion_params.width = args.width  # type: ignore
-    extra = getattr(diffusion_params, "extra_args", {}) or {}
-    extra["cfg_text_scale"] = args.cfg_text_scale
-    extra["timestep_shift"] = args.timestep_shift
-    if args.modality in ("img2text", "video2text"):
-        extra["max_think_tokens"] = args.max_text_tokens
-        extra["do_sample"] = args.do_sample
-        extra["text_temperature"] = args.text_temperature
-    if args.negative_prompt is not None:
-        extra["negative_prompt"] = args.negative_prompt
-    diffusion_params.extra_args = extra  # type: ignore
-
-    outputs = list(omni.generate(prompts=formatted, sampling_params_list=params_list))
-
-    if args.modality in ("img2text", "video2text"):
-        for i, req_output in enumerate(outputs):
-            text = getattr(req_output, "output", None) or getattr(req_output, "text", None)
-            if text is None and hasattr(req_output, "custom_output"):
-                text = (req_output.custom_output or {}).get("text_output")
-            print(f"[Output {i}] {text}")
-        return
-
-    if args.modality in ("text2video", "video2video", "image2video"):
-        for i, req_output in enumerate(outputs):
-            custom = getattr(req_output, "custom_output", None) or {}
-            frames = custom.get("video_frames")
-            if not frames:
-                print(f"[Output {i}] no video frames returned")
-                continue
-            mp4_path = os.path.join(args.output, f"lance_{i}.mp4")
-            try:
-                import imageio.v3 as iio
-                import numpy as np
-
-                arr = np.stack([np.asarray(f) for f in frames], axis=0)
-                iio.imwrite(mp4_path, arr, fps=args.fps, codec="libx264", quality=8)
-                print(f"[Output {i}] Saved MP4 ({len(frames)} frames, {arr.shape[1]}x{arr.shape[2]}) to {mp4_path}")
-            except Exception as exc:
-                print(f"[Output {i}] MP4 encode failed ({exc!r}); writing per-frame ONGs instead")
-                frame_dir = os.path.join(args.output, f"lance_{i}_frames")
-                os.makedirs(frame_dir, exist_ok=True)
-                for j, frame in enumerate(frames):
-                    frame.save(os.path.join(frame_dir, f"{j:04d}.png"))
-                print(f"  → {frame_dir} ({len(frames)} frames)")
-        return
-
-    idx = 0
-    for req_output in outputs:
-        images = getattr(req_output, "images", None)
-        if not images:
-            continue
-        for j, img in enumerate(images):
-            path = os.path.join(args.output, f"lance_{idx}_{j}.png")
-            img.save(path)
-            print(f"[Output] Saved image to {path}")
-        idx += 1
-
-
-if __name__ == "__main__":
-    main()

examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py

@@ -19,6 +19,7 @@ from vllm import SamplingParams
 from vllm.multimodal.image import convert_image_mode
 
 from vllm_omni import Omni
+from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
 
 DEFAULT_SYSTEM = "You are a helpful assistant."
 DEFAULT_QUESTION = "Please summarize the content of this image."
@@ -48,6 +49,7 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Enable diffusion pipeline profiler to display stage durations.",
     )
+    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 

examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_t2i.py

@@ -29,6 +29,7 @@ from PIL import Image
 from vllm.sampling_params import SamplingParams
 
 from vllm_omni import Omni
+from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
 
 logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 logger = logging.getLogger(__name__)
@@ -117,6 +118,7 @@ def parse_args() -> argparse.Namespace:
     )
     p.add_argument("--out", type=str, default="output.png", help="Path to save the generated image.")
     p.add_argument("--trust-remote-code", action="store_true", help="Trust remote code when loading the model.")
+    nullify_stage_engine_defaults(p)
     args = p.parse_args()
     if not args.prompt:
         args.prompt = ["A stylish woman with sunglasses riding a motorcycle in NYC."]

examples/offline_inference/mimo_audio/end2end.py

@@ -21,10 +21,11 @@ from message_convert import (
     to_prompt,
 )
 from vllm import SamplingParams
+from vllm.utils.argparse_utils import FlexibleArgumentParser
 
+from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniTokensPrompt
-from vllm_omni.utils.tracking_parser import TrackingArgumentParser
 
 SEED = 42
 
@@ -342,7 +343,7 @@ def main(args):
 
 
 def parse_args():
-    parser = TrackingArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
+    parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
     parser.add_argument(
         "--model-name",
         "-m",
@@ -438,6 +439,7 @@ def parse_args():
         "vllm_omni/deploy/mimo_audio.yaml based on the HF model_type.",
     )
 
+    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 

examples/offline_inference/ming_flash_omni/README.md

@@ -2,13 +2,12 @@
 
 [Ming-flash-omni-2.0](https://github.com/inclusionAI/Ming) is an omni-modal model supporting text, image, video, and audio understanding, with text and speech outputs.
 
-vLLM-Omni supports three deployment modes:
+vLLM-Omni supports two deployment modes:
 
 | Mode | Deploy config | Output |
 |------|--------------|--------|
 | Thinker + Talker (omni-speech, default) | `vllm_omni/deploy/ming_flash_omni.yaml` | Text + Audio |
 | Thinker only (multimodal understanding) | `vllm_omni/deploy/ming_flash_omni_thinker_only.yaml` | Text |
-| Thinker + Imagegen (text-to-image / img2img) | `vllm_omni/deploy/ming_flash_omni_image.yaml` | Image (online-serving only at the moment) |
 
 For standalone TTS (talker only), see the [Ming-flash-omni-TTS section in the Text-To-Speech hub](../text_to_speech/README.md#ming-flash-omni-tts).
 
@@ -53,7 +52,7 @@ python examples/offline_inference/ming_flash_omni/end2end.py \
     --image-path ./3_0.png
 ```
 
-### Omni-Speech (Thinker + Talker) {#omni-speech-thinker--talker}
+### Omni-Speech (Thinker + Talker)
 
 The default deploy YAML already runs thinker+talker, so spoken output only requires requesting `audio` (or `text,audio`) modalities.
 The thinker processes your multimodal input, generates text, then the talker synthesises the response as speech.
@@ -88,17 +87,6 @@ The default deploy YAML allocates thinker on GPUs 0–3 and talker on GPU 3 for
 
 Pass `--deploy-config /path/to/your_deploy.yaml` to any of the commands above to override the bundled deploy config.
 
-### Image generation (text-to-image / img2img)
-
-The diffusion-stage image-generation path is currently only wired through
-the online OpenAI-compatible chat endpoint (`/v1/chat/completions` with
-`"modalities": ["image"]`). For payloads, optional knobs
-(`steps`/`cfg`/`seed`/`byte5_text`/`negative_prompt`), and the img2img
-reference-image flow, see the
-[image-gen section in the online-serving README](../../online_serving/ming_flash_omni/README.md#image-generation-text-to-image).
-`end2end.py` does not yet exercise the imagegen stage; that is tracked as
-follow-up work.
-
 ## Online serving
 
 For online serving via the OpenAI-compatible API, see [examples/online_serving/ming_flash_omni/README.md](../../online_serving/ming_flash_omni/README.md).

examples/offline_inference/ming_flash_omni/end2end.py

@@ -17,13 +17,14 @@ from vllm.assets.image import ImageAsset
 from vllm.assets.video import VideoAsset, video_to_ndarrays
 from vllm.multimodal.image import convert_image_mode
 from vllm.multimodal.media.audio import load_audio
+from vllm.utils.argparse_utils import FlexibleArgumentParser
 
 import vllm_omni
+from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
 from vllm_omni.entrypoints.omni import Omni
 
 # Imports the processor also registers itself
 from vllm_omni.transformers_utils.processors.ming import MingFlashOmniProcessor  # noqa: F401
-from vllm_omni.utils.tracking_parser import TrackingArgumentParser
 
 SEED = 42
 MODEL_NAME = "Jonathan1909/Ming-flash-omni-2.0"
@@ -396,7 +397,7 @@ def main(args):
 
 
 def parse_args():
-    parser = TrackingArgumentParser(description="Ming-flash-omni 2.0 offline inference example")
+    parser = FlexibleArgumentParser(description="Ming-flash-omni 2.0 offline inference example")
     parser.add_argument(
         "--query-type",
         "-q",
@@ -495,6 +496,7 @@ def parse_args():
         help="Output directory for results.",
     )
 
+    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 

examples/offline_inference/qwen2_5_omni/end2end.py

@@ -19,9 +19,10 @@ from vllm.assets.video import VideoAsset, video_to_ndarrays
 from vllm.multimodal.image import convert_image_mode
 from vllm.multimodal.media.audio import load_audio
 from vllm.sampling_params import SamplingParams
+from vllm.utils.argparse_utils import FlexibleArgumentParser
 
+from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
 from vllm_omni.entrypoints.omni import Omni
-from vllm_omni.utils.tracking_parser import TrackingArgumentParser
 
 SEED = 42
 
@@ -431,7 +432,7 @@ def main(args):
 
 
 def parse_args():
-    parser = TrackingArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
+    parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
     parser.add_argument(
         "--model",
         type=str,
@@ -553,6 +554,7 @@ def parse_args():
         default=False,
         help="Use py_generator mode. The returned type of Omni.generate() is a Python Generator object.",
     )
+    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 

examples/offline_inference/qwen3_omni/end2end.py

@@ -19,9 +19,10 @@ from vllm.assets.image import ImageAsset
 from vllm.assets.video import VideoAsset, video_to_ndarrays
 from vllm.multimodal.image import convert_image_mode
 from vllm.multimodal.media.audio import load_audio
+from vllm.utils.argparse_utils import FlexibleArgumentParser
 
+from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
 from vllm_omni.entrypoints.omni import Omni
-from vllm_omni.utils.tracking_parser import TrackingArgumentParser
 
 SEED = 42
 
@@ -421,7 +422,7 @@ def main(args):
 
 
 def parse_args():
-    parser = TrackingArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
+    parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
     parser.add_argument(
         "--model",
         type=str,
@@ -560,6 +561,7 @@ def parse_args():
         help="Model dtype (auto, half, float16, bfloat16, float, float32).",
     )
 
+    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 

examples/offline_inference/qwen3_omni/end2end_async_chunk.py

@@ -39,9 +39,10 @@ from vllm.assets.image import ImageAsset
 from vllm.assets.video import VideoAsset, video_to_ndarrays
 from vllm.multimodal.image import convert_image_mode
 from vllm.multimodal.media.audio import load_audio
+from vllm.utils.argparse_utils import FlexibleArgumentParser
 
+from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
 from vllm_omni.entrypoints.async_omni import AsyncOmni
-from vllm_omni.utils.tracking_parser import TrackingArgumentParser
 
 logger = logging.getLogger(__name__)
 
@@ -382,9 +383,7 @@ async def run_all(args):
     print(f"[Info] Creating AsyncOmni with deploy_config={args.deploy_config}")
     async_omni = None
     try:
-        # ``from_cli_args`` forwards only explicitly-passed CLI args so
-        # argparse defaults do not silently override deploy YAML values.
-        async_omni = AsyncOmni.from_cli_args(args, model=args.model)
+        async_omni = AsyncOmni(**vars(args))
 
         # Use default sampling params from stage config (they are pre-configured
         # in the YAML for each stage).
@@ -454,7 +453,7 @@ async def run_all(args):
 
 
 def parse_args():
-    parser = TrackingArgumentParser(
+    parser = FlexibleArgumentParser(
         description=(
             "Offline inference with async_chunk enabled via AsyncOmni. "
             "Downstream stages start before upstream stages finish, "
@@ -590,6 +589,7 @@ def parse_args():
         default=16000,
         help="Sampling rate for audio loading.",
     )
+    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 

examples/offline_inference/sensenova_u1/README.md

@@ -37,10 +37,6 @@ python end2end.py --modality text2text \
 # Custom resolution (image generation)
 python end2end.py --prompt "A futuristic cityscape at sunset" \
                   --width 2048 --height 1024 --think
-
-# Text-to-image with Cache-DiT acceleration
-python end2end.py --prompt "A cute cat" \
-                  --cache-backend cache_dit \
 ```
 
 > **Note**: Default configuration works on a single **NVIDIA A100 (80GB)** or **H100** GPU.
@@ -139,8 +135,6 @@ python end2end.py \
 | `--tensor-parallel-size` | 1 | Number of GPUs for tensor parallelism |
 | `--enforce-eager` | False | Disable torch.compile |
 | `--enable-cpu-offload` | False | Enable module-wise (sequential) CPU offload to reduce peak VRAM |
-| `--cache-backend` | None | Set to `cache_dit` for Cache-DiT acceleration |
-| `--enable-cache-dit-summary` | False | Print Cache-DiT cache statistics after generation |
 
 ## Reducing GPU Memory Usage
 

examples/offline_inference/sensenova_u1/end2end.py

@@ -29,6 +29,7 @@ import os
 
 from PIL import Image
 
+from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
@@ -154,19 +155,8 @@ def parse_args():
             "of extra CPU<->GPU transfers."
         ),
     )
-    parser.add_argument(
-        "--cache-backend",
-        type=str,
-        default=None,
-        choices=["cache_dit"],
-        help="Cache backend to use for image generation acceleration.",
-    )
-    parser.add_argument(
-        "--enable-cache-dit-summary",
-        action="store_true",
-        help="Enable cache-dit summary logging after diffusion forward passes.",
-    )
 
+    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 
@@ -192,8 +182,6 @@ def main():
         tensor_parallel_size=args.tensor_parallel_size,
         enforce_eager=args.enforce_eager,
         enable_cpu_offload=args.enable_cpu_offload,
-        cache_backend=args.cache_backend,
-        enable_cache_dit_summary=args.enable_cache_dit_summary,
     )
 
     extra_args = {
@@ -237,7 +225,6 @@ def main():
     print(f"  Seed           : {args.seed}")
     print(f"  Think mode     : {args.think}")
     print(f"  TP size        : {args.tensor_parallel_size}")
-    print(f"  Cache backend  : {args.cache_backend or 'none'}")
     print(f"{'=' * 60}\n")
 
     # Build prompt dict

examples/offline_inference/text_to_image/text_to_image.py

@@ -307,6 +307,9 @@ def parse_args() -> argparse.Namespace:
         help="Supplementary auxiliary text encoder parameters model name or path (especially for Hidream-l1-full).",
     )
     current_omni_platform.pre_register_and_update(parser)
+    from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
+
+    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 

examples/offline_inference/text_to_speech/README.md

@@ -271,14 +271,6 @@ python examples/offline_inference/text_to_speech/omnivoice/end2end.py \
     --lang zh
 ```
 
-### Seed for Reproducibility
-```bash
-python examples/offline_inference/text_to_speech/omnivoice/end2end.py \
-    --model k2-fsa/OmniVoice \
-    --text "Hello, this is a test." \
-    --seed 42
-```
-
 ### Notes
 - Stage 0 (Generator): Qwen3-0.6B with 32-step iterative unmasking.
 - Stage 1 (Decoder): HiggsAudioV2 RVQ + DAC at 24 kHz.

examples/offline_inference/text_to_speech/cosyvoice3/end2end.py

@@ -10,6 +10,7 @@ import soundfile as sf
 from vllm import SamplingParams
 from vllm.multimodal.media.audio import load_audio
 
+from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokenizer
 from vllm_omni.model_executor.models.cosyvoice3.utils import extract_text_token
@@ -64,6 +65,7 @@ def run_e2e():
         required=True,
         help="Path to tokenizer directory (e.g., <model_path>/CosyVoice-BlankEN).",
     )
+    nullify_stage_engine_defaults(parser)
     args = parser.parse_args()
     # Ensure tokenizer directory exists
     if not os.path.exists(args.tokenizer):

examples/offline_inference/text_to_speech/fish_speech/end2end.py

@@ -26,14 +26,16 @@ import torch
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+
 from vllm_omni import AsyncOmni, Omni
+from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
 from vllm_omni.model_executor.models.fish_speech.dac_utils import DAC_HOP_LENGTH, DAC_SAMPLE_RATE
 from vllm_omni.model_executor.models.fish_speech.prompt_utils import (
     build_fish_text_only_prompt_ids,
     estimate_fish_voice_clone_prompt_len_from_normalized,
     normalize_fish_voice_clone_texts,
 )
-from vllm_omni.utils.tracking_parser import TrackingArgumentParser
 
 logger = logging.getLogger(__name__)
 
@@ -209,7 +211,7 @@ async def main_streaming(args):
 
 
 def parse_args():
-    parser = TrackingArgumentParser(
+    parser = FlexibleArgumentParser(
         description="Fish Speech S2 Pro offline inference with vLLM Omni",
     )
     parser.add_argument(
@@ -265,6 +267,7 @@ def parse_args():
         default=False,
         help="Stream audio chunks as they arrive via AsyncOmni.",
     )
+    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 

examples/offline_inference/text_to_speech/higgs_audio_v2/README.md

@@ -1,38 +0,0 @@
-# higgs-audio v2 — offline example
-
-Drives Stage 0 (DualFFN talker) + Stage 1 (HiggsAudio codec) for `bosonai/higgs-audio-v2-generation-3B-base` end-to-end through the vLLM-Omni engine and writes a 24 kHz mono WAV per prompt.
-
-## Prerequisites
-
-Voice clone needs `transformers>=5.3.0` — vllm-omni loads the audio codec via HF's `HiggsAudioV2TokenizerModel`, instantiated from the `k2-fsa/OmniVoice/audio_tokenizer/` subdirectory (only that ~806 MB subdir is downloaded). The boson-ai standalone tokenizer repo's `model.safetensors` is actually a copy of the 3B talker LM, so HF can't load it directly; the k2 bundle ships the same codec weights repackaged with HF-compatible key naming.
-
-```bash
-pip install -U "transformers>=5.3.0"
-```
-
-## Quick start
-
-Plain TTS:
-```bash
-python examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py \
-    --texts "Hello world." "The quick brown fox jumps over the lazy dog." \
-    --output-dir results/wavs
-```
-
-## Voice cloning
-
-Pass both `--ref-audio` and `--ref-text` together:
-```bash
-python examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py \
-    --texts "Hello, this is a cloned voice." \
-    --ref-audio /path/to/reference.wav \
-    --ref-text  "Exact transcript spoken in reference.wav." \
-    --output-dir results/wavs
-```
-
-## Notes
-
-- Output: 24 kHz mono WAV.
-- Deploy config: `vllm_omni/deploy/higgs_audio_v2.yaml` (auto-loaded by HF `model_type`).
-- `--ref-text` must be the real transcript of `--ref-audio`; mismatched text degrades cloned-voice quality.
-- For online serving, see [`examples/online_serving/text_to_speech/higgs_audio_v2/`](../../../online_serving/text_to_speech/higgs_audio_v2/).

examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py

@@ -1,203 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Offline higgs-audio v2 inference example.
-
-Runs Stage 0 (DualFFN talker) + Stage 1 (HiggsAudio codec) end-to-end
-through the vLLM-Omni engine without going through the HTTP server, and
-saves a 24 kHz mono WAV per prompt.
-
-Example:
-
-    python examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py \\
-        --texts "Hello world." \\
-                "The quick brown fox jumps over the lazy dog." \\
-        --output-dir results/wavs
-"""
-
-from __future__ import annotations
-
-import os
-
-# DeepGEMM FP8 kernels require an optional backend that may not be installed.
-# Disable the warmup before importing vLLM so engine startup falls back to the
-# regular gemm path. Users with deep_gemm installed can override these.
-os.environ.setdefault("VLLM_USE_DEEP_GEMM", "0")
-os.environ.setdefault("VLLM_MOE_USE_DEEP_GEMM", "0")
-
-import time
-from pathlib import Path
-
-import numpy as np
-import soundfile as sf
-import torch
-
-from vllm_omni import Omni
-from vllm_omni.utils.tracking_parser import TrackingArgumentParser
-
-SAMPLE_RATE = 24_000
-DEFAULT_TEXTS = (
-    "Hello world.",
-    "The quick brown fox jumps over the lazy dog.",
-)
-
-
-def parse_args():
-    parser = TrackingArgumentParser(description="Offline higgs-audio v2 inference")
-    parser.add_argument(
-        "--model",
-        type=str,
-        default="bosonai/higgs-audio-v2-generation-3B-base",
-        help="Stage-0 talker model id or path.",
-    )
-    parser.add_argument(
-        "--texts",
-        type=str,
-        nargs="+",
-        default=list(DEFAULT_TEXTS),
-        help="One or more plain-text prompts to synthesize.",
-    )
-    parser.add_argument(
-        "--output-dir",
-        type=str,
-        default="results/wavs",
-        help="Directory to write per-prompt WAV files.",
-    )
-    parser.add_argument(
-        "--deploy-config",
-        type=str,
-        default=None,
-        help="Override the deploy config path. Auto-loads "
-        "vllm_omni/deploy/higgs_audio_v2.yaml from the HF model_type by default.",
-    )
-    parser.add_argument(
-        "--max-new-tokens",
-        type=int,
-        default=500,
-        help="Cap on Stage-0 codec frames per request.",
-    )
-    parser.add_argument(
-        "--ref-audio",
-        type=str,
-        default=None,
-        help="Reference clip for voice clone (path to a WAV file). Paired with --ref-text.",
-    )
-    parser.add_argument(
-        "--ref-text",
-        type=str,
-        default=None,
-        help="Transcript of the reference clip. Required when --ref-audio is set.",
-    )
-    return parser.parse_args()
-
-
-def _slugify(text: str) -> str:
-    import re
-
-    slug = re.sub(r"\s+", "_", text.strip().lower())
-    slug = re.sub(r"[^a-z0-9_]+", "", slug)
-    return slug[:48] or "prompt"
-
-
-def _extract_pcm(multimodal_output: dict) -> torch.Tensor:
-    """Pull the final concatenated PCM tensor out of a request's multimodal_output."""
-    audio = multimodal_output.get("model_outputs")
-    if audio is None:
-        audio = multimodal_output.get("audio")
-    if audio is None:
-        raise ValueError(f"no audio key in multimodal_output: {list(multimodal_output.keys())}")
-    if isinstance(audio, list):
-        valid = [torch.as_tensor(a).float().cpu().reshape(-1) for a in audio if a is not None]
-        if not valid:
-            raise ValueError("audio list is empty")
-        return torch.cat(valid, dim=0) if len(valid) > 1 else valid[0]
-    return torch.as_tensor(audio).float().cpu().reshape(-1)
-
-
-def _pcm_to_int16(pcm: torch.Tensor) -> np.ndarray:
-    arr = pcm.numpy()
-    if arr.dtype.kind == "f":
-        arr = np.clip(arr, -1.0, 1.0)
-        arr = (arr * 32767.0).astype(np.int16)
-    else:
-        arr = arr.astype(np.int16)
-    return arr
-
-
-def main():
-    args = parse_args()
-    if (args.ref_audio is None) != (args.ref_text is None):
-        raise SystemExit("--ref-audio and --ref-text must be supplied together")
-    output_dir = Path(args.output_dir)
-    output_dir.mkdir(parents=True, exist_ok=True)
-
-    engine = Omni(model=args.model, deploy_config=args.deploy_config)
-
-    # Build prompt_token_ids using the same path serving_speech.py takes online.
-    from transformers import AutoProcessor
-
-    from vllm_omni.model_executor.models.higgs_audio_v2.higgs_audio_v2_tokenizer import (
-        build_plain_text_prompt,
-        build_voice_clone_prompt,
-        input_ids_to_python_list,
-    )
-
-    processor = AutoProcessor.from_pretrained(args.model, trust_remote_code=True)
-
-    # Voice-clone path: load the reference clip once. The HF processor will
-    # encode it via the bundled HiggsAudioV2TokenizerModel each time we build
-    # a prompt below. This is cheap on CPU for a few-second clip.
-    ref_wav: np.ndarray | None = None
-    ref_sr: int | None = None
-    if args.ref_audio is not None:
-        ref_wav, ref_sr = sf.read(args.ref_audio, always_2d=False)
-        if ref_wav.ndim == 2:
-            ref_wav = ref_wav.mean(axis=1)
-
-    print(f"Model       : {args.model}")
-    print(f"Prompts     : {len(args.texts)}")
-    print(f"Output dir  : {output_dir}")
-    print(f"Voice clone : {'yes' if ref_wav is not None else 'no'}")
-
-    # Run one prompt at a time. The Stage-0 talker's per-slot audio state is
-    # request-scoped; submitting multiple prompts in the same engine.generate()
-    # call would batch them in the AR runner and exercise a code path that is
-    # not validated for this model yet.
-    total_elapsed = 0.0
-    total_dur = 0.0
-    for text in args.texts:
-        if ref_wav is not None:
-            out = build_voice_clone_prompt(processor, text, ref_wav, int(ref_sr), args.ref_text)
-            prompt = {
-                "prompt_token_ids": out["prompt_token_ids"],
-                # Bare tensors (NOT list-wrapped): the msgspec serializer in
-                # vllm_omni.data_entry_keys routes torch.Tensor → tensor_data
-                # and list[Tensor] → list_data (which silently strips tensors).
-                "additional_information": {
-                    "audio_input_ids": out["audio_input_ids"],
-                    "audio_input_ids_mask": out["audio_input_ids_mask"],
-                },
-            }
-        else:
-            inputs = build_plain_text_prompt(processor, text)
-            prompt = {"prompt_token_ids": input_ids_to_python_list(inputs)}
-        t_start = time.perf_counter()
-        outputs = engine.generate([prompt])
-        elapsed = time.perf_counter() - t_start
-        total_elapsed += elapsed
-
-        mm = outputs[0].outputs[0].multimodal_output
-        pcm = _extract_pcm(mm)
-        slug = _slugify(text)
-        out_path = output_dir / f"{slug}.wav"
-        sf.write(str(out_path), _pcm_to_int16(pcm), SAMPLE_RATE, format="WAV", subtype="PCM_16")
-        dur = pcm.numel() / SAMPLE_RATE
-        total_dur += dur
-        print(f"  {slug:<50} dur={dur:5.2f}s  -> {out_path}")
-
-    rtf = total_elapsed / total_dur if total_dur > 0 else float("inf")
-    print(f"Total infer : {total_elapsed:.2f}s  total audio: {total_dur:.2f}s  RTF: {rtf:.3f}")
-
-
-if __name__ == "__main__":
-    os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
-    main()

examples/offline_inference/text_to_speech/ming_flash_omni_tts/end2end.py

@@ -5,11 +5,11 @@ from typing import Any
 
 import soundfile as sf
 import torch
-
-from vllm_omni.utils.tracking_parser import TrackingArgumentParser
+from vllm.utils.argparse_utils import FlexibleArgumentParser
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 
+from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniTokensPrompt
 from vllm_omni.model_executor.models.ming_flash_omni.prompt_utils import (
@@ -68,7 +68,7 @@ def save_audio(mm: dict[str, Any], output_path: str) -> None:
 
 
 def parse_args():
-    parser = TrackingArgumentParser(description="Ming-flash-omni standalone talker offline e2e example")
+    parser = FlexibleArgumentParser(description="Ming-flash-omni standalone talker offline e2e example")
     parser.add_argument("--model", type=str, default=MODEL_NAME, help="Model name or local path.")
     parser.add_argument(
         "--deploy-config",
@@ -89,6 +89,7 @@ def parse_args():
     parser.add_argument("--init-timeout", type=int, default=600, help="Engine init timeout in seconds.")
     parser.add_argument("--stage-init-timeout", type=int, default=300, help="Single stage init timeout in seconds.")
 
+    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 

examples/offline_inference/text_to_speech/moss_tts/end2end.py

@@ -1,377 +0,0 @@
-"""Offline inference example for MOSS-TTS variants.
-
-Covers all five models:
-  - OpenMOSS-Team/MOSS-TTS           (8B, general TTS)
-  - OpenMOSS-Team/MOSS-TTSD-v1.0     (8B, dialogue TTS)
-  - OpenMOSS-Team/MOSS-SoundEffect   (8B, sound effect synthesis)
-  - OpenMOSS-Team/MOSS-TTS-Realtime  (1.7B, low-latency streaming)
-  - OpenMOSS-Team/MOSS-VoiceGenerator(1.7B, voice design)
-
-Usage
------
-# Standard TTS with voice cloning:
-python end2end.py \\
-    --model OpenMOSS-Team/MOSS-TTS \\
-    --text "Hello, this is a MOSS-TTS test." \\
-    --ref-audio path/to/reference.wav \\
-    --output output.wav
-"""
-
-from __future__ import annotations
-
-import argparse
-import gc
-import sys
-
-import soundfile as sf
-import torch
-
-
-def _load_ref_audio(path: str, target_sr: int = 24000) -> torch.Tensor:
-    """Load and resample reference audio to (1, T) at target_sr."""
-    data, sr = sf.read(path, always_2d=True)  # (T, C)
-    wav = torch.from_numpy(data.T.astype("float32"))  # (C, T)
-    if wav.shape[0] > 1:
-        wav = wav.mean(dim=0, keepdim=True)
-    if sr != target_sr:
-        import torchaudio
-
-        wav = torchaudio.functional.resample(wav, sr, target_sr)
-    return wav  # (1, T)
-
-
-def _build_unified_codes(
-    model_id: str,
-    text: str,
-    ref_audio: torch.Tensor | None,
-    *,
-    instruction: str | None = None,
-    tokens: int | None = None,
-    quality: str | None = None,
-    sound_event: str | None = None,
-    ambient_sound: str | None = None,
-    language: str | None = None,
-) -> tuple[list[int], torch.Tensor, int]:
-    """Build the upstream MossTTSDelay unified-codes prompt.
-
-    Returns (text_token_ids, audio_codes_LxNQ, n_vq) where
-    ``text_token_ids`` is the prefill prompt for the talker and
-    ``audio_codes_LxNQ`` is the matching delay-pattern audio code grid
-    (``audio_pad_code`` everywhere except inside the reference-audio block).
-    """
-    from transformers import AutoProcessor
-
-    proc = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
-
-    reference: list[torch.Tensor] | None = None
-    if ref_audio is not None:
-        # Pre-encode so the processor doesn't need torchaudio.load (which
-        # now requires torchcodec on recent torchaudio releases).
-        # encode_audios_from_wav expects (C, T) per item.
-        wav_2d = ref_audio if ref_audio.dim() == 2 else ref_audio.unsqueeze(0)
-        codes_list = proc.encode_audios_from_wav([wav_2d], sampling_rate=24000, n_vq=proc.model_config.n_vq)
-        reference = [codes_list[0]]
-
-    user_msg_kwargs: dict = {"text": text}
-    if reference is not None:
-        user_msg_kwargs["reference"] = reference
-    if instruction is not None:
-        user_msg_kwargs["instruction"] = instruction
-    if tokens is not None:
-        user_msg_kwargs["tokens"] = tokens
-    if quality is not None:
-        user_msg_kwargs["quality"] = quality
-    if sound_event is not None:
-        user_msg_kwargs["sound_event"] = sound_event
-    if ambient_sound is not None:
-        user_msg_kwargs["ambient_sound"] = ambient_sound
-    if language is not None:
-        user_msg_kwargs["language"] = language
-
-    user_msg = proc.build_user_message(**user_msg_kwargs)
-    batch = proc(conversations=[[user_msg]], mode="generation")
-    unified = batch["input_ids"][0]  # (L, 1+n_vq)
-    n_vq = int(unified.shape[-1] - 1)
-    text_ids = unified[:, 0].tolist()
-    audio_codes = unified[:, 1:].contiguous().to(torch.int64)  # (L, n_vq)
-
-    # Free the processor (incl. ~1.6B-param audio_tokenizer on CPU) before
-    # we instantiate Omni — we only needed it for the one-shot encode.
-    del proc
-    gc.collect()
-
-    return text_ids, audio_codes, n_vq
-
-
-def _build_realtime_prompt(
-    model_id: str,
-    text: str,
-    ref_audio: torch.Tensor | None,
-) -> tuple[list[int], torch.Tensor, int]:
-    """Build the prefill prompt for MOSS-TTS-Realtime.
-
-    Realtime uses a different prompt format than the delay variants: a system
-    prompt + (optional) voice-clone context with the encoded reference audio
-    inlined, followed by the user text and an assistant header. The format is
-    a ``(L, channels+1=17)`` int grid where col 0 is text/special tokens and
-    cols 1..16 are RVQ codebook entries (``audio_channel_pad`` outside the
-    ref-audio block).
-
-    The upstream ``MossTTSRealtimeProcessor`` is shipped via
-    ``trust_remote_code`` but is not auto-discovered by ``AutoProcessor``
-    (no ``processor_config.json``) — we import its module directly from the
-    snapshot directory and use its ``make_ensemble``/``make_user_prompt``.
-    """
-    import importlib.util
-    import sys as _sys
-    from pathlib import Path
-
-    import numpy as np
-    import torch
-    from huggingface_hub import snapshot_download
-    from transformers import AutoModel, AutoTokenizer
-
-    snap_dir = Path(snapshot_download(repo_id=model_id))
-    proc_module_path = snap_dir / "processing_mossttsrealtime.py"
-    if not proc_module_path.exists():
-        raise FileNotFoundError(f"realtime processor module missing at {proc_module_path}")
-
-    # Load the processor module without polluting sys.path globally.
-    spec = importlib.util.spec_from_file_location("_moss_tts_realtime_proc", proc_module_path)
-    mod = importlib.util.module_from_spec(spec)  # type: ignore[arg-type]
-    _sys.modules[spec.name] = mod  # type: ignore[union-attr]
-    spec.loader.exec_module(mod)  # type: ignore[union-attr]
-
-    tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
-    realtime_processor = mod.MossTTSRealtimeProcessor(tokenizer=tokenizer)
-
-    # Encode reference audio via the standalone MOSS Audio Tokenizer (CPU is
-    # fine — it runs once per request).
-    audio_tokens = None
-    if ref_audio is not None:
-        codec = AutoModel.from_pretrained("OpenMOSS-Team/MOSS-Audio-Tokenizer", trust_remote_code=True).eval()
-        wav_2d = ref_audio if ref_audio.dim() == 2 else ref_audio.unsqueeze(0)
-        with torch.no_grad():
-            enc = codec.batch_encode([wav_2d.squeeze(0)], num_quantizers=16)
-        # enc.audio_codes: (NQ, B, T); we want (T, NQ) numpy.
-        codes = enc.audio_codes[:, 0, : int(enc.audio_codes_lengths[0].item())]
-        audio_tokens = codes.transpose(0, 1).contiguous().cpu().numpy()  # (T, 16)
-        del codec
-        gc.collect()
-
-    # Build the inference-style prefill (mirrors upstream
-    # ``MossTTSRealtimeInference._build_prefill_batch``):
-    #   system_prompt [+ voice_clone_block]
-    #   <|im_start|>assistant\n
-    #   <text_tokens...>            (col 0 = text, cols 1..16 = audio_channel_pad)
-    #   <last text token>           (col 1 = audio_bos at the last text position)
-    # Generation continues from here; the local transformer emits real audio
-    # codes for codebook 0 directly (no BOS frame at gen_step=0).
-    system_grid = realtime_processor.make_ensemble(prompt_audio_tokens=audio_tokens)  # (L_sys, 17)
-
-    # Append "<|im_start|>assistant\n" header (the snapshot's make_ensemble
-    # omits it; inferencer.py's version includes it. We add it here.)
-    assistant_header_ids = tokenizer.encode("<|im_start|>assistant\n", add_special_tokens=False)
-    assistant_grid = np.full((len(assistant_header_ids), 17), 1024, dtype=np.int64)
-    assistant_grid[:, 0] = assistant_header_ids
-
-    # Append the FIRST 12 text tokens with audio_bos (1025) at the last
-    # included text position (audio column 0 == grid col 1). The remaining
-    # text tokens are fed one-per-step during generation via the talker's
-    # compute_logits → vLLM sampler pipeline (see ``ids.all`` in the
-    # additional_information). This matches upstream
-    # ``MossTTSRealtimeInference._build_prefill_batch`` (prefill_max_text=12).
-    text_ids_only = tokenizer.encode(text, add_special_tokens=False)
-    if not text_ids_only:
-        raise ValueError("Empty text after tokenisation")
-    PREFILL_MAX_TEXT = 12
-    cur_len = min(len(text_ids_only), PREFILL_MAX_TEXT)
-    text_grid = np.full((cur_len, 17), 1024, dtype=np.int64)
-    text_grid[:, 0] = text_ids_only[:cur_len]
-    text_grid[-1, 1] = 1025  # audio_bos_token at last *prefilled* text token
-
-    grid = np.concatenate([system_grid, assistant_grid, text_grid], axis=0)
-    text_ids = grid[:, 0].tolist()
-    audio_codes = torch.from_numpy(grid[:, 1:].astype(np.int64).copy())  # (L, 16)
-
-    # ``remaining_text_ids`` will be fed one-per-step by the talker.
-    remaining_text_ids = list(text_ids_only[cur_len:])
-    return text_ids, audio_codes, 16, remaining_text_ids
-
-
-def run_tts(args: argparse.Namespace) -> None:
-    from vllm import SamplingParams
-
-    from vllm_omni import Omni
-
-    deploy_map = {
-        "OpenMOSS-Team/MOSS-TTS": "moss_tts",
-        "OpenMOSS-Team/MOSS-TTSD-v1.0": "moss_ttsd",
-        "OpenMOSS-Team/MOSS-SoundEffect": "moss_sound_effect",
-        "OpenMOSS-Team/MOSS-TTS-Realtime": "moss_tts_realtime",
-        "OpenMOSS-Team/MOSS-VoiceGenerator": "moss_voice_generator",
-    }
-    deploy_config = deploy_map.get(args.model, "moss_tts")
-
-    is_sound_effect = "SoundEffect" in args.model
-    is_voice_gen = "VoiceGenerator" in args.model
-    is_realtime = "Realtime" in args.model
-
-    ref_audio = None
-    builder_kwargs: dict = {}
-
-    remaining_text_ids: list[int] | None = None
-    if is_realtime:
-        # MOSS-TTS-Realtime uses a different prompt format and has its own
-        # per-step depth transformer; the delay-style ``MossTTSDelayProcessor``
-        # path doesn't apply.
-        if not args.ref_audio:
-            sys.exit("MOSS-TTS-Realtime currently requires --ref-audio (voice clone).")
-        ref_audio = _load_ref_audio(args.ref_audio)
-        text = args.text or ""
-        text_ids, audio_codes, n_vq, remaining_text_ids = _build_realtime_prompt(args.model, text, ref_audio)
-        print(
-            f"Prefill prompt: {len(text_ids)} tokens, {n_vq}-quantizer audio block (realtime); "
-            f"{len(remaining_text_ids)} text tokens to stream"
-        )
-    elif is_sound_effect:
-        if not args.ambient_sound:
-            sys.exit("MOSS-SoundEffect requires --ambient-sound")
-        builder_kwargs["ambient_sound"] = args.ambient_sound
-        # MOSS-SoundEffect emits ~12.5 codec frames per second; ``tokens``
-        # tells the AR loop the target duration.
-        if args.duration_tokens:
-            builder_kwargs["tokens"] = args.duration_tokens
-        elif args.duration_seconds:
-            builder_kwargs["tokens"] = max(1, int(float(args.duration_seconds) * 12.5))
-        text = args.text or ""
-    elif is_voice_gen:
-        # MOSS-VoiceGenerator needs an instruction (voice description) plus
-        # the text to synthesise; no reference audio.
-        if not args.instruction:
-            sys.exit("MOSS-VoiceGenerator requires --instruction (voice description)")
-        builder_kwargs["instruction"] = args.instruction
-        text = args.text or ""
-    else:
-        if not args.ref_audio:
-            sys.exit("Voice cloning requires --ref-audio")
-        ref_audio = _load_ref_audio(args.ref_audio)
-        text = args.text or ""
-
-    if not is_realtime:
-        text_ids, audio_codes, n_vq = _build_unified_codes(args.model, text, ref_audio, **builder_kwargs)
-        print(f"Prefill prompt: {len(text_ids)} tokens, {n_vq}-quantizer audio block")
-
-    omni = Omni(model=args.model, deploy_config=deploy_config, stage_init_timeout=600)
-
-    talker_sp = SamplingParams(
-        temperature=1.7,
-        top_p=0.8,
-        top_k=25,
-        max_tokens=2048,
-        seed=42,
-        detokenize=False,
-    )
-    codec_sp = SamplingParams(
-        temperature=0.0,
-        top_p=1.0,
-        top_k=-1,
-        max_tokens=65536,
-        seed=42,
-        detokenize=True,
-    )
-
-    additional_info: dict = {
-        # Reference-audio codes consumed by the talker prefill (additive
-        # embedding); shape is (L, n_vq) where pad_code marks non-audio
-        # positions. The OmniPayload schema reserves ``codes.ref`` for
-        # exactly this kind of conditioning input.
-        "codes": {"ref": audio_codes},
-    }
-    if remaining_text_ids:
-        # Realtime variant: feed the remaining text tokens (everything past
-        # the first ``PREFILL_MAX_TEXT``) one-per-step during decode via the
-        # talker's compute_logits → vLLM sampler one-hot trick.
-        additional_info["ids"] = {"all": remaining_text_ids}
-
-    request = {
-        "prompt_token_ids": text_ids,
-        "additional_information": additional_info,
-    }
-
-    import time as _time
-
-    chunks: list[torch.Tensor] = []
-    t_gen_start = _time.monotonic()
-    t_first_audio: float | None = None
-    for stage_outputs in omni.generate(request, [talker_sp, codec_sp]):
-        # Each yielded ``stage_outputs`` is an OmniRequestOutput; the codec
-        # stage attaches its waveform under ``multimodal_output["model_outputs"]``.
-        # Avoid ``a or b`` when ``a`` may be a tensor (truthiness raises).
-        mm = stage_outputs.multimodal_output or {}
-        wav = mm.get("audio")
-        if wav is None:
-            wav = mm.get("model_outputs")
-        if wav is None:
-            continue
-        if isinstance(wav, list):
-            non_empty = [w.reshape(-1) for w in wav if isinstance(w, torch.Tensor) and w.numel() > 0]
-            if non_empty and t_first_audio is None:
-                t_first_audio = _time.monotonic()
-            chunks.extend(non_empty)
-        elif isinstance(wav, torch.Tensor) and wav.numel() > 0:
-            if t_first_audio is None:
-                t_first_audio = _time.monotonic()
-            chunks.append(wav.reshape(-1))
-
-    t_gen_end = _time.monotonic()
-    if not chunks:
-        print("No audio in output — check model loading logs.")
-        return
-
-    wav_t = torch.cat(chunks).float().cpu()
-    sf.write(args.output, wav_t.numpy(), 24000)
-    print(f"Saved {len(wav_t) / 24000:.2f}s of audio to {args.output}")
-    ttfa_ms = (t_first_audio - t_gen_start) * 1000.0 if t_first_audio is not None else float("nan")
-    gen_total_ms = (t_gen_end - t_gen_start) * 1000.0
-    print(f"TTFA_MS={ttfa_ms:.1f} GEN_TOTAL_MS={gen_total_ms:.1f} AUDIO_S={len(wav_t) / 24000:.2f}")
-
-
-def main() -> None:
-    parser = argparse.ArgumentParser(description="MOSS-TTS offline inference example")
-    parser.add_argument(
-        "--model",
-        default="OpenMOSS-Team/MOSS-TTS",
-        help="HuggingFace model ID",
-    )
-    parser.add_argument("--text", default="", help="Text to synthesise (TTS variants)")
-    parser.add_argument("--ref-audio", default=None, help="Reference audio for voice cloning")
-    parser.add_argument(
-        "--ambient-sound",
-        default=None,
-        help="Sound description for MOSS-SoundEffect",
-    )
-    parser.add_argument(
-        "--duration-tokens",
-        type=int,
-        default=None,
-        help="Target duration in tokens (1 s ≈ 12.5 tokens) for MOSS-SoundEffect",
-    )
-    parser.add_argument(
-        "--duration-seconds",
-        type=float,
-        default=None,
-        help="Target duration in seconds (converted via ~12.5 tokens/s) for MOSS-SoundEffect",
-    )
-    parser.add_argument(
-        "--instruction",
-        default=None,
-        help="Voice description for MOSS-VoiceGenerator (e.g. 'a young woman with a clear voice')",
-    )
-    parser.add_argument("--output", default="output.wav", help="Output WAV path")
-    args = parser.parse_args()
-    run_tts(args)
-
-
-if __name__ == "__main__":
-    main()

examples/offline_inference/text_to_speech/moss_tts_nano/end2end.py

@@ -38,8 +38,7 @@ from pathlib import Path
 import soundfile as sf
 import torch
 from vllm import SamplingParams
-
-from vllm_omni.utils.tracking_parser import TrackingArgumentParser
+from vllm.utils.argparse_utils import FlexibleArgumentParser
 
 # Prevent multiprocessing from re-importing CUDA in the wrong context.
 os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
@@ -156,7 +155,7 @@ def main(args) -> None:
 
 
 def parse_args():
-    parser = TrackingArgumentParser(description="MOSS-TTS-Nano offline inference")
+    parser = FlexibleArgumentParser(description="MOSS-TTS-Nano offline inference")
     parser.add_argument("--text", default="Hello, this is MOSS-TTS-Nano speaking.", help="Text to synthesize.")
     parser.add_argument(
         "--ref-audio",

examples/offline_inference/text_to_speech/omnivoice/end2end.py

@@ -21,6 +21,7 @@ import os
 import numpy as np
 import soundfile as sf
 
+from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
@@ -79,12 +80,7 @@ def run_e2e():
         default=600,
         help="Stage initialization timeout in seconds",
     )
-    parser.add_argument(
-        "--seed",
-        type=int,
-        default=None,
-        help="Random seed for generation",
-    )
+    nullify_stage_engine_defaults(parser)
     args = parser.parse_args()
 
     if not os.path.exists(args.stage_config):
@@ -126,7 +122,7 @@ def run_e2e():
     if mm_processor_kwargs:
         prompts["mm_processor_kwargs"] = mm_processor_kwargs
 
-    sampling_params_list = [OmniDiffusionSamplingParams(extra_args={"seed": args.seed})]
+    sampling_params_list = [OmniDiffusionSamplingParams()]
 
     print(f"Generating speech for: {args.text}")
 

examples/offline_inference/text_to_speech/qwen3_tts/end2end.py

@@ -15,8 +15,10 @@ import torch
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+
 from vllm_omni import AsyncOmni, Omni
-from vllm_omni.utils.tracking_parser import TrackingArgumentParser
+from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
 
 logger = logging.getLogger(__name__)
 
@@ -41,8 +43,8 @@ def _estimate_prompt_len(
     """
     try:
         from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import Qwen3TTSConfig
-        from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import (
-            Qwen3TTSPromptEmbedsBuilder,
+        from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import (
+            Qwen3TTSTalkerForConditionalGeneration,
         )
 
         if model_name not in _cache:
@@ -122,7 +124,7 @@ def _estimate_prompt_len(
             except Exception:
                 return None
 
-        return Qwen3TTSPromptEmbedsBuilder.estimate_prompt_len_from_additional_information(
+        return Qwen3TTSTalkerForConditionalGeneration.estimate_prompt_len_from_additional_information(
             additional_information=additional_information,
             task_type=task_type,
             tokenize_prompt=lambda t: tok(t, padding=False)["input_ids"],
@@ -433,7 +435,7 @@ async def main_streaming(args):
 
 
 def parse_args():
-    parser = TrackingArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
+    parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
     parser.add_argument(
         "--query-type",
         "-q",
@@ -546,6 +548,7 @@ def parse_args():
         help="Number of prompts per batch (default: 1, sequential).",
     )
 
+    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 

examples/offline_inference/text_to_speech/voxcpm2/end2end.py

@@ -12,16 +12,17 @@ from pathlib import Path
 
 import soundfile as sf
 import torch
+from vllm.utils.argparse_utils import FlexibleArgumentParser
 
 from vllm_omni import Omni
-from vllm_omni.utils.tracking_parser import TrackingArgumentParser
+from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
 
 REPO_ROOT = Path(__file__).resolve().parents[4]
 SAMPLE_RATE = 48_000
 
 
 def parse_args():
-    parser = TrackingArgumentParser(description="Offline VoxCPM2 native AR inference")
+    parser = FlexibleArgumentParser(description="Offline VoxCPM2 native AR inference")
     parser.add_argument(
         "--model",
         type=str,
@@ -59,6 +60,7 @@ def parse_args():
         default=None,
         help="Optional transcript of --ref-audio (enables continuation mode).",
     )
+    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 

examples/offline_inference/text_to_speech/voxtral_tts/end2end.py

@@ -27,10 +27,11 @@ except ImportError:
         "pip install -e ."
     )
 from vllm import SamplingParams
+from vllm.utils.argparse_utils import FlexibleArgumentParser
 
 from vllm_omni import AsyncOmni
+from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
 from vllm_omni.entrypoints.omni import Omni
-from vllm_omni.utils.tracking_parser import TrackingArgumentParser
 
 logger = logging.getLogger(__name__)
 
@@ -229,7 +230,7 @@ def run_non_streaming(inputs, sampling_params_list, model_name, args, output_dir
 
 
 def parse_args() -> Namespace:
-    parser = TrackingArgumentParser(description="Demo on using vLLM for offline inference with Voxtral TTS")
+    parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with Voxtral TTS")
     parser.add_argument(
         "--model",
         type=str,
@@ -312,6 +313,7 @@ def parse_args() -> Namespace:
         default=None,
         help="Quantization method (e.g. 'fp8'). Applied to the language model only.",
     )
+    nullify_stage_engine_defaults(parser)
     return parser.parse_args()
 
 

examples/offline_inference/vace/vace_video_generation.md

@@ -1,105 +1,24 @@
 # VACE Video Generation
 
-Generate videos from text prompts, images, or video conditions using vLLM-Omni's VACE diffusion pipeline.
+[VACE](https://github.com/ali-vilab/VACE) (Video All-in-one Creation Engine) supports multiple video tasks through a single model.
 
-- `vace_video_generation.py`: command-line script for multi-mode video generation with advanced options.
+| Model | Architecture | Model Weights (bf16) | HuggingFace |
+|-------|-------------|----------------------|-------------|
+| Wan2.1-VACE (1.3B) | Wan2.1 | ~10 GB | [Wan-AI/Wan2.1-VACE-1.3B-diffusers](https://huggingface.co/Wan-AI/Wan2.1-VACE-1.3B-diffusers) |
+| Wan2.1-VACE (14B) | Wan2.1 | ~38 GB | [Wan-AI/Wan2.1-VACE-14B-diffusers](https://huggingface.co/Wan-AI/Wan2.1-VACE-14B-diffusers) |
 
-## Table of Contents
-
-- [Overview](#overview)
-- [Quick Start](#quick-start)
-- [Key Arguments](#key-arguments)
-- [More CLI Examples](#more-cli-examples)
-
-## Overview
-
-[VACE](https://github.com/ali-vilab/VACE) (Video All-in-one Creation Engine) supports multiple video tasks through a single unified model, including text-to-video, image-to-video, first-last-frame interpolation, inpainting, and reference image-guided generation.
-
-### Supported Models
-
-| Model | Architecture | Peak VRAM (GiB) * | Model Weights (GiB) | HuggingFace |
-|-------|-------------|-------------------|---------------------|-------------|
-| Wan2.1-VACE (1.3B) | Wan2.1 | TBD | ~10 | [Wan-AI/Wan2.1-VACE-1.3B-diffusers](https://huggingface.co/Wan-AI/Wan2.1-VACE-1.3B-diffusers) |
-| Wan2.1-VACE (14B) | Wan2.1 | TBD | ~38 | [Wan-AI/Wan2.1-VACE-14B-diffusers](https://huggingface.co/Wan-AI/Wan2.1-VACE-14B-diffusers) |
-
-!!! info
-*Peak VRAM: based on basic single-card usage, 480x832 resolution, 81 frames, 30 inference steps, without any acceleration/optimization features.
-
-Default model: `Wan-AI/Wan2.1-VACE-14B-diffusers`
-
-## Quick Start
-
-### Python API
-
-Text-to-video generation:
-
-```python
-import torch
-
-from vllm_omni.entrypoints.omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-if __name__ == "__main__":
-    omni = Omni(model="Wan-AI/Wan2.1-VACE-1.3B-diffusers")
-    prompt = "A sleek robot stands in a vast warehouse filled with boxes"
-    outputs = omni.generate(
-        prompt,
-        OmniDiffusionSamplingParams(
-            height=480,
-            width=832,
-            num_frames=81,
-            num_inference_steps=30,
-            guidance_scale=5.0,
-        ),
-    )
-    video = outputs[0].images
-
-    from diffusers.utils import export_to_video
-    export_to_video(list(video[0]), "t2v_output.mp4", fps=16)
-    omni.close()
-```
-
-### Local CLI Usage
+## Text-to-Video (T2V)
 
 ```bash
 python vace_video_generation.py \
   --mode t2v \
   --prompt "A sleek robot stands in a vast warehouse filled with boxes" \
   --height 480 --width 832 --num-frames 81 \
+  --num-inference-steps 30 --guidance-scale 5.0 --flow-shift 5.0 \
   --output t2v_output.mp4
 ```
 
-## Key Arguments
-
-| Argument | Type | Default | Description |
-| -------- | ---- | ------- | ----------- |
-| `--mode` | str | `"t2v"` | VACE task mode: `t2v`, `i2v`, `v2lf`, `flf2v`, `inpaint`, `r2v` |
-| `--model` | str | `"Wan-AI/Wan2.1-VACE-14B-diffusers"` | Model ID or local path |
-| `--prompt` | str | `"A cat walking in a garden"` | Text description of desired video |
-| `--negative-prompt` | str | `""` | Negative prompt for classifier-free guidance |
-| `--image` | str | `None` | Input image path (for I2V, R2V, FLF2V, inpaint modes) |
-| `--last-image` | str | `None` | Last frame image path (for FLF2V mode) |
-| `--height` | int | `480` | Output video height in pixels (should be a multiple of 16) |
-| `--width` | int | `832` | Output video width in pixels (should be a multiple of 16) |
-| `--num-frames` | int | `81` | Number of video frames to generate |
-| `--num-inference-steps` | int | `30` | Number of denoising steps (more steps = higher quality, slower) |
-| `--guidance-scale` | float | `5.0` | Classifier-free guidance scale |
-| `--flow-shift` | float | `5.0` | Scheduler flow shift parameter |
-| `--seed` | int | `42` | Random seed for deterministic sampling |
-| `--fps` | int | `16` | Frames per second for the saved MP4 |
-| `--output` | str | `"vace_output.mp4"` | Path to save the generated video |
-| `--vae-use-tiling` | flag | on | Enable VAE tiling for memory optimization |
-| `--cfg-parallel-size` | int | `1` | Set to `2` to enable CFG Parallel |
-| `--ulysses-degree` | int | `1` | Ulysses sequence parallel degree for multi-GPU inference |
-| `--ring-degree` | int | `1` | Ring sequence parallel degree for hybrid Ulysses + Ring inference |
-| `--tensor-parallel-size` | int | — | Tensor parallel size |
-| `--enforce-eager` | flag | off | Disable torch.compile |
-
-> If you encounter OOM errors, try `--vae-use-tiling` or multi-GPU parallelism options (`--ulysses-degree`, `--cfg-parallel-size`).
-
-## More CLI Examples
-
-### Image-to-Video (I2V)
+## Image-to-Video (I2V)
 
 First frame is kept, remaining frames are generated:
 
@@ -112,7 +31,7 @@ python vace_video_generation.py \
   --output i2v_output.mp4
 ```
 
-### First-Last-Frame Interpolation (FLF2V)
+## First-Last-Frame Interpolation (FLF2V)
 
 ```bash
 python vace_video_generation.py \
@@ -123,7 +42,7 @@ python vace_video_generation.py \
   --output flf2v_output.mp4
 ```
 
-### Inpainting
+## Inpainting
 
 Center vertical stripe is masked and regenerated:
 
@@ -136,7 +55,7 @@ python vace_video_generation.py \
   --output inpaint_output.mp4
 ```
 
-### Reference Image-guided (R2V)
+## Reference Image-guided (R2V)
 
 ```bash
 python vace_video_generation.py \
@@ -146,3 +65,24 @@ python vace_video_generation.py \
   --height 480 --width 832 --num-frames 81 \
   --output r2v_output.mp4
 ```
+
+## Key Arguments
+
+- `--mode`: VACE task mode (`t2v`, `i2v`, `flf2v`, `inpaint`, `r2v`).
+- `--model`: Model ID (default: `Wan-AI/Wan2.1-VACE-1.3B-diffusers`).
+- `--image`: Input image for I2V, inpainting, and R2V modes.
+- `--last-image`: Last frame image for FLF2V mode.
+- `--prompt`: Text description of desired video.
+- `--height/--width`: Output resolution (default 480x832). Dimensions should be multiples of 16.
+- `--num-frames`: Number of frames (default 81).
+- `--guidance-scale`: CFG scale (default 5.0).
+- `--flow-shift`: Scheduler flow shift (default 5.0).
+- `--num-inference-steps`: Number of denoising steps (default 30).
+- `--fps`: Frames per second for the saved MP4 (default 16).
+- `--output`: Path to save the generated video.
+- `--vae-use-tiling`: Enable VAE tiling for memory optimization.
+- `--ulysses-degree`: Ulysses sequence parallelism degree for multi-GPU.
+- `--cfg-parallel-size`: CFG parallel size for multi-GPU.
+- `--tensor-parallel-size`: Tensor parallel size.
+
+> If you encounter OOM errors, try `--vae-use-tiling` or multi-GPU parallelism options.

examples/offline_inference/x_to_video_audio/x_to_video_audio.md

@@ -102,5 +102,3 @@ Key arguments:
 - `--audio-negative-prompt`: negative prompt for audio generation.
 - `--enable-cpu-offload`: enable CPU offload (defaults False).
 - `--cache-backend`: enable `cache_dit` for acceleration.
-- `--quantization`: online (dynamic) quantization method — `fp8` or `int8`.
-  (VAEs, the T5 text encoder, norms and modulation stay bf16.)

examples/offline_inference/x_to_video_audio/x_to_video_audio.py

@@ -4,7 +4,6 @@
 import argparse
 import re
 import time
-from typing import Any
 
 import numpy as np
 from PIL import Image
@@ -42,12 +41,6 @@ def parse_args() -> argparse.Namespace:
         choices=[1, 2, 3, 4],
         help="Number of GPUs used for classifier free guidance parallel size (max 4 branches).",
     )
-    parser.add_argument(
-        "--tensor-parallel-size",
-        type=int,
-        default=1,
-        help=("Number of GPUs used for tensor parallelism (TP) inside the DiT."),
-    )
     parser.add_argument(
         "--video-negative-prompt",
         default="jitter, bad hands, blur, distortion",
@@ -59,13 +52,6 @@ def parse_args() -> argparse.Namespace:
         help="Negative prompt for audio.",
     )
     parser.add_argument("--output", default="dreamid_output.mp4", help="Output video path.")
-    parser.add_argument(
-        "--quantization",
-        type=str,
-        default=None,
-        choices=["fp8", "int8"],
-        help=("Online (dynamic) quantization method for the DreamID-Omni transformer. "),
-    )
     parser.add_argument(
         "--enable-cpu-offload",
         action="store_true",
@@ -114,27 +100,6 @@ def load_image_and_audio(image_paths, audio_paths):
     return image, audio
 
 
-def _extract_peak_memory_mb(result: Any) -> float:
-    """Pull worker-reported peak VRAM (MiB) generation result.
-
-    Mirrors vllm_omni/entrypoints/openai/serving_video.py:_extract_peak_memory_mb.
-    """
-    if isinstance(result, list):
-        result = result[0] if result else None
-    if result is None:
-        return 0.0
-    val = getattr(result, "peak_memory_mb", 0.0)
-    if not val:
-        inner = getattr(result, "request_output", None)
-        if isinstance(inner, list):
-            inner = inner[0] if inner else None
-        val = getattr(inner, "peak_memory_mb", 0.0)
-    try:
-        return float(val or 0.0)
-    except (TypeError, ValueError):
-        return 0.0
-
-
 def main() -> None:
     args = parse_args()
     if args.prompt is None and args.prompt_file is None:
@@ -177,7 +142,6 @@ def main() -> None:
 
     parallel_config = DiffusionParallelConfig(
         cfg_parallel_size=args.cfg_parallel_size,
-        tensor_parallel_size=args.tensor_parallel_size,
         use_hsdp=args.use_hsdp,
         hsdp_shard_size=args.hsdp_shard_size,
         hsdp_replicate_size=args.hsdp_replicate_size,
@@ -198,7 +162,7 @@ def main() -> None:
             "scm_steps_policy": "dynamic",
         }
 
-    omni_kwargs = dict(
+    omni = Omni(
         model=args.model,
         parallel_config=parallel_config,
         model_type=args.model_type,
@@ -207,9 +171,6 @@ def main() -> None:
         cache_backend=args.cache_backend,
         cache_config=cache_config,
     )
-    if args.quantization is not None:
-        omni_kwargs["quantization"] = args.quantization
-    omni = Omni(**omni_kwargs)
     start = time.perf_counter()
     outputs = omni.generate(prompt, sampling_params)
     elapsed = time.perf_counter() - start
@@ -247,9 +208,6 @@ def main() -> None:
         f.write(video_bytes)
     print(f"Saved generated video to {output_path}")
     print(f"Total time: {elapsed:.2f}s")
-    peak_mb = _extract_peak_memory_mb(outputs)
-    if peak_mb:
-        print(f"Worker peak GPU memory (reserved): {peak_mb:.2f} MiB ({peak_mb / 1024:.2f} GiB)")
 
 
 if __name__ == "__main__":

examples/online_serving/bagel/README.md

@@ -68,67 +68,6 @@ vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 --tensor-parallel-size
 
 Or set `tensor_parallel_size` per stage in a custom deploy YAML.
 
-### VAE Patch Parallelism
-
-[VAE Patch Parallelism](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/parallelism/vae_patch_parallel.html) distributes Bagel VAE **decode/encode** across multiple GPUs by splitting latent tiles. It lowers **per-GPU peak memory during VAE decode**, which helps high-resolution `text2img` / `img2img` when VAE becomes a bottleneck.
-
-**Scope for Bagel:**
-
-| Topology | VAE patch parallel |
-| :------- | :----------------- |
-| **Single-stage** (DiT only) | Supported on stage 0 (`BagelPipeline` + `DistributedAutoEncoder`) |
-| **Two-stage** | Supported on **stage 1 (DiT)** only; stage 0 (Thinker) uses encoder-only VAE and is unrelated |
-
-**Requirements:**
-
-- `vae_patch_parallel_size > 1` and a distributed VAE (`DistributedAutoEncoder` on the DiT pipeline).
-- The DiT process group must have at least `vae_patch_parallel_size` ranks. In practice this means the diffusion stage `world_size` must be ≥ 2 (commonly `tensor_parallel_size=2` on that stage).
-- `vae_use_tiling` must be enabled. If you set `vae_patch_parallel_size > 1` and omit tiling, the registry auto-enables `vae_use_tiling` at startup.
-
-VAE patch parallel **reuses the DiT process group** (`dit_group`); it does not create a separate VAE-only worker pool. It is not a substitute for single-GPU VAE tiling (`vae_pp=1`).
-
-**Online serving (single-stage, 2 GPUs):**
-
-```bash
-CUDA_VISIBLE_DEVICES=0,1 vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 \
-    --deploy-config vllm_omni/deploy/bagel_single_stage.yaml \
-    --tensor-parallel-size 2 \
-    --vae-patch-parallel-size 2 \
-    --vae-use-tiling
-```
-
-**Online serving (two-stage, VAE PP on DiT stage 1):** use a custom deploy YAML, for example:
-
-```yaml
-stages:
-  - stage_id: 0
-    devices: "0"
-    # Thinker (AR) — no VAE patch parallel here
-
-  - stage_id: 1
-    devices: "0,1"
-    vae_use_tiling: true
-    parallel_config:
-      tensor_parallel_size: 2
-      vae_patch_parallel_size: 2
-```
-
-```bash
-vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 \
-    --deploy-config /path/to/bagel_vae_pp.yaml
-```
-
-**Verify it is active** (check server logs at startup):
-
-```text
-INFO ... vae_patch_parallel_size=2 requires vae_use_tiling; automatically enabling it.
-```
-
-| CLI flag | Default | Description |
-| :------- | :------ | :---------- |
-| `--vae-patch-parallel-size` | `1` | Number of DiT ranks used for VAE tile parallelism. Set to `2` or higher to enable. Should be ≤ DiT process group size (typically match `--tensor-parallel-size` on the diffusion stage). |
-| `--vae-use-tiling` | off | Enable VAE spatial tiling. Required for VAE patch parallel (auto-enabled when `vae_patch_parallel_size > 1`). |
-
 #### Hybrid Sharded Data Parallel (HSDP)
 
 For larger Bagel deployments on multiple GPUs, you can enable HSDP (Hybrid Sharded Data Parallel) by modifying the stage configuration (for example, [`bagel.yaml`](../../../vllm_omni/deploy/bagel.yaml)). HSDP shards transformer weights across GPUs to reduce per-GPU memory usage.
@@ -384,8 +323,6 @@ python openai_chat_client.py \
 | `stages[].gpu_memory_utilization` | per-stage | Fraction of GPU memory to use |
 | `stages[].enforce_eager` | per-stage | Disable CUDA graphs |
 | `stages[].tensor_parallel_size` | per-stage | TP degree for this stage |
-| `stages[].parallel_config.vae_patch_parallel_size` | per-stage (DiT) | VAE tile parallelism degree (DiT stage only) |
-| `stages[].vae_use_tiling` | per-stage (DiT) | Enable VAE tiling (required for VAE patch parallel) |
 | `connectors` | top-level | Define available connector instances (SHM, Mooncake) |
 | `platforms` | top-level | Platform-specific overrides (e.g. `xpu`) |
 

examples/online_serving/dreamzero/README.md

@@ -1,132 +0,0 @@
-# DreamZero
-
-DreamZero is the robot-policy serving example for the OpenPI-compatible `/v1/realtime/robot/openpi` endpoint.
-
-## Supported checkpoints
-
-| Model | HuggingFace repo |
-|---|---|
-| DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` |
-
-## Quick start
-
-### Start the server
-
-```bash
-bash examples/online_serving/dreamzero/run_server.sh
-```
-
-By default this launches:
-
-```bash
-vllm serve GEAR-Dreams/DreamZero-DROID --omni --port 8000 \
-    --served-model-name dreamzero-droid \
-    --deploy-config vllm_omni/deploy/dreamzero_tp1_cfg2.yaml \
-    --enforce-eager --disable-log-stats
-```
-
-Override `MODEL`, `PORT`, `HOST`, `DEPLOY_CONFIG`, or `SERVED_MODEL_NAME` through the script environment if needed.
-
-### Install example dependencies
-
-The core `pip install -e .` setup does not include the extra packages used by the DreamZero example scripts.
-
-- `openpi_client.py`:
-  `openpi-client`, `websockets`, `opencv-python`
-- `droid_sim_eval_client.py`:
-  `mediapy`, `websockets`, `openpi-client`
-
-The DROID sim-eval script also needs an Isaac Lab environment that provides `isaaclab`, `isaaclab_tasks`, `sim_evals`, and `gymnasium`.
-
-If you run the DROID client on Python < 3.12, also install `typing-extensions`.
-
-### Configure TP and CFG parallelism
-
-The bundled DreamZero configs intentionally keep only:
-
-| Config | Purpose |
-|---|---|
-| `vllm_omni/deploy/dreamzero.yaml` | Default TP=1, CFG parallel disabled |
-| `vllm_omni/deploy/dreamzero_tp1_cfg2.yaml` | TP=1, CFG parallel size=2 |
-
-For other topologies, use CLI parallelism flags and update stage 0 `devices` with `--stage-overrides`. The number of listed devices must match `tensor_parallel_size * cfg_parallel_size`.
-
-TP=2 with CFG parallel disabled:
-
-```bash
-vllm serve GEAR-Dreams/DreamZero-DROID --omni --port 8000 \
-    --served-model-name dreamzero-droid \
-    --deploy-config vllm_omni/deploy/dreamzero.yaml \
-    --tensor-parallel-size 2 \
-    --stage-overrides '{"0": {"devices": "0,1"}}' \
-    --enforce-eager --disable-log-stats
-```
-
-TP=2 with CFG parallel size=2:
-
-```bash
-vllm serve GEAR-Dreams/DreamZero-DROID --omni --port 8000 \
-    --served-model-name dreamzero-droid \
-    --deploy-config vllm_omni/deploy/dreamzero.yaml \
-    --tensor-parallel-size 2 \
-    --cfg-parallel-size 2 \
-    --stage-overrides '{"0": {"devices": "0,1,2,3"}}' \
-    --enforce-eager --disable-log-stats
-```
-
-### Download example assets
-
-The OpenPI client and DROID sim-eval example expect the three camera MP4 files in `outputs/dreamzero/assets`.
-
-```bash
-hf download YangshenDeng/vllm-omni-dreamzero-assets --repo-type dataset --local-dir outputs/dreamzero/assets
-```
-
-### Run the OpenPI client
-
-```bash
-python examples/online_serving/dreamzero/openpi_client.py \
-    --host 127.0.0.1 \
-    --port 8000 \
-    --video-dir outputs/dreamzero/assets
-```
-
-This client uses downloaded example videos and talks to the OpenPI websocket server.
-
-### Run DROID sim eval
-
-```bash
-${ISAACLAB_LAUNCHER} -p examples/online_serving/dreamzero/droid_sim_eval_client.py \
-    --host 127.0.0.1 \
-    --port 8000 \
-    --scene 1 \
-    --episodes 1 \
-    --headless \
-    --device cuda:0
-```
-
-Set `ISAACLAB_LAUNCHER=path/to/isaaclab.sh` from the vLLM-Omni repository root before running the command.
-This launches Isaac Lab / sim-evals and runs the DROID benchmark loop against the same websocket endpoint.
-
-### Export comparison videos offline
-
-```bash
-python examples/offline_inference/dreamzero/export_prediction_video.py \
-    --deploy-config vllm_omni/deploy/dreamzero_tp1_cfg2.yaml \
-    --save-input-video \
-    --save-gif
-```
-
-The export script uses local `Omni` inference, not the websocket server. It writes the input rollout video and the predicted output video artifacts for side-by-side inspection.
-
-### MolmoSpace demo
-
-```bash
-python examples/online_serving/dreamzero/molmospace_dreamzero_eval_demo.py \
-    --host 127.0.0.1 \
-    --port 8000 \
-    --benchmark_dir /path/to/benchmark \
-    --output_dir /path/to/output
-```
-
-This demo adapts DreamZero to the MolmoSpace-style remote policy eval loop.

examples/online_serving/dreamzero/droid_sim_eval_client.py

@@ -1,819 +0,0 @@
-#!/usr/bin/env python3
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""Run one or more DROID sim-eval rollouts against the vLLM DreamZero server.
-
-This script is the vLLM/OpenPI adaptation of the upstream DreamZero
-`eval_utils/run_sim_eval.py` sim-eval client.
-
-Behavior intentionally kept close to upstream:
-- same DROID observation extraction (`external_cam`, `external_cam_2`,
-  `wrist_cam`, joint position, gripper position)
-- same resize-with-pad preprocessing to `(180, 320)`
-- same `open_loop_horizon=8`
-- same gripper binarization rule (`> 0.5 -> 1`, else `0`)
-- same per-scene language prompts
-
-Unlike upstream DreamZero, vLLM serves the compatible websocket policy endpoint
-at `/v1/realtime/robot/openpi`, so this script includes the path suffix in the
-client URI.
-
-Run this script through Isaac Lab's launcher from the vLLM-Omni repository
-root, for example:
-
-    "${ISAACLAB_LAUNCHER}" -p \
-      examples/online_serving/dreamzero/droid_sim_eval_client.py \
-      --host 127.0.0.1 \
-      --port 8000 \
-      --scene 1 \
-      --episodes 1 \
-      --headless \
-      --device cuda:1
-"""
-
-from __future__ import annotations
-
-import argparse
-import json
-import logging
-import sys
-import time
-import uuid
-from dataclasses import dataclass
-from datetime import UTC, datetime
-from pathlib import Path
-from typing import Any
-
-import numpy as np
-import torch
-
-try:
-    import mediapy
-except ImportError as exc:  # pragma: no cover - runtime dependency guard
-    raise ImportError("Optional DROID sim-eval client requires `mediapy`.") from exc
-
-try:
-    from typing import override
-except ImportError:
-    try:
-        from typing_extensions import override
-    except ImportError as exc:  # pragma: no cover - runtime dependency guard
-        raise ImportError("Optional DROID sim-eval client requires `typing-extensions` on Python < 3.12.") from exc
-
-try:
-    import websockets.sync.client
-except ImportError as exc:  # pragma: no cover - runtime dependency guard
-    raise ImportError("Optional DROID sim-eval client requires `websockets`.") from exc
-
-# NOTE:
-# This directory already contains a local file named `openpi_client.py`.
-# However, what we want here is the *installed* `openpi_client` package from
-# upstream OpenPI, not the sibling example file. When a script is executed
-# directly, Python often puts the script directory at `sys.path[0]`, which
-# would cause `import openpi_client` to resolve to the local example file and
-# create a circular import.
-#
-# To avoid that ambiguity, temporarily remove the current example directory
-# from the front of `sys.path`, import the real package, and then restore the
-# path afterwards.
-example_dir = str(Path(__file__).resolve().parent)
-removed_path = False
-if sys.path and sys.path[0] == example_dir:
-    sys.path.pop(0)
-    removed_path = True
-try:
-    from openpi_client import image_tools, msgpack_numpy
-    from openpi_client.base_policy import BasePolicy
-except ImportError as exc:  # pragma: no cover - runtime dependency guard
-    raise ImportError("Optional DROID sim-eval client requires the `openpi-client` package.") from exc
-finally:
-    if removed_path:
-        sys.path.insert(0, example_dir)
-
-# -----------------------------------------------------------------------------
-# Constants
-# -----------------------------------------------------------------------------
-#
-# These values intentionally mirror the upstream DreamZero sim-eval client
-# where possible. The important distinction is:
-#
-# - ACTION_HORIZON = 24
-#     The model returns 24 future actions per inference call.
-# - DEFAULT_OPEN_LOOP_HORIZON = 8
-#     The sim client only executes the first 8 actions locally before asking
-#     the server to replan from a fresh observation.
-#
-# So a single server call predicts 24x8 actions, but the rollout consumes only
-# 8 of them before replanning.
-PING_INTERVAL_SECS = 60
-PING_TIMEOUT_SECS = 600
-DEFAULT_PATH = "/v1/realtime/robot/openpi"
-DEFAULT_OPEN_LOOP_HORIZON = 8
-ACTION_HORIZON = 24
-ACTION_DIM = 8
-DEFAULT_OUTPUT_ROOT = Path("runs") / "dreamzero_sim_eval"
-SCENE_PROMPTS = {
-    1: "put the cube in the bowl",
-    2: "pick up the can and put it in the mug",
-    3: "put the banana in the bin",
-}
-
-
-def _decode_action_response(response: bytes | str) -> np.ndarray:
-    if isinstance(response, str):
-        raise RuntimeError(f"Error in inference server:\n{response}")
-    decoded = msgpack_numpy.unpackb(response)
-    if isinstance(decoded, dict) and decoded.get("type") == "error":
-        message = decoded.get("message", decoded)
-        raise RuntimeError(f"Error in inference server:\n{message}")
-    return np.asarray(decoded, dtype=np.float32)
-
-
-@dataclass(frozen=True)
-class StepRecord:
-    """One fully materialized rollout step for later JSON export.
-
-    The `episode_00.json` artifact is intended to be human-readable and
-    post-process-friendly. Instead of keeping raw tensors around, each step is
-    flattened into plain Python types so it can be serialized directly.
-    """
-
-    # Index of this control step within the episode.
-    step_index: int
-    # Whether this step triggered a fresh model call (as opposed to reusing the
-    # cached open-loop chunk from the previous server response).
-    used_server_call: bool
-    # End-to-end latency of the server call that produced the current chunk.
-    # This is `None` on steps that only reuse cached actions.
-    chunk_latency_s: float | None
-    # The concrete 8-D action sent into the simulator at this step.
-    action: list[float]
-    # Observed 7-DoF arm joint positions before the next environment step.
-    joint_position: list[float]
-    # Observed gripper scalar before the next environment step.
-    gripper_position: list[float]
-    # Reward and termination signals directly returned by the simulator.
-    reward: float
-    terminated: bool
-    truncated: bool
-    # Optional scene object positions for downstream debugging / success
-    # heuristics. This may be empty if the environment does not expose them.
-    scene_objects: dict[str, list[float]]
-
-
-class OpenPIWebsocketClientPolicy(BasePolicy):
-    """Minimal websocket client for the DreamZero/OpenPI policy protocol.
-
-    Protocol shape:
-    - connect -> server immediately sends a metadata dict
-    - infer   -> send msgpack observation, receive action chunk
-    - reset   -> send msgpack reset command, receive confirmation string
-
-    This class intentionally stays very small because the more interesting
-    DreamZero-specific behavior lives one layer above, in
-    `DreamZeroJointPosClient`.
-    """
-
-    def __init__(
-        self,
-        host: str = "127.0.0.1",
-        port: int = 8000,
-        path: str = DEFAULT_PATH,
-    ) -> None:
-        # vLLM serves the robot endpoint under `/v1/realtime/robot/openpi`.
-        self._uri = f"ws://{host}:{port}{path}"
-        # Upstream protocol uses msgpack with numpy support, not JSON.
-        self._packer = msgpack_numpy.Packer()
-        # Connect immediately and cache the server handshake metadata.
-        self._ws, self._server_metadata = self._wait_for_server()
-
-    def get_server_metadata(self) -> dict[str, Any]:
-        """Return a copy of the server handshake metadata."""
-
-        return dict(self._server_metadata)
-
-    def _wait_for_server(self):
-        """Connect to the websocket server and read the initial metadata frame."""
-
-        logging.info("Connecting to %s", self._uri)
-        conn = websockets.sync.client.connect(
-            self._uri,
-            compression=None,
-            max_size=None,
-            ping_interval=PING_INTERVAL_SECS,
-            ping_timeout=PING_TIMEOUT_SECS,
-        )
-        metadata = msgpack_numpy.unpackb(conn.recv())
-        if not isinstance(metadata, dict):
-            raise TypeError(f"Expected dict metadata from server, got {type(metadata)!r}")
-        return conn, metadata
-
-    @override
-    def infer(self, obs: dict[str, Any]) -> np.ndarray:
-        """Send an inference request and return the decoded action chunk."""
-
-        # Keep the upstream DreamZero/OpenPI convention that the request itself
-        # tells the server which logical endpoint is being called.
-        payload = dict(obs)
-        payload["endpoint"] = "infer"
-        self._ws.send(self._packer.pack(payload))
-        response = self._ws.recv()
-        return _decode_action_response(response)
-
-    @override
-    def reset(self, reset_info: dict[str, Any] | None = None) -> str:
-        """Tell the server to reset its session-side state."""
-
-        payload = dict(reset_info or {})
-        payload["endpoint"] = "reset"
-        self._ws.send(self._packer.pack(payload))
-        response = self._ws.recv()
-        if isinstance(response, str):
-            return response
-        decoded = msgpack_numpy.unpackb(response)
-        if not isinstance(decoded, dict) or decoded.get("status") != "reset successful":
-            raise RuntimeError(f"Unexpected reset response: {decoded!r}")
-        return str(decoded["status"])
-
-    def close(self) -> None:
-        """Close the websocket connection explicitly."""
-
-        self._ws.close()
-
-
-class DreamZeroJointPosClient:
-    """DROID sim-eval client that talks to the vLLM OpenPI websocket server.
-
-    This is the main compatibility layer between:
-    - Isaac Lab DROID observations (`obs["policy"][...]`)
-    - DreamZero/OpenPI websocket payloads
-    - local open-loop action reuse across several simulator steps
-
-    In other words:
-        simulator obs -> websocket request -> action chunk -> one action per step
-    """
-
-    def __init__(
-        self,
-        remote_host: str = "127.0.0.1",
-        remote_port: int = 8000,
-        path: str = DEFAULT_PATH,
-        open_loop_horizon: int = DEFAULT_OPEN_LOOP_HORIZON,
-    ) -> None:
-        # Low-level transport client.
-        self.client = OpenPIWebsocketClientPolicy(remote_host, remote_port, path=path)
-        # Number of actions to execute locally before replanning.
-        self.open_loop_horizon = open_loop_horizon
-        # Cursor into the currently cached action chunk.
-        self.actions_from_chunk_completed = 0
-        # Most recent `(ACTION_HORIZON, ACTION_DIM)` server response.
-        self.pred_action_chunk: np.ndarray | None = None
-        # Session id is part of the DreamZero serving contract. Changing it
-        # causes the server side to treat the rollout as a fresh episode.
-        self.session_id = str(uuid.uuid4())
-        # Simple runtime stats for reporting.
-        self.server_calls = 0
-        self.last_chunk_latency_s: float | None = None
-        self.last_used_server_call = False
-
-    def metadata(self) -> dict[str, Any]:
-        """Expose the server metadata to callers / logs."""
-
-        return self.client.get_server_metadata()
-
-    def reset(self) -> str:
-        """Reset local chunk state and remote session state.
-
-        Local reset:
-        - drop cached action chunk
-        - rewind chunk cursor
-        - allocate a fresh session id
-
-        Remote reset:
-        - send a websocket `reset` message so the server can clear any
-          request/session-side state it associates with this client
-        """
-
-        self.actions_from_chunk_completed = 0
-        self.pred_action_chunk = None
-        self.session_id = str(uuid.uuid4())
-        self.last_chunk_latency_s = None
-        self.last_used_server_call = False
-        return self.client.reset({})
-
-    def infer(self, obs: dict[str, Any], instruction: str) -> dict[str, Any]:
-        """Turn one simulator observation into one executable 8-D action.
-
-        Key behavior:
-        - call the server only when the local chunk cache is empty/exhausted
-        - otherwise, keep consuming the cached chunk open-loop
-        - always return exactly one 8-D action for the current simulator step
-        """
-
-        # Convert Isaac Lab observation structure into a plain numpy-friendly
-        # record that is easier to serialize and visualize.
-        curr_obs = self._extract_observation(obs)
-        self.last_used_server_call = False
-
-        # Replan if:
-        # 1. this is the first step of a rollout / chunk
-        # 2. we already consumed `open_loop_horizon` actions from the current chunk
-        # 3. no cached chunk is currently available
-        if (
-            self.actions_from_chunk_completed == 0
-            or self.actions_from_chunk_completed >= self.open_loop_horizon
-            or self.pred_action_chunk is None
-        ):
-            self.actions_from_chunk_completed = 0
-            # Build the exact DreamZero/OpenPI payload expected by the server.
-            #
-            # Notes:
-            # - images are resized/padded to the serving contract's 180x320
-            # - proprio is cast to float64 to match upstream client behavior
-            # - cartesian_position is currently unused by DreamZero DROID, so
-            #   a dummy zero vector is sent for protocol completeness
-            request_data = {
-                "observation/exterior_image_0_left": image_tools.resize_with_pad(curr_obs["right_image"], 180, 320),
-                "observation/exterior_image_1_left": image_tools.resize_with_pad(curr_obs["left_image"], 180, 320),
-                "observation/wrist_image_left": image_tools.resize_with_pad(curr_obs["wrist_image"], 180, 320),
-                "observation/joint_position": curr_obs["joint_position"].astype(np.float64),
-                "observation/cartesian_position": np.zeros((6,), dtype=np.float64),
-                "observation/gripper_position": curr_obs["gripper_position"].astype(np.float64),
-                "prompt": instruction,
-                "session_id": self.session_id,
-            }
-
-            # Measure end-to-end server latency for this chunk request.
-            start = time.perf_counter()
-            actions = self.client.infer(request_data)
-            self.last_chunk_latency_s = time.perf_counter() - start
-            self.last_used_server_call = True
-            self.server_calls += 1
-
-            # DreamZero DROID serving is expected to return an action chunk with
-            # 24 future actions, each action being 8-D.
-            if actions.ndim != 2:
-                raise AssertionError(f"Expected 2D action array, got shape {actions.shape}")
-            if actions.shape != (ACTION_HORIZON, ACTION_DIM):
-                raise AssertionError(f"Expected action shape {(ACTION_HORIZON, ACTION_DIM)}, got {actions.shape}")
-            self.pred_action_chunk = actions
-
-        # Consume exactly one action row from the cached chunk for this
-        # simulator step.
-        action = np.array(self.pred_action_chunk[self.actions_from_chunk_completed], copy=True)
-        self.actions_from_chunk_completed += 1
-
-        # Upstream DreamZero sim-eval binarizes the gripper command.
-        action[-1] = 1.0 if action[-1].item() > 0.5 else 0.0
-
-        # Produce a human-friendly visualization strip for videos:
-        # right external | wrist | left external
-        img1 = image_tools.resize_with_pad(curr_obs["right_image"], 224, 224)
-        img2 = image_tools.resize_with_pad(curr_obs["wrist_image"], 224, 224)
-        img3 = image_tools.resize_with_pad(curr_obs["left_image"], 224, 224)
-        viz = np.concatenate([img1, img2, img3], axis=1)
-
-        # Return both the executable action and auxiliary debug info.
-        return {
-            "action": action,
-            "viz": viz,
-            "joint_position": curr_obs["joint_position"],
-            "gripper_position": curr_obs["gripper_position"],
-            "used_server_call": self.last_used_server_call,
-            "chunk_latency_s": self.last_chunk_latency_s if self.last_used_server_call else None,
-        }
-
-    @staticmethod
-    def _extract_observation(obs_dict: dict[str, Any]) -> dict[str, np.ndarray]:
-        """Extract the pieces DreamZero cares about from Isaac Lab observations.
-
-        `sim-evals` exposes camera frames and robot state inside the
-        `obs["policy"]` group. This helper converts those tensors into numpy
-        arrays so they can be fed into image preprocessing / websocket packing.
-        """
-
-        policy = obs_dict["policy"]
-        # Isaac Lab stores camera observations as batched tensors; use env 0.
-        right_image = policy["external_cam"][0].clone().detach().cpu().numpy()
-        left_image = policy["external_cam_2"][0].clone().detach().cpu().numpy()
-        wrist_image = policy["wrist_cam"][0].clone().detach().cpu().numpy()
-        # Robot proprioception.
-        joint_position = policy["arm_joint_pos"].clone().detach().cpu().numpy()
-        gripper_position = policy["gripper_pos"].clone().detach().cpu().numpy()
-
-        return {
-            "right_image": right_image,
-            "left_image": left_image,
-            "wrist_image": wrist_image,
-            "joint_position": joint_position,
-            "gripper_position": gripper_position,
-        }
-
-
-def _scene_instruction(scene: int) -> str:
-    """Map a numeric scene id onto the fixed language prompt used for rollout."""
-
-    try:
-        return SCENE_PROMPTS[scene]
-    except KeyError as exc:
-        raise ValueError(f"Unsupported scene {scene}. Available scenes: {sorted(SCENE_PROMPTS)}") from exc
-
-
-def _capture_scene_objects(env: Any) -> dict[str, list[float]]:
-    """Best-effort extraction of scene object root positions.
-
-    This is only for debugging / reporting. The rollout logic itself does not
-    depend on these positions.
-    """
-
-    objects: dict[str, list[float]] = {}
-    scene = getattr(env, "scene", None)
-    if scene is None:
-        return objects
-
-    # Skip non-task entities such as cameras, the robot, and lighting.
-    for name in scene.keys():
-        if name in {"robot", "external_cam", "external_cam_2", "wrist_cam", "sphere_light", "scene"}:
-            continue
-        entity = scene[name]
-        data = getattr(entity, "data", None)
-        root_pos_w = getattr(data, "root_pos_w", None)
-        if root_pos_w is None:
-            continue
-        # Convert tensors to plain lists for JSON serialization.
-        value = root_pos_w[0].detach().cpu().to(torch.float32).tolist()
-        objects[str(name)] = [float(x) for x in value]
-    return objects
-
-
-def _maybe_infer_success(scene: int, final_objects: dict[str, list[float]]) -> dict[str, Any]:
-    """Best-effort geometric heuristic.
-
-    The simulator itself does not expose a built-in success term; this function
-    provides a transparent fallback for human-readable reporting only.
-    """
-
-    task_pairs = {
-        1: ("cube", "bowl"),
-        2: ("can", "mug"),
-        3: ("banana", "bin"),
-    }
-    source_name, target_name = task_pairs.get(scene, (None, None))
-    if source_name not in final_objects or target_name not in final_objects:
-        return {
-            "has_builtin_success": False,
-            "heuristic_success": None,
-            "reason": "scene object names unavailable for heuristic",
-        }
-
-    source = np.asarray(final_objects[source_name], dtype=np.float32)
-    target = np.asarray(final_objects[target_name], dtype=np.float32)
-    xy_distance = float(np.linalg.norm(source[:2] - target[:2]))
-    z_delta = float(source[2] - target[2])
-    heuristic_success = bool(xy_distance < 0.12 and z_delta > -0.08)
-    return {
-        "has_builtin_success": False,
-        "heuristic_success": heuristic_success,
-        "xy_distance": xy_distance,
-        "z_delta": z_delta,
-        "source_object": source_name,
-        "target_object": target_name,
-        "reason": "no built-in env success flag; using final object pose heuristic",
-    }
-
-
-def _dump_json(path: Path, payload: dict[str, Any]) -> None:
-    """Write a JSON file with stable UTF-8 formatting."""
-
-    path.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
-
-
-def _scalar_from_env_value(value: Any) -> float:
-    """Normalize simulator scalar outputs into a plain Python float.
-
-    Isaac Lab / Gym values may come back as tensors, numpy arrays, tuples, or
-    direct Python scalars depending on the wrapper stack. Centralizing the
-    conversion here makes the rollout loop cleaner and more robust.
-    """
-
-    if isinstance(value, torch.Tensor):
-        return float(value.reshape(-1)[0].detach().cpu().item())
-    if isinstance(value, np.ndarray):
-        return float(value.reshape(-1)[0])
-    if isinstance(value, (list, tuple)):
-        return float(np.asarray(value).reshape(-1)[0])
-    return float(value)
-
-
-def _bool_from_env_value(value: Any) -> bool:
-    """Normalize simulator boolean-like outputs into a plain Python bool."""
-
-    if isinstance(value, torch.Tensor):
-        return bool(value.reshape(-1)[0].detach().cpu().item())
-    if isinstance(value, np.ndarray):
-        return bool(value.reshape(-1)[0])
-    if isinstance(value, (list, tuple)):
-        return bool(np.asarray(value).reshape(-1)[0])
-    return bool(value)
-
-
-def _make_output_dir(output_root: Path, scene: int) -> Path:
-    """Create a timestamped output directory for one sim-eval run."""
-
-    timestamp = datetime.now(UTC).strftime("%Y%m%dT%H%M%SZ")
-    output_dir = output_root / f"scene{scene}_{timestamp}"
-    output_dir.mkdir(parents=True, exist_ok=True)
-    return output_dir
-
-
-def main() -> None:
-    """Entry point for one or more DROID simulation rollouts.
-
-    High-level flow:
-    1. parse command-line flags
-    2. bootstrap Isaac Lab / sim-evals imports
-    3. create the DROID environment
-    4. connect the DreamZero websocket client
-    5. run `episodes` rollouts
-    6. export videos + JSON summaries
-    """
-
-    # Script-level arguments.
-    parser = argparse.ArgumentParser(description=__doc__)
-    parser.add_argument("--episodes", type=int, default=1, help="Number of episodes to run.")
-    parser.add_argument("--scene", type=int, default=1, help="DROID scene id (1/2/3).")
-    parser.add_argument("--host", type=str, default="127.0.0.1", help="vLLM DreamZero server host.")
-    parser.add_argument("--port", type=int, default=8000, help="vLLM DreamZero server port.")
-    parser.add_argument("--path", type=str, default=DEFAULT_PATH, help="Websocket path suffix.")
-    parser.add_argument(
-        "--open-loop-horizon",
-        type=int,
-        default=DEFAULT_OPEN_LOOP_HORIZON,
-        help="How many actions to consume locally before requesting the next chunk.",
-    )
-    parser.add_argument(
-        "--output-root",
-        type=Path,
-        default=DEFAULT_OUTPUT_ROOT,
-        help="Directory where videos and trajectory logs are stored.",
-    )
-
-    try:
-        from isaaclab.app import AppLauncher
-    except ImportError as exc:  # pragma: no cover - runtime dependency guard
-        raise ImportError(
-            "Optional DROID sim-eval client requires Isaac Lab (`isaaclab`). "
-            "Launch it from an Isaac Lab environment, e.g. via `isaaclab.sh -p`."
-        ) from exc
-
-    # Let Isaac Lab inject its own runtime flags (e.g. headless, device).
-    AppLauncher.add_app_launcher_args(parser)
-    args = parser.parse_args()
-
-    # DreamZero sim-eval always needs camera observations enabled.
-    args.enable_cameras = True
-    # Boot Isaac Sim / Isaac Lab.
-    app_launcher = AppLauncher(args)
-    simulation_app = app_launcher.app
-    # Set defaults so the `finally` block can clean up safely even if an
-    # earlier step fails.
-    env = None
-    client = None
-
-    # Import simulator modules only *after* the app is launched. This matches
-    # Isaac Lab's required import ordering.
-    try:
-        import gymnasium as gym
-    except ImportError as exc:  # pragma: no cover - runtime dependency guard
-        raise ImportError("Optional DROID sim-eval client requires `gymnasium`.") from exc
-
-    try:
-        import sim_evals.environments  # noqa: F401
-    except ImportError as exc:  # pragma: no cover - runtime dependency guard
-        raise ImportError(
-            "Optional DROID sim-eval client requires the external `sim-evals` package or checkout to be importable."
-        ) from exc
-
-    try:
-        from isaaclab_tasks.utils import parse_env_cfg
-    except ImportError as exc:  # pragma: no cover - runtime dependency guard
-        raise ImportError("Optional DROID sim-eval client requires `isaaclab_tasks`.") from exc
-
-    # Resolve output location and scene prompt.
-    output_dir = _make_output_dir(args.output_root.expanduser().resolve(), args.scene)
-    instruction = _scene_instruction(args.scene)
-
-    # Build the DROID environment configuration from `sim-evals`.
-    env_cfg = parse_env_cfg(
-        "DROID",
-        device=args.device,
-        num_envs=1,
-        use_fabric=True,
-    )
-    # Select one of the pre-authored scenes/tasks.
-    env_cfg.set_scene(args.scene)
-    env = gym.make("DROID", cfg=env_cfg)
-
-    # Upstream sim-evals resets twice so materials / cameras are fully ready.
-    obs, _ = env.reset()
-    obs, _ = env.reset()
-
-    # Connect the websocket policy client.
-    client = DreamZeroJointPosClient(
-        remote_host=args.host,
-        remote_port=args.port,
-        path=args.path,
-        open_loop_horizon=args.open_loop_horizon,
-    )
-
-    # Aggregated per-run results.
-    all_episode_summaries: list[dict[str, Any]] = []
-    max_steps = int(env.env.max_episode_length)
-    logging.info("DreamZero metadata: %s", client.metadata())
-    logging.info("Scene %s prompt: %s", args.scene, instruction)
-    logging.info("Writing outputs to %s", output_dir)
-
-    try:
-        # No gradients are needed in inference-only rollout mode.
-        with torch.no_grad():
-            for episode_index in range(args.episodes):
-                # Per-episode collectors.
-                frames: list[np.ndarray] = []
-                step_records: list[StepRecord] = []
-                episode_start = time.perf_counter()
-                server_time_s = 0.0
-                final_reward = 0.0
-                terminated = False
-                truncated = False
-
-                for step_index in range(max_steps):
-                    # Ask the policy for the next action. Internally this may or
-                    # may not trigger a real server request depending on whether
-                    # the local chunk cache has been exhausted.
-                    logging.debug("Episode %d step %d: requesting action", episode_index, step_index)
-                    result = client.infer(obs, instruction)
-                    logging.debug(
-                        "Episode %d step %d: got action (server_call=%s latency=%s)",
-                        episode_index,
-                        step_index,
-                        result["used_server_call"],
-                        result["chunk_latency_s"],
-                    )
-                    # Save one visualization frame per simulator step.
-                    frames.append(result["viz"])
-
-                    # Isaac Lab expects batched actions, hence `[None]`.
-                    action_tensor = torch.tensor(result["action"], dtype=torch.float32)[None]
-                    logging.debug("Episode %d step %d: stepping env", episode_index, step_index)
-                    obs, reward, term, trunc, info = env.step(action_tensor)
-                    logging.debug("Episode %d step %d: env.step returned", episode_index, step_index)
-                    logging.debug("Episode %d step %d: parsing reward/flags", episode_index, step_index)
-                    logging.debug(
-                        "Episode %d step %d: raw types reward=%s term=%s trunc=%s",
-                        episode_index,
-                        step_index,
-                        type(reward).__name__,
-                        type(term).__name__,
-                        type(trunc).__name__,
-                    )
-                    # Normalize environment outputs into plain Python scalars so
-                    # the rest of the code does not depend on wrapper-specific types.
-                    reward_value = _scalar_from_env_value(reward)
-                    term_value = _bool_from_env_value(term)
-                    trunc_value = _bool_from_env_value(trunc)
-                    logging.debug(
-                        "Episode %d step %d: parsed reward=%s term=%s trunc=%s",
-                        episode_index,
-                        step_index,
-                        reward_value,
-                        term_value,
-                        trunc_value,
-                    )
-                    # Keep scene-object capture optional. It is useful for
-                    # debugging / success heuristics, but the rollout should not
-                    # fail if the environment does not expose object roots.
-                    scene_objects = _capture_scene_objects(env)
-
-                    # Accumulate total server-side time only on steps that
-                    # triggered a fresh chunk inference.
-                    if result["chunk_latency_s"] is not None:
-                        server_time_s += float(result["chunk_latency_s"])
-
-                    # Materialize one JSON-serializable trajectory record.
-                    logging.debug("Episode %d step %d: appending trajectory", episode_index, step_index)
-                    step_records.append(
-                        StepRecord(
-                            step_index=step_index,
-                            used_server_call=bool(result["used_server_call"]),
-                            chunk_latency_s=(
-                                float(result["chunk_latency_s"]) if result["chunk_latency_s"] is not None else None
-                            ),
-                            action=[float(x) for x in np.asarray(result["action"], dtype=np.float32).tolist()],
-                            joint_position=[
-                                float(x) for x in np.asarray(result["joint_position"], dtype=np.float32).tolist()
-                            ],
-                            gripper_position=[
-                                float(x) for x in np.asarray(result["gripper_position"], dtype=np.float32).tolist()
-                            ],
-                            reward=reward_value,
-                            terminated=term_value,
-                            truncated=trunc_value,
-                            scene_objects=scene_objects,
-                        )
-                    )
-                    logging.debug("Episode %d step %d: trajectory appended", episode_index, step_index)
-
-                    # Track final status for summary export.
-                    final_reward = reward_value
-                    terminated = term_value
-                    truncated = trunc_value
-                    if term_value or trunc_value:
-                        # End the rollout early if the environment terminates.
-                        break
-
-                # Episode-level timing and video export.
-                episode_wall_time_s = time.perf_counter() - episode_start
-                video_path = output_dir / f"episode_{episode_index:02d}.mp4"
-                logging.info("Episode %d: writing video to %s", episode_index, video_path)
-                mediapy.write_video(video_path, frames, fps=15)
-
-                # Reset the policy server between episodes.
-                logging.info("Episode %d: sending reset", episode_index)
-                reset_response = client.reset()
-                final_objects = step_records[-1].scene_objects if step_records else {}
-                success_report = _maybe_infer_success(args.scene, final_objects)
-
-                # Assemble the per-episode summary that is written to
-                # `episode_XX.json`.
-                episode_summary = {
-                    "episode_index": episode_index,
-                    "prompt": instruction,
-                    "video_path": str(video_path),
-                    "steps_executed": len(step_records),
-                    "max_steps": max_steps,
-                    "terminated": terminated,
-                    "truncated": truncated,
-                    "final_reward": final_reward,
-                    "server_calls": client.server_calls,
-                    "server_time_s": server_time_s,
-                    "episode_wall_time_s": episode_wall_time_s,
-                    "avg_server_time_per_call_s": (
-                        server_time_s / client.server_calls if client.server_calls else None
-                    ),
-                    "reset_response": reset_response,
-                    "success_report": success_report,
-                    "server_metadata": client.metadata(),
-                    "trajectory": [record.__dict__ for record in step_records],
-                }
-                _dump_json(output_dir / f"episode_{episode_index:02d}.json", episode_summary)
-                all_episode_summaries.append(episode_summary)
-
-                logging.info(
-                    "Episode %d done: steps=%d wall=%.2fs server_calls=%d heuristic_success=%s",
-                    episode_index,
-                    len(step_records),
-                    episode_wall_time_s,
-                    client.server_calls,
-                    success_report.get("heuristic_success"),
-                )
-
-                # Reset per-episode counters while keeping the client alive.
-                client.server_calls = 0
-
-        # Top-level run summary across all episodes.
-        summary = {
-            "scene": args.scene,
-            "prompt": instruction,
-            "episodes": args.episodes,
-            "host": args.host,
-            "port": args.port,
-            "path": args.path,
-            "device": args.device,
-            "output_dir": str(output_dir),
-            "server_metadata": client.metadata(),
-            "episodes_summary": all_episode_summaries,
-        }
-        _dump_json(output_dir / "summary.json", summary)
-        # Also print the summary to stdout so the caller can capture it in logs.
-        print(json.dumps(summary, ensure_ascii=False, indent=2))
-    finally:
-        # Best-effort cleanup. Avoid masking the main error if cleanup fails.
-        try:
-            if client is not None:
-                client.client.close()
-        except Exception:
-            pass
-        if env is not None:
-            env.close()
-        simulation_app.close()
-
-
-if __name__ == "__main__":
-    # Keep the script-level logs readable. Per-step rollout details are still
-    # available via `DEBUG` if needed, but websocket / asyncio internals are
-    # usually too noisy for normal usage.
-    logging.basicConfig(level=logging.INFO)
-    logging.getLogger("websockets").setLevel(logging.WARNING)
-    logging.getLogger("asyncio").setLevel(logging.WARNING)
-    main()

examples/online_serving/dreamzero/molmospace_dreamzero_eval_demo.py

@@ -1,86 +0,0 @@
-from __future__ import annotations
-
-import argparse
-import os
-import sys
-from pathlib import Path
-
-os.environ.setdefault("MUJOCO_GL", "egl")
-os.environ.setdefault("PYOPENGL_PLATFORM", "egl")
-
-# Import base configs at module top level so the subclasses below are pickle-
-# resolvable (worker processes import this module fresh via __main__).
-from molmo_spaces.configs.policy_configs_baselines import (  # noqa: E402
-    DreamZeroPolicyConfig,
-)
-from molmo_spaces.evaluation.configs.evaluation_configs import (  # noqa: E402
-    DreamZeroPolicyEvalConfig,
-)
-
-
-# We only need to change the backend host and port to the vllm-host!
-class DreamZeroVllmOmniPolicyConfig(DreamZeroPolicyConfig):
-    remote_config: dict = dict(host="127.0.0.1", port=8000)
-
-
-class DreamZeroVllmOmniEvalConfig(DreamZeroPolicyEvalConfig):
-    policy_config: DreamZeroVllmOmniPolicyConfig = DreamZeroVllmOmniPolicyConfig()
-
-
-def main() -> int:
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--host", default="127.0.0.1")
-    parser.add_argument("--port", type=int, default=8000)
-    parser.add_argument(
-        "--benchmark_dir",
-        required=True,
-        help=(
-            "Path to a MolmoSpaces benchmark directory, for example "
-            "$MOLMOSPACES_BENCHMARK_DIR/20260327/ithor/FrankaCloseHardBench/"
-            "FrankaCloseHardBench_20260206_json_benchmark"
-        ),
-    )
-    parser.add_argument("--max_episodes", type=int, default=1)
-    parser.add_argument("--task_horizon_steps", type=int, default=80)
-    parser.add_argument(
-        "--output_dir",
-        required=True,
-        help="Directory to write evaluation outputs (created if missing).",
-    )
-    parser.add_argument("--episode_idx", type=int, default=None)
-    args = parser.parse_args()
-
-    policy_config = DreamZeroVllmOmniPolicyConfig(remote_config=dict(host=args.host, port=args.port))
-    DreamZeroVllmOmniPolicyConfig.model_fields["remote_config"].default = policy_config.remote_config
-    DreamZeroVllmOmniEvalConfig.model_fields["policy_config"].default = policy_config
-
-    # Import after env vars are set so MuJoCo picks EGL.
-    from molmo_spaces.evaluation import run_evaluation
-
-    cfg_cls = DreamZeroVllmOmniEvalConfig
-
-    output_dir = args.output_dir
-    Path(output_dir).mkdir(parents=True, exist_ok=True)
-
-    print(f"[eval] benchmark_dir={args.benchmark_dir}")
-    print(f"[eval] max_episodes={args.max_episodes} task_horizon_steps={args.task_horizon_steps}")
-    print(f"[eval] remote policy: ws://{args.host}:{args.port}/v1/realtime/robot/openpi")
-
-    results = run_evaluation(
-        eval_config_cls=cfg_cls,
-        benchmark_dir=Path(args.benchmark_dir),
-        max_episodes=args.max_episodes,
-        task_horizon_steps=args.task_horizon_steps,
-        num_workers=1,
-        use_wandb=False,
-        output_dir=output_dir,
-        episode_idx=args.episode_idx,
-    )
-
-    print(f"[eval] success={results.success_count}/{results.total_count} ({results.success_rate:.1%})")
-    print(f"[eval] output_dir={results.output_dir}")
-    return 0
-
-
-if __name__ == "__main__":
-    sys.exit(main())

examples/online_serving/dreamzero/openpi_client.py

@@ -1,311 +0,0 @@
-#!/usr/bin/env python3
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from __future__ import annotations
-
-import argparse
-import json
-import logging
-import sys
-import uuid
-from dataclasses import dataclass
-from pathlib import Path
-from typing import Any
-
-import numpy as np
-
-try:
-    import cv2
-except ImportError as exc:  # pragma: no cover - runtime dependency guard
-    raise ImportError("DreamZero OpenPI example requires `opencv-python`.") from exc
-
-try:
-    example_dir = str(Path(__file__).resolve().parent)
-    removed_path = False
-    if sys.path and sys.path[0] == example_dir:
-        sys.path.pop(0)
-        removed_path = True
-    try:
-        from openpi_client.websocket_client_policy import WebsocketClientPolicy
-    finally:
-        if removed_path:
-            sys.path.insert(0, example_dir)
-except ImportError as exc:  # pragma: no cover - runtime dependency guard
-    raise ImportError("DreamZero OpenPI example requires `openpi-client`.") from exc
-
-DEFAULT_HOST = "127.0.0.1"
-DEFAULT_PORT = 8000
-DEFAULT_PATH = "/v1/realtime/robot/openpi"
-DEFAULT_PROMPT = "Move the pan forward and use the brush in the middle of the plates to brush the inside of the pan"
-ACTION_HORIZON = 24
-DEFAULT_ACTION_DIM = 8
-RELATIVE_OFFSETS = [-23, -16, -8, 0]
-REPO_ROOT = Path(__file__).resolve().parents[3]
-ASSET_REPO_ID = "YangshenDeng/vllm-omni-dreamzero-assets"
-DEFAULT_VIDEO_DIR = REPO_ROOT / "outputs" / "dreamzero" / "assets"
-CAMERA_FILES = {
-    "observation/exterior_image_0_left": "exterior_image_1_left.mp4",
-    "observation/exterior_image_1_left": "exterior_image_2_left.mp4",
-    "observation/wrist_image_left": "wrist_image_left.mp4",
-}
-
-
-def _as_action_array(response: Any) -> np.ndarray:
-    if isinstance(response, dict) and response.get("type") == "error":
-        message = response.get("message", response)
-        raise RuntimeError(f"Inference failed: {message}")
-    return np.asarray(response, dtype=np.float32)
-
-
-@dataclass(frozen=True)
-class DreamZeroServerMetadata:
-    image_resolution: tuple[int, int]
-    n_external_cameras: int
-    needs_wrist_camera: bool
-    needs_stereo_camera: bool
-    needs_session_id: bool
-    action_space: str
-
-    @classmethod
-    def from_dict(cls, payload: dict[str, Any]) -> DreamZeroServerMetadata:
-        required_keys = (
-            "image_resolution",
-            "n_external_cameras",
-            "needs_wrist_camera",
-            "needs_stereo_camera",
-            "needs_session_id",
-            "action_space",
-        )
-        missing_keys = [key for key in required_keys if key not in payload]
-        if missing_keys:
-            raise ValueError(f"Missing DreamZero metadata keys: {missing_keys}")
-
-        image_resolution = payload["image_resolution"]
-        if not isinstance(image_resolution, (list, tuple)) or len(image_resolution) != 2:
-            raise ValueError(f"Invalid image_resolution: {image_resolution!r}")
-
-        return cls(
-            image_resolution=(int(image_resolution[0]), int(image_resolution[1])),
-            n_external_cameras=int(payload["n_external_cameras"]),
-            needs_wrist_camera=bool(payload["needs_wrist_camera"]),
-            needs_stereo_camera=bool(payload["needs_stereo_camera"]),
-            needs_session_id=bool(payload["needs_session_id"]),
-            action_space=str(payload["action_space"]),
-        )
-
-
-def _openpi_uri(host: str, port: int, path: str) -> str:
-    if host.startswith(("ws://", "wss://")):
-        return host
-    return f"ws://{host}:{port}{path}"
-
-
-def _close_policy_client(client: WebsocketClientPolicy) -> None:
-    ws = getattr(client, "_ws", None)
-    if ws is not None:
-        ws.close()
-
-
-def load_all_frames(video_path: Path) -> np.ndarray:
-    cap = cv2.VideoCapture(str(video_path))
-    frames = []
-    while True:
-        ok, frame = cap.read()
-        if not ok:
-            break
-        frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
-    cap.release()
-    if not frames:
-        raise RuntimeError(f"No frames loaded from {video_path}")
-    return np.stack(frames, axis=0)
-
-
-def load_camera_frames(video_dir: Path) -> dict[str, np.ndarray]:
-    camera_frames: dict[str, np.ndarray] = {}
-    for camera_key, file_name in CAMERA_FILES.items():
-        video_path = video_dir / file_name
-        if not video_path.exists():
-            raise FileNotFoundError(
-                f"Missing DreamZero example asset: {video_path}. "
-                "Download the example videos with: "
-                f"`hf download {ASSET_REPO_ID} --repo-type dataset --local-dir {video_dir}`"
-            )
-        camera_frames[camera_key] = load_all_frames(video_path)
-    return camera_frames
-
-
-def build_frame_schedule(total_frames: int, num_chunks: int) -> list[list[int]]:
-    chunks: list[list[int]] = []
-    current_frame = 23
-    for _ in range(num_chunks):
-        indices = [max(current_frame + offset, 0) for offset in RELATIVE_OFFSETS]
-        if indices[-1] >= total_frames:
-            break
-        chunks.append(indices)
-        current_frame += ACTION_HORIZON
-    return chunks
-
-
-def make_obs_from_video(
-    camera_frames: dict[str, np.ndarray],
-    frame_indices: list[int],
-    *,
-    prompt: str,
-    session_id: str,
-) -> dict[str, Any]:
-    obs: dict[str, Any] = {}
-    for camera_key, all_frames in camera_frames.items():
-        selected = all_frames[frame_indices]
-        obs[camera_key] = selected[0] if len(frame_indices) == 1 else selected
-
-    obs["observation/joint_position"] = np.zeros(7, dtype=np.float32)
-    obs["observation/cartesian_position"] = np.zeros(6, dtype=np.float32)
-    obs["observation/gripper_position"] = np.zeros(1, dtype=np.float32)
-    obs["prompt"] = prompt
-    obs["session_id"] = session_id
-    return obs
-
-
-def build_demo_observations(
-    camera_frames: dict[str, np.ndarray],
-    *,
-    prompt: str,
-    session_id: str,
-    num_chunks: int = 2,
-) -> list[dict[str, Any]]:
-    if num_chunks < 1:
-        raise ValueError("num_chunks must be at least 1")
-
-    total_frames = min(frames.shape[0] for frames in camera_frames.values())
-    observations = [
-        make_obs_from_video(
-            camera_frames,
-            [0],
-            prompt=prompt,
-            session_id=session_id,
-        )
-    ]
-    for indices in build_frame_schedule(total_frames, num_chunks - 1):
-        observations.append(
-            make_obs_from_video(
-                camera_frames,
-                indices,
-                prompt=prompt,
-                session_id=session_id,
-            )
-        )
-    return observations
-
-
-def validate_session_result(
-    result: dict[str, Any],
-    *,
-    expected_action_horizon: int = ACTION_HORIZON,
-    expected_action_dim: int = DEFAULT_ACTION_DIM,
-    expected_action_count: int | None = None,
-) -> None:
-    metadata = DreamZeroServerMetadata.from_dict(result["metadata"])
-    if metadata.image_resolution != (180, 320):
-        raise AssertionError(f"Unexpected image_resolution: {metadata.image_resolution}")
-    if metadata.n_external_cameras != 2:
-        raise AssertionError(f"Unexpected n_external_cameras: {metadata.n_external_cameras}")
-    if not metadata.needs_wrist_camera:
-        raise AssertionError("DreamZero example expects wrist camera metadata")
-    if metadata.action_space != "joint_position":
-        raise AssertionError(f"Unexpected action_space: {metadata.action_space}")
-
-    actions = result["actions"]
-    if expected_action_count is not None and len(actions) != expected_action_count:
-        raise AssertionError(f"Expected {expected_action_count} action tensors, got {len(actions)}")
-    if not actions:
-        raise AssertionError("Expected at least one action tensor")
-    for index, action in enumerate(actions):
-        if action.shape != (expected_action_horizon, expected_action_dim):
-            raise AssertionError(
-                f"Action {index} shape mismatch: expected "
-                f"{(expected_action_horizon, expected_action_dim)}, got {action.shape}"
-            )
-        if not np.isfinite(action).all():
-            raise AssertionError(f"Action {index} contains non-finite values")
-
-
-def run_policy_session(
-    *,
-    host: str = DEFAULT_HOST,
-    port: int = DEFAULT_PORT,
-    path: str = DEFAULT_PATH,
-    video_dir: Path = DEFAULT_VIDEO_DIR,
-    prompt: str = DEFAULT_PROMPT,
-    session_id: str | None = None,
-    num_chunks: int = 2,
-) -> dict[str, Any]:
-    session_id = session_id or str(uuid.uuid4())
-    camera_frames = load_camera_frames(video_dir)
-    observations = build_demo_observations(
-        camera_frames,
-        prompt=prompt,
-        session_id=session_id,
-        num_chunks=num_chunks,
-    )
-
-    uri = _openpi_uri(host, port, path)
-    logging.info("Connecting to %s", uri)
-    client = WebsocketClientPolicy(host=uri)
-    try:
-        metadata = dict(client.get_server_metadata())
-        actions = [_as_action_array(client.infer(obs)) for obs in observations]
-        return {
-            "metadata": metadata,
-            "actions": actions,
-            "session_id": session_id,
-        }
-    finally:
-        _close_policy_client(client)
-
-
-def format_action_summary(index: int, action: np.ndarray) -> str:
-    return (
-        f"Action {index}: shape={tuple(action.shape)} dtype={action.dtype} "
-        f"min={action.min():.6f} max={action.max():.6f}"
-    )
-
-
-def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description="DreamZero OpenPI client example with downloaded real videos.")
-    parser.add_argument("--host", default=DEFAULT_HOST)
-    parser.add_argument("--port", type=int, default=DEFAULT_PORT)
-    parser.add_argument("--path", default=DEFAULT_PATH)
-    parser.add_argument(
-        "--video-dir", type=Path, default=DEFAULT_VIDEO_DIR, help="Directory containing the three camera MP4 files."
-    )
-    parser.add_argument("--prompt", default=DEFAULT_PROMPT)
-    parser.add_argument("--session-id", default=None)
-    parser.add_argument("--num-chunks", type=int, default=2)
-    return parser.parse_args()
-
-
-def main() -> int:
-    args = parse_args()
-    logging.basicConfig(level=logging.INFO)
-
-    result = run_policy_session(
-        host=args.host,
-        port=args.port,
-        path=args.path,
-        video_dir=args.video_dir,
-        prompt=args.prompt,
-        session_id=args.session_id,
-        num_chunks=args.num_chunks,
-    )
-    validate_session_result(result, expected_action_count=args.num_chunks)
-
-    print("Server metadata:", json.dumps(result["metadata"], sort_keys=True))
-    for index, action in enumerate(result["actions"]):
-        print(format_action_summary(index, action))
-    print("Session ID:", result["session_id"])
-    return 0
-
-
-if __name__ == "__main__":
-    raise SystemExit(main())

examples/online_serving/dreamzero/run_server.sh

@@ -1,24 +0,0 @@
-#!/usr/bin/env bash
-set -euo pipefail
-
-MODEL="${MODEL:-GEAR-Dreams/DreamZero-DROID}"
-HOST="${HOST:-127.0.0.1}"
-PORT="${PORT:-8000}"
-DEPLOY_CONFIG="${DEPLOY_CONFIG:-vllm_omni/deploy/dreamzero_tp1_cfg2.yaml}"
-SERVED_MODEL_NAME="${SERVED_MODEL_NAME:-dreamzero-droid}"
-
-args=(
-  serve
-  "$MODEL"
-  --omni
-  --host "$HOST"
-  --port "$PORT"
-  --served-model-name "$SERVED_MODEL_NAME"
-  --deploy-config "$DEPLOY_CONFIG"
-  --enforce-eager
-  --disable-log-stats
-)
-
-ATTENTION_BACKEND="${ATTENTION_BACKEND:-torch}" \
-DIFFUSION_ATTENTION_BACKEND="${DIFFUSION_ATTENTION_BACKEND:-TORCH_SDPA}" \
-vllm "${args[@]}"

examples/online_serving/dynin_omni/README.md

@@ -22,14 +22,13 @@ export PYTHONPATH=<transformers_modules_path>:$PYTHONPATH
 export HF_MODULES_CACHE=<transformers_modules_path>
 ```
 
-Run from repository root. The 3-stage pipeline is resolved automatically from
-the model type (`vllm_omni/deploy/dynin_omni.yaml`), so no config path is
-required. Pass `--deploy-config` only to override the defaults:
+Run from repository root:
 
 ```bash
 vllm-omni serve snu-aidas/Dynin-Omni \
   --omni \
-  --port 8091
+  --port 8091 \
+  --stage-configs-path "$(pwd)/vllm_omni/model_executor/stage_configs/dynin_omni.yaml"
 ```
 
 If `vllm-omni` is not in PATH, run:
@@ -37,7 +36,8 @@ If `vllm-omni` is not in PATH, run:
 ```bash
 PYTHONPATH="$(pwd)" python -m vllm_omni.entrypoints.cli.main serve snu-aidas/Dynin-Omni \
   --omni \
-  --port 8091
+  --port 8091 \
+  --stage-configs-path "$(pwd)/vllm_omni/model_executor/stage_configs/dynin_omni.yaml"
 ```
 
 Wait until the server logs show both `All stages initialized successfully` and

examples/online_serving/lance/README.md

@@ -1,40 +0,0 @@
-# Lance — online serving
-
-OpenAI-compatible chat completions API for Lance, served by `vllm-omni`.
-
-## Start the server
-
-```bash
-bash examples/online_serving/lance/run_server.sh
-# or, with overrides
-MODEL=bytedance-research/Lance \
-DEPLOY_CONFIG=vllm_omni/deploy/lance.yaml \
-PORT=8091 \
-    bash examples/online_serving/lance/run_server.sh
-```
-
-The default deploy config is the single-stage `lance.yaml`; for `video_edit`
-or `text2video` workloads, point ``DEPLOY_CONFIG`` at a video-checkpoint
-config (or pass ``--model bytedance-research/Lance/Lance_3B_Video`` to
-`run_server.sh`).
-
-## Send a request
-
-```bash
-# text-to-image
-python examples/online_serving/lance/openai_chat_client.py \
-    --prompt "A cute corgi astronaut on the moon, cinematic" \
-    --modality text2img \
-    --output corgi.png
-
-# image edit
-python examples/online_serving/lance/openai_chat_client.py \
-    --prompt "Convert this into a vibrant cartoon-style illustration" \
-    --modality img2img \
-    --image-url path/to/photo.png \
-    --output edited.png
-```
-
-The client is shared with BAGEL — same OpenAI message format, same
-``modalities`` and ``num_inference_steps`` / ``seed`` / ``height`` /
-``width`` knobs.

examples/online_serving/lance/openai_chat_client.py

@@ -1,184 +0,0 @@
-#!/usr/bin/env python3
-"""
-Lance OpenAI-compatible chat client for image generation and multimodal tasks.
-
-Usage:
-    python examples/online_serving/lance/openai_chat_client.py --prompt "A cute cat" --output output.png
-    python examples/online_serving/lance/openai_chat_client.py --prompt "Describe this image" --image-url https://example.com/image.png
-"""
-
-import argparse
-import base64
-from pathlib import Path
-
-import requests
-
-
-def generate_image(
-    prompt: str,
-    server_url: str = "http://localhost:8091",
-    image_url: str | None = None,
-    height: int | None = None,
-    width: int | None = None,
-    steps: int | None = None,
-    seed: int | None = None,
-    negative_prompt: str | None = None,
-    modality: str = "text2img",  # "text2img" (default), "img2img", "img2text", "text2text"
-) -> bytes | str | None:
-    """Generate an image or text using the chat completions API.
-
-    Args:
-        prompt: Text description or prompt
-        server_url: Server URL
-        image_url: URL or path to input image (for img2img/img2text)
-        height: Image height in pixels
-        width: Image width in pixels
-        steps: Number of inference steps
-        seed: Random seed
-        negative_prompt: Negative prompt
-        modality: Task modality hint
-
-    Returns:
-        Image bytes (for image outputs) or Text string (for text outputs) or None if failed
-    """
-
-    # Construct Message Content
-    content = [{"type": "text", "text": f"<|im_start|>{prompt}<|im_end|>"}]
-
-    if image_url:
-        # Check if local file
-        if Path(image_url).exists():
-            with open(image_url, "rb") as f:
-                b64_data = base64.b64encode(f.read()).decode("utf-8")
-                final_image_url = f"data:image/jpeg;base64,{b64_data}"
-        else:
-            final_image_url = image_url
-
-        content.append({"type": "image_url", "image_url": {"url": final_image_url}})
-
-    messages = [{"role": "user", "content": content}]
-
-    # Build request payload with all parameters at top level
-    # Note: vLLM ignores "extra_body", so we put parameters directly in the payload
-    payload = {"messages": messages}
-
-    # Set output modalities at top level
-    if modality == "text2img" or modality == "img2img":
-        payload["modalities"] = ["image"]
-    elif modality == "img2text" or modality == "text2text":
-        payload["modalities"] = ["text"]
-
-    # Add generation parameters directly to payload
-    if height is not None:
-        payload["height"] = height
-    if width is not None:
-        payload["width"] = width
-    if steps is not None:
-        payload["num_inference_steps"] = steps
-    if seed is not None:
-        payload["seed"] = seed
-    if negative_prompt:
-        payload["negative_prompt"] = negative_prompt
-
-    # Send request
-    try:
-        print(f"Sending request to {server_url} with modality {modality}...")
-        response = requests.post(
-            f"{server_url}/v1/chat/completions",
-            headers={"Content-Type": "application/json"},
-            json=payload,
-            timeout=300,
-        )
-        response.raise_for_status()
-        data = response.json()
-
-        # Extract content - check ALL choices since server may return multiple
-        # (e.g., text in choices[0], image in choices[1])
-        choices = data.get("choices", [])
-
-        # First pass: look for image output in any choice
-        for choice in choices:
-            choice_content = choice.get("message", {}).get("content")
-
-            # Handle Image Output
-            if isinstance(choice_content, list) and len(choice_content) > 0:
-                first_item = choice_content[0]
-                if isinstance(first_item, dict) and "image_url" in first_item:
-                    img_url_str = first_item["image_url"].get("url", "")
-                    if img_url_str.startswith("data:image"):
-                        _, b64_data = img_url_str.split(",", 1)
-                        return base64.b64decode(b64_data)
-
-        # Second pass: look for text output if no image found
-        for choice in choices:
-            choice_content = choice.get("message", {}).get("content")
-            if isinstance(choice_content, str) and choice_content:
-                return choice_content
-
-        print(f"Unexpected response format: {choices}")
-        return None
-
-    except Exception as e:
-        print(f"Error: {e}")
-        return None
-
-
-def main():
-    parser = argparse.ArgumentParser(description="Bagel multimodal chat client")
-    parser.add_argument("--prompt", "-p", default="A cute cat", help="Text prompt")
-    parser.add_argument("--output", "-o", default="bagel_output.png", help="Output file (for image results)")
-    parser.add_argument("--server", "-s", default="http://localhost:8091", help="Server URL")
-
-    # Modality Control
-    parser.add_argument("--image-url", "-i", type=str, help="Input image URL or local path")
-    parser.add_argument(
-        "--modality",
-        "-m",
-        default="text2img",
-        choices=["text2img", "img2img", "img2text", "text2text"],
-        help="Task modality",
-    )
-
-    # Generation Params
-    parser.add_argument("--height", type=int, default=512, help="Image height")
-    parser.add_argument("--width", type=int, default=512, help="Image width")
-    parser.add_argument("--steps", type=int, default=25, help="Inference steps")
-    parser.add_argument("--seed", type=int, default=42, help="Random seed")
-    parser.add_argument("--negative", help="Negative prompt")
-
-    args = parser.parse_args()
-
-    print(f"Mode: {args.modality}")
-    if args.image_url:
-        print(f"Input Image: {args.image_url}")
-
-    result = generate_image(
-        prompt=args.prompt,
-        server_url=args.server,
-        image_url=args.image_url,
-        height=args.height,
-        width=args.width,
-        steps=args.steps,
-        seed=args.seed,
-        negative_prompt=args.negative,
-        modality=args.modality,
-    )
-
-    if result:
-        if isinstance(result, bytes):
-            # It's an image
-            output_path = Path(args.output)
-            output_path.write_bytes(result)
-            print(f"Image saved to: {output_path}")
-            print(f"Size: {len(result) / 1024:.1f} KB")
-        elif isinstance(result, str):
-            # It's text
-            print("Response:")
-            print(result)
-    else:
-        print("Failed to generate response")
-        exit(1)
-
-
-if __name__ == "__main__":
-    main()

examples/online_serving/lance/run_server.sh

@@ -1,15 +0,0 @@
-#!/bin/bash
-# Lance online serving startup script
-
-MODEL="${MODEL:-bytedance-research/Lance}"
-PORT="${PORT:-8091}"
-DEPLOY_CONFIG="${DEPLOY_CONFIG:-vllm_omni/deploy/lance.yaml}"
-
-echo "Starting Lance server..."
-echo "Model:         $MODEL"
-echo "Deploy config: $DEPLOY_CONFIG"
-echo "Port:          $PORT"
-
-vllm serve "$MODEL" --omni \
-    --deploy-config "$DEPLOY_CONFIG" \
-    --port "$PORT"

examples/online_serving/mimo_audio/openai_chat_completion_client_for_multimodal_generation.py

@@ -12,8 +12,7 @@ from typing import Any
 import requests
 from openai import OpenAI
 from vllm.assets.audio import AudioAsset
-
-from vllm_omni.utils.tracking_parser import TrackingArgumentParser
+from vllm.utils.argparse_utils import FlexibleArgumentParser
 
 # Modify OpenAI's API key and API base to use vLLM's API server.
 openai_api_key = "EMPTY"
@@ -494,7 +493,7 @@ def run_multimodal_generation(args) -> None:
 
 
 def parse_args():
-    parser = TrackingArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
+    parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
     parser.add_argument(
         "--query-type",
         "-q",

examples/online_serving/ming_flash_omni/README.md

@@ -10,7 +10,6 @@ Please refer to [README.md](../../../README.md)
 |------|---------------|--------|
 | Thinker + Talker (omni-speech, default) | `vllm serve ... --omni` | Text + Audio |
 | Thinker only (multimodal understanding) | `vllm serve ... --omni --deploy-config vllm_omni/deploy/ming_flash_omni_thinker_only.yaml` | Text |
-| Thinker + Imagegen (text-to-image / img2img) | `vllm serve ... --omni --deploy-config vllm_omni/deploy/ming_flash_omni_image.yaml` | Image |
 
 For standalone TTS (talker only), see the [Ming-flash-omni-TTS section in the Text-To-Speech hub](../text_to_speech/README.md#ming-flash-omni-tts).
 
@@ -57,91 +56,8 @@ bash run_curl_multimodal_generation.sh use_image
 bash run_curl_multimodal_generation.sh use_audio
 bash run_curl_multimodal_generation.sh use_video
 bash run_curl_multimodal_generation.sh use_mixed_modalities
-bash run_curl_multimodal_generation.sh use_image_gen
 ```
 
-## Image generation (text-to-image)
-
-Ming-flash-omni-2.0 also exposes an image-generation (diffusion) stage.
-Launch with the image deploy YAML, which adds an image-gen stage behind
-the thinker:
-
-```bash
-vllm serve Jonathan1909/Ming-flash-omni-2.0 --omni \
-    --deploy-config vllm_omni/deploy/ming_flash_omni_image.yaml \
-    --stage-init-timeout 1800 \
-    --init-timeout 1800 \
-    --port 8091
-```
-
-Then request image output by passing `"modalities": ["image"]`:
-
-```bash
-curl -s http://127.0.0.1:8091/v1/chat/completions \
-    -H "Content-Type: application/json" \
-    -d '{
-      "model": "Jonathan1909/Ming-flash-omni-2.0",
-      "messages": [{"role": "user", "content": "Please draw a cute cat."}],
-      "modalities": ["image"]
-    }' | jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' | base64 -d > ming_imagegen.png
-```
-
-### Optional knobs
-
-Pass image-gen overrides via the diffusion-stage `sampling_params_list[1].extra_args.image_gen`:
-
-| Key | Default | Description |
-| --- | --- | --- |
-| `height` / `width` | from config (1024) | Output resolution (multiples of `vae_scale_factor * 2`, currently 16). |
-| `steps` | 30 | Number of FlowMatchEuler denoise steps. |
-| `cfg` | 2.0 | Classifier-free guidance scale. |
-| `seed` | 42 | Per-request RNG seed (deterministic when ByT5 is also seed-stable). |
-| `byte5_text` | (auto) | Override the glyph text for ByT5 enhancement; raw strings are auto-wrapped to Ming's `Text "...". ` format. |
-| `negative_prompt` | empty | Real CFG negative conditioning (must be set on **stage-0 thinker** `sampling_params` so `expand_cfg_prompts` spawns the companion). |
-
-Example with all the knobs:
-
-```bash
-curl http://127.0.0.1:8091/v1/chat/completions \
-    -H "Content-Type: application/json" \
-    -d '{
-      "model": "Jonathan1909/Ming-flash-omni-2.0",
-      "modalities": ["image"],
-      "sampling_params_list": [
-        {"temperature":0.4,"top_p":0.9,"top_k":1,"max_tokens":1,"seed":42,
-         "extra_args":{"image_gen":{"negative_prompt":"ugly, blurry, distorted"}}},
-        {"seed":42,"extra_args":{"image_gen":{
-            "steps":6,"cfg":1.5,"height":512,"width":512,"seed":123,
-            "byte5_text":["理解与生成统一"]}}}
-      ],
-      "messages": [{"role": "user", "content": "Draw a poster."}]
-    }' | jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' | base64 -d > ming_imagegen_knobs.png
-```
-
-### img2img (reference image)
-
-Add an `image_url` content part to the user message; the parser routes it
-into the diffusion stage as `extra[reference_image]`:
-
-```jsonc
-"messages": [{
-  "role": "user",
-  "content": [
-    {"type": "text", "text": "Change the background to a sandy beach at sunset."},
-    {"type": "image_url", "image_url": {"url": "data:image/png;base64,<base64>"}}
-  ]
-}]
-```
-
-### GPU layout
-
-The shipped `ming_flash_omni_image.yaml` allocates the thinker on GPUs 0–3
-(TP=4) and the diffusion stage on GPU 4 (TP=1). Copy the YAML and edit
-`devices` per stage to relocate; with fewer GPUs available, drop
-the thinker TP to 2 and run the diffusion stage on a free card. Image-gen
-warmup takes roughly an extra 30–60 s on top of the thinker — set
-`--stage-init-timeout 1800` if the default 300 s is too tight.
-
 ## Modality control
 
 | `modalities` | Server config | Output |
@@ -149,7 +65,6 @@ warmup takes roughly an extra 30–60 s on top of the thinker — set
 | `["text"]` or omitted | Thinker only | Text |
 | `["audio"]` | Thinker + Talker | Audio (speech) |
 | `["text", "audio"]` | Thinker + Talker | Text + Audio |
-| `["image"]` | Thinker + Imagegen (image deploy YAML) | Image (PNG, base64 in `choices[0].message.content`) |
 
 For ready-to-copy curl examples (text / audio / multimodal input, SSE
 streaming, reasoning mode), see the recipe at